def abs_vort_dt_at_pcent(run, rot_fac=1., lev=150.,lat_bound=45., res=0.01, interp=True):
'''Calculate the (normalised) vorticity tendency at the precipitation centroid'''
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc')
# Interpolate in time if wanted
if interp:
times_new = np.arange(1., 72.2, 0.2)
precip_new = time_interp(data.precipitation, times_new)
u_new = time_interp(data.ucomp.sel(pfull=lev), times_new)
v_new = time_interp(data.vcomp.sel(pfull=lev), times_new)
else:
precip_new = data.precipitation
u_new = data.ucomp.sel(pfull=lev)
v_new = data.vcomp.sel(pfull=lev)
# Find precipitation centroid
precip_new = make_sym(precip_new)
data = xr.Dataset({'precipitation': precip_new}, coords=precip_new.coords)
precip_centroid(data, lat_bound=lat_bound, res=res)
# Calculate vorticity
v_dx = gr.ddx(v_new) # dvdx
u_dy = gr.ddy(u_new) # dudy
omega = 7.2921150e-5 * rot_fac
f = 2 * omega * np.sin(data.lat *np.pi/180)
vor = (v_dx - u_dy + f)*86400.
div = gr.ddx(u_new) + gr.ddy(v_new)
stretching_mean = (-86400. * vor * div).mean('lon')
vor = make_sym(vor, asym=True)
stretching_mean = make_sym(stretching_mean, asym=True)
# Take time derivative of absolute vorticity
if interp:
dvordt = gr.ddt(vor.mean('lon'))*86400.*5.
else:
dvordt = gr.ddt(vor.mean('lon'))*86400.
# Also normalise this by the value of vorticity
dvordtvor = dvordt/vor.mean('lon')
#dvordtvor = stretching_mean/vor.mean('lon')
# Interpolate vorticity in latitude to match precipitation centroid lats
lats = [data.lat[i] for i in range(len(data.lat)) if data.lat[i] >= -lat_bound and data.lat[i] <= lat_bound]
lats_new = np.arange(-lat_bound, lat_bound+res, res)
stretching_mean = lat_interp(stretching_mean.sel(lat=lats), lats_new)
dvordt = lat_interp(dvordt.sel(lat=lats), lats_new)
dvordtvor = lat_interp(dvordtvor.sel(lat=lats), lats_new)
# Get and return dvordt and dvordtvor at the precipitation centroid, as well as the precipitation centroid itself
st_pcent = [float(stretching_mean[i,:].sel(lat=data.p_cent.values[i]).values) for i in range(len(data.xofyear))]
st_pcent = xr.DataArray(np.asarray(st_pcent), coords=[stretching_mean.xofyear.values], dims=['xofyear'])
dvordt_pcent = [float(dvordt[i,:].sel(lat=data.p_cent.values[i]).values) for i in range(len(data.xofyear))]
dvordt_pcent = xr.DataArray(np.asarray(dvordt_pcent), coords=[dvordt.xofyear.values], dims=['xofyear'])
dvordtvor_pcent = [float(dvordtvor[i,:].sel(lat=data.p_cent.values[i]).values) for i in range(len(data.xofyear))]
dvordtvor_pcent = xr.DataArray(np.asarray(dvordtvor_pcent), coords=[dvordtvor.xofyear.values], dims=['xofyear'])
return dvordt_pcent, dvordtvor_pcent, data.p_cent, st_pcent
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))
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 div_at_pcent(run,
rot_fac=1.,
lev=150.,
lat_bound=45.,
res=0.01,
interp=True):
'''Calculate the (normalised) vorticity tendency at the precipitation centroid'''
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
# Interpolate in time if wanted
if interp:
times_new = np.arange(1., 72.2, 0.2)
precip_new = time_interp(data.precipitation, times_new)
u_new = time_interp(data.ucomp.sel(pfull=lev), times_new)
v_new = time_interp(data.vcomp.sel(pfull=lev), times_new)
else:
precip_new = data.precipitation
u_new = data.ucomp.sel(pfull=lev)
v_new = data.vcomp.sel(pfull=lev)
# Find precipitation centroid
precip_new = make_sym(precip_new)
data = xr.Dataset({'precipitation': precip_new}, coords=precip_new.coords)
precip_centroid(data, lat_bound=lat_bound, res=res)
# Calculate divergence
v_dx = gr.ddx(v_new) # dvdx
u_dy = gr.ddy(u_new) # dudy
div = (gr.ddx(u_new) + gr.ddy(v_new)).mean('lon') * 86400.
omega = 7.2921150e-5 * rot_fac
f = 2 * omega * np.sin(data.lat * np.pi / 180)
vor = (-1. * gr.ddy(u_new) * 0. + f).mean('lon') * 86400.
div = make_sym(div)
vor = make_sym(vor, asym=True)
div = div * vor
div = vor
# Interpolate divergence in latitude to match precipitation centroid lats
lats = [
data.lat[i] for i in range(len(data.lat))
if data.lat[i] >= -lat_bound and data.lat[i] <= lat_bound
]
lats_new = np.arange(-lat_bound, lat_bound + res, res)
div = lat_interp(div.sel(lat=lats), lats_new)
# Get and return dvordt and dvordtvor at the precipitation centroid, as well as the precipitation centroid itself
div_pcent = [
float(div[i, :].sel(lat=data.p_cent.values[i]).values)
for i in range(len(data.xofyear))
]
div_pcent = xr.DataArray(np.asarray(div_pcent),
coords=[div.xofyear.values],
dims=['xofyear'])
return div_pcent, data.p_cent
def abs_vort_mean(run, rot_fac=1., lev=150.):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
precip_centroid(data)
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
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)
# Take gradients of vorticity
dvordx = gr.ddx(vor)
dvordy = gr.ddy(vor, vector=False)
# Horizontal material derivative
horiz_md_mean = (-86400.**2. *
(data.ucomp.sel(pfull=lev) * dvordx +
data.vcomp.sel(pfull=lev) * dvordy)).mean('lon')
# Calculate divergence and stretching term
div = gr.ddx(data.ucomp.sel(pfull=lev)) + gr.ddy(data.vcomp.sel(pfull=lev))
stretching_mean = (-86400.**2. * vor * div).mean('lon')
vor = vor.mean('lon')
vor_mag = np.abs(vor)
stretching_mag = stretching_mean
stretching_mag[:, 0:32] = stretching_mag[:, 0:32] * -1.
lat_itcz, lat_nh, lat_sh = get_streamfun_zeros(data, lev=lev)
def take_cell_average(var, lats_min, lats_max):
cell_av = np.zeros(len(data.xofyear.values))
for i in range(len(data.xofyear.values)):
lats = [
data.lat[j] for j in range(len(data.lat))
if data.lat[j] > lats_min[i] and data.lat[j] < lats_max[i]
]
cell_av[i] = var[i].sel(lat=lats).mean('lat')
return cell_av
vor_av_sh = take_cell_average(vor_mag, lat_sh, lat_itcz)
stretching_av_sh = take_cell_average(stretching_mag, lat_sh, lat_itcz)
plt.plot(stretching_av_sh, 'k')
plt.figure(2)
plt.plot(data.p_cent, stretching_av_sh, 'kx')
plt.show()
def plumb_flux(data, plev=200.):
prefac_a = plev/1000. * np.cos(np.pi*data.lat/180.)
prefac_b = 1/(2*mc.omega*mc.a*np.sin(np.pi*data.lat/180.))
u_anom = (data.ucomp - data.ucomp.mean('lon')).sel(pfull=plev)
v_anom = (data.vcomp - data.vcomp.mean('lon')).sel(pfull=plev)
z_anom = mc.grav*(data.height - data.height.mean('lon')).sel(pfull=plev)
f_x = prefac_a * (v_anom**2 - prefac_b * gr.ddx(v_anom * z_anom))
f_y = prefac_a * (-u_anom * v_anom + prefac_b * gr.ddx(u_anom * z_anom))
return f_x, f_y
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 abs_vort_dt_at_pcent(run, rot_fac=1., lev=150., lat_bound=45.):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
data['precipitation'] = make_sym(data.precipitation)
precip_centroid(data, lat_bound=lat_bound)
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
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) * 86400.
vor = make_sym(vor, asym=True)
# Take time derivative of absolute vorticity
dvordt = gr.ddt(vor.mean('lon')) * 86400.
# Select latitudes over which to evaluate precip centroid
lats = [
data.lat[i] for i in range(len(data.lat))
if data.lat[i] >= -lat_bound and data.lat[i] <= lat_bound
]
dvordt_lats = dvordt.sel(lat=lats).values
#f_lats = f.sel(lat=lats).values
# Interpolate vorticity tendency and f in latitude
lats_new = np.arange(-lat_bound, lat_bound + 0.1, 0.1)
fun = spint.interp1d(lats,
dvordt_lats,
axis=-1,
fill_value='extrapolate',
kind='quadratic')
dvordt_new = fun(lats_new)
dvordt = xr.DataArray(dvordt_new,
coords=[data.xofyear.values, lats_new],
dims=['xofyear', 'lat'])
#fun = spint.interp1d(lats, f_lats, axis=-1, fill_value='extrapolate', kind='quadratic')
#f_new = fun(lats_new)
#f = xr.DataArray(f_new, coords=[lats_new], dims=['lat'])
dvordt_pcent = [
float(dvordt[i, :].sel(lat=data.p_cent.values[i]).values)
for i in range(len(data.xofyear))
]
#f_pcent = [float(f.sel(lat=data.p_cent.values[i]).values*86400.) for i in range(len(data.xofyear))]
dvordt_pcent = xr.DataArray(np.asarray(dvordt_pcent),
coords=[data.xofyear.values],
dims=['xofyear'])
return dvordt_pcent
def vor_plot(run, ax, tf, linestyle='-', rot_fac=1., lev=150.):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
# Calculate vorticity
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
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) * 86400.
div = (gr.ddx(data.ucomp) + gr.ddy(data.vcomp)).sel(pfull=lev)
vor = (make_sym(vor, asym=True)).mean('lon')
div = (make_sym(div)).mean('lon') * 8640.
m = ((mc.omega * mc.a**2. * np.cos(data.lat * np.pi / 180.)**2. +
data.ucomp.mean('lon') * mc.a * np.cos(data.lat * np.pi / 180.)).sel(
pfull=lev)) / (mc.omega * mc.a**2.)
#dvordt = gr.ddt(vor.mean('lon'))*86400.
#stretching = (-86400. * vor * div).mean('lon')
#adv = -86400. * (data.vcomp.sel(pfull=lev) * gr.ddy(vor, vector=False)).mean('lon')
#dvordy = gr.ddy(vor, vector=False).mean('lon')*8640.
#dvordy[tf[0]:tf[1],:].mean('xofyear').plot(ax=ax, color='r', linestyle=linestyle)
#m[:,tf[0]:tf[1]].mean('xofyear').plot(ax=ax, color='r', linestyle=linestyle)
vor[tf[0]:tf[1], :].mean('xofyear').plot(ax=ax,
color='k',
linestyle=linestyle)
#stretching[tf[0]:tf[1],:].mean('xofyear').plot(ax=ax, color='b', linestyle=linestyle)
#adv[tf[0]:tf[1],:].mean('xofyear').plot(ax=ax, color='r', linestyle=linestyle)
#(stretching + adv)[tf[0]:tf[1],:].mean('xofyear').plot(ax=ax, color='r', linestyle=linestyle)
ax.set_title('')
ax.set_xlabel('')
#ax.set_ylim(0.5,1.)
ax.set_ylim(-10, 10)
ax.grid(True, linestyle=':')
ax.set_ylabel('Absolute vorticity, day$^{-1}$')
def abs_vort_dt_plot(run, rot_fac=1., lev=150.):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
psi /= 1.e9
data['precipitation'] = make_sym(data.precipitation)
precip_centroid(data)
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
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) * 86400.
vor = make_sym(vor, asym=True)
psi = make_sym(psi, asym=True)
# Take time derivative of absolute vorticity
dvordt = gr.ddt(vor.mean('lon')) * 86400.
plot_dir = '/scratch/rg419/plots/paper_2_figs/abs_vort_dt/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
rcParams['figure.figsize'] = 5.5, 4.3
rcParams['font.size'] = 16
fig = plt.figure()
ax1 = fig.add_subplot(111)
f1 = dvordt.plot.contourf(ax=ax1,
x='xofyear',
y='lat',
levels=np.arange(-0.06, 0.07, 0.01),
add_colorbar=False,
add_labels=False)
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(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.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('Absolute vorticity tendency, day$^{-2}$')
plt.savefig(plot_dir + 'abs_vort_dt_' + run + '.pdf', format='pdf')
plt.close()
def div_plot(run, ax, 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)
# Calculate divergence
u_dx = gr.ddx(data.ucomp) # dudx
v_dy = gr.ddy(data.vcomp) # dvdy
div = (u_dx + v_dy).sel(pfull=lev) * 86400.
div = make_sym(div)
div = div.mean('lon')
f1 = div.plot.contourf(ax=ax,
x='xofyear',
y='lat',
levels=np.arange(-1.2, 1.3, 0.2),
add_colorbar=False,
add_labels=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_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.set_xlabel('Pentad')
ax.grid(True, linestyle=':')
box = ax.get_position()
axColor = plt.axes(
[box.x0 + box.width * 0.92, box.y0 + box.y0 * 0.12, 0.015, box.height])
cb1 = fig.colorbar(f1, cax=axColor, orientation='vertical')
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()
data_u = data_u['var33'].load().loc['1958-01':'2016-12']
# v has different time coord to u, presumably due to how Stephen has downloaded/averaged. I think the two are equivalent, so just substitute the time dimension into v
data_v_temp = xr.open_dataset(
'/disca/share/reanalysis_links/jra_55/1958_2016/vcomp_monthly/atmos_monthly_together.nc'
)
data_v = xr.DataArray(data_v_temp.var34.values,
coords={
'time': data_u.time,
'lev': data_v_temp.lev,
'lat': data_v_temp.lat,
'lon': data_v_temp.lon
},
dims=['time', 'lev', 'lat', 'lon'])
dvdx = gr.ddx(data_v, a=6371.0e3)
dudy = gr.ddy(data_u, a=6371.0e3)
data_vo = (dvdx - dudy) * 86400.
data_t = xr.open_dataset(
'/disca/share/reanalysis_links/jra_55/1958_2016/temp_monthly/atmos_monthly_together.nc'
)
data_q = xr.open_dataset(
'/disca/share/rg419/JRA_55/sphum_monthly/atmos_monthly_together.nc')
data_z = xr.open_dataset(
'/disca/share/reanalysis_links/jra_55/1958_2016/height_monthly/atmos_monthly_together.nc'
)
data_mse = (mc.cp_air * data_t.var11 + mc.L * data_q.var51 +
9.81 * data_z.var7) / 1000.
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 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 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 precip_u_hms_isca(run, lonin=[110., 120.], rotfac=1.):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
lons = [
data.lon[i].values for i in range(len(data.lon))
if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]
]
fig, (ax1, ax2, ax3) = plt.subplots(3, sharex=True)
(data.precipitation * 86400.).sel(lon=lons).mean('lon').plot.contourf(
ax=ax1,
x='xofyear',
y='lat',
levels=np.arange(2., 15., 2.),
add_colorbar=False,
add_labels=False,
cmap='Blues')
(data.precipitation * 86400.).sel(lon=lons).mean('lon').plot.contour(
ax=ax1,
x='xofyear',
y='lat',
levels=np.arange(6., 1006., 1000.),
colors='k',
add_labels=False)
u_850 = data.ucomp.sel(pfull=850.).sel(lon=lons).mean('lon')
v_850 = data.vcomp.sel(pfull=850.).sel(lon=lons).mean('lon')
u_200 = data.ucomp.sel(pfull=200.).sel(lon=lons).mean('lon')
v_200 = data.vcomp.sel(pfull=200.).sel(lon=lons).mean('lon')
u_850.plot.contourf(ax=ax2,
x='xofyear',
y='lat',
cmap='Greys',
levels=np.arange(0., 1001., 1000.),
add_labels=False,
add_colorbar=False,
extend='neither')
b = ax2.quiver(data.xofyear, data.lat, u_850.T, v_850.T, scale=400.)
ax2.add_patch(patches.Rectangle((0, -15), 7.5, 10, facecolor='0.9'))
ax2.quiverkey(b,
0.05,
0.05,
10,
'10 m/s',
fontproperties={
'weight': 'bold',
'size': 10
},
color='k',
labelcolor='k',
labelsep=0.03)
omega = 7.2921150e-5 * rotfac
f = 2 * omega * np.sin(data.lat * np.pi / 180)
vort = ((gr.ddx(data.vcomp.sel(pfull=200.)) - gr.ddy(
data.ucomp.sel(pfull=200.))).sel(lon=lons).mean('lon') + f) * 86400.
#vort = ((- gr.ddy(data.ucomp.sel(pfull=200.))).sel(lon=lons).mean('lon') + f) * 86400.
vort.plot.contourf(ax=ax3,
x='xofyear',
y='lat',
levels=np.arange(-12., 12., 2.),
add_labels=False,
add_colorbar=False)
b = ax3.quiver(data.xofyear, data.lat, u_200.T, v_200.T, scale=1000.)
ax3.add_patch(patches.Rectangle((0, -15), 7.5, 10, facecolor='0.9'))
ax3.quiverkey(b,
0.05,
0.05,
50,
'50 m/s',
fontproperties={
'weight': 'bold',
'size': 10
},
color='k',
labelcolor='k',
labelsep=0.03)
ax3.set_xlabel('Pentad')
ax1.set_ylabel('Latitude')
ax2.set_ylabel('Latitude')
ax3.set_ylabel('Latitude')
ax1.set_ylim([-15, 45])
ax2.set_ylim([-15, 45])
ax3.set_ylim([-15, 45])
ax1.grid(True, linestyle=':', linewidth=2, color='k')
ax2.grid(True, linestyle=':', linewidth=2, color='k')
ax3.grid(True, linestyle=':', linewidth=2, color='k')
plt.savefig(plot_dir + 'precip_u_hms_' + run + '.pdf', format='pdf')
plt.close()
def h_w_mass_flux_monthly(run, lev=500., dp=5000.):
data = xr.open_dataset(
'/scratch/rg419/obs_and_reanalysis/era_v_clim_alllevs.nc')
data_u = xr.open_dataset(
'/scratch/rg419/obs_and_reanalysis/era_u_clim_alllevs.nc')
plot_dir = '/scratch/rg419/plots/overturning_monthly/era/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data.coords['month'] = (data.xofyear - 1) // 6 + 1
data = data.groupby('month').mean(('xofyear'))
# Create a VectorWind instance to handle the computation
w = VectorWind(data.ucomp.sel(pfull=np.arange(50., 950., 50.)),
data.vcomp.sel(pfull=np.arange(50., 950., 50.)))
# Compute variables
streamfun, vel_pot = w.sfvp()
uchi, vchi, upsi, vpsi = w.helmholtz()
coslat = np.cos(data.lat * np.pi / 180)
# Evaluate mass fluxes for the zonal and meridional components (Walker and Hadley) following Schwendike et al. 2014
mass_flux_zon = (gr.ddx(uchi)).cumsum('pfull') * dp * coslat / mc.grav
mass_flux_merid = (gr.ddy(vchi)).cumsum('pfull') * dp * coslat / mc.grav
# Set figure parameters
rcParams['figure.figsize'] = 15, 11
rcParams['font.size'] = 14
# Start figure with 12 subplots
fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8),
(ax9, ax10, ax11, ax12)) = plt.subplots(3,
4,
sharex='col',
sharey='row')
axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12]
i = 0
for ax in axes:
f1 = mass_flux_merid.sel(pfull=lev)[i, :, :].plot.contourf(
ax=ax,
x='lon',
y='lat',
add_labels=False,
add_colorbar=False,
levels=np.arange(-0.0065, 0.0066, 0.001),
extend='both')
mass_flux_zon.sel(pfull=lev)[i, :, :].plot.contour(ax=ax,
x='lon',
y='lat',
add_labels=False,
colors='k',
levels=np.arange(
0.0005, 0.0066,
0.001))
mass_flux_zon.sel(pfull=lev)[i, :, :].plot.contour(
ax=ax,
x='lon',
y='lat',
add_labels=False,
colors='0.5',
levels=np.arange(-0.0065, -0.00049, 0.001))
i = i + 1
ax.set_ylim(-60, 60)
ax.set_xticks(np.arange(0, 361, 90))
ax.set_yticks(np.arange(-60, 61, 30))
ax.grid(True, linestyle=':')
plt.subplots_adjust(left=0.05,
right=0.97,
top=0.95,
bottom=0.1,
hspace=0.2,
wspace=0.2)
cb1 = fig.colorbar(f1,
ax=axes,
use_gridspec=True,
orientation='horizontal',
fraction=0.05,
pad=0.1,
aspect=30,
shrink=0.5)
cb1.set_label(
'Vertical mass flux associated with meridional circulation, kgm$^{-2}$s$^{-1}$'
)
figname = plot_dir + 'h_w_' + run + '.pdf'
plt.savefig(figname, format='pdf')
plt.close()
def 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 vort_budg_terms(run, lonin=[-1.,361.], rot_fac=1.):
'''Evaluate pentad mean vorticity budget and differences from daily snapshot budget RG 3/11/2017
Imputs: run = run_name
lonin = longitude range to average over
do_ss = run is steady state
rot_fac = scale factor for Earth's rotation rate
planetary_only = only plot planetary vorticity terms
no_eddies = don't plot transient eddies too
ll = keep longitude dimension for a lat-lon plot'''
#Load in vorticity budget term means
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/vort_eq_'+run+'.nc')
if run in ['ap_2', 'full_qflux']:
data_vort = data
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/vort_eq_uv'+run+'.nc')
print('vorticity budget data loaded')
if lonin[1]>lonin[0]:
lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]]
else:
lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]]
# Calculate vertical component of pentad mean vorticity parts: f and dv/dx - du/dy
omega = 7.2921150e-5 * rot_fac
f = 2 * omega * np.sin(data.lat *np.pi/180)
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
# Evaluate relative vorticity
vor = v_dx - u_dy
# Take divergencc of u*vor
dvorudx = gr.ddx(vor*data.ucomp)
dvorvdy = gr.ddy(vor*data.vcomp)
div_uvor = -86400.**2 * (dvorudx + dvorvdy)
dfudx = gr.ddx(data.ucomp * f)
dfvdy = gr.ddy(data.vcomp * f)
div_uf = -86400.**2 * (dfudx + dfvdy)
# If run is ap_2 or full_qflux, then vorticity budget and velocities were saved separately. Load the budget up now
if run in ['ap_2','full_qflux']:
data = data_vort
# Another glitch with ap_2 and full_qflux - these are only on 150 hPa.
coord_dict = {'xofyear': ('xofyear', data.pentad),
'lat': ('lat', data.lat)}
dim_list = ['xofyear', 'lat']
# For a Hovmoller, keep time, pressure, and lat
else:
coord_dict = {'xofyear': ('xofyear', data.pentad),
'pfull': ('pfull', data.pfull),
'lat': ('lat', data.lat)}
dim_list = ['xofyear', 'pfull', 'lat']
# Specify output dictionary to be written
output_dict = {'div_uvor': (dim_list, div_uvor.sel(lon=lons).mean('lon')),
'div_uf': (dim_list, div_uf.sel(lon=lons).mean('lon'))}
if run in ['ap_2','full_qflux']:
# Single level transients from ap_2, full_qflux mistakes
transient_hm = ((data.stretching.sel(pfull=150).values + data.horiz_md.sel(pfull=150).values) * 86400.**2.
- (div_uvor + div_uf).sel(pfull=150))
else:
# Average out longitude for other plot types
transient_hm = ((data.stretching + data.horiz_md)*86400.**2. - div_uvor - div_uf).sel(lon=lons).mean('lon')
output_dict.update({'transient_hm': (dim_list, transient_hm)})
# Create a dataset of the terms to be plotted and return
ds = xr.Dataset(output_dict, coords=coord_dict)
return ds
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 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.]):
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()
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 vort_budg_terms(run, lonin=[-1.,361.], do_ss=False, rot_fac=1., planetary_only=False, no_eddies=False, ll=False):
'''Evaluate pentad mean vorticity budget and differences from daily snapshot budget RG 3/11/2017
Imputs: run = run_name
lonin = longitude range to average over
do_ss = run is steady state
rot_fac = scale factor for Earth's rotation rate
planetary_only = only plot planetary vorticity terms
no_eddies = don't plot transient eddies too
ll = keep longitude dimension for a lat-lon plot'''
#Load in vorticity budget term means
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_'+run+'.nc')
if run in ['ap_2', 'full_qflux']:
data_vort = data
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_uv'+run+'.nc')
print('vorticity budget data loaded')
if lonin[1]>lonin[0]:
lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]]
else:
lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]]
# Calculate vertical component of pentad mean absolute vorticity = f + dv/dx - du/dy
omega = 7.2921150e-5 * rot_fac
f = 2 * omega * np.sin(data.lat *np.pi/180)
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
# If only want to look at parts relating to planetary vorticity, substract relative part (this way turns f into a 4d array)
if planetary_only:
vor = v_dx - u_dy + f - v_dx + u_dy
else:
vor = v_dx - u_dy + f
# Take gradients of vorticity
dvordx = gr.ddx(vor)
dvordy = gr.ddy(vor, vector=False)
# Horizontal material derivative
horiz_md_mean = -86400.**2. * (data.ucomp * dvordx + data.vcomp * dvordy)
# Calculate divergence and stretching term
div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
stretching_mean = -86400.**2. * vor * div
if ll:
#Keep lon dimension if want a steady state lat-lon plot
horiz_md_av = horiz_md_mean
stretching_av = stretching_mean
dvordy_av = dvordy
vor_av = vor*86400.
v_av = data.vcomp
div_av = div*86400.
else:
# Take zonal mean of all over specified longitude range
horiz_md_av = horiz_md_mean.sel(lon=lons).mean('lon')
stretching_av = stretching_mean.sel(lon=lons).mean('lon')
dvordy_av = dvordy.sel(lon=lons).mean('lon')
vor_av = vor.sel(lon=lons).mean('lon')*86400.
v_av = data.vcomp.sel(lon=lons).mean('lon')
div_av = div.sel(lon=lons).mean('lon')*86400.
# If run is ap_2 or full_qflux, then vorticity budget and velocities were saved separately. Load the budget up now
if run in ['ap_2','full_qflux']:
data = data_vort
# For a steady state lat-pressure run, only need pfull and lat
if do_ss:
coord_dict = {'pfull': ('pfull', data.pfull),
'lat': ('lat', data.lat)}
dim_list = ['pfull', 'lat']
# For a steady state lat-lon run, also need lon. Keep pfull to allow a choice of plotting level
elif ll:
if 'pentad' in horiz_md_av.coords:
coord_dict = {'xofyear': ('xofyear', data.pentad),
'pfull': ('pfull', data.pfull),
'lat': ('lat', data.lat),
'lon': ('lon', data.lon)}
dim_list = ['xofyear', 'pfull', 'lat', 'lon']
else:
coord_dict = {'pfull': ('pfull', data.pfull),
'lat': ('lat', data.lat),
'lon': ('lon', data.lon)}
dim_list = ['pfull', 'lat', 'lon']
# Another glitch with ap_2 and full_qflux - these are only on 150 hPa.
elif run in ['ap_2','full_qflux']:
coord_dict = {'xofyear': ('xofyear', data.pentad),
'lat': ('lat', data.lat)}
dim_list = ['xofyear', 'lat']
# For a Hovmoller, keep time, pressure, and lat
else:
coord_dict = {'xofyear': ('xofyear', data.pentad),
'pfull': ('pfull', data.pfull),
'lat': ('lat', data.lat)}
dim_list = ['xofyear', 'pfull', 'lat']
# Specify output dictionary to be written
output_dict = {'horiz_md': (dim_list, horiz_md_av),
'stretching': (dim_list, stretching_av),
'dvordy': (dim_list, dvordy_av),
'v': (dim_list, v_av),
'vor': (dim_list, vor_av),
'div': (dim_list, div_av)}
if not (planetary_only or no_eddies):
if run in ['ap_2','full_qflux']:
# Single level transients from ap_2, full_qflux mistakes
transient_s_hm = (data.stretching.sel(pfull=150).values * 86400.**2. - stretching_mean).sel(lon=lons).mean('lon')
transient_h_hm = (data.horiz_md.sel(pfull=150).values * 86400.**2. - horiz_md_mean).sel(lon=lons).mean('lon')
elif ll:
# Lat lon transients
transient_s_hm = data.stretching.values * 86400.**2. - stretching_mean
transient_h_hm = data.horiz_md.values * 86400.**2. - horiz_md_mean
else:
# Average out longitude for other plot types
transient_s_hm = (data.stretching.values * 86400.**2. - stretching_mean).sel(lon=lons).mean('lon')
transient_h_hm = (data.horiz_md.values * 86400.**2. - horiz_md_mean).sel(lon=lons).mean('lon')
transient_hm = transient_s_hm + transient_h_hm
output_dict.update({'transient_hm': (dim_list, transient_hm),
'transient_s_hm': (dim_list, transient_s_hm),
'transient_h_hm': (dim_list, transient_h_hm)})
# Create a dataset of the terms to be plotted and return
ds = xr.Dataset(output_dict, coords=coord_dict)
return ds
data_v = data_v['vwnd'].load().loc['1948-01':'2016-12']
print('v loaded')
data_t = xr.open_mfdataset(
'/disca/share/reanalysis_links/NCEP_NCAR/air.mon.mean.nc')
data_t = data_t['air'].load().loc['1948-01':'2016-12'] + 273.15
print('t loaded')
data_q = xr.open_mfdataset(
'/disca/share/reanalysis_links/NCEP_NCAR/shum.mon.mean.nc')
data_q = data_q['shum'].load().loc['1948-01':'2016-12'] / 1000.
print('q loaded')
data_z = xr.open_mfdataset(
'/disca/share/reanalysis_links/NCEP_NCAR/hgt.mon.mean.nc')
data_z = data_z['hgt'].load().loc['1948-01':'2016-12'] * 9.81
print('z loaded')
dvdx = gr.ddx(data_v, a=6371.0e3)
dudy = gr.ddy(data_u, a=6371.0e3)
data_vo = (dvdx.values - dudy.values) * 86400.
data_vo = xr.DataArray(data_vo, dims=data_u.dims, coords=data_u.coords)
lh = 2.5e6
cp = 1005.
data_mse = (data_q * lh + data_t * cp + data_z) / 1000.
print('Data loaded ok')
land_mask = '/scratch/rg419/python_scripts/land_era/ERA-I_Invariant_0125.nc'
land = xr.open_dataset(land_mask)
def plot_winter_climate(data,
title,
levels_clim=np.arange(-50., 51., 5.),
data_mse = (data_q * lh + data_t * cp + data_z) / 1000.
data_vo = data_vo * 86400.
# Create a VectorWind instance to handle the computation
w = VectorWind(data_u.sel(level=np.arange(50., 950., 50.)),
data_v.sel(level=np.arange(50., 950., 50.)))
# Compute variables
streamfun, vel_pot = w.sfvp()
uchi, vchi, upsi, vpsi = w.helmholtz()
coslat = np.cos(data_u.latitude * np.pi / 180)
dp = 5000.
# Evaluate mass fluxes for the zonal and meridional components (Walker and Hadley) following Schwendike et al. 2014
mass_flux_zon = (gr.ddx(
uchi, latname='latitude',
lonname='longitude')).cumsum('level') * dp * coslat / 9.81
mass_flux_merid = (gr.ddy(
vchi, latname='latitude')).cumsum('level') * dp * coslat / 9.81
print('Data loaded ok')
land_mask = '/scratch/rg419/python_scripts/land_era/ERA-I_Invariant_0125.nc'
land = xr.open_dataset(land_mask)
def plot_winter_climate(data,
title,
levels_clim=np.arange(-50., 51., 5.),
levels=np.arange(-5., 5.1, 0.5),
local=False,
def stretching_plot(run, ax_h, ax_s, ax_v, pentad, tibet=True):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_' +
run + '.nc')
land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc'
land = xr.open_dataset(land_file)
lons = [
data.lon[i] for i in range(len(data.lon))
if data.lon[i] >= 60. and data.lon[i] < 150.
]
lats = [
data.lat[i] for i in range(len(data.lat))
if data.lat[i] >= -30. and data.lat[i] < 60.
]
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
v_dx = gr.ddx(data.vcomp.sel(pfull=150)) # dvdx
u_dy = gr.ddy(data.ucomp.sel(pfull=150)) # dudy
vor = v_dx - u_dy + f
div = gr.ddx(data.ucomp.sel(pfull=150)) + gr.ddy(data.vcomp.sel(pfull=150))
dvordx = gr.ddx(vor)
dvordy = gr.ddy(vor, vector=False)
horiz_md_mean = -86400.**2. * (data.ucomp.sel(pfull=150) * dvordx +
data.vcomp.sel(pfull=150) * dvordy)
stretching_mean = -86400.**2. * vor * div
levels = np.arange(-1.5, 1.6, 0.5)
fh = horiz_md_mean[pentad, :, :].plot.contourf(x='lon',
y='lat',
levels=levels,
ax=ax_h,
add_labels=False,
extend='both',
add_colorbar=False)
land.land_mask.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
ax=ax_h,
colors='k',
add_colorbar=False,
add_labels=False)
if tibet:
land.zsurf.plot.contourf(x='lon',
y='lat',
ax=ax_h,
levels=np.arange(2500., 100001., 97500.),
add_labels=False,
extend='neither',
add_colorbar=False,
alpha=0.5,
cmap='Greys_r')
fs = stretching_mean[pentad, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(
-1.5, 1.6, 0.5),
ax=ax_s,
add_labels=False,
extend='both',
add_colorbar=False)
land.land_mask.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
ax=ax_s,
colors='k',
add_colorbar=False,
add_labels=False)
if tibet:
land.zsurf.plot.contourf(x='lon',
y='lat',
ax=ax_s,
levels=np.arange(2500., 100001., 97500.),
add_labels=False,
extend='neither',
add_colorbar=False,
alpha=0.5,
cmap='Greys_r')
fv = (vor[pentad, :, :] * 86400.).plot.contourf(x='lon',
y='lat',
levels=np.arange(
-12., 13., 2.),
ax=ax_v,
add_labels=False,
extend='both',
add_colorbar=False)
land.land_mask.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
ax=ax_v,
colors='k',
add_colorbar=False,
add_labels=False)
if tibet:
land.zsurf.plot.contourf(x='lon',
y='lat',
ax=ax_v,
levels=np.arange(2500., 100001., 97500.),
add_labels=False,
extend='neither',
add_colorbar=False,
alpha=0.5,
cmap='Greys_r')
for ax in [ax_h, ax_s, ax_v]:
ax.set_xlim(60, 150)
ax.set_ylim(-30, 60)
ax.set_xticks(np.arange(60., 151., 30.))
ax.set_yticks(np.arange(-30., 61., 30.))
ax.grid(True, linestyle=':')
ax_h.set_xticklabels('')
ax_s.set_xticklabels('')
ax_v.set_xticklabels(['60', '90', '120', '150'])
return fh, fs, fv
def abs_vort_dt_plot(run, ax, rot_fac=1., lev=150., dvordt_flag=False):
'''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)
# Calculate vorticity
v_dx = gr.ddx(data.vcomp) # dvdx
u_dy = gr.ddy(data.ucomp) # dudy
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)*86400.
vor = make_sym(vor, asym=True)
# Take time derivative of absolute vorticity
dvordt = gr.ddt(vor.mean('lon'))*86400.
# Also normalise this by the value of vorticity
dvordtvor = dvordt/vor.mean('lon')
# Plot!
plot_dir = '/scratch/rg419/plots/paper_2_figs/abs_vort_dt/'
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)
if dvordt_flag:
f1 = dvordt.plot.contourf(ax=ax, x='xofyear', y='lat', levels=np.arange(-0.06,0.07,0.01), add_colorbar=False, add_labels=False)
else:
f1 = dvordtvor.plot.contourf(ax=ax, x='xofyear', y='lat', levels=np.arange(-0.05,0.055,0.005), add_colorbar=False, add_labels=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_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.set_xlabel('Pentad')
ax.grid(True,linestyle=':')
#originalSize = get(gca, 'Position')
#set(f1, 'Position', originalSize)
#plt.subplots_adjust(left=0.15, right=0.95, top=0.95, bottom=0.05)
box = ax.get_position()
#ax.set_position([box.x0*1.05, box.y0, box.width, box.height])
axColor = plt.axes([box.x0 + box.width*0.92, box.y0+box.y0*0.12, 0.015, box.height])
#cb1=fig.colorbar(f1, ax=ax, use_gridspec=True, orientation = 'horizontal',fraction=0.15, pad=0.2, aspect=40)
#cb1=plt.colorbar(f1, ax=axColor, use_gridspec=True, orientation = 'vertical',fraction=0.15, pad=0.05, aspect=30)
cb1=fig.colorbar(f1, cax=axColor, orientation = 'vertical')
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 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')
