def partition_advection(data, lons, lev=150):
#Do vv terms
vv_trans_dy = -86400. * gr.ddy( (data.vcomp_sq - data.vcomp * data.vcomp).sel(pfull=lev) , uv=True)
v = data.vcomp.sel(pfull=lev).load() # v
v_dy = -86400. * gr.ddy( v ) # dudy
v_dvdy_zav = v.sel(lon=lons).mean('lon') * v_dy.sel(lon=lons).mean('lon') # [v][dudy]
v_dvdy_stat = (v * v_dy).sel(lon=lons).mean('lon') - v_dvdy_zav # v*dudy* = [vdudy] - [v][dudy]
data['vv_trans_dy'] = (('xofyear','lat'), vv_trans_dy.sel(lon=lons).mean('lon'))
data['v_dvdy_stat'] = (('xofyear','lat'), v_dvdy_stat)
data['v_dvdy_zav'] = (('xofyear','lat'), v_dvdy_zav )
print('vv terms done')
#Do vw terms
vw_trans_dp = -86400. * gr.ddp( (data.vcomp_omega - data.vcomp * data.omega).sel(lon=lons).mean('lon') )
w = data.omega.sel(pfull=lev).load() # w
v_dp = -86400. * (gr.ddp(data.vcomp)).sel(pfull=lev) # dudp
w_dvdp_zav = w.sel(lon=lons).mean('lon') * v_dp.sel(lon=lons).mean('lon')
w_dvdp_stat = (w * v_dp).sel(lon=lons).mean('lon') - w_dvdp_zav
data['vw_trans_dp'] = (('xofyear','lat'), vw_trans_dp.sel(pfull=lev))
data['w_dvdp_stat'] = (('xofyear','lat'), w_dvdp_stat)
data['w_dvdp_zav'] = (('xofyear','lat'), w_dvdp_zav )
print('vw terms done')
def fig_9(run, ax, pentad=40, lev=850.):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
heating = (data.dt_tg_condensation + data.dt_tg_convection +
data.dt_tg_diffusion + data.tdt_rad).mean('lon') * 86400.
rho = data.pfull * 100. / mc.rdgas / data.temp
dtdy = gr.ddy(data.temp.mean('lon'), vector=False)
dtdp = gr.ddp(data.temp.mean('lon'))
dtdt = gr.ddt(data.temp.mean('lon')) * 86400.
vt_eddy = data.vcomp_temp.mean(
'lon') - data.vcomp.mean('lon') * data.temp.mean('lon')
wt_eddy = data.omega_temp.mean(
'lon') - data.omega.mean('lon') * data.temp.mean('lon')
vdtdy_mean = -1. * data.vcomp.mean('lon') * dtdy * 86400.
wdtdp_mean = -1. * data.omega.mean('lon') * dtdp * 86400.
expansion_term = (data.omega / rho / mc.cp_air).mean('lon') * 86400.
div_vt_eddy = -1. * gr.ddy(vt_eddy, vector=True) * 86400.
div_wt_eddy = -1. * gr.ddp(wt_eddy) * 86400.
lats = [
data.lat[i] for i in range(len(data.lat))
if data.lat[i] >= -40. and data.lat[i] <= 40.
]
heating.sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax, color='k')
dtdt.sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax,
color='k',
linestyle=':')
#vdtdy_mean.sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax, color='C1')
#(wdtdp_mean + heating).sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax, color='C2')
expansion_term.sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax,
color='C1')
(vdtdy_mean + wdtdp_mean + expansion_term).sel(pfull=lev,
xofyear=pentad,
lat=lats).plot(
ax=ax,
color='k',
linestyle='--')
#vdtdy_mean.sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax, color='C2', linestyle='--')
#wdtdp_mean.sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax, color='C3', linestyle='--')
#(vdtdy_mean + wdtdp_mean + expansion_term).sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='--')
#(-dthetadt + vdthetady_mean + wdthetadp_mean + heating_theta + div_vt_eddy + div_wt_eddy).sel(pfull=lev, xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle=':')
(div_vt_eddy + div_wt_eddy).sel(pfull=lev, xofyear=pentad,
lat=lats).plot(ax=ax,
color='k',
linestyle='-.')
ax.set_title('')
ax.set_xlabel('')
#ax.set_ylim(-0.25,0.25)
ax.set_ylim(-5., 5.)
ax.set_xlim(-30., 30.)
ax.grid(True, linestyle=':')
ax.set_ylabel('Heating rate, K/day')
def get_budget_terms(energy, uenergy, venergy, wenergy, diab):
u_e_eddy = uenergy - energy * data.ucomp
v_e_eddy = venergy - energy * data.vcomp
w_e_eddy = wenergy - energy * data.omega
eddy_conv = -1. * (gr.ddx(u_e_eddy) + gr.ddy(v_e_eddy) +
gr.ddp(w_e_eddy))
vert_adv = -1. * (data.omega * gr.ddp(energy))
horiz_adv = -1. * (data.ucomp * gr.ddx(energy) +
data.vcomp * gr.ddy(energy, vector=False))
denergydt = gr.ddt(energy) * 20.
total = (diab + horiz_adv + vert_adv + eddy_conv) * 10.
return diab, horiz_adv, vert_adv, eddy_conv, denergydt, total
def gross_stability(run, moist=False, i=1):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
convTtotheta = (1000. / data.pfull)**mc.kappa
theta = data.temp * convTtotheta
theta_equiv = (data.temp + mc.L / mc.cp_air * data.sphum /
(1 - data.sphum)) * convTtotheta
dthetady = gr.ddy(theta.mean('lon'), vector=False)
dthetadp = gr.ddp(theta.mean('lon'))
dthetady_equiv = gr.ddy(theta_equiv.mean('lon'), vector=False)
dthetadp_equiv = gr.ddp(theta_equiv.mean('lon'))
vdthetady_mean = data.vcomp.mean('lon') * dthetady
wdthetadp_mean = data.omega.mean('lon') * dthetadp
vdthetady_mean_equiv = data.vcomp.mean('lon') * dthetady_equiv
wdthetadp_mean_equiv = data.omega.mean('lon') * dthetadp_equiv
def column_int(var_in):
var_int = mc.cp_air * var_in.sum('pfull') * 5000. / mc.grav
return var_int
div_vt_mean_int = -1. * column_int(wdthetadp_mean + vdthetady_mean)
div_vt_mean_int_equiv = -1. * column_int(wdthetadp_mean_equiv +
vdthetady_mean_equiv)
vt_mean_int = column_int(data.vcomp.mean('lon') * theta.mean('lon'))
vt_mean_int_equiv = column_int(
data.vcomp.mean('lon') * theta_equiv.mean('lon'))
#vt_mean_int.plot.contourf(x='xofyear', y='lat')
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
#psi /= 1.e9
psi = np.abs(psi).max('pfull')
#plt.figure(2)
#psi.plot.contourf(x='xofyear', y='lat')
gross_stab = (2. * np.pi * mc.a * np.cos(data.lat * np.pi / 180.) *
np.abs(vt_mean_int)) / psi
gross_moist_stab = (2. * np.pi * mc.a * np.cos(data.lat * np.pi / 180.) *
np.abs(vt_mean_int_equiv)) / psi
plt.figure(i)
gross_moist_stab.plot.contourf(x='xofyear',
y='lat',
levels=np.arange(0., 2.e5, 2.e4))
def incompressibility(run, lev=850.):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
rho = data.pfull * 100. / data.temp / mc.rdgas
rho = data.pfull * 100. / mc.rdgas / data.temp / (
1 + (mc.rvgas - mc.rdgas) / mc.rdgas * data.sphum)
dudx = gr.ddx(data.ucomp * rho)
dvdy = gr.ddy(data.vcomp * rho)
dwdp = gr.ddp(data.omega * rho)
divU = (dudx + dvdy + dwdp).mean('lon').sel(pfull=lev) * 86400.
rcParams['figure.figsize'] = 5.5, 3.
rcParams['font.size'] = 14
fig, ax = plt.subplots(1, sharex=True)
divU.plot.contourf(ax=ax, x='xofyear', y='lat', add_labels=False)
ax.set_xlabel('Pentad')
ax.set_ylabel('Latitude')
plt.subplots_adjust(left=0.15, right=0.95, top=0.97, bottom=0.2)
plt.savefig(plot_dir + 'incompressibility_rho_' + run + '.pdf',
format='pdf')
plt.close()
def fig_8(run, ax, pentad=45, rotfac=1., dayfac=1.):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc')
heating = data.dt_tg_condensation + data.dt_tg_convection + data.dt_tg_diffusion + data.tdt_rad
convTtotheta=(1000./data.pfull)**mc.kappa
theta = data.temp*convTtotheta
heating_theta = heating*convTtotheta
dthetady = gr.ddy(theta.mean('lon'), vector=False)
dthetadp = gr.ddp(theta.mean('lon'))
dthetadt = gr.ddt(theta.mean('lon'))
vdthetady_mean = data.vcomp.mean('lon') * dthetady
wdthetadp_mean = data.omega.mean('lon') * dthetadp
adv_heating = -1. *(vdthetady_mean + wdthetadp_mean)
sinphi = np.sin(data.lat * np.pi/180.)
cosphi = np.cos(data.lat * np.pi/180.)
w_850 = data.omega.sel(pfull=850.).mean('lon')
w_max, w_max_lat = get_latmax(-1.*w_850)
theta_850 = theta.sel(pfull=850.).mean('lon')
theta_max, theta_max_lat = get_latmax(theta_850)
sinphimax = np.sin(theta_max_lat * np.pi/180.)
theta_m = theta_max - 300.*(mc.omega * rotfac)**2.*mc.a**2./(2.*mc.grav*14000.) * (sinphi**2.-sinphimax**2.)**2./cosphi**2.
def adj_theta(theta_in, adj_by, dayfac=dayfac):
if 'lon' in adj_by.coords:
return theta_in + adj_by.sel(pfull=850.).mean('lon') * 86400. * dayfac
else:
return theta_in + adj_by.sel(pfull=850.) * 86400. * dayfac
theta_rc = adj_theta(theta_850, heating_theta)
theta_adv = adj_theta(theta_850, adv_heating)
theta_net = adj_theta(theta_850, dthetadt, dayfac=20.)
lats = [data.lat[i] for i in range(len(data.lat)) if data.lat[i] >= -60. and data.lat[i] <= 60.]
theta_m.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='0.7')
theta_rc.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='--')
theta_adv.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='-.')
theta_850.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C0')
theta_net.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C2')
ax.plot([w_max_lat.sel(xofyear=pentad),w_max_lat.sel(xofyear=pentad)], [295,315], color='0.7', linestyle=':')
ax.set_title('')
ax.set_xlabel('')
ax.set_ylim(295.,315.)
ax.set_xlim(-35.,35.)
ax.grid(True,linestyle=':')
ax.set_ylabel('$\Theta$, K')
def fig_9(run, ax, pentad=40):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc')
convTtotheta=(1000./data.pfull)**mc.kappa
heating = data.dt_tg_condensation + data.dt_tg_convection + data.dt_tg_diffusion + data.tdt_rad
rho = data.pfull*100./mc.rdgas/data.temp
heating_theta = (heating*convTtotheta).mean('lon')
theta = data.temp*convTtotheta
dthetady = gr.ddy(theta.mean('lon'), vector=False)
dthetadp = gr.ddp(theta.mean('lon'))
dthetadt = gr.ddt(theta.mean('lon'))
vcomp_theta = data.vcomp_temp*convTtotheta
vcomp_theta_eddy = vcomp_theta.mean('lon') - data.vcomp.mean('lon')*theta.mean('lon')
vdthetady_mean = data.vcomp.mean('lon') * dthetady
wdthetadp_mean = data.omega.mean('lon') * dthetadp
def column_int(var_in):
var_int = mc.cp_air * var_in.sum('pfull')*5000./mc.grav
return var_int
vdtdy_mean_int = -1. * column_int(vdthetady_mean)
wdtdp_mean_int = -1. * column_int(wdthetadp_mean)
vt_eddy_int = -1. * column_int(vcomp_theta_eddy)
div_vt_eddy_int = gr.ddy(vt_eddy_int, vector=True)
w_850 = data.omega.sel(pfull=850.).mean('lon')
w_max, w_max_lat = get_latmax(-1.*w_850)
heating_theta_int = column_int(heating_theta)
dthetadt_int = column_int(dthetadt)
lats = [data.lat[i] for i in range(len(data.lat)) if data.lat[i] >= -40. and data.lat[i] <= 40.]
dthetadt_int.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C2')
heating_theta_int.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='--')
#wdtdp_mean_int.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C1')
#vdtdy_mean_int.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='--')
(vdtdy_mean_int + wdtdp_mean_int).sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='-.')
#(-dthetadt_int + vdtdy_mean_int + wdtdp_mean_int + heating_theta_int + div_vt_eddy_int).sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle=':')
#(vdtdy_mean_int + wdtdp_mean_int + heating_theta_int + div_vt_eddy_int).sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle=':')
div_vt_eddy_int.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle=':')
ax.plot([w_max_lat.sel(xofyear=pentad),w_max_lat.sel(xofyear=pentad)], [-500.,500.], color='0.7', linestyle=':')
ax.set_title('')
ax.set_xlabel('')
ax.set_ylim(-500.,500.)
ax.set_xlim(-35.,35.)
ax.grid(True,linestyle=':')
ax.set_ylabel('Heating rate, W/m$^2$')
def fig_9_moist(run, ax, pentad=40):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
convTtotheta = (1000. / data.pfull)**mc.kappa
heating = data.dt_tg_condensation + data.dt_tg_convection + data.dt_tg_diffusion + data.tdt_rad
hum_tend = data.dt_qg_condensation + data.dt_qg_convection + data.dt_qg_diffusion
heating_theta_equiv = ((heating + mc.L / mc.cp_air * hum_tend /
(1 - data.sphum)**2.) * convTtotheta).mean('lon')
theta_equiv = (data.temp + mc.L / mc.cp_air * data.sphum /
(1 - data.sphum)) * convTtotheta
dthetady_equiv = gr.ddy(theta_equiv.mean('lon'), vector=False)
dthetadp_equiv = gr.ddp(theta_equiv.mean('lon'))
dthetadt_equiv = gr.ddt(theta_equiv.mean('lon'))
vcomp_theta_equiv = (data.vcomp_temp +
mc.L / mc.cp_air * data.sphum_v) * convTtotheta
vcomp_theta_eddy_equiv = vcomp_theta_equiv.mean(
'lon') - data.vcomp.mean('lon') * theta_equiv.mean('lon')
vdthetady_mean_equiv = data.vcomp.mean('lon') * dthetady_equiv
wdthetadp_mean_equiv = data.omega.mean('lon') * dthetadp_equiv
def column_int(var_in):
var_int = mc.cp_air * var_in.sum('pfull') * 5000. / mc.grav
return var_int
vdtdy_mean_int_equiv = -1. * column_int(vdthetady_mean_equiv)
wdtdp_mean_int_equiv = -1. * column_int(wdthetadp_mean_equiv)
vt_eddy_int_equiv = -1. * column_int(vcomp_theta_eddy_equiv)
div_vt_eddy_int_equiv = gr.ddy(vt_eddy_int_equiv, vector=True)
heating_theta_int_equiv = column_int(heating_theta_equiv)
dthetadt_int_equiv = column_int(dthetadt_equiv)
lats = [
data.lat[i] for i in range(len(data.lat))
if data.lat[i] >= -60. and data.lat[i] <= 60.
]
w_850 = data.omega.sel(pfull=850.).mean('lon')
w_max, w_max_lat = get_latmax(-1. * w_850)
dthetadt_int_equiv.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C2')
heating_theta_int_equiv.sel(xofyear=pentad, lat=lats).plot(ax=ax,
color='k',
linestyle='--')
(vdtdy_mean_int_equiv + wdtdp_mean_int_equiv).sel(
xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='-.')
#(vdtdy_mean_int_equiv + wdtdp_mean_int_equiv + heating_theta_int_equiv + div_vt_eddy_int_equiv).sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle=':')
div_vt_eddy_int_equiv.sel(xofyear=pentad, lat=lats).plot(ax=ax,
color='k',
linestyle=':')
ax.plot([w_max_lat.sel(xofyear=pentad),
w_max_lat.sel(xofyear=pentad)], [-250., 250.],
color='0.7',
linestyle=':')
ax.set_title('')
ax.set_xlabel('')
ax.set_ylim(-250., 250.)
ax.set_xlim(-35., 35.)
ax.grid(True, linestyle=':')
ax.set_ylabel('Heating rate, W/m$^2$')
def heat_budg_hm_fluxform(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')
uT_dx = -86400. * gr.ddx(data.ucomp_temp.sel(pfull=lev)).mean('lon')
vT_dy = -86400. * gr.ddy(data.vcomp_temp.sel(pfull=lev)).mean('lon')
wT_dp = -86400. * gr.ddp(data.omega_temp).sel(pfull=lev).mean('lon')
uT_eddy = data.ucomp_temp - data.ucomp * data.temp
vT_eddy = data.vcomp_temp - data.vcomp * data.temp
wT_eddy = data.omega_temp - data.omega * data.temp
#uT_dx = -86400. * (data.ucomp * gr.ddx(data.temp)).sel(pfull=lev).mean('lon')
#vT_dy = -86400. * (data.vcomp * gr.ddy(data.temp, vector=False)).sel(pfull=lev).mean('lon')
#wT_dp = -86400. * (data.omega * gr.ddp(data.temp)).sel(pfull=lev).mean('lon')
uT_dx_eddy = -86400. * gr.ddx(uT_eddy.sel(pfull=lev)).mean('lon')
vT_dy_eddy = -86400. * gr.ddy(vT_eddy.sel(pfull=lev)).mean('lon')
wT_dp_eddy = -86400. * gr.ddp(wT_eddy).sel(pfull=lev).mean('lon')
tdt_rad = data.tdt_rad.sel(pfull=lev).mean('lon')*86400.
tdt_conv = data.dt_tg_convection.sel(pfull=lev).mean('lon')*86400.
tdt_cond = data.dt_tg_condensation.sel(pfull=lev).mean('lon')*86400.
tdt_diff = data.dt_tg_diffusion.sel(pfull=lev).mean('lon')*86400.
diabatic = tdt_rad + tdt_conv + tdt_cond + tdt_diff
eddies = uT_dx_eddy + vT_dy_eddy + wT_dp_eddy
rho = data.pfull*100./data.temp/mc.rdgas
asc_cool = (data.omega /(mc.cp_air * rho)).sel(pfull=lev).mean('lon') *86400.
heat_sum = uT_dx + vT_dy + wT_dp + diabatic + asc_cool #+ eddies
Tdt = gr.ddt(data.temp.sel(pfull=lev)).mean('lon') *86400.
horiz_adv = uT_dx + vT_dy
vertical_term = wT_dp + asc_cool
levels = np.arange(-100,101.1,10.)
#levels = np.arange(-20,21.,2.)
#levels = np.arange(-3,3.1,0.2)
#levels = np.arange(-1,1.1,0.1)
# Nine subplots
fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2, 3, sharex='col', sharey='row')
plt.set_cmap('RdBu_r')
plot_vars = [horiz_adv, vertical_term, diabatic, heat_sum, Tdt, eddies]
axes = [ax1,ax2,ax3,ax4,ax5,ax6]
labels = ['a)','b)','c)','d)','e)','f)']
for i in range(6):
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)
axes[i].set_ylim(-60,60)
axes[i].grid(True,linestyle=':')
axes[i].set_yticks(np.arange(-60.,61.,30.))
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('Horizontal advection', fontsize=17)
ax2.set_title('Vertical advection and expansion', fontsize=17)
ax3.set_title('Diabatic', fontsize=17)
ax4.set_title('Residual', fontsize=17)
ax5.set_title('Temperature tendency', fontsize=17)
ax6.set_title('Eddy convergence', 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)
figname = 'heat_budg_fluxform_' +run+ '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def vort_eq(run,
month,
filename='plev_daily',
period_fac=1.,
rot_fac=1.,
do_ss=False,
day_fac=1.):
'''Evaluate terms in the vorticity budget using daily data from simulations. RG 2/11/17
Inputs: run = run name
month = month to evaluate
filename = name of file to look for, default is pressure interpolated daily data
period_fac = ratio of the orbital period to Earth's
rot_fac = ratio of the rotation rate to Earth's
do_ss = evaluate for a steady state simulation, do not average daily data, average all saved data together
'''
#Load in dataset
try:
name_temp = '/scratch/rg419/Data_moist/' + run + '/run%03d/' + filename + '.nc'
name = name_temp % month
#read data into xarr
data = xr.open_dataset(name, decode_times=False)
except:
name_temp = '/scratch/rg419/Data_moist/' + run + '/run%04d/' + filename + '.nc'
name = name_temp % month
#read data into xarr
data = xr.open_dataset(name, decode_times=False)
# Calculate planetary vorticity
omega = 7.2921150e-5 * rot_fac
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
# Take gradients of vorticity
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)
})
# If a steady state run, leave time dimension in to average out later. Otherwise take pentad means for climatology
if do_ss:
dsout = ds
else:
ds.coords['xofyear'] = np.mod(ds.time / day_fac - 1.,
360. * period_fac) // 5 + 1.
dsout = ds.groupby('xofyear').mean(('time'))
try:
fileout = '/scratch/rg419/Data_moist/' + run + '/run%03d/vort_eq.nc'
fileout = fileout % month
dsout.to_netcdf(path=fileout)
except:
fileout = '/scratch/rg419/Data_moist/' + run + '/run%04d/vort_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.sel(pfull=lev)**2)) # <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).sel(lon=lons).mean(
'lon') # [u]dudx*
u_dudx_cross2 = (u_ed * u_dx.mean('lon')).sel(lon=lons).mean(
'lon') # u*[dudx]
u_dudx_stat = (u_ed * u_dx_ed).sel(lon=lons).mean('lon') # u*dudx*
data['uu_trans_dx'] = (('pentad', 'lat'),
uu_trans_dx.sel(lon=lons).mean('lon'))
data['u_dudx_cross1'] = (('pentad', 'lat'), u_dudx_cross1)
data['u_dudx_cross2'] = (('pentad', 'lat'), u_dudx_cross2)
data['u_dudx_stat'] = (('pentad', 'lat'), u_dudx_stat)
data['u_dudx_zav'] = (('pentad', 'lat'), u_dudx_zav)
print('uu terms done')
#Next do uv terms
uv_trans_dy = -86400. * gr.ddy(
(data.ucomp_vcomp - data.ucomp.sel(pfull=lev) * data.vcomp), uv=True)
v = data.vcomp.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).sel(lon=lons).mean(
'lon') # [v]dudy*
v_dudy_cross2 = (v_ed * u_dy.mean('lon')).sel(lon=lons).mean(
'lon') # v*[dudy]
v_dudy_stat = (v_ed * u_dy_ed).sel(lon=lons).mean('lon') # v*dudy*
data['uv_trans_dy'] = (('pentad', 'lat'),
uv_trans_dy.sel(lon=lons).mean('lon'))
data['v_dudy_cross1'] = (('pentad', 'lat'), v_dudy_cross1)
data['v_dudy_cross2'] = (('pentad', 'lat'), v_dudy_cross2)
data['v_dudy_stat'] = (('pentad', 'lat'), v_dudy_stat)
data['v_dudy_zav'] = (('pentad', '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_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).sel(lon=lons).mean(
'lon') # [w]dudp*
w_dudp_cross2 = (w_ed * u_dp.mean('lon')).sel(lon=lons).mean(
'lon') # w*[dudp]
w_dudp_stat = (w_ed * u_dp_ed).sel(lon=lons).mean('lon') # w*dudp*
data['uw_trans_dp'] = (('pentad', 'lat'), uw_trans_dp.sel(pfull=lev))
data['w_dudp_cross1'] = (('pentad', 'lat'), w_dudp_cross1)
data['w_dudp_cross2'] = (('pentad', 'lat'), w_dudp_cross2)
data['w_dudp_stat'] = (('pentad', 'lat'), w_dudp_stat)
data['w_dudp_zav'] = (('pentad', 'lat'), w_dudp_zav)
print('uw terms done')
def bs_mombudg(run, ax, pentad=45, lev=150., dayfac=3., rotfac=1.):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
damping = data.dt_ug_diffusion # + data.tdt_diss_rdamp
dudy = gr.ddy(data.ucomp.mean('lon'))
duvdy = gr.ddy(data.ucomp_vcomp.mean('lon'), uv=True)
dudp = gr.ddp(data.ucomp.mean('lon'))
duwdp = gr.ddp(data.ucomp_omega.mean('lon'))
dudt = gr.ddt(data.ucomp.mean('lon'))
vdudy_mean = data.vcomp.mean('lon') * dudy
wdudp_mean = data.omega.mean('lon') * dudp
adv_tend = -1. * (vdudy_mean + wdudp_mean)
vert_adv = -1. * wdudp_mean
horiz_adv = -1. * vdudy_mean
eddy_tend = -1. * duvdy - 1. * duwdp - adv_tend
vert_eddy = -1. * duwdp - vert_adv
horiz_eddy = -1. * duvdy - horiz_adv
sinphi = np.sin(data.lat * np.pi / 180.)
cosphi = np.cos(data.lat * np.pi / 180.)
coriolis = 2. * rotfac * mc.omega * sinphi * data.vcomp
def get_latmax(data_in):
data_max = data_in.max('lat')
data_max_lat = data_in.lat.values[data_in.argmax('lat').values]
data_max_lat = xr.DataArray(data_max_lat,
coords=[data_in.xofyear],
dims=['xofyear'])
return data_max, data_max_lat
w_850 = data.omega.sel(pfull=850.).mean('lon')
w_max, w_max_lat = get_latmax(-1. * w_850)
convTtotheta = (1000. / data.pfull)**mc.kappa
theta_equiv_850 = ((data.temp + mc.L / mc.cp_air * data.sphum /
(1 - data.sphum)) *
convTtotheta).sel(pfull=850.).mean('lon')
theta_max, theta_max_lat = get_latmax(theta_equiv_850)
cosphimax = np.cos(theta_max_lat * np.pi / 180.)
u_150 = data.ucomp.sel(pfull=lev).mean('lon')
u_m = rotfac * mc.omega * mc.a * (cosphimax**2. - cosphi**2.) / cosphi
def adj_u(u_in, adj_by, dayfac=dayfac):
if 'lon' in adj_by.coords:
return u_in + adj_by.sel(pfull=lev).mean('lon') * 86400. * dayfac
else:
return u_in + adj_by.sel(pfull=lev) * 86400. * dayfac
u_damp = adj_u(u_150, damping)
u_cor = adj_u(u_150, coriolis)
u_adv = adj_u(u_150, adv_tend)
u_vadv = adj_u(u_150, horiz_adv)
u_hadv = adj_u(u_150, vert_adv)
u_eddy = adj_u(u_150, eddy_tend)
u_veddy = adj_u(u_150, horiz_eddy)
u_heddy = adj_u(u_150, vert_eddy)
resid = coriolis + damping + adv_tend + eddy_tend
u_resid = adj_u(u_150, resid, dayfac=10. * dayfac)
u_net = adj_u(u_150, dudt, dayfac=10. * dayfac)
lats = [
data.lat[i] for i in range(len(data.lat))
if data.lat[i] >= -60. and data.lat[i] <= 60.
]
u_m.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='0.7')
u_cor.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='--')
u_adv.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='-.')
u_eddy.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle=':')
u_150.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C0')
u_net.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C2')
ax.plot([w_max_lat.sel(xofyear=pentad),
w_max_lat.sel(xofyear=pentad)], [-75., 75.],
color='0.7',
linestyle=':')
ax.set_title('')
ax.set_xlabel('')
ax.set_xlim(-35., 35.)
ax.set_ylim(-75., 75.)
ax.grid(True, linestyle=':')
ax.set_ylabel('u, m/s')
def plot_vort_dev(video=False, threed=True):
data_w = xr.open_dataset('/disca/share/rg419/jra_omega_pentad_clim.nc')
data_u = xr.open_dataset('/disca/share/rg419/jra_ucomp_pentad_clim.nc')
data_v_temp = xr.open_dataset(
'/disca/share/rg419/jra_vcomp_pentad_clim.nc')
# v has different time coord to u, presumably due to how Stephen has downloaded/averaged. I think the two are equivalent, so just substitute the time dimension into v
data_v = xr.DataArray(data_v_temp.var34.values,
coords={
'pentad': data_u.pentad,
'lat': data_u.lat,
'lon': data_u.lon
},
dims=('pentad', 'lat', 'lon'))
print('files opened')
data_u = data_u.var33
data_w = data_w.var39
zon_adv = data_u.sel(lev=20000.) * gr.ddx(data_u.sel(lev=20000.))
merid_adv = data_v * gr.ddy(data_u.sel(lev=20000.))
vert_adv = data_w * (gr.ddp(data_u, pname='lev')).sel(lev=20000.) * 100.
sinphi = np.sin(data_u.lat * np.pi / 180.)
f = 2. * mc.omega * sinphi
if threed:
rossby = (zon_adv + merid_adv + vert_adv) / (f * data_v)
else:
rossby = merid_adv / (f * data_v)
# Start figure with 1 subplots
rcParams['figure.figsize'] = 10, 5
rcParams['font.size'] = 14
for i in range(72):
fig, ax1 = plt.subplots()
title = 'Pentad ' + str(int(data_u.pentad[i]))
f1 = rossby.sel(pentad=i + 1).plot.contourf(x='lon',
y='lat',
ax=ax1,
add_labels=False,
add_colorbar=False,
extend='both',
zorder=1,
levels=np.arange(
0., 1.1, 0.1))
# data_p.precipitation.sel(pentad=i+1).plot.contour(x='lon', y='lat', ax=ax1, add_labels=False, extend='both', zorder=1, levels=np.arange(2.,21.,2.), colors='k')
ax1.grid(True, linestyle=':')
ax1.set_ylim(-60., 60.)
ax1.set_yticks(np.arange(-60., 61., 30.))
ax1.set_xticks(np.arange(0., 361., 90.))
ax1.set_title(title)
land_mask = '/scratch/rg419/python_scripts/land_era/ERA-I_Invariant_0125.nc'
land = xr.open_dataset(land_mask)
land.lsm[0, :, :].plot.contour(ax=ax1,
x='longitude',
y='latitude',
levels=np.arange(-1., 2., 1.),
add_labels=False,
colors='k')
ax1.set_ylabel('Latitude')
ax1.set_xlabel('Longitude')
plt.subplots_adjust(left=0.1,
right=0.97,
top=0.93,
bottom=0.05,
hspace=0.25,
wspace=0.2)
cb1 = fig.colorbar(f1,
ax=ax1,
use_gridspec=True,
orientation='horizontal',
fraction=0.05,
pad=0.15,
aspect=60,
shrink=0.5)
vidstr = ''
if video:
vidstr = 'video/'
if threed:
plot_dir = '/scratch/rg419/plots/zonal_asym_runs/gill_development/jra/' + vidstr + '/rossby3d/'
else:
plot_dir = '/scratch/rg419/plots/zonal_asym_runs/gill_development/jra/' + vidstr + '/rossby/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
if video:
plt.savefig(plot_dir + 'rossby_and_precip_' +
str(int(data_u.pentad[i])) + '.png',
format='png')
else:
plt.savefig(plot_dir + 'rossby_and_precip_' +
str(int(data_u.pentad[i])) + '.pdf',
format='pdf')
plt.close()
def mse_budg_vert_int(run, land_mask=None):
plot_dir = '/scratch/rg419/plots/zonal_asym_runs/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
data[
'rflux_surf'] = data.t_surf**4. * mc.stefan - data.flux_sw - data.flux_lw
data['rflux_toa'] = data.toa_sw - data.olr
mse = (mc.cp_air * data.temp + mc.L * data.sphum +
mc.grav * data.height) / 1000.
e = ((mc.cp_air - mc.rdgas) * data.temp + mc.L * data.sphum +
mc.grav * data.height) / 1000.
u_mse = (mc.cp_air * data.ucomp_temp + mc.L * data.sphum_u +
mc.grav * data.ucomp_height) / 1000.
v_mse = (mc.cp_air * data.vcomp_temp + mc.L * data.sphum_v +
mc.grav * data.vcomp_height) / 1000.
w_mse = (mc.cp_air * data.omega_temp + mc.L * data.sphum_w +
mc.grav * data.omega_height) / 1000.
u_mse_eddy = u_mse - mse * data.ucomp
v_mse_eddy = v_mse - mse * data.vcomp
w_mse_eddy = w_mse - mse * data.omega
u_dmsedx = data.ucomp * gr.ddx(mse)
v_dmsedy = data.vcomp * gr.ddy(mse, vector=False)
w_dmsedp = data.omega * gr.ddp(mse)
dumsedx_eddy = gr.ddx(u_mse_eddy)
dvmsedy_eddy = gr.ddy(v_mse_eddy)
dwmsedp_eddy = gr.ddp(w_mse_eddy)
dedt = gr.ddt(e) * 5.
Fnet = data.flux_lhe + data.flux_t + data.rflux_surf + data.rflux_toa
def column_int(var_in):
var_int = mc.cp_air * var_in.sum('pfull') * 5000. / mc.grav
return var_int
u_dmsedx = column_int(u_dmsedx)
v_dmsedy = column_int(v_dmsedy)
w_dmsedp = -1. * column_int(w_dmsedp)
dumsedx_eddy = column_int(dumsedx_eddy)
dvmsedy_eddy = column_int(dvmsedy_eddy)
dwmsedp_eddy = column_int(dwmsedp_eddy)
dedt = column_int(dedt)
horiz_adv = -1. * (u_dmsedx + v_dmsedy)
eddy_conv = -1. * (dumsedx_eddy + dvmsedy_eddy + dwmsedp_eddy)
total = (Fnet + horiz_adv + eddy_conv + w_dmsedp) * 5.
for pentad in [32, 38, 44, 50, 56]:
# Start figure with 4 subplots
rcParams['figure.figsize'] = 14, 5
rcParams['font.size'] = 11
fig, ((ax1, ax2, ax3, ax4, ax5), (ax6, ax7, ax8, ax9,
ax10)) = plt.subplots(2,
5,
sharex='col',
sharey='row')
axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10]
f1 = Fnet.sel(xofyear=pentad).plot.contourf(ax=ax1,
x='lon',
y='lat',
add_labels=False,
extend='both',
levels=np.arange(
-200., 201., 20.),
add_colorbar=False)
horiz_adv.sel(xofyear=pentad).plot.contourf(ax=ax2,
x='lon',
y='lat',
add_labels=False,
extend='both',
levels=np.arange(
-200., 201., 20.),
add_colorbar=False)
w_dmsedp.sel(xofyear=pentad).plot.contourf(ax=ax3,
x='lon',
y='lat',
add_labels=False,
extend='both',
levels=np.arange(
-200., 201., 20.),
add_colorbar=False)
eddy_conv.sel(xofyear=pentad).plot.contourf(ax=ax4,
x='lon',
y='lat',
add_labels=False,
extend='both',
levels=np.arange(
-200., 201., 20.),
add_colorbar=False)
dedt.sel(xofyear=pentad).plot.contourf(ax=ax5,
x='lon',
y='lat',
add_labels=False,
extend='both',
levels=np.arange(
-200., 201., 20.),
add_colorbar=False)
data.flux_lhe.sel(xofyear=pentad).plot.contourf(ax=ax6,
x='lon',
y='lat',
add_labels=False,
extend='both',
levels=np.arange(
-200., 201., 20.),
add_colorbar=False)
data.flux_t.sel(xofyear=pentad).plot.contourf(ax=ax7,
x='lon',
y='lat',
add_labels=False,
extend='both',
levels=np.arange(
-200., 201., 20.),
add_colorbar=False)
data.rflux_surf.sel(xofyear=pentad).plot.contourf(ax=ax8,
x='lon',
y='lat',
add_labels=False,
extend='both',
levels=np.arange(
-200., 201.,
20.),
add_colorbar=False)
data.rflux_toa.sel(xofyear=pentad).plot.contourf(ax=ax9,
x='lon',
y='lat',
add_labels=False,
extend='both',
levels=np.arange(
-200., 201., 20.),
add_colorbar=False)
total.sel(xofyear=pentad).plot.contourf(ax=ax10,
x='lon',
y='lat',
add_labels=False,
extend='both',
levels=np.arange(
-200., 201., 20.),
add_colorbar=False)
for ax in axes:
ax.set_ylim(-15., 45.)
ax.set_xlim(90., 225.)
ax.set_yticks(np.arange(-15., 45., 15.))
ax.set_xticks(np.arange(90., 226., 45.))
ax.grid(True, linestyle=':')
if not land_mask == None:
land = xr.open_dataset(land_mask)
for ax in axes:
land.land_mask.plot.contour(x='lon',
y='lat',
ax=ax,
levels=np.arange(-1., 2., 1.),
add_labels=False,
colors='k')
plt.subplots_adjust(left=0.05,
right=0.97,
top=0.97,
bottom=0.03,
hspace=0.25,
wspace=0.2)
cb1 = fig.colorbar(f1,
ax=axes,
use_gridspec=True,
orientation='horizontal',
fraction=0.05,
pad=0.05,
aspect=60,
shrink=1.)
plt.savefig(plot_dir + 'mse_budg_' + run + '_' + str(pentad) + '.pdf',
format='pdf')
plt.close()
def ang_mom_grad_plot(run, rot_fac=1., lev=150.):
'''Plot dvordt or 1/vor * dvordt'''
#Load data
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
#Calculate psi to overplot
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
psi /= 1.e9
psi = make_sym(psi, asym=True)
# Make precip symmetric and find the precip centroid
data['precipitation'] = make_sym(data.precipitation)
precip_centroid(data)
data['ucomp'] = make_sym(data.ucomp)
data['vcomp'] = make_sym(data.vcomp, asym=True)
data['omega'] = make_sym(data.omega)
# Calculate vorticity
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
u_dp = gr.ddp(data.ucomp) # dudp
omega = 7.2921150e-5 * rot_fac
f = 2 * omega * np.sin(data.lat * np.pi / 180)
vor = (v_dx - u_dy + f).sel(pfull=lev).mean('lon') * 86400.
vertical_term = (-1. * data.omega / data.vcomp *
u_dp).sel(pfull=lev).mean('lon') * 86400.
ang_mom_grad = vor + vertical_term
# Plot!
plot_dir = '/scratch/rg419/plots/paper_2_figs/ang_mom_grad/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
rcParams['figure.figsize'] = 5.5, 4.3
rcParams['font.size'] = 14
fig = plt.figure()
ax1 = fig.add_subplot(111)
f1 = ang_mom_grad.plot.contourf(ax=ax1,
x='xofyear',
y='lat',
levels=np.arange(-6., 6.1, 1.),
add_colorbar=False,
add_labels=False,
extend='both')
psi.sel(pfull=lev).plot.contour(ax=ax1,
x='xofyear',
y='lat',
levels=np.arange(-500., 0., 100.),
add_labels=False,
colors='0.7',
linewidths=2,
linestyles='--')
psi.sel(pfull=lev).plot.contour(ax=ax1,
x='xofyear',
y='lat',
levels=np.arange(0., 510., 100.),
add_labels=False,
colors='0.7',
linewidths=2)
psi.sel(pfull=lev).plot.contour(ax=ax1,
x='xofyear',
y='lat',
levels=np.arange(-1000., 1010., 1000.),
add_labels=False,
colors='0.5',
linewidths=2)
data.p_cent.plot.line(ax=ax1, color='k', linewidth=2)
ax1.set_ylim([-60, 60])
ax1.set_ylabel('Latitude')
ax1.set_xticks([12, 24, 36, 48, 60, 72])
ax1.set_yticks([-60, -30, 0, 30, 60])
ax1.set_xlabel('Pentad')
ax1.grid(True, linestyle=':')
plt.subplots_adjust(left=0.17,
right=0.9,
top=0.97,
bottom=0.07,
hspace=0.3)
cb1 = fig.colorbar(f1,
ax=ax1,
use_gridspec=True,
orientation='horizontal',
fraction=0.15,
pad=0.2,
aspect=40)
cb1.set_label('Angular momentum gradient')
plt.savefig(plot_dir + 'ang_mom_grad_' + run + '.pdf', format='pdf')
plt.close()
def plot_heatbudg_dev(run,
land_mask=None,
lev=500.,
qscale=100.,
windtype='',
ref_arrow=5):
plot_dir = '/scratch/rg419/plots/zonal_asym_runs/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
convTtotheta = (1000. / data.pfull)**mc.kappa
data['theta'] = data.temp * convTtotheta
data['dthetadt'] = gr.ddt(data.theta) * 86400.
# Take monthly averages
data.coords['month'] = (data.xofyear - 1) // 6 + 1
data = data.groupby('month').mean(('xofyear'))
heating = (data.dt_tg_condensation + data.dt_tg_convection +
data.dt_tg_diffusion + data.tdt_rad) * 86400.
data['precipitation'] = (data.precipitation * 86400.)
heating_theta = heating * convTtotheta
dthetadx = gr.ddx(data.theta)
dthetady = gr.ddy(data.theta, vector=False)
dthetadp = gr.ddp(data.theta)
udthetadx_mean = data.ucomp * dthetadx
vdthetady_mean = data.vcomp * dthetady
wdthetadp_mean = data.omega * dthetadp
horiz_adv_heating = -1. * (udthetadx_mean + vdthetady_mean) * 86400.
vert_adv_heating = -86400. * wdthetadp_mean
# Start figure with 12 subplots
rcParams['figure.figsize'] = 10, 5
rcParams['font.size'] = 14
fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2,
3,
sharex='col',
sharey='row')
axes = [ax1, ax2, ax3, ax4, ax5, ax6]
titles = ['May', 'June', 'July', 'August', 'September', 'October']
blist = []
for i in range(6):
#f1 = data.dthetadt.sel(pfull=lev)[i+4,:,:].plot.contourf(x='lon', y='lat', ax=axes[i], levels = np.arange(-0.3,0.31,0.02), add_labels=False, add_colorbar=False, extend='both', zorder=1)
f1 = heating_theta.sel(pfull=lev)[i + 4, :, :].plot.contourf(
x='lon',
y='lat',
ax=axes[i],
levels=np.arange(-10., 10.1, 2.),
add_labels=False,
add_colorbar=False,
extend='both',
zorder=1)
#data_zanom.slp[i+4,:,:].plot.contour(x='lon', y='lat', ax=axes[i], levels = np.arange(-15.,16.,3.), add_labels=False, colors='0.5', alpha=0.5)
axes[i].contour(data.lon,
data.lat,
horiz_adv_heating.sel(pfull=lev)[i + 4, :, :],
levels=np.arange(0., 10.1, 2.),
colors='m',
alpha=0.5,
zorder=2)
axes[i].contour(data.lon,
data.lat,
horiz_adv_heating.sel(pfull=lev)[i + 4, :, :],
levels=np.arange(-10., 0., 1.),
colors='m',
alpha=0.5,
linestyle='--',
zorder=2)
axes[i].contour(data.lon,
data.lat,
vert_adv_heating.sel(pfull=lev)[i + 4, :, :],
levels=np.arange(0., 10.1, 1.),
colors='k',
alpha=0.5,
zorder=2)
axes[i].contour(data.lon,
data.lat,
vert_adv_heating.sel(pfull=lev)[i + 4, :, :],
levels=np.arange(-10., 0., 2.),
colors='k',
alpha=0.5,
linestyle='--',
zorder=2)
#axes[i].contour(data.lon, data.lat, heating_theta.sel(pfull=lev)[i+4,:,:], levels = np.arange(0.,5.1,1.), colors='k', alpha=0.5, zorder=2)
#axes[i].contour(data.lon, data.lat, heating_theta.sel(pfull=lev)[i+4,:,:], levels = np.arange(-5.,0.,1.), colors='k', alpha=0.5, linestyle='--', zorder=2)
#if windtype=='div':
# b = axes[i].quiver(data.lon[::6], data.lat[::3], uchi_zanom[i+4,::3,::6], vchi_zanom[i+4,::3,::6], scale=qscale, angles='xy', width=0.005, headwidth=3., headlength=5., zorder=3)
#elif windtype=='rot':
# b = axes[i].quiver(data.lon[::6], data.lat[::3], upsi_zanom[i+4,::3,::6], vpsi_zanom[i+4,::3,::6], scale=qscale, angles='xy', width=0.005, headwidth=3., headlength=5., zorder=3)
#else:
# b = axes[i].quiver(data.lon[::6], data.lat[::3], data_zanom.ucomp.sel(pfull=lev)[i+4,::3,::6], data_zanom.vcomp.sel(pfull=lev)[i+4,::3,::6], scale=qscale, angles='xy', width=0.005, headwidth=3., headlength=5., zorder=3)
#blist.append(b)
axes[i].grid(True, linestyle=':')
axes[i].set_ylim(-60., 60.)
axes[i].set_yticks(np.arange(-60., 61., 30.))
axes[i].set_xticks(np.arange(0., 361., 90.))
axes[i].set_title(titles[i])
if not land_mask == None:
land = xr.open_dataset(land_mask)
land.land_mask.plot.contour(x='lon',
y='lat',
ax=axes[i],
levels=np.arange(-1., 2., 1.),
add_labels=False,
colors='k')
for ax in [ax1, ax4]:
ax.set_ylabel('Latitude')
for ax in [ax4, ax5, ax6]:
ax.set_xlabel('Longitude')
#ax4.quiverkey(blist[3], 0.,-0.5, ref_arrow, str(ref_arrow) + ' m/s', fontproperties={'weight': 'bold', 'size': 10}, color='k', labelcolor='k', labelsep=0.03, zorder=10)
plt.subplots_adjust(left=0.1,
right=0.97,
top=0.93,
bottom=0.05,
hspace=0.25,
wspace=0.2)
cb1 = fig.colorbar(f1,
ax=axes,
use_gridspec=True,
orientation='horizontal',
fraction=0.05,
pad=0.15,
aspect=60,
shrink=0.5)
levstr = ''
windtypestr = ''
if lev != 850:
levstr = '_' + str(lev)
if windtype != '':
windtypestr = '_' + windtype
plt.savefig(plot_dir + 'heat_budg_' + run + '.pdf', format='pdf')
plt.close()
def partition_advection(data, lons, lev=850, do_theta=False):
if do_theta:
convTtotheta=(1000./data.pfull)**(2./7.)
theta = data.temp*convTtotheta
data['ucomp_temp'] = data['ucomp_temp']*convTtotheta
data['vcomp_temp'] = data['vcomp_temp']*convTtotheta
data['omega_temp'] = data['omega_temp']*convTtotheta
data['temp'] = data['temp']*convTtotheta
#First do uT terms
uT_trans_dx = -86400. * gr.ddx( (data.ucomp_temp - data.ucomp*data.temp).sel(pfull=lev) ) # <u'T'> = <uT> - <u><T>
u = data.ucomp.sel(pfull=lev) # u
T_dx = -86400. * gr.ddx( data.temp.sel(pfull=lev) ) # dTdx
u_ed = u - u.mean('lon')
T_dx_ed = T_dx - T_dx.mean('lon')
u_dTdx_zav = u.mean('lon') * T_dx.mean('lon') # [u][dTdx], where brackets denote mean over all longitudes
u_dTdx_cross1 = (u.mean('lon') * T_dx_ed).sel(lon=lons).mean('lon') # [u]dTdx*
u_dTdx_cross2 = (u_ed * T_dx.mean('lon')).sel(lon=lons).mean('lon') # u*[dTdx]
u_dTdx_stat = (u_ed * T_dx_ed).sel(lon=lons).mean('lon') # u*dTdx*
data['uT_trans_dx'] = (('xofyear','lat'), uT_trans_dx.sel(lon=lons).mean('lon') )
data['u_dTdx_cross1'] = (('xofyear','lat'), u_dTdx_cross1 )
data['u_dTdx_cross2'] = (('xofyear','lat'), u_dTdx_cross2 )
data['u_dTdx_stat'] = (('xofyear','lat'), u_dTdx_stat )
data['u_dTdx_zav'] = (('xofyear','lat'), u_dTdx_zav )
print('uT terms done')
#Next do vT terms
vT_trans_dy = -86400. * gr.ddy( (data.vcomp_temp - data.vcomp*data.temp).sel(pfull=lev))
v = data.vcomp.sel(pfull=lev).load() # v
T_dy = -86400. * gr.ddy( data.temp.sel(pfull=lev), vector=False) # dTdy
v_ed = v - v.mean('lon')
T_dy_ed = T_dy - T_dy.mean('lon')
v_dTdy_zav = v.mean('lon') * T_dy.mean('lon') # [v][dTdy]
v_dTdy_cross1 = (v.mean('lon') * T_dy_ed).sel(lon=lons).mean('lon') # [v]dTdy*
v_dTdy_cross2 = (v_ed * T_dy.mean('lon')).sel(lon=lons).mean('lon') # v*[dTdy]
v_dTdy_stat = (v_ed * T_dy_ed).sel(lon=lons).mean('lon') # v*dTdy*
data['vT_trans_dy'] = (('xofyear','lat'), vT_trans_dy.sel(lon=lons).mean('lon'))
data['v_dTdy_cross1'] = (('xofyear','lat'), v_dTdy_cross1 )
data['v_dTdy_cross2'] = (('xofyear','lat'), v_dTdy_cross2 )
data['v_dTdy_stat'] = (('xofyear','lat'), v_dTdy_stat)
data['v_dTdy_zav'] = (('xofyear','lat'), v_dTdy_zav )
print('vT terms done')
#Finally do uw terms
wT_trans_dp = -86400. * gr.ddp( (data.omega_temp - data.omega * data.temp).sel(lon=lons).mean('lon') )
w = data.omega.sel(pfull=lev).load() # w
T_dp = -86400. * (gr.ddp(data.temp)).sel(pfull=lev) # dTdp
w_ed = w - w.mean('lon')
T_dp_ed = T_dp - T_dp.mean('lon')
w_dTdp_zav = w.mean('lon') * T_dp.mean('lon') # [w][dTdp]
w_dTdp_cross1 = (w.mean('lon') * T_dp_ed).sel(lon=lons).mean('lon') # [w]dTdp*
w_dTdp_cross2 = (w_ed * T_dp.mean('lon')).sel(lon=lons).mean('lon') # w*[dTdp]
w_dTdp_stat = (w_ed * T_dp_ed).sel(lon=lons).mean('lon') # w*dTdp*
data['wT_trans_dp'] = (('xofyear','lat'), wT_trans_dp.sel(pfull=lev))
data['w_dTdp_cross1'] = (('xofyear','lat'), w_dTdp_cross1 )
data['w_dTdp_cross2'] = (('xofyear','lat'), w_dTdp_cross2 )
data['w_dTdp_stat'] = (('xofyear','lat'), w_dTdp_stat)
data['w_dTdp_zav'] = (('xofyear','lat'), w_dTdp_zav )
print('wT terms done')
def elliptical_equation(run, rotfac=1.):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc')
data = data.mean('lon')
data['precipitation'] = make_sym(data.precipitation)
precip_centroid(data) # Locate precipitation centroid
dpcentdt = gr.ddt(data.p_cent) * 86400.
p_cent_pos = np.abs(data.p_cent.where(dpcentdt>=0.))
eq_time = p_cent_pos.xofyear[p_cent_pos.argmin('xofyear').values]
peak_rate_time = dpcentdt.xofyear[dpcentdt.argmax('xofyear').values]
max_lat_time = data.p_cent.xofyear[data.p_cent.argmax('xofyear').values]
convTtotheta=(1000./data.pfull)**mc.kappa
data['theta'] = data.temp*convTtotheta
#data['vcomp_theta'] = data.vcomp_temp*convTtotheta
#data['omega_theta'] = data.omega_temp*convTtotheta
#heating = data.dt_tg_condensation + data.dt_tg_convection + data.dt_tg_diffusion + data.tdt_rad
#data['heating_theta'] = heating*convTtotheta
sinphi = np.sin(data.lat * np.pi/180.)
cosphi = np.cos(data.lat * np.pi/180.)
#Calculate psi
#psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
dpsidydy = gr.ddy(data.omega, vector=False)
rho = data.pfull*100./mc.rdgas/data.temp
dthetadp = gr.ddp(data.theta)
brunt_fac = -1./rho/data.theta * dthetadp
psi_yy = (brunt_fac * dpsidydy)
dpsidpdp = gr.ddp(-1.*data.vcomp)
coriolis = 2.* rotfac * mc.omega * sinphi * data.vcomp
dudy = gr.ddy(data.ucomp)
cor_fac = -1.*dudy*coriolis + coriolis**2.
psi_pp = (cor_fac * dpsidydy)
dpsidpdy = gr.ddy(-1.*data.vcomp)
dudp = gr.ddp(data.ucomp)
psi_py = (2.*dudp*dpsidpdy*coriolis)
duvdy = gr.ddy(data.ucomp_vcomp, uv=True)
duwdp = gr.ddp(data.ucomp_omega)
vdudy_mean = data.vcomp * dudy
wdudp_mean = data.omega * dudp
eddy_tend = -1.*duvdy -1.*duwdp + vdudy_mean + wdudp_mean
M = data.dt_ug_diffusion + eddy_tend
dMdp = (gr.ddp(M) * coriolis)
dthetady = gr.ddy(data.theta, vector=False)
#dvtdy = gr.ddy(data.vcomp_theta)
#dwtdp = gr.ddp(data.omega_theta)
vdtdy_mean = data.vcomp * dthetady
wdtdp_mean = data.omega * dthetadp
#eddy_tend_th = -1.*dvtdy -1.*dwtdp + vdtdy_mean + wdtdp_mean
#J = data.heating_theta #+ eddy_tend_th
#dJdy = (-1.* gr.ddy(J, vector=False)/rho/data.theta)
plot_dir = '/scratch/rg419/plots/paper_2_figs/revisions/elliptical_eq/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
rcParams['figure.figsize'] = 20., 12.
levels = np.arange(-1.e-12,1.1e-12,0.1e-12)
fig, ((ax1, ax2, ax3, ax4, ax5), (ax6, ax7, ax8, ax9, ax10), (ax11, ax12, ax13, ax14, ax15)) = plt.subplots(3, 5, sharey='row', sharex='col')
axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12, ax13, ax14, ax15]
i=0
for time in [eq_time, peak_rate_time, max_lat_time]:
f1 = psi_yy.sel(xofyear=time).plot.contourf(ax=axes[5*i], x='lat', y='pfull', yincrease=False, levels=levels, add_labels=False, extend='both', add_colorbar=False)
psi_pp.sel(xofyear=time).plot.contourf(ax=axes[5*i+1], x='lat', y='pfull', yincrease=False, levels=levels, add_labels=False, extend='both', add_colorbar=False)
psi_py.sel(xofyear=time).plot.contourf(ax=axes[5*i+2], x='lat', y='pfull', yincrease=False, levels=levels, add_labels=False, extend='both', add_colorbar=False)
dMdp.sel(xofyear=time).plot.contourf(ax=axes[5*i+3], x='lat', y='pfull', yincrease=False, levels=levels, add_labels=False, extend='both', add_colorbar=False)
#dJdy.sel(xofyear=time).plot.contourf(ax=axes[5*i+4], x='lat', y='pfull', yincrease=False, levels=levels, add_labels=False, extend='both', add_colorbar=False)
i=i+1
for ax in axes:
ax.set_xlim(-30.,30.)
ax1.set_title('N$^2d^2psi/dy^2$')
ax2.set_title('$f(f-du/dy)d^2psi/dp^2$')
ax3.set_title('Cross term')
ax4.set_title('dM/dp')
ax5.set_title('dJ/dy')
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)
figname = 'elliptical_eq_' +run+ '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def temp_grad_tend(run, lev=850., moist=False):
rcParams['figure.figsize'] = 12, 5.5
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/paper_2_figs/revisions/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc')
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
psi /= 1.e9
psi = np.abs(psi).max('pfull')
data = data.mean('lon')
convTtotheta=(1000./data.pfull)**mc.kappa
heating = (data.dt_tg_condensation + data.dt_tg_convection + data.dt_tg_diffusion + data.tdt_rad)
hum_tend = data.dt_qg_condensation + data.dt_qg_convection + data.dt_qg_diffusion
if moist:
heating = (heating + mc.L/mc.cp_air * hum_tend/(1-data.sphum)**2.) * convTtotheta
theta = (data.temp + mc.L/mc.cp_air * data.sphum/(1-data.sphum)) * convTtotheta
else:
theta = data.temp*convTtotheta
#vcomp_theta = data.vcomp_temp*convTtotheta
#omega_theta = data.omega_temp*convTtotheta
heating = heating*convTtotheta
dthetady = gr.ddy(theta, vector=False)
dthetadp = gr.ddp(theta)
def column_int(var_in):
var_int = mc.cp_air * var_in.sum('pfull')*5000./mc.grav * 10.**5.
return var_int
dthetadydt = column_int(gr.ddt(dthetady))
div_term = column_int(-1.* gr.ddy(data.vcomp) * dthetady)
horiz_adv = column_int(-1.* data.vcomp * gr.ddy(dthetady))
tilt_term = column_int(-1.* gr.ddy(data.omega, vector=False) * dthetadp)
vert_adv = column_int(-1.* data.omega * gr.ddy(dthetadp, vector=False))
heating_grad = column_int(gr.ddy(heating, vector=False))
#eddies = column_int(-1.*gr.ddy(gr.ddy(vcomp_theta) + gr.ddp(omega_theta))) - (div_term + horiz_adv + tilt_term + vert_adv)
#eddies = dthetadydt - div_term - horiz_adv - tilt_term - vert_adv - heating_grad
term_sum = div_term + horiz_adv + tilt_term + vert_adv + heating_grad
# 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=dthetadydt.plot.contourf(ax=ax1, x='xofyear', y='lat', extend='both', levels=np.arange(-2.,2.1,0.2), add_labels=False)
psi.plot.contour(ax=ax1, x='xofyear', y='lat', extend='both', add_labels=False, colors='k', alpha=0.2, levels=np.arange(0,600.,100.))
ax1.set_ylabel('Latitude')
ax1.set_title('$d\Theta/dtdy$', 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
div_term.plot.contourf(ax=ax2, x='xofyear', y='lat', extend='both', levels=np.arange(-9.,9.1,0.5), add_labels=False)
psi.plot.contour(ax=ax2, x='xofyear', y='lat', extend='both', add_labels=False, colors='k', alpha=0.2, levels=np.arange(0,600.,100.))
ax2.grid(True,linestyle=':')
ax2.set_title('$-dv/dy*d\Theta/dy$', fontsize=17)
ax2.set_ylim(-60,60)
ax2.text(-5, 60, 'b)')
#Third plot
horiz_adv.plot.contourf(ax=ax3, x='xofyear', y='lat', extend='both', levels=np.arange(-9.,9.1,0.5), add_labels=False)
psi.plot.contour(ax=ax3, x='xofyear', y='lat', extend='both', add_labels=False, colors='k', alpha=0.2, levels=np.arange(0,600.,100.))
ax3.grid(True,linestyle=':')
ax3.set_ylim(-60,60)
ax3.set_title('$-vd\Theta/dydy$', fontsize=17)
ax3.text(-5, 60, 'c)')
#Fourth plot
tilt_term.plot.contourf(ax=ax4, x='xofyear', y='lat', extend='both', levels=np.arange(-9.,9.1,0.5), add_labels=False)
psi.plot.contour(ax=ax4, x='xofyear', y='lat', extend='both', add_labels=False, colors='k', alpha=0.2, levels=np.arange(0,600.,100.))
ax4.grid(True,linestyle=':')
ax4.set_ylabel('Latitude')
ax4.set_ylim(-60,60)
ax4.set_title('$-dw/dy*d\Theta/dp$', fontsize=17)
ax4.set_yticks(np.arange(-60.,61.,30.))
ax4.text(-15, 60, 'd)')
#Fifth plot
vert_adv.plot.contourf(ax=ax5, x='xofyear', y='lat', extend='both', levels=np.arange(-9.,9.1,0.5), add_labels=False)
psi.plot.contour(ax=ax5, x='xofyear', y='lat', extend='both', add_labels=False, colors='k', alpha=0.2, levels=np.arange(0,600.,100.))
ax5.grid(True,linestyle=':')
ax5.set_title('$-wd\Theta/dydp$', fontsize=17)
ax5.set_ylim(-60,60)
ax5.text(-5, 60, 'e)')
#Sixth plot
heating_grad.plot.contourf(ax=ax6, x='xofyear', y='lat', extend='both', levels=np.arange(-9.,9.1,0.5), add_labels=False)
psi.plot.contour(ax=ax6, x='xofyear', y='lat', extend='both', add_labels=False, colors='k', alpha=0.2, levels=np.arange(0,600.,100.))
ax6.grid(True,linestyle=':')
ax6.set_ylim(-60,60)
ax6.set_title('Diabatic heating gradient', fontsize=17)
ax6.text(-5, 60, 'f)')
plt.subplots_adjust(right=0.97, left=0.1, top=0.95, bottom=0.1, 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 moist:
figname = 'temp_grad_tend_' +run+ '_moist.pdf'
else:
figname = 'temp_grad_tend_' +run+ '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def fig_8_moist(run, ax, pentad=45, rotfac=1., dayfac=5.):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
heating = data.dt_tg_condensation + data.dt_tg_convection + data.dt_tg_diffusion + data.tdt_rad
hum_tend = data.dt_qg_condensation + data.dt_qg_convection + data.dt_qg_diffusion
convTtotheta = (1000. / data.pfull)**mc.kappa
theta = data.temp * convTtotheta
heating_theta = heating * convTtotheta
heating_equiv_theta = heating_theta + mc.L / mc.cp_air * hum_tend / (
1 - data.sphum)**2. * convTtotheta
sinphi = np.sin(data.lat * np.pi / 180.)
cosphi = np.cos(data.lat * np.pi / 180.)
theta_850 = theta.sel(pfull=850.).mean('lon')
theta_equiv = (data.temp + mc.L / mc.cp_air * data.sphum /
(1 - data.sphum)) * convTtotheta
theta_equiv_850 = theta_equiv.sel(pfull=850.).mean('lon')
lats = [
data.lat[i] for i in range(len(data.lat))
if data.lat[i] >= -10. and data.lat[i] <= 10.
]
Tt = (data.temp.sel(pfull=150.).mean(('lon')) *
cosphi).sel(lat=lats).sum('lat') / cosphi.sel(lat=lats).sum('lat')
Ts = (data.t_surf.mean(('lon')) *
cosphi).sel(lat=lats).sum('lat') / cosphi.sel(lat=lats).sum('lat')
dthetady_equiv = gr.ddy(theta_equiv.mean('lon'), vector=False)
dthetadp_equiv = gr.ddp(theta_equiv.mean('lon'))
dthetadt_equiv = gr.ddt(theta_equiv.mean('lon'))
vdthetady_mean = data.vcomp.mean('lon') * dthetady_equiv
wdthetadp_mean = data.omega.mean('lon') * dthetadp_equiv
adv_heating_equiv = -1. * (vdthetady_mean + wdthetadp_mean)
w_850 = data.omega.sel(pfull=850.).mean('lon')
w_max, w_max_lat = get_latmax(-1. * w_850)
theta_max, theta_max_lat = get_latmax(theta_equiv_850)
cosphimax = np.cos(theta_max_lat * np.pi / 180.)
chi = (rotfac * mc.omega)**2 * mc.a**2. / mc.cp_air / (Ts - Tt)
theta_m_equiv = theta_max * np.exp(
-1. * chi * (cosphimax**2. - cosphi**2.)**2. / cosphi**2.)
def adj_theta(theta_in, adj_by, dayfac=dayfac):
if 'lon' in adj_by.coords:
return theta_in + adj_by.sel(
pfull=850.).mean('lon') * 86400. * dayfac
else:
return theta_in + adj_by.sel(pfull=850.) * 86400. * dayfac
theta_equiv_rc = adj_theta(theta_equiv_850,
heating_equiv_theta,
dayfac=dayfac)
theta_equiv_adv = adj_theta(theta_equiv_850,
adv_heating_equiv,
dayfac=dayfac)
theta_equiv_net = adj_theta(theta_equiv_850,
dthetadt_equiv,
dayfac=dayfac * 3.)
lats = [
data.lat[i] for i in range(len(data.lat))
if data.lat[i] >= -60. and data.lat[i] <= 60.
]
theta_m_equiv.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='0.7')
theta_equiv_rc.sel(xofyear=pentad, lat=lats).plot(ax=ax,
color='k',
linestyle='--')
theta_equiv_adv.sel(xofyear=pentad, lat=lats).plot(ax=ax,
color='k',
linestyle='-.')
theta_equiv_850.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C0')
theta_equiv_net.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C2')
ax.plot([w_max_lat.sel(xofyear=pentad),
w_max_lat.sel(xofyear=pentad)], [300., 380.],
color='0.7',
linestyle=':')
ax.set_title('')
ax.set_xlabel('')
ax.set_ylim(300., 380.)
ax.set_xlim(-35., 35.)
ax.grid(True, linestyle=':')
ax.set_ylabel('$\Theta$, K')
def plot_vort_dev(run, land_mask=None, lev=200, video=False, threed=True):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
data['precipitation'] = (data.precipitation * 86400.)
zon_adv = data.ucomp * gr.ddx(data.ucomp)
merid_adv = data.vcomp * gr.ddy(data.ucomp)
vert_adv = data.omega * gr.ddp(data.ucomp)
sinphi = np.sin(data.lat * np.pi / 180.)
f = 2. * mc.omega * sinphi
if threed:
rossby = (zon_adv + merid_adv + vert_adv) / (f * data.vcomp)
else:
rossby = merid_adv / (f * data.vcomp)
# Start figure with 1 subplots
rcParams['figure.figsize'] = 10, 5
rcParams['font.size'] = 14
for i in range(72):
fig, ax1 = plt.subplots()
title = 'Pentad ' + str(int(data.xofyear[i]))
f1 = rossby.sel(xofyear=i + 1,
pfull=lev).plot.contourf(x='lon',
y='lat',
ax=ax1,
add_labels=False,
add_colorbar=False,
extend='both',
zorder=1,
levels=np.arange(
0., 1.1, 0.1))
data.precipitation.sel(xofyear=i + 1).plot.contour(x='lon',
y='lat',
ax=ax1,
add_labels=False,
extend='both',
zorder=1,
levels=np.arange(
2., 21., 2.),
colors='k')
ax1.grid(True, linestyle=':')
ax1.set_ylim(-60., 60.)
ax1.set_yticks(np.arange(-60., 61., 30.))
ax1.set_xticks(np.arange(0., 361., 90.))
ax1.set_title(title)
if not land_mask == None:
land = xr.open_dataset(land_mask)
land.land_mask.plot.contour(x='lon',
y='lat',
ax=ax1,
levels=np.arange(-1., 2., 1.),
add_labels=False,
colors='k')
ax1.set_ylabel('Latitude')
ax1.set_xlabel('Longitude')
plt.subplots_adjust(left=0.1,
right=0.97,
top=0.93,
bottom=0.05,
hspace=0.25,
wspace=0.2)
cb1 = fig.colorbar(f1,
ax=ax1,
use_gridspec=True,
orientation='horizontal',
fraction=0.05,
pad=0.15,
aspect=60,
shrink=0.5)
vidstr = ''
if video:
vidstr = 'video/'
if threed:
plot_dir = '/scratch/rg419/plots/zonal_asym_runs/gill_development/' + run + '/' + vidstr + '/rossby3d/'
else:
plot_dir = '/scratch/rg419/plots/zonal_asym_runs/gill_development/' + run + '/' + vidstr + '/rossby/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
if video:
plt.savefig(plot_dir + 'rossby_and_precip_' +
str(int(data.xofyear[i])) + '.png',
format='png')
else:
plt.savefig(plot_dir + 'rossby_and_precip_' +
str(int(data.xofyear[i])) + '.pdf',
format='pdf')
plt.close()
def vdtdy_plot(run, ax, lev=850.):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
psi /= 1.e9
try:
data['precipitation'] = make_sym(data.precipitation)
except:
data['precipitation'] = data.convection_rain + data.condensation_rain
data['precipitation'] = make_sym(data.precipitation)
precip_centroid(data)
convTtotheta = (1000. / data.pfull)**(2. / 7.)
theta = data.temp * convTtotheta
rho = lev * 100. / Rd / data.temp.sel(
pfull=lev) / (1 + (Rv - Rd) / Rd * data.sphum.sel(pfull=lev))
expansion_term = (1. / rho / cp *
data.omega.sel(pfull=lev)).mean('lon') * 86400.
#expansion_term = (1./rho/cp ).mean('lon') * 86400.
v = data.vcomp.sel(pfull=lev) # v
T_dy = -86400. * gr.ddy(data.temp.sel(pfull=lev), vector=False) # dTdy
vdTdy = v.mean('lon') * T_dy.mean('lon') # [v][dTdy]
w = data.omega.sel(pfull=lev) # w
T_dp = -86400. * (gr.ddp(data.temp)).sel(pfull=lev) # dTdp
wdTdp = w.mean('lon') * T_dp.mean('lon') # [w][dTdp]
dTdt = gr.ddt(data.temp).sel(pfull=lev).mean('lon') * 86400.
#dvtdy = -1.*(gr.ddy(data.vcomp * data.temp)).mean('lon').sel(pfull=lev)*86400.
#dwtdp = -1.*(gr.ddp(data.omega * data.temp)).mean('lon').sel(pfull=lev)*86400.
#dvtdy_colav = -1.*gr.ddy((data.vcomp * theta).mean('lon').sel(pfull=np.arange(50.,501.,50.)).mean('pfull')*5000./9.8)*1004.64
#dvHdy_colav = -1.*gr.ddy((data.vcomp * (data.temp*1004.64 + data.sphum*2.500e6)).mean('lon').mean('pfull')*5000./9.8)
#f1 = (dvtdy_colav).plot.contourf(ax=ax, x='xofyear', y='lat', add_labels=False, levels=np.arange(-30.,33.,3.), extend='both', add_colorbar=False)
#f1 = (vdTdy).plot.contourf(ax=ax, x='xofyear', y='lat', add_labels=False, levels=np.arange(-5.,5.5,0.5), extend='both', add_colorbar=False)
f1 = (vdTdy + wdTdp).plot.contourf(ax=ax,
x='xofyear',
y='lat',
add_labels=False,
levels=np.arange(-3., 3.1, 0.5),
extend='both',
add_colorbar=False)
#f1 = (dTdt).plot.contourf(ax=ax, x='xofyear', y='lat', add_labels=False, levels=np.arange(-0.2,0.21,0.02), extend='both', add_colorbar=False)
#f1 = (vdTdy + wdTdp + expansion_term).plot.contourf(ax=ax, x='xofyear', y='lat', add_labels=False, extend='both', add_colorbar=False)
#f1 = (-1.*expansion_term-T_dp.mean('lon')).plot.contourf(ax=ax, x='xofyear', y='lat', add_labels=False, extend='both', add_colorbar=False)
psi.sel(pfull=500).plot.contour(ax=ax,
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=ax,
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=ax,
x='xofyear',
y='lat',
levels=np.arange(-1000., 1010., 1000.),
add_labels=False,
colors='0.5',
linewidths=2)
data.p_cent.plot.line(ax=ax, color='k', linewidth=2)
ax.set_xlabel('')
ax.set_ylim([-60, 60])
ax.set_ylabel('Latitude')
ax.set_xticks([12, 24, 36, 48, 60, 72])
ax.set_yticks([-60, -30, 0, 30, 60])
ax.grid(True, linestyle=':')
return f1
def fig_8(run, ax, pentad=45, lev=200., dayfac=3., rotfac=1.):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc')
sinphi = np.sin(data.lat * np.pi/180.)
cosphi = np.cos(data.lat * np.pi/180.)
zeta = 2.* rotfac * mc.omega*sinphi -1.* gr.ddy(data.ucomp.mean('lon')) #* 86400.
nu = gr.ddp(data.ucomp.mean('lon')) #* 86400.
dzetady = gr.ddy(zeta, vector=False)
dzetadp = gr.ddp(zeta)
dwdy = gr.ddy(data.omega, vector=False)
dvdy = gr.ddy(data.vcomp)
dwdp = gr.ddp(data.omega)
dzetadt = gr.ddt(zeta)
vdzetady_mean = data.vcomp.mean('lon') * dzetady
wdzetadp_mean = data.omega.mean('lon') * dzetadp
adv_tend = -1. *(vdzetady_mean + wdzetadp_mean)
vert_adv = -1.*wdzetadp_mean
horiz_adv = -1.*vdzetady_mean
stretching = -1.* zeta * dvdy
tilting = nu * dwdy
eddy_tend = dzetadt - adv_tend - stretching - tilting
w_850 = data.omega.sel(pfull=850.).mean('lon')
w_max = w_850.min('lat')
w_max_lat = data.lat.values[w_850.argmin('lat').values]
w_max_lat = xr.DataArray(w_max_lat, coords=[data.xofyear], dims=['xofyear'])
cosphimax = np.cos(w_max_lat * np.pi/180.)
u_m = rotfac * mc.omega * mc.a * (cosphimax**2. - cosphi**2.)/cosphi
zeta_m = 2.*mc.omega*sinphi -1.* gr.ddy(u_m)
def adj_zeta(zeta_in, adj_by, dayfac=dayfac):
if 'lon' in adj_by.coords:
return zeta_in + adj_by.sel(pfull=lev).mean('lon') * 86400.**2. * dayfac
else:
return zeta_in + adj_by.sel(pfull=lev) * 86400.**2. * dayfac
zeta_150 = zeta.sel(pfull=150.)*86400.
zeta_adv = adj_zeta(zeta_150, adv_tend); zeta_vadv = adj_zeta(zeta_150, horiz_adv); zeta_hadv = adj_zeta(zeta_150, vert_adv)
zeta_stretching = adj_zeta(zeta_150, stretching)
zeta_tilting = adj_zeta(zeta_150, tilting)
zeta_eddy = adj_zeta(zeta_150, eddy_tend);
zeta_net = adj_zeta(zeta_150, dzetadt, dayfac=10.*dayfac)
lats = [data.lat[i] for i in range(len(data.lat)) if data.lat[i] >= -60. and data.lat[i] <= 60.]
zeta_m.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='0.7', linestyle='-')
zeta_stretching.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='--')
zeta_tilting.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='0.7', linestyle='--')
zeta_adv.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle='-.')
#u_hadv.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C1', linestyle='-.')
#u_vadv.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C2', linestyle='-.')
zeta_eddy.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='k', linestyle=':')
zeta_150.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C0')
zeta_net.sel(xofyear=pentad, lat=lats).plot(ax=ax, color='C2')
ax.plot([w_max_lat.sel(xofyear=pentad),w_max_lat.sel(xofyear=pentad)], [-7.5,7.5], color='0.7', linestyle=':')
ax.set_title('')
ax.set_xlabel('')
ax.set_xlim(-35.,35.)
ax.set_ylim(-7.5,7.5)
ax.grid(True,linestyle=':')
ax.set_ylabel('zeta, /s')
def vdtdy_plot(run, lev=850.):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc')
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
psi /= 1.e9
try:
data['precipitation'] = make_sym(data.precipitation)
except:
data['precipitation'] = data.convection_rain + data.condensation_rain
data['precipitation'] = make_sym(data.precipitation)
precip_centroid(data)
convTtotheta=(1000./data.pfull)**(2./7.)
theta = data.temp*convTtotheta
#rcParams['figure.figsize'] = 5.5, 4.3
#theta.mean('lon').sel(xofyear=40).plot.contourf(x='lat', y='pfull', yincrease=False)
#plt.show()
dvtdy = -1.*(gr.ddy(data.vcomp * theta)).mean('lon').sel(pfull=lev)*86400.
dwtdp = -1.*(gr.ddp(data.omega * theta)).mean('lon').sel(pfull=lev)*86400.
adv_tend_grad = gr.ddy(dvtdy+dwtdp, vector=False)*10.**5.
dTdt = (gr.ddt(theta)).mean('lon').sel(pfull=lev)*86400./theta.mean('lon').sel(pfull=lev)
dvtdy_colav = -1.*gr.ddy((data.vcomp * theta).mean('lon').mean('pfull')*5000./9.8)*1004.64
#dwtdp = -1.*gr.ddp((data.omega * theta).mean('lon').sum('pfull')*50./9.8)
dTdt_colav = gr.ddt(theta.mean('lon').sum('pfull')*5000./9.8)*1004.64
plot_dir = '/scratch/rg419/plots/paper_2_figs/dvtdy/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
#rcParams['figure.figsize'] = 5.5, 4.3
rcParams['figure.figsize'] = 5, 10.5
rcParams['font.size'] = 14
#rcParams['font.size'] = 16
#fig = plt.figure()
#ax1 = fig.add_subplot(111)
fig, ((ax1), (ax2), (ax3)) = plt.subplots(3, 1)
f1 = (dvtdy+dwtdp).plot.contourf(ax=ax1, x='xofyear', y='lat', add_labels=False, levels=np.arange(-50.,55.,5.), extend='both')
psi.sel(pfull=500).plot.contour(ax=ax1, 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=ax1, 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=ax1, x='xofyear', y='lat', levels=np.arange(-1000.,1010.,1000.), add_labels=False, colors='0.5', linewidths=2)
data.p_cent.plot.line(ax=ax1,color='k', linewidth=2)
ax1.set_ylim([-60,60])
ax1.set_ylabel('Latitude')
ax1.set_xticks([12,24,36,48,60,72])
ax1.set_yticks([-60,-30,0,30,60])
ax1.set_xlabel('Pentad')
ax1.grid(True,linestyle=':')
f1 = adv_tend_grad.plot.contourf(ax=ax2, x='xofyear', y='lat', add_labels=False, levels=np.arange(-10.,11.,1.), extend='both')
psi.sel(pfull=500).plot.contour(ax=ax2, 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=ax2, 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=ax2, x='xofyear', y='lat', levels=np.arange(-1000.,1010.,1000.), add_labels=False, colors='0.5', linewidths=2)
data.p_cent.plot.line(ax=ax2,color='k', linewidth=2)
ax2.set_ylim([-60,60])
ax2.set_ylabel('Latitude')
ax2.set_xticks([12,24,36,48,60,72])
ax2.set_yticks([-60,-30,0,30,60])
ax2.set_xlabel('Pentad')
ax2.grid(True,linestyle=':')
#f1 = dvtdy_colav.plot.contourf(ax=ax3, x='xofyear', y='lat', add_labels=False, levels=np.arange(-50.,51.,5.), extend='both')
#psi.sel(pfull=500).plot.contour(ax=ax3, 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=ax3, 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=ax3, x='xofyear', y='lat', levels=np.arange(-1000.,1010.,1000.), add_labels=False, colors='0.5', linewidths=2)
#data.p_cent.plot.line(ax=ax3,color='k', linewidth=2)
#ax3.set_ylim([-60,60])
#ax3.set_ylabel('Latitude')
#ax3.set_xticks([12,24,36,48,60,72])
#ax3.set_yticks([-60,-30,0,30,60])
#ax3.set_xlabel('Pentad')
#ax3.grid(True,linestyle=':')
plt.subplots_adjust(left=0.15, right=0.95, top=0.95, bottom=0.05)
# cb1=fig.colorbar(f1, ax=ax1, use_gridspec=True, orientation = 'horizontal',fraction=0.15, pad=0.2, aspect=40)
# cb1.set_label('Advective temperature tendency, Kday$^{-1}$')
plt.savefig(plot_dir+'adv_tend_' + run + '.pdf', format='pdf')
plt.close()
def mom_budg_fun(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
#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 rossby_plot(run,
ax,
rot_fac=1.,
lev=200.,
type='vor_only',
plottype='rossby'):
'''Plot dvordt or 1/vor * dvordt'''
#Load data
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
#Calculate psi to overplot
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
psi /= 1.e9
psi = make_sym(psi, asym=True)
# Make precip symmetric and find the precip centroid
data['precipitation'] = make_sym(data.precipitation)
data['ucomp'] = make_sym(data.ucomp)
data['vcomp'] = make_sym(data.vcomp, asym=True)
data['omega'] = make_sym(data.omega)
precip_centroid(data)
# Calculate vorticity
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
u_dp = gr.ddp(data.ucomp) # dudp
omega = 7.2921150e-5 * rot_fac
f = 2 * omega * np.sin(data.lat * np.pi / 180)
vor = (v_dx - u_dy).sel(pfull=lev)
metric = (data.ucomp / mc.a *
np.tan(data.lat * np.pi / 180)).sel(pfull=lev)
vertical_term = (-1. * data.omega / data.vcomp * u_dp).sel(pfull=lev)
#vor = make_sym(vor, asym=True)
#metric = make_sym(metric, asym=True)
#vertical_term = make_sym(vertical_term, asym=True)
if type == 'vor_only':
rossby = (-1. * vor / f).mean('lon')
elif type == 'metric':
rossby = (-1. * (metric + vor) / f).mean('lon')
elif type == 'vertical':
rossby = (-1. * (vor + vertical_term) / f).mean('lon')
elif type == 'full':
rossby = (-1. * (metric + vor + vertical_term) / f).mean('lon')
levels = np.arange(-0.9, 1.0, 0.3)
if plottype == 'drodt':
rossby = gr.ddt(rossby) * 84600.
levels = np.arange(-0.1, 0.1, 0.01)
f1 = rossby.plot.contourf(ax=ax,
x='xofyear',
y='lat',
levels=levels,
add_colorbar=False,
add_labels=False,
extend='both')
psi.sel(pfull=lev).plot.contour(ax=ax,
x='xofyear',
y='lat',
levels=np.arange(-500., 0., 100.),
add_labels=False,
colors='0.7',
linewidths=2,
linestyles='--')
psi.sel(pfull=lev).plot.contour(ax=ax,
x='xofyear',
y='lat',
levels=np.arange(0., 510., 100.),
add_labels=False,
colors='0.7',
linewidths=2)
psi.sel(pfull=lev).plot.contour(ax=ax,
x='xofyear',
y='lat',
levels=np.arange(-1000., 1010., 1000.),
add_labels=False,
colors='0.5',
linewidths=2)
data.p_cent.plot.line(ax=ax, color='k', linewidth=2)
ax.set_ylim([-30, 30])
ax.set_ylabel('Latitude')
ax.set_xticks([12, 24, 36, 48, 60, 72])
ax.set_yticks([-30, -15, 0, 15, 30])
ax.set_xlabel('Pentad')
ax.grid(True, linestyle=':')
#plt.subplots_adjust(left=0.15, right=0.95, top=0.95, bottom=0.05)
cb1 = fig.colorbar(f1,
ax=ax,
use_gridspec=True,
orientation='horizontal',
fraction=0.15,
pad=0.2,
aspect=40)
cb1.set_label('Rossby number')
