def precip_psi_plot(run, ax):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
f1 = precip_mse_plot(data, ax, plot_type='precip', precip_contour=None)
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
psi /= 1.e9
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(-500., 510., 500.),
add_labels=False,
colors='0.5',
linewidths=2)
return f1
def plot_overturning(data_u, data_v, ax, lonin=[-1.,361.], region='All', sanity_check=False):
lons = get_lons(lonin,data_u)
levs = np.arange(50.,1000.,50.)
ds_summer = xr.Dataset({'ucomp': (['pfull', 'lat', 'lon'], data_u.var33.values), 'vcomp': (['pfull', 'lat', 'lon'], data_v.var34.values)},
coords={'pfull': ('pfull', v_summer_a.lev/100.),
'lat': ('lat', v_summer_a.lat),
'lon': ('lon', v_summer_a.lon)})
ds_summer = ds_summer.sel(pfull=levs)
psi = mass_streamfunction(ds_summer, lons=lons, dp_in=50.)#, use_v_locally=True)
psi /= 1.e9
u = ds_summer.ucomp
f1 = u.sel(lon=lons).mean('lon').plot.contourf(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(-50.,50.1,5.), extend='both', add_labels=False, add_colorbar=False)
psi.plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(0.,601,50.), colors='k', add_labels=False)
psi.plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(-600.,0.,50.), colors='k', linestyles='dashed', add_labels=False)
m = mc.omega * mc.a**2. * np.cos(u.lat*np.pi/180.)**2. + u.sel(lon=lons).mean('lon') * mc.a * np.cos(u.lat*np.pi/180.)
m_levs = mc.omega * mc.a**2. * np.cos(np.arange(-60.,1.,5.)*np.pi/180.)**2.
m.plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=m_levs, colors='k', alpha=0.7, add_labels=False)
ax.set_xlim(-35,35)
ax.set_xticks(np.arange(-30,31,15))
ax.grid(True,linestyle=':')
ax.set_title(region)
plt.subplots_adjust(left=0.05, right=0.97, top=0.95, bottom=0.05, hspace=0.3, wspace=0.2)
return f1
def psi_plot(run, ax, title=''):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
mse = (mc.cp_air * data.temp + mc.L * data.sphum +
mc.grav * data.height).mean('lon').sel(pfull=850.) / 1000.
dmsedy = gr.ddy(mse, vector=False) * 100000.
dtsurfdy = gr.ddy(data.t_surf.mean('lon'), vector=False) * 100000.
psi /= 1.e9
psi = np.abs(psi).max('pfull')
#mse.plot.contour(ax=ax, x='xofyear', y='lat', add_labels=False, levels=np.arange(200.,401.,10.), colors='w')
dmsedy.plot.contourf(ax=ax,
x='xofyear',
y='lat',
add_labels=False,
levels=np.arange(-3., 3.1, 0.25),
add_colorbar=False)
#dtsurfdy.plot.contourf(ax=ax, x='xofyear', y='lat', add_labels=False, levels=np.arange(-1.5,1.6,0.1), add_colorbar=False)
psi.plot.contour(ax=ax,
x='xofyear',
y='lat',
extend='both',
add_labels=False,
levels=np.arange(0, 101., 100.),
add_colorbar=False,
alpha=0.4)
ax.set_title(title, fontsize=17)
ax.set_ylim(-60, 60)
ax.grid(True, linestyle=':')
ax.set_yticks(np.arange(-60., 61., 30.))
ax.set_xticks(np.arange(0., 72., 18.))
def psi_mean_clim(run):
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
area = cell_area(42, '/scratch/rg419/Isca/')
area_xr = xr.DataArray(area, [('lat', data.lat ), ('lon', data.lon)])
area_xr = area_xr.mean('lon')
psi_mean = ((psi*area_xr).sum(('lat'))/area_xr.sum(('lat'))).mean('pfull')
#psi_mean = psi_mean*-1.
return psi_mean
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 get_streamfun_zeros(data, lev=150., sanity_check=False):
# Calculate mass streamfunction
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
psi /= 1.e9
psi = psi.sel(pfull=lev)
n = len(psi.xofyear.values)
psi_mask = np.ones(psi.shape)
psi_mask[psi <= 0] = 0.
psi_mask = xr.DataArray(psi_mask, coords=psi.coords, dims=['lat', 'xofyear'])
borders = psi_mask.diff('lat')
# Identify Hadley cell edges and ITCZ by:
# ITCZ: positive border closest to the equator
# cell edges: negative border poleward of ITCZ
lat_itcz = np.zeros(n,)
lat_nh = np.zeros(n,)
lat_sh = np.zeros(n,)
for i in range(n):
lats_pos = data.latb[1:-1].values[np.where(borders[:,i] == 1)[0]]
lats_pos = lats_pos[(lats_pos>-30.) & (lats_pos<30.)]
if len(lats_pos) > 1:
if np.all(lats_pos>0.):
lats_pos = np.max(lats_pos)
elif np.all(lats_pos<0.):
lats_pos = np.min(lats_pos)
else:
lats_pos = lats_pos[ np.abs(lats_pos)==np.max(np.abs(lats_pos))]
lat_itcz[i] = lats_pos
lats_neg = data.latb[1:-1].values[np.where(borders[:,i] == -1)[0]]
lat_nh[i] = np.min(lats_neg[(lats_neg - lat_itcz[i] > 0) & (lats_neg > 0)])
lat_sh[i] = np.max(lats_neg[(lats_neg - lat_itcz[i] < 0) & (lats_neg < 0)])
if sanity_check == True:
psi.plot.contourf()
plt.plot(data.xofyear, lat_itcz,'k')
plt.plot(data.xofyear, lat_nh,'k')
plt.plot(data.xofyear, lat_sh,'k')
plt.show()
return lat_itcz, lat_nh, lat_sh
def overturning_plot(data,
ax_in,
lonin=[-1., 361.],
do_xlabels=False,
month_labels=True,
thresh=0.,
nh=False):
lons = pick_lons(data, lonin)
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50., lons=lons)
psi /= 1.e9
edge_loc, psi_max, psi_max_loc = get_edge_psi(data,
lonin=lonin,
lev=500.,
thresh=thresh,
nh=nh)
f1 = psi.sel(pfull=500.).plot.contourf(ax=ax_in,
x='xofyear',
y='lat',
levels=np.arange(-500., 510., 50.),
add_colorbar=False,
add_labels=False,
extend='both')
edge_loc.plot(color='k', ax=ax_in)
psi_max_loc.plot(color='k', ax=ax_in)
ax_in.set_ylabel('Latitude')
ax_in.set_ylim(-60, 60)
ax_in.set_yticks(np.arange(-60., 61., 30.))
ax_in.grid(True, linestyle=':')
if month_labels:
mn_dic = month_dic(1)
tickspace = list(range(13, 72, 18))
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
ax_in.set_xlim((1, 72))
ax_in.set_xticks(tickspace)
if do_xlabels:
ax_in.set_xlabel('')
ax_in.set_xticklabels(labels, rotation=25)
return f1
def psi_u_plot(data, tf, ax):
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
psi /= 1.e9
n = len(data.xofyear.values) // 2
psi_temp = np.zeros(psi.values.shape)
for i in range(0, n):
psi_temp[:, i, :] = (psi[:, i, :].values -
psi[::-1, i + n, :].values) / 2.
psi_temp[:, i + n, :] = -1. * psi_temp[::-1, i, :]
psi = xr.DataArray(psi_temp,
coords=[data.lat, data.xofyear.values, psi.pfull],
dims=['lat', 'xofyear', 'pfull'])
f1 = data.ucomp[tf[0]:tf[1], :, :].mean(
('xofyear', 'lon')).plot.contourf(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(-50., 50.1, 5.),
extend='both',
add_labels=False,
add_colorbar=False)
psi[:, tf[0]:tf[1], :].mean('xofyear').plot.contour(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(
0., 301, 60.),
colors='k',
add_labels=False)
psi[:, tf[0]:tf[1], :].mean('xofyear').plot.contour(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(
-300., 0., 60.),
colors='k',
linestyles='dashed',
add_labels=False)
ax.set_xlim(-60, 60)
ax.grid(True, linestyle=':')
return f1
def sf_spinup(run, months_list, filenames=['plev_pentad']):
# Function to open files for a specfied month range and filename.
def open_files(run, months, filename):
name_temp = '/scratch/rg419/Data_moist/' + run + '/run%04d/'+filename+'.nc'
names = [name_temp % m for m in range( months[0], months[1]) ]
data = xr.open_mfdataset( names, decode_times=False, chunks={'time': 30})
# Reduce dataset so that only vcomp is retained
data = xr.Dataset({'vcomp': data.vcomp}, coords=data.vcomp.coords)
#data.coords['month'] = data.time//30 + 1
#data = data.groupby('month').mean(('time'))
return data
arrays = []
i=0
for filename in filenames:
data = open_files(run, months_list[i], filename)
arrays.append(data)
i=i+1
data = xr.concat(arrays, dim='time')
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
psi /= 1.e9
area = cell_area(42, '/scratch/rg419/Isca/')
area_xr = xr.DataArray(area, [('lat', data.lat ), ('lon', data.lon)])
area_xr = area_xr.mean('lon')
psi_mean = ((psi*area_xr).sum(('lat'))/area_xr.sum(('lat'))).mean('pfull')
psi_mean = psi_mean*-1.
# The rolling function will bring up alot of true_divide errors, but still works.
# Can silence these by writing them to a file (although maybe this is slower)
temp = sys.stderr # store original stdout object for later
sys.stderr = open('log.txt', 'w') # redirect all prints to this log file
#psi_mean = psi_mean.rolling(time=60).mean()
#adj_time_70 = np.min(psi_mean.time.where(psi_mean >= 70.,drop=True).values)
adj_time = np.min(psi_mean.time.where(psi_mean >= 74.,drop=True).values)
#adj_time_75 = np.min(psi_mean.time.where(psi_mean >= 75.,drop=True).values)
sys.stderr.close() # ordinary file object
sys.stderr = temp
return psi_mean, adj_time
def precip_psi_plot(run, ax, label='a)', p_cent=True):
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 = make_sym(psi, asym=True)
data['precipitation'] = make_sym(data.precipitation)
f1 = precip_mse_plot(data,
ax,
plot_type='precip',
precip_contour=None,
p_cent=p_cent)
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='--',
alpha=0.7)
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,
alpha=0.7)
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,
alpha=0.7)
ax.text(-10, 60, label)
return f1
def psi_u_plot(data, tf, ax):
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
psi /= 1.e9
f1 = data.ucomp[tf[0]:tf[1], :, :].mean(
('xofyear', 'lon')).plot.contourf(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(-50., 50.1, 5.),
extend='both',
add_labels=False,
add_colorbar=False)
psi[:, tf[0]:tf[1], :].mean('xofyear').plot.contour(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(
0., 301, 60.),
colors='k',
add_labels=False)
psi[:, tf[0]:tf[1], :].mean('xofyear').plot.contour(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(
-300., 0., 60.),
colors='k',
linestyles='dashed',
add_labels=False)
ax.set_xlim(-60, 60)
ax.grid(True, linestyle=':')
return f1
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 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 overturning_monthly(run, lonin=[-1.,361.]):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc')
plot_dir = '/scratch/rg419/plots/overturning_monthly_ageo/' + run + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data.coords['month'] = (data.xofyear - 1) //6 + 1
data = data.groupby('month').mean(('xofyear'))
#Coriolis
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat *np.pi/180)
dphidx = gr.ddx(data.height)
dphidx = 9.8 * dphidx
v_ageo = data.vcomp - dphidx/f
v_geo = dphidx/f
dphidy = gr.ddy(data.height, vector=False)
dphidy = 9.8 * dphidy
u_ageo = data.ucomp - dphidy/f
u_geo = dphidx/f
# 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.)))
w = VectorWind(u_ageo.sel(pfull=np.arange(50.,950.,50.)), v_ageo.sel(pfull=np.arange(50.,950.,50.)))
# Compute variables
streamfun, vel_pot = w.sfvp()
uchi, vchi, upsi, vpsi = w.helmholtz()
ds_chi = xr.Dataset({'vcomp': (vchi)},
coords={'month': ('month', vchi.month),
'pfull': ('pfull', vchi.pfull),
'lat': ('lat', vchi.lat),
'lon': ('lon', vchi.lon)})
w = VectorWind(u_geo.sel(pfull=np.arange(50.,950.,50.)), v_geo.sel(pfull=np.arange(50.,950.,50.)))
# Compute variables
streamfun, vel_pot = w.sfvp()
uchi, vchi, upsi, vpsi = w.helmholtz()
ds_chi_geo = xr.Dataset({'vcomp': (vchi)},
coords={'month': ('month', vchi.month),
'pfull': ('pfull', vchi.pfull),
'lat': ('lat', vchi.lat),
'lon': ('lon', vchi.lon)})
def get_lons(lonin, data):
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]]
return lons
lons = get_lons(lonin,data)
psi = mass_streamfunction(data, lons=lons, dp_in=50.)
psi /= 1.e9
psi_chi = mass_streamfunction(ds_chi, lons=lons, dp_in=-50., intdown=False)
psi_chi /= 1.e9
psi_chi_geo = mass_streamfunction(ds_chi_geo, lons=lons, dp_in=-50., intdown=False)
psi_chi_geo /= 1.e9
# Set figure parameters
rcParams['figure.figsize'] = 10, 7
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)
axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12]
i=0
for ax in axes:
psi_chi[:,i,:].plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(0.,601,100.), colors='k', add_labels=False)
psi_chi[:,i,:].plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(-600.,0.,100.), colors='k', linestyles='dashed', add_labels=False)
i=i+1
ax.set_xlim(-35,35)
ax.set_xticks(np.arange(-30,31,15))
ax.grid(True,linestyle=':')
plt.subplots_adjust(left=0.1, right=0.97, top=0.95, bottom=0.1, hspace=0.3, wspace=0.3)
if lonin == [-1.,361.]:
plt.savefig(plot_dir + 'psi_chi_' + run + '.pdf', format='pdf')
else:
figname = plot_dir + 'psi_chi_' + run + '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(figname, format='pdf')
plt.close()
# Start figure with 12 subplots
fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8), (ax9, ax10, ax11, ax12)) = plt.subplots(3, 4)
axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12]
i=0
for ax in axes:
psi_chi_geo[:,i,:].plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(0.,601,100.), colors='k', add_labels=False)
psi_chi_geo[:,i,:].plot.contour(ax=ax, x='lat', y='pfull', yincrease=False, levels=np.arange(-600.,0.,100.), colors='k', linestyles='dashed', add_labels=False)
i=i+1
ax.set_xlim(-35,35)
ax.set_xticks(np.arange(-30,31,15))
ax.grid(True,linestyle=':')
plt.subplots_adjust(left=0.1, right=0.97, top=0.95, bottom=0.1, hspace=0.3, wspace=0.3)
if lonin == [-1.,361.]:
plt.savefig(plot_dir + 'psi_chi_geo_' + run + '.pdf', format='pdf')
else:
figname = plot_dir + 'psi_chi_geo_' + run + '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(figname, format='pdf')
plt.close()
def vort_eq_hm(run, lev=150., lonin=[-1.,361.], planetary_only=False, month_labels=True, rot_fac=1., no_eddies=False, add_psi=True):
'''Plot vorticity budget Hovmoller. RG 3/11/2017
Imputs: run = run_name
lev = level to plot
lonin = longitude range to average over
planetary_only = only plot planetary vorticity terms
month_labels = label x axis with month of year (no good for non earth year length)
rot_fac = scale factor for Earth's rotation rate
no_eddies = don't plot transient eddies too'''
# Call the above function to get fields to plot. For a Hovmoller, run is by definition not steady state. Pass other inputs on.
ds = vort_budg_terms(run, lonin=lonin, rot_fac=rot_fac, planetary_only=planetary_only, no_eddies=no_eddies)
if add_psi:
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run + '.nc')
if lonin[1]>lonin[0]:
lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]]
else:
lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]]
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50., lons=lons)
psi /= 1.e9
# Determine figure size based on number of panels
if planetary_only or no_eddies:
rcParams['figure.figsize'] = 15, 6.25
else:
rcParams['figure.figsize'] = 15, 8
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
levels = np.arange(-1.5,1.6,0.25)
print('starting plotting')
# Set number of panels and declare which are the left column, bottom row.
if planetary_only or no_eddies:
f, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2, 3, sharex='col', sharey='row')
left_column = [ax1,ax4]
bottom_row = [ax4,ax5,ax6]
all_plots = [ax1,ax2,ax3,ax4,ax5,ax6]
else:
f, ((ax1, ax2, ax3), (ax4, ax5, ax6), (ax7, ax8, ax9)) = plt.subplots(3, 3, sharex='col', sharey='row')
left_column = [ax1,ax4,ax7]
bottom_row = [ax7,ax8,ax9]
all_plots = [ax1,ax2,ax3,ax4,ax5,ax6,ax7,ax8,ax9]
plt.set_cmap('RdBu_r')
# Select the pressure level, if needed
if run in ['ap_2','full_qflux']:
ds = ds
else:
ds = ds.sel(pfull=lev)
# Plot!
f2 = ds.horiz_md.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax1, extend = 'both', add_labels=False)
if add_psi:
ax1.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax1.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax1.set_title('Horizontal advection', fontsize=17)
ax1.text(-15, 60, 'a)')
ds.dvordy.plot.contourf(x='xofyear', y='lat', levels=np.arange(-1.e-10,1.05e-10,0.1e-10), ax=ax2, extend = 'both', add_labels=False)
if planetary_only:
ax2.set_title('$\partial f /\partial y$', fontsize=17)
else:
ax2.set_title('$\partial (\overline{\zeta} + f) /\partial y$', fontsize=17)
ax2.text(-7, 60, 'b)')
ds.v.plot.contourf(x='xofyear', y='lat', levels=np.arange(-12.,13,2.), ax=ax3, extend = 'both', add_labels=False)
ax3.set_title('$\overline{v}$', fontsize=17)
ax3.text(-7, 60, 'c)')
ds.stretching.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax4, extend = 'both', add_labels=False)
if add_psi:
ax4.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax4.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax4.set_title('Vortex stretching', fontsize=17)
ax4.text(-15, 60, 'd)')
ds.div.plot.contourf(x='xofyear', y='lat', levels=np.arange(-0.6,0.7,0.1), ax=ax5, extend = 'both', add_labels=False)
ax5.set_title('Divergence', fontsize=17)
ax5.text(-7, 60, 'e)')
ds.vor.plot.contourf(x='xofyear', y='lat', levels=np.arange(-12.,13.,2.), ax=ax6, extend = 'both', add_labels=False)
if planetary_only:
ax6.set_title('Planetary vorticity', fontsize=17)
else:
ax6.set_title('Absolute vorticity', fontsize=17)
ax6.text(-7, 60, 'f)')
# Plot transients if needed
if not (planetary_only or no_eddies):
ds.transient_hm.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax7, extend = 'both', add_labels=False)
if add_psi:
ax7.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax7.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax7.set_title('Transient eddy vorticity tendency', fontsize=17)
ax7.text(-15, 60, 'g)')
ds.transient_h_hm.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax8, extend = 'both', add_labels=False)
ax8.set_title('Transient horizontal adv.', fontsize=17)
ax8.text(-7, 60, 'h)')
ds.transient_s_hm.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax9, extend = 'both', add_labels=False)
ax9.set_title('Transient vortex stretching', fontsize=17)
ax9.text(-7, 60, 'i)')
# Add grid, set y limits and ticks
for ax in all_plots:
ax.grid(True,linestyle=':')
ax.set_ylim(-60,60)
ax.set_yticks(np.arange(-60,61,30))
# Label left column
for ax in left_column:
ax.set_ylabel('Latitude')
# If month labels are being used, load in the month dictionary and set up labelling, then label bottom row
if month_labels:
mn_dic = month_dic(1)
tickspace = list(range(13,72,18))
labels = [mn_dic[(k+5)/6 ] for k in tickspace]
for ax in bottom_row:
ax.set_xticks(tickspace)
ax.set_xticklabels(labels,rotation=25)
plt.subplots_adjust(right=0.97, left=0.08, top=0.93, bottom=0.1, hspace=0.25, wspace=0.15)
# Set plot name based on type of plot.
if lonin == [-1.,361.]:
lon_tag = ''
else:
lon_tag = '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1]))
if planetary_only:
f_tag = '_fonly'
else:
f_tag = ''
figname = 'vort_breakdown_' + run + lon_tag + f_tag + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
# Return to xarray
field_out = xr.DataArray(field_out, dims=field.dims, coords=field.coords)
field_out = field_out.transpose(*dims_in)
return field_out
if __name__ == '__main__':
"""Examples/sanity check"""
import matplotlib.pyplot as plt
from hadley_cell import mass_streamfunction
data = xr.open_dataset(
'/scratch/rg419/Data_moist/climatologies/sn_1.000.nc')
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
psi /= 1.e9
precip_test = flip_field(data.precipitation)
plt.figure(1)
data.precipitation.mean('lon').plot.contourf(x='xofyear', y='lat')
plt.figure(2)
precip_test.mean('lon').plot.contourf(x='xofyear', y='lat')
plt.show()
precip_test = make_sym(data.precipitation)
print('precip done')
psi_test = make_sym(psi, asym=True)
print('psi done')
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 monthly_overturning_jra(lonin=[-1., 361.],
region='All',
sanity_check=False):
plot_dir = '/scratch/rg419/plots/monsoon_review_figs/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
# Load in data
data_u = xr.open_dataset(
'/disca/share/reanalysis_links/jra_55/1958_2016/ucomp_monthly/atmos_monthly_together.nc'
)
data_v = xr.open_dataset(
'/disca/share/reanalysis_links/jra_55/1958_2016/vcomp_monthly/atmos_monthly_together.nc'
)
# Make climatologies
u_clim = data_u.groupby('time.month').mean('time')
v_clim = data_v.groupby('time.month').mean('time')
print('means taken')
def get_lons(lonin, data):
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]
]
return lons
lons = get_lons(lonin, data_u)
levs = np.arange(50., 1000., 50.)
ds_clim = xr.Dataset(
{
'ucomp': (['month', 'pfull', 'lat', 'lon'], u_clim.var33.values),
'vcomp': (['month', 'pfull', 'lat', 'lon'], v_clim.var34.values)
},
coords={
'month': ('month', v_clim.month),
'pfull': ('pfull', v_clim.lev / 100.),
'lat': ('lat', v_clim.lat),
'lon': ('lon', v_clim.lon)
})
ds_clim = ds_clim.sel(pfull=levs)
psi = mass_streamfunction(ds_clim, lons=lons, dp_in=50.)
psi /= 1.e9
# Set figure parameters
rcParams['figure.figsize'] = 14, 9
rcParams['font.size'] = 14
# Start figure with 12 subplots
# Start figure with 12 subplots
fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8),
(ax9, ax10, ax11, ax12)) = plt.subplots(3, 4)
axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12]
months = [
'Jan', 'Feb', 'Mar', 'Apr', 'May', 'June', 'Jul', 'Aug', 'Sept', 'Oct',
'Nov', 'Dec'
]
def plot_psi_u(ax, psi, u, month):
f1 = u.sel(lon=lons,
month=month).mean('lon').plot.contourf(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(
-50., 50.1, 5.),
extend='both',
add_labels=False,
add_colorbar=False)
psi.sel(month=month).plot.contour(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(0., 601, 50.),
colors='k',
add_labels=False)
psi.sel(month=month).plot.contour(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(-600., 0., 50.),
colors='k',
linestyles='dashed',
add_labels=False)
m = mc.omega * mc.a**2. * np.cos(u.lat * np.pi / 180.)**2. + u.sel(
lon=lons, month=month).mean('lon') * mc.a * np.cos(
u.lat * np.pi / 180.)
m_levs = mc.omega * mc.a**2. * np.cos(
np.arange(-60., 1., 5.) * np.pi / 180.)**2.
m.plot.contour(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=m_levs,
colors='0.7',
add_labels=False)
ax.set_xlim(-35, 35)
ax.set_xticks(np.arange(-30, 31, 15))
ax.grid(True, linestyle=':')
ax.set_title(months[i - 1])
return f1
for i in range(1, 13):
f1 = plot_psi_u(axes[i - 1], psi, ds_clim.ucomp, i)
for i in range(8, 12):
axes[i].set_xlabel('Latitude')
for i in range(0, 12, 4):
axes[i].set_ylabel('Pressure, hPa')
plt.subplots_adjust(left=0.1,
right=0.97,
top=0.95,
bottom=0.1,
hspace=0.3,
wspace=0.3)
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('Zonal wind speed, m/s')
if lonin == [-1., 361.]:
plt.savefig(plot_dir + 'monthly_psi_u_jra.pdf', format='pdf')
else:
figname = plot_dir + 'monthly_psi_u_jra_' + region + '.pdf'
plt.savefig(figname, format='pdf')
plt.close()
def vort_eq_hm(run, lev=150., lonin=[-1.,361.], month_labels=True, rot_fac=1., add_psi=True):
'''Plot vorticity budget Hovmoller. RG 3/11/2017
Imputs: run = run_name
lev = level to plot
lonin = longitude range to average over
planetary_only = only plot planetary vorticity terms
month_labels = label x axis with month of year (no good for non earth year length)
rot_fac = scale factor for Earth's rotation rate
no_eddies = don't plot transient eddies too'''
# Call the above function to get fields to plot. For a Hovmoller, run is by definition not steady state. Pass other inputs on.
ds = vort_budg_terms(run, lonin=lonin, rot_fac=rot_fac)
if add_psi:
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc')
if lonin[1]>lonin[0]:
lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]]
else:
lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]]
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50., lons=lons)
psi /= 1.e9
# Determine figure size based on number of panels
rcParams['figure.figsize'] = 15, 4
rcParams['font.size'] = 14
rcParams['text.usetex'] = True
levels = np.arange(-1.5,1.6,0.25)
print('starting plotting')
# Set number of panels and declare which are the left column, bottom row.
f, (ax1, ax2, ax3) = plt.subplots(1, 3, sharey='row')
left_column = [ax1]
all_plots = [ax1,ax2,ax3]
plt.set_cmap('RdBu_r')
# Select the pressure level, if needed
if run in ['ap_2','full_qflux']:
ds = ds
else:
ds = ds.sel(pfull=lev)
# Plot!
f2 = ds.div_uvor.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax1, extend = 'both', add_labels=False)
if add_psi:
ax1.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=500), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax1.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=500), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax1.set_title('Relative vorticty part', fontsize=17)
ax1.text(-15, 60, 'a)')
ds.div_uf.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax2, extend = 'both', add_labels=False)
if add_psi:
ax2.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=500), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax2.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=500), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax2.set_title('Planetary vorticity part', fontsize=17)
ax2.text(-7, 60, 'b)')
ds.transient_hm.plot.contourf(x='xofyear', y='lat', levels=levels, ax=ax3, extend = 'both', add_labels=False)
if add_psi:
ax3.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,100.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax3.contour(data.xofyear, data.lat, -1.*psi.sel(pfull=150), levels=np.arange(-500.,510.,500.), add_labels=False, alpha=0.15, colors='k', linewidths=2)
ax3.set_title('Transient eddy vorticity tendency', fontsize=17)
ax3.text(-7, 60, 'c)')
# Add grid, set y limits and ticks
for ax in all_plots:
ax.grid(True,linestyle=':')
ax.set_ylim(-60,60)
ax.set_yticks(np.arange(-60,61,30))
if month_labels:
mn_dic = month_dic(1)
tickspace = list(range(13,72,18))
labels = [mn_dic[(k+5)/6 ] for k in tickspace]
ax.set_xticks(tickspace)
ax.set_xticklabels(labels,rotation=25)
# Label left column
for ax in left_column:
ax.set_ylabel('Latitude')
plt.subplots_adjust(right=0.97, left=0.08, top=0.93, bottom=0.1, hspace=0.25, wspace=0.15)
# Set plot name based on type of plot.
if lonin == [-1.,361.]:
lon_tag = ''
else:
lon_tag = '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1]))
figname = 'alt_breakdown_' + run + lon_tag + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def overturning_monthly(run, lonin=[-1., 361.]):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
plot_dir = '/scratch/rg419/plots/overturning_monthly/' + run + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data.coords['month'] = (data.xofyear - 1) // 6 + 1
data = data.groupby('month').mean(('xofyear'))
#data['vcomp'] = data.vcomp.fillna(0.)
#data['ucomp'] = data.ucomp.fillna(0.)
# 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.)))
#w = VectorWind(data.ucomp, data.vcomp)
# Compute variables
streamfun, vel_pot = w.sfvp()
uchi, vchi, upsi, vpsi = w.helmholtz()
#print(vchi.pfull)
#print(data.pfull)
#data.vcomp.mean('lon')[0,:,:].plot.contourf(x='lat', y='pfull', yincrease=False, add_labels=False)
#plt.figure(2)
#vchi.mean('lon')[0,:,:].plot.contourf(x='lat', y='pfull', yincrease=False, add_labels=False)
#plt.show()
ds_chi = xr.Dataset({'vcomp': (vchi)},
coords={
'month': ('month', vchi.month),
'pfull': ('pfull', vchi.pfull),
'lat': ('lat', vchi.lat),
'lon': ('lon', vchi.lon)
})
ds_psi = xr.Dataset({'vcomp': (vpsi)},
coords={
'month': ('month', vchi.month),
'pfull': ('pfull', vchi.pfull),
'lat': ('lat', vchi.lat),
'lon': ('lon', vchi.lon)
})
def get_lons(lonin, data):
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]
]
return lons
lons = get_lons(lonin, data)
psi = mass_streamfunction(data, lons=lons, dp_in=50., use_v_locally=True)
psi /= 1.e9
psi_chi = mass_streamfunction(ds_chi, lons=lons, dp_in=50.)
psi_chi /= 1.e9
# Set figure parameters
rcParams['figure.figsize'] = 10, 7
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)
axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12]
i = 0
for ax in axes:
psi[:, i, :].plot.contour(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(0., 601, 100.),
colors='k',
add_labels=False)
psi[:, i, :].plot.contour(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(-600., 0., 100.),
colors='k',
linestyles='dashed',
add_labels=False)
f1 = data.ucomp.sel(month=i +
1).sel(lon=lons).mean('lon').plot.contourf(
ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(-50., 50.1, 5.),
extend='both',
add_labels=False,
add_colorbar=False)
m = mc.omega * mc.a**2. * np.cos(
psi.lat * np.pi / 180.)**2. + data.ucomp.sel(
lon=lons).mean('lon') * mc.a * np.cos(psi.lat * np.pi / 180.)
m_levs = mc.omega * mc.a**2. * np.cos(
np.arange(-60., 1., 5.) * np.pi / 180.)**2.
m.sel(month=i + 1).plot.contour(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=m_levs,
colors='0.7',
add_labels=False)
i = i + 1
ax.set_xlim(-35, 35)
ax.set_xticks(np.arange(-30, 31, 15))
ax.grid(True, linestyle=':')
plt.subplots_adjust(left=0.1,
right=0.97,
top=0.95,
bottom=0.1,
hspace=0.3,
wspace=0.3)
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('Zonal wind speed, m/s')
if lonin == [-1., 361.]:
plt.savefig(plot_dir + 'psi_u_' + run + '.pdf', format='pdf')
else:
figname = plot_dir + 'psi_u_' + run + '_' + str(int(
lonin[0])) + '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(figname, format='pdf')
plt.close()
# Start figure with 12 subplots
fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8),
(ax9, ax10, ax11, ax12)) = plt.subplots(3, 4)
axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12]
i = 0
for ax in axes:
psi_chi[:, i, :].plot.contour(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(0., 601, 100.),
colors='k',
add_labels=False)
psi_chi[:, i, :].plot.contour(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(-600., 0., 100.),
colors='k',
linestyles='dashed',
add_labels=False)
f1 = uchi.sel(month=i + 1).sel(lon=lons).mean('lon').plot.contourf(
ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(-3., 3.1, 0.5),
extend='both',
add_labels=False,
add_colorbar=False)
i = i + 1
ax.set_xlim(-35, 35)
ax.set_xticks(np.arange(-30, 31, 15))
ax.grid(True, linestyle=':')
plt.subplots_adjust(left=0.1,
right=0.97,
top=0.95,
bottom=0.1,
hspace=0.3,
wspace=0.3)
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('Zonal wind speed, m/s')
if lonin == [-1., 361.]:
plt.savefig(plot_dir + 'psi_chi_u_' + run + '.pdf', format='pdf')
else:
figname = plot_dir + 'psi_chi_u_' + run + '_' + str(int(
lonin[0])) + '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(figname, format='pdf')
plt.close()
def overturning_hm(run,
regions=[[350, 10], [80, 100], [170, 190], [260, 280]]):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
plot_dir = '/scratch/rg419/plots/overturning_monthly/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
# 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()
ds_chi = xr.Dataset({'vcomp': (vchi)},
coords={
'xofyear': ('xofyear', vchi.xofyear),
'pfull': ('pfull', vchi.pfull),
'lat': ('lat', vchi.lat),
'lon': ('lon', vchi.lon)
})
def get_lons(lonin, data):
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]
]
return lons
# Set figure parameters
rcParams['figure.figsize'] = 10, 7
rcParams['font.size'] = 14
# Start figure with 4 subplots
fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2,
2,
sharex='col',
sharey='row')
axes = [ax1, ax2, ax3, ax4]
i = 0
for ax in axes:
lons = get_lons(regions[i], data)
psi_chi = mass_streamfunction(ds_chi, lons=lons, dp_in=50.)
psi_chi /= 1.e9
i = i + 1
f1 = psi_chi.sel(pfull=500).plot.contourf(ax=ax,
x='xofyear',
y='lat',
add_labels=False,
add_colorbar=False,
levels=np.arange(
-500., 501., 100.),
extend='both')
ax1.set_title('West coast')
ax2.set_title('Land')
ax3.set_title('East coast')
ax4.set_title('Ocean')
for ax in [ax1, ax2, ax3, ax4]:
ax.grid(True, linestyle=':')
ax.set_ylim(-60, 60)
ax.set_yticks(np.arange(-60., 61., 30.))
ax.set_xticks([0, 18, 36, 54, 72])
ax3.set_xlabel('Pentad')
ax4.set_xlabel('Pentad')
ax1.set_ylabel('Latitude')
ax3.set_ylabel('Latitude')
plt.subplots_adjust(left=0.1,
right=0.97,
top=0.95,
bottom=0.1,
hspace=0.3,
wspace=0.3)
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('Overturning streamfunction')
# Save as a pdf
plt.savefig(plot_dir + 'regional_overturning_hm_' + run + '.pdf',
format='pdf')
plt.close()
data.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 get_edge_psi_time(data,
lons,
sanity_check=False,
lev=500.,
thresh=0.,
intdown=True,
nh=False):
"""Get cell edge for climatological data
Inputs: data = input data file
lons = longitudes to average over
sanity_check = plot up overturning and latitudes of max and division to check
lev = level to look at
thresh = threshold to use for boundary
intdown = direction to integrate for mass streamfunction
nh = Only calculate for northern hemisphere """
# Calculate mass streamfunction
psi_coarse = mass_streamfunction(data,
a=6376.0e3,
dp_in=50.,
lons=lons,
intdown=intdown)
psi_coarse /= 1.e9
# Interpolate between 35S and 35N (Should help to avoid mistakes with Ferrel cell etc)
f = spint.interp1d(psi_coarse.lat,
psi_coarse.sel(pfull=lev),
axis=0,
fill_value='extrapolate')
psi = f(lats)
psi = xr.DataArray(psi,
coords=[lats, psi_coarse.xofyear],
dims=['lat', 'xofyear'])
# Create a mask that puts NaNs where abs(psi) is below a threshold
psi_mask = np.ones(psi.values.shape)
if thresh == 0.:
psi_mask[psi.values <= thresh] = np.nan
elif thresh <= 1.:
# If a fraction of the max streamfunction is specified, use this to define the mask
if nh:
psi_max = psi.min('lat')
thresh_frac = psi_max.values * -thresh
else:
psi_max = np.abs(psi).max('lat')
thresh_frac = psi_max.values * thresh
for i in range(psi.xofyear.shape[0]):
psi_mask[np.abs(psi[:, i]).values <= thresh_frac[i], i] = np.nan
else:
psi_mask[np.abs(psi).values <= thresh] = np.nan
psi_mask = xr.DataArray(psi_mask, psi.coords)
#plt.figure(1)
#psi.plot.contourf(x='xofyear', y='lat')
#plt.figure(2)
#psi_mask.plot.contourf(x='xofyear', y='lat')
#plt.show()
# Use mask to mask psi at 500hPa
psi_masked = psi * psi_mask
# Get a list of times where abs(psi) has values above the threshold
times_defd = []
for i in range(0, len(psi_masked.xofyear)):
if np.any(np.isfinite(psi_masked[:, i])):
times_defd.append(np.float(psi_masked.xofyear[i]))
# Reduce psi to only include times where it goes above the threshold
if thresh == 0.:
psi_red = psi.sel(xofyear=times_defd)
else:
psi_red = psi_masked.sel(xofyear=times_defd)
psi_mask_red = np.nan_to_num(psi_mask.sel(xofyear=times_defd))
if nh:
# If only looking at Northern Hemisphere, now take psi_red and look for min instead of taking abs and finding max
psi_max = psi_red.min('lat')
psi_max_loc = psi_red.lat.values[psi_red.argmin('lat').values]
psi_max_loc = xr.DataArray(psi_max_loc,
coords=[psi_red.xofyear],
dims=['xofyear'])
ispos = []
#Create a list of which times the min of psi_red is positive or negative
for i in range(0, len(psi_max_loc)):
ispos.append(
psi_red.values[psi_red.argmin('lat').values[i], i] >= 0.)
# Multiply by -1 to give a positive sign for the SH cell
psi_max = -1. * psi_max
else:
psi_max = np.abs(psi_red).max('lat')
psi_max_loc = psi_red.lat.values[np.abs(psi_red).argmax('lat').values]
psi_max_loc = xr.DataArray(psi_max_loc,
coords=[psi_red.xofyear],
dims=['xofyear'])
# Create list of which times the max of abs(psi_red) is positive vs negative
ispos = []
for i in range(0, len(psi_max_loc)):
ispos.append(
psi_red.values[np.abs(psi_red).argmax('lat').values[i],
i] >= 0.)
# Use this first to get the sign of psi_max right
#NB this line not working properly on 26/01/2018. Use caution when this option is used.
psi_max[ispos] = -1. * psi_max[ispos]
# Now locate cell edge finally!
edge_loc = np.zeros([
len(psi_max_loc),
])
for i in range(0, len(psi_max_loc)):
# Locate places where the psi mask changes sign
edges_are_at = psi_red.lat[np.where(
psi_mask_red[:-1, i] != psi_mask_red[1:, i])[0]]
try:
# Use a try statement to avoid errors where edge is poorly defined
if ispos[i]:
# If the overturning is positive, look for the first edge to the south of the max
# If only looking at the Northern hemisphere don't need this
if nh:
edge_loc[i] = np.nan
else:
edge_loc[i] = np.max(
edges_are_at[edges_are_at < psi_max_loc[i]])
else:
# If the overturning is negative, look for the first edge to the north of the max
edge_loc[i] = np.min(
edges_are_at[edges_are_at > psi_max_loc[i]])
except:
# If we can't find an edge, set the value to nan for this time
edge_loc[i] = np.nan
# Convert to an xarray dataarray for ease of plotting etc.
edge_loc = xr.DataArray(edge_loc,
coords=[psi_red.xofyear],
dims=['xofyear'])
# Option to plot up psi as a sanity check that all points are ok
if sanity_check:
psi_masked.plot.contourf(levels=np.arange(-500., 510., 50.))
psi_max_loc.plot()
edge_loc.plot()
plt.figure(2)
plt.plot(edge_loc, psi_max, 'x')
#plt.figure(3)
#plt.plot(psi_max_loc.xofyear, psi_max_loc, 'x')
#plt.show()
return edge_loc, psi_max, psi_max_loc
def get_edge_psi_ss(data,
lons,
sanity_check=False,
lev=500.,
thresh=0.,
intdown=True,
nh=False):
"""Get cell edge for steady state data
Inputs: data = input data file
lons = longitudes to average over
sanity_check = plot up overturning and latitudes of max and division to check
lev = level to look at
thresh = threshold to use for boundary
intdown = direction to integrate for mass streamfunction
nh = Only calculate for northern hemisphere """
# Calculate mass streamfunction
psi_coarse = mass_streamfunction(data,
a=6376.0e3,
dp_in=50.,
lons=lons,
intdown=intdown)
psi_coarse /= 1.e9
# Interpolate between 35S and 35N (Should help to avoid mistakes with Ferrel cell etc)
f = spint.interp1d(psi_coarse.lat,
psi_coarse.sel(pfull=lev),
axis=0,
fill_value='extrapolate')
psi = f(lats)
psi = xr.DataArray(psi, coords=[lats], dims=['lat'])
# Create a mask that puts NaNs where abs(psi) is below a threshold
psi_mask = np.ones(psi.values.shape)
if thresh == 0.:
psi_mask[psi.values <= thresh] = np.nan
elif thresh == None:
if nh:
psi_max = psi.min('lat')
thresh_frac = psi_max.values * -0.05
else:
psi_max = np.abs(psi).max('lat')
thresh_frac = psi_max.values * 0.05
psi_mask[np.abs(psi).values <= thresh_frac] = np.nan
else:
psi_mask[np.abs(psi).values <= thresh] = np.nan
psi_mask = xr.DataArray(psi_mask, psi.coords)
# Use mask to mask psi at 500hPa
psi_masked = psi * psi_mask
# Reduce psi to only where it goes above the threshold
if thresh == 0.:
psi_red = psi
else:
psi_red = psi_masked
psi_mask_red = np.nan_to_num(psi_mask)
if nh:
psi_max = psi_red.min('lat')
psi_max_loc = psi_red.lat.values[psi_red.argmin('lat').values]
# Create list of which times the max of abs(psi_rad) is positive vs negative
ispos = psi_red.values[psi_red.argmin('lat').values] >= 0.
else:
psi_max = np.abs(psi_red).max('lat')
psi_max_loc = psi_red.lat.values[np.abs(psi_red).argmax('lat').values]
# Create list of which times the max of abs(psi_rad) is positive vs negative
ispos = psi_red.values[np.abs(psi_red).argmax('lat').values] >= 0.
# Use this first to get the sign of psi_max right
if not ispos: psi_max = -1. * psi_max
# Locate places where the psi mask changes sign
edges_are_at = psi_red.lat[np.where(
psi_mask_red[:-1] != psi_mask_red[1:])[0]]
try:
# Use a try statement to avoid errors where edge is poorly defined
if ispos:
# If the overturning is positive, look for the first edge to the south of the max
edge_loc = np.max(edges_are_at[edges_are_at < psi_max_loc])
else:
# If the overturning is negative, look for the first edge to the north of the max
edge_loc = np.min(edges_are_at[edges_are_at > psi_max_loc])
except:
# If we can't find an edge, set the value to nan
edge_loc = np.nan
if sanity_check:
psi_coarse.plot.contourf(x='lat',
y='pfull',
yincrease=False,
levels=np.arange(-500., 501., 50.))
plt.plot(edge_loc, 500., 'kx', mew=2)
plt.plot(psi_max_loc, 500., 'kx', mew=2)
plt.show()
return edge_loc, psi_max, psi_max_loc
def heat_budg_hm(run, lev=850, filename='plev_pentad', timeav='pentad', period_fac=1.,lonin=[-1.,361.], plot_precip=True, do_theta=False):
rcParams['figure.figsize'] = 12, 8
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/heat_budg/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/'+run+'.nc')
if lonin[1]>lonin[0]:
lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]]
else:
lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]]
#advective terms
partition_advection(data, lons, lev=lev)
try:
precip_centroid(data)
except:
data['precipitation'] = data.convection_rain + data.condensation_rain
precip_centroid(data)
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
psi /= 1.e9
heat_mean = data.u_dTdx_zav #+ data.w_dTdp_zav #+data.v_dTdy_zav #
heat_cross = data.u_dTdx_cross1 + data.v_dTdy_cross1 + data.w_dTdp_cross1 + data.u_dTdx_cross2 + data.v_dTdy_cross2 + data.w_dTdp_cross2
heat_cross1 = data.u_dTdx_cross1 + data.v_dTdy_cross1 + data.w_dTdp_cross1
heat_cross2 = data.u_dTdx_cross2 + data.v_dTdy_cross2 + data.w_dTdp_cross2
heat_trans = data.uT_trans_dx + data.vT_trans_dy + data.wT_trans_dp
heat_stat = data.u_dTdx_stat + data.v_dTdy_stat + data.w_dTdp_stat
heat_crossu = data.u_dTdx_cross1 + data.u_dTdx_cross2
heat_crossv = data.v_dTdy_cross1 + data.v_dTdy_cross2
heat_crossw = data.w_dTdp_cross1 + data.w_dTdp_cross2
tdt_rad = data.tdt_rad.sel(pfull=lev).sel(lon=lons).mean('lon')*86400.
tdt_conv = data.dt_tg_convection.sel(pfull=lev).sel(lon=lons).mean('lon')*86400.
tdt_cond = data.dt_tg_condensation.sel(pfull=lev).sel(lon=lons).mean('lon')*86400.
tdt_diff = data.dt_tg_diffusion.sel(pfull=lev).sel(lon=lons).mean('lon')*86400.
if do_theta:
convTtotheta=(1000./data.pfull)**(2./7.)
diabatic = (tdt_rad + tdt_conv + tdt_cond + tdt_diff)*convTtotheta.sel(pfull=lev)
heat_sum = heat_mean + heat_trans + heat_stat + heat_cross + diabatic
asc_cool = heat_sum * 0.
Tdt = gr.ddt(data.temp.sel(pfull=lev)).sel(lon=lons).mean('lon') *86400.
else:
diabatic = tdt_rad + tdt_conv + tdt_cond + tdt_diff
rho = data.pfull*100./data.temp/287.04
asc_cool = (data.omega /(1004.64 * rho)).sel(pfull=lev).sel(lon=lons).mean('lon') *86400.
heat_sum = heat_mean + heat_trans + heat_stat + heat_cross + diabatic + asc_cool
Tdt = gr.ddt(data.temp.sel(pfull=lev)).sel(lon=lons).mean('lon') *86400.
#levels = np.arange(-20,21.1,2.)
#levels = np.arange(-10,10.1,1.)
levels = np.arange(-1,1.1,0.1)
mn_dic = month_dic(1)
tickspace = list(range(13,72,18))
ticklabels = [mn_dic[(k+5)/6 ] for k in tickspace]
# Nine subplots
fig, ((ax1, ax2, ax3), (ax4, ax5, ax6), (ax7, ax8, ax9)) = plt.subplots(3, 3, sharex='col', sharey='row')
plt.set_cmap('RdBu_r')
plot_vars = [heat_mean, heat_trans, heat_sum, asc_cool, diabatic,
tdt_rad, tdt_diff, tdt_conv, tdt_cond]
axes = [ax1,ax2,ax3,ax4,ax5,ax6,ax7,ax8,ax9]
labels = ['a)','b)','c)','d)','e)','f)','g)','h)','i)']
for i in range(9):
f1=plot_vars[i].plot.contourf(ax=axes[i], x='xofyear', y='lat', extend='both', add_labels=False, cmap='RdBu_r', levels = levels, add_colorbar=False)
if plot_precip:
psi.sel(pfull=500).plot.contour(ax=axes[i], x='xofyear', y='lat', levels=np.arange(-500.,0.,100.), add_labels=False, colors='0.7', linewidths=2, linestyles='--')
psi.sel(pfull=500).plot.contour(ax=axes[i], x='xofyear', y='lat', levels=np.arange(0.,510.,100.), add_labels=False, colors='0.7', linewidths=2)
psi.sel(pfull=500).plot.contour(ax=axes[i], x='xofyear', y='lat', levels=np.arange(-1000.,1010.,1000.), add_labels=False, colors='0.5', linewidths=2)
data.p_cent.plot.line(ax=axes[i],color='k', linewidth=2)
axes[i].set_xlabel(' ')
axes[i].set_ylabel(' ')
#totp.plot.contour(ax=axes[i], x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2)
axes[i].set_ylim(-60,60)
axes[i].grid(True,linestyle=':')
axes[i].set_yticks(np.arange(-60.,61.,30.))
axes[i].set_xticks(tickspace)
for i in [0,3,6]:
axes[i].set_ylabel('Latitude')
axes[i].text(-15, 60, labels[i])
for i in [1,2,4,5,7,8]:
axes[i].text(-5, 60, labels[i])
for i in [6,7,8]:
axes[i].set_xticklabels(ticklabels,rotation=25)
# set titles
ax1.set_title('Mean state advection', fontsize=17)
ax2.set_title('Transient eddy flux conv.', fontsize=17)
ax3.set_title('Residual', fontsize=17)
ax4.set_title('w cooling', fontsize=17)
ax5.set_title('Total diabatic heating', fontsize=17)
ax6.set_title('Radiative heating', fontsize=17)
ax7.set_title('Diffusive heating', fontsize=17)
ax8.set_title('Convective heating', fontsize=17)
ax9.set_title('Condensational heating', fontsize=17)
plt.subplots_adjust(right=0.97, left=0.1, top=0.95, bottom=0., hspace=0.25, wspace=0.12)
#Colorbar
cb1=fig.colorbar(f1, ax=axes, use_gridspec=True, orientation = 'horizontal',fraction=0.15, pad=0.15, aspect=30, shrink=0.5)
if lonin == [-1.,361.]:
figname = 'zon_heat_budg_' +run+ '_u_850.pdf'
else:
figname = 'zon_heat_budg_' + run + '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def abs_vort_dt_plot(run, rot_fac=1., lev=150.):
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
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.
# Take gradients of vorticity
dvordx = gr.ddx(vor)
dvordy = gr.ddy(vor, vector=False)
# Horizontal material derivative
horiz_md_mean = -86400. * (data.ucomp.sel(pfull=lev) * dvordx +
data.vcomp.sel(pfull=lev) * dvordy)
# Calculate divergence and stretching term
div = gr.ddx(data.ucomp.sel(pfull=lev)) + gr.ddy(data.vcomp.sel(pfull=lev))
stretching_mean = -86400. * vor * div
#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.
stretching_mean = stretching_mean.mean('lon')
horiz_md_mean = horiz_md_mean.mean('lon')
plot_dir = '/scratch/rg419/plots/paper_2_figs/abs_vort_dt/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
rcParams['figure.figsize'] = 5, 12
rcParams['font.size'] = 14
fig, (ax1, ax2, ax3, ax4) = plt.subplots(4, 1)
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(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 = stretching_mean.plot.contourf(ax=ax2,
x='xofyear',
y='lat',
levels=np.arange(-1.5, 1.6, 0.25),
add_colorbar=False,
add_labels=False)
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 = horiz_md_mean.plot.contourf(ax=ax3,
x='xofyear',
y='lat',
levels=np.arange(-1.5, 1.6, 0.25),
add_colorbar=False,
add_labels=False)
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=':')
f1 = (stretching_mean + horiz_md_mean).plot.contourf(ax=ax4,
x='xofyear',
y='lat',
levels=np.arange(
-1.5, 1.6, 0.25),
add_colorbar=False,
add_labels=False)
psi.sel(pfull=500).plot.contour(ax=ax4,
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=ax4,
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=ax4,
x='xofyear',
y='lat',
levels=np.arange(-1000., 1010., 1000.),
add_labels=False,
colors='0.5',
linewidths=2)
data.p_cent.plot.line(ax=ax4, color='k', linewidth=2)
ax4.set_ylim([-60, 60])
ax4.set_ylabel('Latitude')
ax4.set_xticks([12, 24, 36, 48, 60, 72])
ax4.set_yticks([-60, -30, 0, 30, 60])
ax4.set_xlabel('Pentad')
ax4.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 + 'vort_terms_' + run + '.pdf', format='pdf')
plt.close()
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 surface_plot(run,
lonin=[-1., 361.],
diff_run=None,
scale_fac=1.,
do_make_sym=True):
print('what?')
rcParams['figure.figsize'] = 15, 10
rcParams['font.size'] = 14
plot_dir = '/scratch/rg419/plots/surface_fluxes/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
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['flux_lw_up'] = data.t_surf**4. * mc.stefan
data['dTdt'] = gr.ddt(data.t_surf) * 86400. * scale_fac
if do_make_sym:
for field in [
't_surf', 'dTdt', 'flux_sw', 'flux_lw', 'flux_lw_up', 'flux_t',
'flux_lhe'
]:
data[field] = make_sym(data[field])
lons = pick_lons(data, lonin)
if not diff_run == None:
diff_data = xr.open_dataset(
'/disca/share/rg419/Data_moist/climatologies/' + diff_run + '.nc')
data = (data - diff_data).sel(lon=lons).mean('lon')
levels_t = np.arange(-10., 11., 1.)
levels_dt = np.arange(-0.5, 0.51, 0.05)
levels_flux = np.arange(-80., 81., 5.)
else:
data = data.sel(lon=lons).mean('lon')
levels_t = np.arange(270., 306., 2.5)
levels_dt = np.arange(-0.3, 0.31, 0.05)
levels_flux = np.arange(-300., 305., 20.)
f, ((ax1, ax2), (ax3, ax4), (ax5, ax6)) = plt.subplots(3,
2,
sharex='col',
sharey='row')
left_column = [ax1, ax3, ax5]
bottom_row = [ax5, ax6]
all_plots = [ax1, ax2, ax3, ax4, ax5, ax6]
#check = data.flux_sw + (data.flux_lw - flux_lw_up) - data.flux_t - data.flux_lhe
#check = data.flux_sw - data.flux_lhe
data.t_surf.plot.contourf(x='xofyear',
y='lat',
levels=levels_t,
ax=ax1,
extend='both',
add_labels=False)
ax1.set_title('SST')
data.dTdt.plot.contourf(x='xofyear',
y='lat',
ax=ax2,
levels=levels_dt,
extend='both',
add_labels=False,
cmap='RdBu_r')
ax2.set_title('d(SST)/dt')
data.flux_sw.plot.contourf(x='xofyear',
y='lat',
levels=levels_flux,
ax=ax3,
extend='both',
add_labels=False,
cmap='RdBu_r')
ax3.set_title('Net downward SW flux')
(data.flux_lw - data.flux_lw_up).plot.contourf(x='xofyear',
y='lat',
levels=levels_flux,
ax=ax4,
extend='both',
add_labels=False,
cmap='RdBu_r')
ax4.set_title('Net downward LW flux')
(-1. * data.flux_t).plot.contourf(x='xofyear',
y='lat',
levels=levels_flux,
ax=ax5,
extend='both',
add_labels=False,
cmap='RdBu_r')
#check.plot.contourf(x='xofyear', y='lat', levels=np.arange(-300.,305.,20.), ax=ax5, extend = 'both', add_labels=False, cmap='RdBu_r')
ax5.set_title('Downward sensible heat flux')
(-1. * data.flux_lhe).plot.contourf(x='xofyear',
y='lat',
levels=levels_flux,
ax=ax6,
extend='both',
add_labels=False,
cmap='RdBu_r')
ax6.set_title('Downward latent heat flux')
i = 0
labels = ['a)', 'b)', 'c)', 'd)', 'e)', 'f)']
for ax in all_plots:
ax.grid(True, linestyle=':')
ax.set_xticks([12, 24, 36, 48, 60, 72])
ax.set_ylim([-60, 60])
ax.set_yticks([-60, -30, 0, 30, 60])
ax.text(-8, 60., labels[i])
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_ylabel('')
ax.set_xlabel('')
i = i + 1
for ax in left_column:
ax.set_ylabel('Latitude')
for ax in bottom_row:
ax.set_xlabel('Pentad')
plt.subplots_adjust(left=0.07,
right=0.97,
top=0.97,
bottom=0.05,
hspace=0.2,
wspace=0.1)
if lonin == [-1., 361.]:
if not diff_run == None:
plt.savefig(plot_dir + 'surface_fluxes_' + run + '_' + diff_run +
'.pdf',
format='pdf')
else:
plt.savefig(plot_dir + 'surface_fluxes_' + run + '.pdf',
format='pdf')
else:
figname = plot_dir + 'surface_fluxes_' + run + '_' + str(int(
lonin[0])) + '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(figname, format='pdf')
plt.close()
def evap_parts(run, lonin=[-1., 361.], do_make_sym=True):
rcParams['figure.figsize'] = 6, 9
rcParams['font.size'] = 14
plot_dir = '/scratch/rg419/plots/surface_fluxes/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
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)
lons = pick_lons(data, lonin)
# Density of lowest level air
#rho_a = 95000./Rd/data.temp.sel(pfull=950.) / (1 + (Rv - Rd)/Rd*data.sphum.sel(pfull=950.))
# Look at an average over a zonal region
v_a = data.wind_speed.sel(lon=lons).mean('lon')
q_a = data.q_atm.sel(lon=lons).mean('lon')
q_s = data.q_surf.sel(lon=lons).mean('lon')
qdiff = (q_a - q_s) * 1000.
rho_a = data.rho.sel(lon=lons).mean('lon')
#rho_a.plot.contourf()
#plt.colorbar()
#plt.show()
#(v_a*qdiff).plot.contourf(x='xofyear', y='lat', levels=np.arange(-0.05,0.05,0.005))
#plt.show()
#q_a.plot.contourf(x='xofyear', y='lat', levels=np.arange(0,0.02,0.001))
#plt.figure(2)
#q_s.plot.contourf(x='xofyear', y='lat', levels=np.arange(0,0.02,0.001))
#plt.show()
if do_make_sym:
qdiff = make_sym(qdiff)
rho_a = make_sym(rho_a)
v_a = make_sym(v_a)
f, (ax1, ax2, ax3) = plt.subplots(3, sharex=True)
qdiff.plot.contourf(x='xofyear',
y='lat',
ax=ax1,
extend='both',
add_labels=False,
levels=np.arange(-10., 10.1, 1.))
ax1.set_title('q$_{a}$ - q$_{s}$')
rho_a.plot.contourf(x='xofyear',
y='lat',
ax=ax2,
extend='both',
add_labels=False,
levels=np.arange(1.12, 1.26, 0.01))
ax2.set_title(r'$\rho_{a}$')
v_a.plot.contourf(x='xofyear',
y='lat',
levels=np.arange(0., 20., 2.),
ax=ax3,
extend='both',
add_labels=False)
ax3.set_title('v$_{a}$')
labels = ['a)', 'b)', 'c)']
i = 0
for ax in [ax1, ax2, ax3]:
#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_yticks([-60, -30, 0, 30, 60])
ax.set_ylabel('Latitude')
ax.text(-10, 60., labels[i])
ax.grid(True, linestyle=':')
i = i + 1
ax3.set_xticks([12, 24, 36, 48, 60, 72])
ax3.set_xlabel('Pentad')
plt.subplots_adjust(left=0.12, right=0.99, top=0.95, bottom=0.06)
plt.savefig(plot_dir + 'evap_parts_revisions_' + run + '.pdf',
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')
