def aht_vs_psi_eq(run, period_fac=1.):
#Load in data, add area to dataset
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
psi = mass_streamfunction(data, mc.a, dp_in=50.)
psi /= 1.e9
psi_eq = psi.sel(pfull=500.).isel(lat=[31, 32]).mean('lat')
ahteq = aht_eq(run)[0]
ahteq.coords['month'] = np.mod(ahteq.xofyear - 1,
72. * period_fac) // (6 * period_fac) + 1
ahteq_month = ahteq.groupby('month').mean(('xofyear'))
psi_eq.coords['month'] = np.mod(psi_eq.xofyear - 1,
72. * period_fac) // (6 * period_fac) + 1
psi_eq_month = psi_eq.groupby('month').mean(('xofyear'))
plt.plot(psi_eq, ahteq, 'xk', alpha=0.7)
plt.plot(psi_eq_month, ahteq_month, 'xk', ms=7, mew=2)
for i in range(0, len(month_list)):
plt.text(psi_eq_month[i] + 0.1,
ahteq_month[i] + 0.1,
month_list[i],
fontsize=14)
plt.ylabel('Atmospheric heat transport at the Equator (PW)')
plt.xlabel('500 hPa Mass Streamfunction at the Equator')
plt.grid(True, linestyle=':')
plt.xlim([-600, 600])
plt.ylim([-10, 10])
plt.savefig('/scratch/rg419/plots/aht_work/psi_aht_' + run + '.pdf',
format='pdf')
plt.close()
def get_edge_psi(run, lonin=[-1.,361.]):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/'+run+'.nc')
omega = 7.2921150e-5
a= 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi/180)
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]]
lats = data.lat[np.abs(data.lat) <= 30.]
psi = mass_streamfunction(data, a=6376.0e3, lons=lons, dp_in=50.)
psi /= 1.e9
cp = 287.04/2*7
L = 2.500e6
#data['mse'] = (cp*data.temp + L*data.sphum).sel(lon=lons).mean('lon')
data['mse'] = (data.convection_rain + data.condensation_rain).sel(lon=lons).mean('lon')
#mse_max_loc = data.lat[data.mse.sel(pfull=850.).argmax('lat').values].values
mse_max_loc = data.lat[data.mse.argmax('lat').values].values
psi_min = -1.*psi.sel(pfull=500.).sel(lat=lats).min('lat')
#psi.sel(pfull=500.).plot.contourf(x='xofyear', y='lat')
#plt.figure(2)
#plt.plot( psi_min)
#plt.figure(3)
#plt.plot(mse_max_loc)
#plt.show()
return mse_max_loc, psi_min
def get_edge_psi(run, lonin=[-1., 361.], sanity_check=True, lev=lev):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
omega = 7.2921150e-5
a = 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi / 180.)
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]
]
#lats = data.lat[np.abs(data.lat) <= 35.]
#Calculate mass streamfunction
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50., lons=lons)
psi /= 1.e9
f = spint.interp1d(psi.lat, psi, axis=0, fill_value='extrapolate')
lats_new = np.arange(-90, 90.1, 0.5)
psi_new = f(lats_new)
psi_new = xr.DataArray(psi_new,
coords=[lats_new, psi.xofyear, psi.pfull],
dims=['lat', 'xofyear', 'pfull'])
lats = psi_new.lat[np.abs(psi_new.lat) <= 35.]
#Create array of zeros in lat-time shape
edge_loc = np.zeros(psi_new.sel(pfull=lev).values.shape)
# Mask where overturning magnitude is less than -50.
edge_loc[psi_new.sel(pfull=lev).values <= -1. * thresh] = 1.
edge_loc = xr.DataArray(edge_loc, psi_new.sel(pfull=lev).coords)
#Take diff along lat axis between 35 N and S and find location of first maximum in range
eli = edge_loc.sel(lat=lats).diff('lat').argmin('lat')
edge_loc_lat = psi_new.lat.sel(lat=lats)[eli]
psi_min = -1. * psi_new.sel(pfull=lev).sel(lat=lats).min('lat')
if sanity_check:
#Plot as sanity check
psi.sel(pfull=lev).plot.contourf(levels=np.arange(-300., 301., 50.))
plt.plot(edge_loc_lat, 'r')
plt.figure(2)
plt.plot(psi_min)
plt.figure(3)
plt.plot(edge_loc_lat[21:36], psi_min[21:36], 'b')
plt.plot(edge_loc_lat[50:70], psi_min[50:70], 'g')
plt.xlim([0, 30])
plt.show()
return edge_loc_lat, psi_min
def get_edge_psi_ss(data, lons, sanity_check=False, lev=500., thresh=0., intdown=True):
"""Get cell edge for steady state data"""
# 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'])
psi_mask = np.ones(psi.values.shape)
if thresh == 0.:
psi_mask[psi.values <= thresh] = np.nan
else:
psi_mask[np.abs(psi).values <= thresh] = np.nan
psi_mask = xr.DataArray( psi_mask, psi.coords)
psi_masked = psi * psi_mask
if thresh ==0.:
psi_red = psi
else:
psi_red = psi_masked
psi_mask_red = np.nan_to_num(psi_mask)
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
def upsi(run,
months,
before_after,
filename='plev_pentad',
timeav='pentad',
period_fac=1.,
lonin=[-1., 361.]):
rcParams['figure.figsize'] = 8, 7
rcParams['font.size'] = 25
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
#data = time_means(run, months, filename=filename, timeav=timeav,period_fac=period_fac)
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
a = 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi / 180)
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, lons=lons, dp_in=50.)
psi /= 1.e9
psi_before = psi[::-1, before_after[0]:before_after[0] +
4, ::-1].mean('xofyear').transpose()
psi_after = psi[::-1, before_after[1]:before_after[1] +
4, ::-1].mean('xofyear').transpose()
plt.contour(data.lat,
data.pfull,
psi_before,
colors='k',
levels=np.arange(-300., 301., 50.),
linewidths=2)
plt.gca().invert_yaxis()
plt.ylabel('Pressure, hPa')
plt.xlabel('Latitude')
plt.xlim(-60, 60)
plt.grid(True, linestyle=':')
plt.tight_layout()
plt.savefig(plot_dir + 'psi_below.pdf', format='pdf')
plt.close()
def psi_max_500(run, lonin=[-1., 361.], lev=500., intdown=True):
"""Front end so same function can be used for both climatological and steady state data"""
#Load data
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
# Choose lons to average over
if lonin[1] > lonin[0]:
lons = [
data.lon[i] for i in range(len(data.lon))
if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]
]
else:
lons = [
data.lon[i] for i in range(len(data.lon))
if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]
]
# 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)
# Check if data has a time dimension
if 'xofyear' in data.coords:
psi = xr.DataArray(psi,
coords=[lats, psi_coarse.xofyear],
dims=['lat', 'xofyear'])
psi_max = xr.DataArray(np.zeros(len(psi.xofyear), ),
coords=[psi.xofyear],
dims=['xofyear'])
for i in range(len(psi.xofyear)):
psi_max[i] = max(psi.isel(xofyear=i).max('lat'),
psi.isel(xofyear=i).min('lat'),
key=abs)
else:
psi = xr.DataArray(psi, coords=[lats], dims=['lat'])
psi_max = max(psi.max('lat'), psi.min('lat'), key=abs)
return psi_max
def u_and_psi_zmean(run,
months,
filename='atmos_pentad',
period_fac=1.,
lonin=[-1., 361.]):
rcParams['figure.figsize'] = 12, 10
plot_dir = '/scratch/rg419/plots/climatology/' + run + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
sn_dic = season_dic(1)
data = time_means(run,
months,
filename=filename,
timeav='season',
period_fac=period_fac)
lons = pick_lons(data, lonin)
uwnd = data.ucomp[:, :, :, lons].mean('lon')
psi = mass_streamfunction(data, a=6376.0e3, lons=lons, dp_in=50.)
psi /= 1.e9
for i in range(0, 4):
print i
ax = uwnd[i, :, :].plot.contourf(x='lat',
y='pfull',
levels=np.arange(-60., 60., 5.),
add_label=False,
add_colorbar=False,
yincrease=False,
extend='both')
cs = psi[:, i, :].plot.contour(x='lat',
y='pfull',
levels=np.arange(-350., 350., 50.),
colors='k',
add_label=False,
add_colorbar=False,
yincrease=False,
extend='both')
plt.clabel(cs, fontsize=15, inline_spacing=-1, fmt='%1.0f')
cb1 = plt.colorbar(ax)
cb1.set_label('Zonal wind speed, m/s')
plt.ylabel('Pressure, hPa')
plt.xlabel('Latitude')
plt.tight_layout()
plt.savefig(plot_dir + 'u_and_psi_zmean_alllons_plev' + sn_dic[i] +
'.png')
plt.close()
def u_psi(run, bef_aft, lonin=[-1., 361.]):
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]
]
print data.vcomp.sel(pfull=900.).sel(xofyear=40.)[50, 33].values
vtest = np.nan_to_num(data.vcomp)
vtest = xr.DataArray(vtest, [('xofyear', data.xofyear),
('pfull', data.pfull), ('lat', data.lat),
('lon', data.lon)])
#data['vcomp'] = (('xofyear','pfull','lat','lon'), vtest)
print data.vcomp.sel(pfull=900.).sel(xofyear=40.)[50, 33].values
psi = mass_streamfunction(data, a=6376.0e3, lons=lons, dp_in=50.)
psi /= 1.e9
psi_before = psi[:, bef_aft[0]:bef_aft[0] +
4, ::-1].mean('xofyear').transpose()
psi_after = psi[:, bef_aft[1]:bef_aft[1] +
4, ::-1].mean('xofyear').transpose()
uv_conv_before, u_ztav_before = uv_partitioning(data, bef_aft[0], lons)
uv_conv_after, u_ztav_after = uv_partitioning(data, bef_aft[1], lons)
ds = xr.Dataset(
{
'psi_before': (['pfull', 'lat'], psi_before),
'psi_after': (['pfull', 'lat'], psi_after),
'uv_conv_before': (['pfull', 'lat'], uv_conv_before),
'uv_conv_after': (['pfull', 'lat'], uv_conv_after),
'u_ztav_before': (['pfull', 'lat'], u_ztav_before),
'u_ztav_after': (['pfull', 'lat'], u_ztav_after)
},
coords={
'pfull': ('pfull', data.pfull),
'lat': ('lat', data.lat)
})
return ds
def get_edge_psi(run, lonin=[-1., 361.], sanity_check=False, lev=500.):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
omega = 7.2921150e-5
a = 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi / 180)
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]
]
lats = data.lat[np.abs(data.lat) <= 35.]
#Calculate mass streamfunction
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50., lons=lons)
psi /= 1.e9
levs = data.pfull[data.pfull <= 950.]
#Create array of zeros in lat-time shape
edge_loc = np.zeros(psi.sel(pfull=lev).values.shape)
# Mask where overturning magnitude is less than -50.
#edge_loc[psi.sel(pfull=lev).values <= -150.] = 1.
edge_loc[psi.sel(pfull=levs).mean('pfull').values <= -120.] = 1.
edge_loc = xr.DataArray(edge_loc, psi.sel(pfull=lev).coords)
#Take diff along lat axis between 35 N and S and find location of first maximum in range
eli = edge_loc.sel(lat=lats).diff('lat').argmin('lat')
edge_loc_lat = data.lat.sel(lat=lats)[eli]
#psi_min = -1.*psi.sel(pfull=levs).sel(lat=lats).min(('lat','pfull'))
psi_min = -1. * psi.sel(pfull=levs).sel(lat=lats).mean('pfull').min('lat')
if sanity_check:
#Plot as sanity check
psi.sel(pfull=levs).mean('pfull').plot.contourf(
levels=np.arange(-300., 301., 50.))
plt.plot(edge_loc_lat, 'r')
plt.figure(2)
plt.plot(psi_min)
plt.show()
return edge_loc_lat, psi_min
def pcent_vs_psi_eq(run, period_fac=1., lonin=[-1., 361.], do_plot=False):
#Load in data, add area to dataset
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, mc.a, dp_in=50., lons=lons)
psi /= 1.e9
psi_eq = psi.sel(pfull=500.).isel(lat=[31, 32]).mean('lat')
psi_eq.coords['month'] = np.mod(psi_eq.xofyear - 1,
72. * period_fac) // (6 * period_fac) + 1
psi_eq_month = psi_eq.groupby('month').mean(('xofyear'))
p_cent = precip_centroid(run, lat_bound=45., lonin=lonin)
p_cent.coords['month'] = np.mod(p_cent.xofyear - 1,
72. * period_fac) // (6 * period_fac) + 1
p_cent_month = p_cent.groupby('month').mean(('xofyear'))
if do_plot:
plt.plot(p_cent, psi_eq, 'xk', alpha=0.7)
plt.plot(p_cent_month, psi_eq_month, 'xk', ms=7, mew=2)
for i in range(0, len(month_list)):
plt.text(p_cent_month[i] + 0.1,
psi_eq_month[i] + 0.1,
month_list[i],
fontsize=14)
plt.xlabel('Precipitation centroid')
plt.ylabel('500 hPa Mass Streamfunction at the Equator')
plt.grid(True, linestyle=':')
plt.ylim([-600, 600])
plt.xlim([-30, 30])
plt.savefig('/scratch/rg419/plots/aht_work/pcent_psi_' + run + '.pdf',
format='pdf')
plt.close()
return p_cent, p_cent_month, psi_eq, psi_eq_month
def u_psi(run, lonin=[-1.,361.]):
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, lons=lons, dp_in=50.)
psi /= 1.e9
psi = psi.T
uv_conv, u_ztav = uv_partitioning(data, lons)
ds = xr.Dataset({'psi': (['pfull', 'lat'], psi),
'uv_conv': (['pfull', 'lat'], uv_conv),
'u_ztav': (['pfull', 'lat'], u_ztav)},
coords={'pfull': ('pfull', data.pfull),
'lat': ('lat', data.lat)})
lwid=2
ds.uv_conv.plot.contourf(x='lat', y='pfull', add_labels = False, levels = np.arange(-10.,10.1,1.), extend = 'both')
plt.contour(ds.lat, ds.pfull, ds.u_ztav, colors='0.7', levels=np.arange(-60.,61.,10.), linewidths=lwid)
psi.plot.contour(x='lat', y='pfull', colors='k', levels=np.arange(0.,301.,60.), linewidths=lwid)
psi.plot.contour(x='lat', y='pfull', colors='k', linestyles='--', levels=np.arange(-300.,0.,60.), linewidths=lwid)
plt.ylabel('Pressure, hPa')
plt.xlabel('Latitude')
plt.xlim(-60,60)
plt.grid(True,linestyle=':')
plt.gca().invert_yaxis()
if lonin == [-1.,361.]:
figname = 'psi_u_' + run + '.pdf'
else:
figname = 'psi_u_' + run + '_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def psi_max_lat(run,months,filename='atmos_pentad',period_fac=1.):
data = time_means(run, months, filename=filename, timeav='pentad',period_fac=period_fac)
#t_surf_max = data.lat[np.argmax(data.t_surf.mean('lon').values, axis=1)]
#data.t_surf.mean('lon').plot.contourf(x='xofyear',y='lat')
omega_max_loc = data.lat[np.argmin(data.omega[36:71,27,:,:].mean('lon').values, axis=1)]
#plt.ylim(-45,45)
#plt.plot(t_surf_max)
plt.figure(2)
psi = mass_streamfunction(data, a=6376.0e3)/1e9
print 'psi evaluated'
#psi_max[0,:] = data.lat[np.argmin( psi.values[:,30:71,27], axis=0)]
psi_max = np.amin( psi[:,36:71,27], axis=0)
plt.plot(omega_max_loc[0:11],psi_max[0:11],'x')
plt.plot(omega_max_loc[12:23],psi_max[12:23],'xr')
plt.plot(omega_max_loc[24:71],psi_max[24:71],'xg')
plt.show()
def psi_regime_plots(run, months, period_fac=1.):
data = time_means(run,
months,
filename='atmos_daily',
timeav='pentad',
period_fac=period_fac)
#find index where omega crosses the equator
omega_max_loc = data.lat.values[np.argmin(
data.omega[:, 27, :, :].mean('lon').values, axis=1)].tolist()
oei = (np.diff([1 if i > 0 else 0
for i in omega_max_loc]).tolist()).index(1) + 1
psi = mass_streamfunction(data.isel(xofyear=[oei - 4, oei, oei + 4]),
a=6376.0e3) / 1e9
print 'psi evaluated'
u = data.ucomp[[oei - 5, oei, oei + 5], :, :, :].mean('lon')
a_cos = a * np.cos(data.lat * np.pi / 180.)
m = (u + omega * a_cos) * a_cos
print 'AM evaluated'
for i in [0, 1, 2]:
psi[:, i, :].plot.contourf(x='lat',
y='pfull',
yincrease=False,
add_labels=False,
levels=range(-350, 351, 50))
m[i, :, :].plot.contour(x='lat',
y='pfull',
yincrease=False,
add_labels=False,
colors='k',
add_colorbar=False,
levels=np.arange(0, np.max(m),
omega * a * a / 10))
plt.xlim(-45, 45)
plt.xlabel('Latitude')
plt.ylabel('Pressure, hPa')
plt.savefig('/scratch/rg419/plots/seasons_and_rotation/psi_m_' + run +
'_' + str(i) + '.png')
plt.close()
print str(i) + ' done'
def u_psi(run, bef_aft, lonin=[-1., 361.]):
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, lons=lons, dp_in=50.)
psi /= 1.e9
psi_before = psi[::-1, bef_aft[0]:bef_aft[0] +
4, ::-1].mean('xofyear').transpose()
psi_after = psi[::-1, bef_aft[1]:bef_aft[1] +
4, ::-1].mean('xofyear').transpose()
uv_conv_before, u_ztav_before = uv_partitioning(data, bef_aft[0], lons)
uv_conv_after, u_ztav_after = uv_partitioning(data, bef_aft[1], lons)
ds = xr.Dataset(
{
'psi_before': (['pfull', 'lat'], psi_before),
'psi_after': (['pfull', 'lat'], psi_after),
'uv_conv_before': (['pfull', 'lat'], uv_conv_before),
'uv_conv_after': (['pfull', 'lat'], uv_conv_after),
'u_ztav_before': (['pfull', 'lat'], u_ztav_before),
'u_ztav_after': (['pfull', 'lat'], u_ztav_after)
},
coords={
'pfull': ('pfull', data.pfull),
'lat': ('lat', data.lat)
})
return ds
def regime_diag(run, months, period_fac=1., rmean=6):
data = time_means(run,
months,
filename='atmos_daily',
timeav='pentad',
period_fac=period_fac)
psi = mass_streamfunction(data, a=6376.0e3) / 1e9
print 'psi evaluated'
psi_max = np.argmin(psi.values[0:40, :, :], axis=0)
psi_rate = np.gradient(psi_max[:, 27])
def running_mean(x, N):
cumsum = np.cumsum(np.insert(x, 0, 0))
return (cumsum[N:] - cumsum[:-N]) / N
psi_rate_out = running_mean(psi_rate, rmean)
return psi_rate_out
def peak_ascent_psi(run, months, period_fac=1.):
data = time_means(run,
months,
filename='atmos_daily',
timeav='pentad',
period_fac=period_fac)
psi = mass_streamfunction(data, a=6376.0e3) / 1e9
print 'psi evaluated'
omega_max_loc = data.lat[np.argmin(data.omega[:,
27, :, :].mean('lon').values,
axis=1)]
print 'peak ascent located'
psi[:, :, 36].plot.contourf(x='xofyear', y='lat', add_labels=False)
plt.plot(data.xofyear, omega_max_loc, 'k')
plt.ylim(-45, 45)
plt.xlabel('Pentad')
plt.ylabel('Latitude')
plt.savefig('/scratch/rg419/plots/seasons_and_rotation/wpsi_' + run +
'.png')
plt.close()
def get_edge_psi(run):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
omega = 7.2921150e-5
a = 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi / 180)
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
psi /= 1.e9
edge_loc = np.zeros(psi.sel(pfull=500).values.shape)
edge_loc[psi.sel(pfull=500).values <= -50.] = 1.
edge_loc = xr.DataArray(edge_loc, psi.sel(pfull=500).coords)
eli = edge_loc.diff('lat')[::-1, :].argmax('lat')
edge_loc_lat = data.lat[eli]
edge_ll_summer = edge_loc_lat[abs(data.lat[eli]) < 30.]
psi_min = psi.sel(pfull=500.).min('lat')
psi_min_summer = -1. * psi_min[abs(data.lat[eli]) < 30.]
return edge_ll_summer, psi_min_summer
from physics import mass_streamfunction
from pylab import rcParams
import sh
import matplotlib.pyplot as plt
rcParams['figure.figsize'] = 7, 7
rcParams['font.size'] = 15
#rcParams['text.usetex'] = True
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/zs_sst.nc')
plot_dir = '/scratch/rg419/plots/axisymmetric_runs/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
psi = mass_streamfunction(data, a=6376.0e3, dp_in=50.)
psi /= 1.e9
#levels = np.arange(-2.,2.1,0.25)
f, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8),
(ax9, ax10, ax11, ax12)) = plt.subplots(3, 4, sharex='col', sharey='row')
plt.set_cmap('RdBu_r')
ax_list = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12]
levels = np.arange(-300., 301., 60.)
for i in range(12):
# Evaluate the mass streamfunction as a function of time, look at progression of peak throughout year
from physics import mass_streamfunction
from data_handling import time_means
import matplotlib.pyplot as plt
import xarray as xr
import numpy as np
data = time_means('sn_3.000', [145, 469],
filename='atmos_pentad',
timeav='pentad',
period_fac=3.)
psi_30_3 = mass_streamfunction(data, a=6376.0e3) / 1e9
print 'psi evaluated'
data_ap10 = time_means('aquaplanet_10m', [109, 145],
filename='atmos_pentad',
timeav='pentad',
period_fac=1.)
psi_10_1 = mass_streamfunction(data_ap10, a=6376.0e3) / 1e9
#u = data.ucomp
#print 'u loaded'
psi_max_30_3 = np.argmin(psi_30_3.values[0:40, :, :], axis=0)
psi_max_10_1 = np.argmin(psi_10_1.values[0:40, :, :], axis=0)
#36
#psi_30_3[:,:,36].plot.contourf(x='xofyear', y='lat', levels = np.arange(-450,451,150), extend='neither')
plt.plot(data.xofyear / 3., data.lat[psi_max_30_3[:, 27]], 'k')
plt.plot(data_ap10.xofyear, data.lat[psi_max_10_1[:, 27]], 'r')
def upsi(run,
months,
before_after,
filename='plev_pentad',
timeav='pentad',
period_fac=1.,
lonin=[-1., 361.]):
rcParams['figure.figsize'] = 15, 10
rcParams['font.size'] = 25
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
#data = time_means(run, months, filename=filename, timeav=timeav,period_fac=period_fac)
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
a = 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi / 180)
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, lons=lons, dp_in=50.)
psi /= 1.e9
psi_before = psi[::-1, before_after[0]:before_after[0] +
4, ::-1].mean('xofyear').transpose()
psi_after = psi[::-1, before_after[1]:before_after[1] +
4, ::-1].mean('xofyear').transpose()
def uv_partitioning(data, start_index):
u_tav = data.ucomp[start_index:start_index +
4, :, :, :].mean('xofyear')
u_ztav = u_tav.sel(lon=lons).mean('lon')
v_tav = data.vcomp[start_index:start_index +
4, :, :, :].mean('xofyear')
v_ztav = v_tav.sel(lon=lons).mean('lon')
uv_trans = (data.ucomp_vcomp[start_index:start_index +
4, :, :, :].mean('xofyear') -
u_tav * v_tav).sel(lon=lons).mean('lon')
uv_stat = (u_tav * v_tav).sel(lon=lons).mean('lon') - u_ztav * v_ztav
uv_conv_trans = xr.DataArray(
cfd((uv_trans * coslat * coslat).values, data.lat * np.pi / 180,
1), [('pfull', data.pfull), ('lat', data.lat)])
uv_conv_trans = -86400. * uv_conv_trans / coslat / coslat / a
uv_conv_stat = xr.DataArray(
cfd((uv_stat * coslat * coslat).values, data.lat * np.pi / 180, 1),
[('pfull', data.pfull), ('lat', data.lat)])
uv_conv_stat = -86400. * uv_conv_stat / coslat / coslat / a
return uv_conv_trans, uv_conv_stat, u_ztav
uv_conv_trans_before, uv_conv_stat_before, u_ztav_before = uv_partitioning(
data, before_after[0])
uv_conv_trans_after, uv_conv_stat_after, u_ztav_after = uv_partitioning(
data, before_after[1])
# Four subplots
f, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2,
2,
sharex='col',
sharey='row')
plt.set_cmap('RdBu_r')
#First plot
f1 = ax1.contourf(data.lat,
data.pfull,
uv_conv_trans_before,
extend='both',
levels=np.arange(-5., 5.1, 1.))
ax1.contour(data.lat,
data.pfull,
u_ztav_before,
colors='0.75',
levels=np.arange(-60., 61., 10.),
linewidths=2)
ax1.contour(data.lat,
data.pfull,
psi_before,
colors='k',
levels=np.arange(-300., 301., 50.),
linewidths=2)
ax1.invert_yaxis()
ax1.set_ylabel('Pressure, hPa')
ax1.set_xlim(-60, 60)
ax1.grid(True, linestyle=':')
#Second plot
f2 = ax2.contourf(data.lat,
data.pfull,
uv_conv_trans_after,
extend='both',
levels=np.arange(-5., 5.1, 1.))
ax2.contour(data.lat,
data.pfull,
u_ztav_after,
colors='0.75',
levels=np.arange(-60., 61., 10.))
ax2.contour(data.lat,
data.pfull,
psi_after,
colors='k',
levels=np.arange(-300., 301., 50.),
linewidths=2)
ax2.grid(True, linestyle=':')
ax2.set_xlim(-60, 60)
#Third plot
f2 = ax3.contourf(data.lat,
data.pfull,
uv_conv_stat_before,
extend='both',
levels=np.arange(-5., 5.1, 1.))
ax3.contour(data.lat,
data.pfull,
u_ztav_before,
colors='0.75',
levels=np.arange(-60., 61., 10.),
linewidths=2)
ax3.contour(data.lat,
data.pfull,
psi_before,
colors='k',
levels=np.arange(-300., 301., 50.),
linewidths=2)
ax3.grid(True, linestyle=':')
ax3.invert_yaxis()
ax3.set_ylabel('Pressure, hPa')
ax3.set_xlabel('Latitude')
#Fourth plot
f2 = ax4.contourf(data.lat,
data.pfull,
uv_conv_stat_after,
extend='both',
levels=np.arange(-5., 5.1, 1.))
ax4.contour(data.lat,
data.pfull,
u_ztav_after,
colors='0.75',
levels=np.arange(-60., 61., 10.),
linewidths=2)
ax4.contour(data.lat,
data.pfull,
psi_after,
colors='k',
levels=np.arange(-300., 301., 50.),
linewidths=2)
ax4.grid(True, linestyle=':')
ax4.set_xlabel('Latitude')
plt.subplots_adjust(right=0.95, top=0.95, bottom=0., hspace=0.1)
#Colorbar
cb1 = f.colorbar(f2,
ax=[ax1, ax2, ax3, ax4],
use_gridspec=True,
orientation='horizontal',
fraction=0.15,
pad=0.1,
aspect=30)
cb1.set_label('Eddy momentum flux convergence, $ms^{-1}day^{-1}$')
plt.savefig(plot_dir + 'upsi_eddies.pdf', format='pdf')
plt.close()
def psi_plot(run, period_fac=1., lonin=[-1., 361.], sanity_check=False):
rcParams['figure.figsize'] = 10, 7
rcParams['font.size'] = 25
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
omega = 7.2921150e-5
a = 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi / 180)
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, lons=lons, dp_in=50.)
psi /= 1.e9
edge_loc = np.zeros(psi.sel(pfull=500).values.shape)
edge_loc[psi.sel(pfull=500).values <= -50.] = 1.
edge_loc = xr.DataArray(edge_loc, psi.sel(pfull=500).coords)
eli = edge_loc.diff('lat')[::-1, :].argmax('lat')
edge_loc_lat = data.lat[eli]
edge_ll_summer = edge_loc_lat[abs(data.lat[eli]) < 30.]
psi_min = psi.sel(pfull=500.).min('lat')
psi_min_summer = -1. * psi_min[abs(data.lat[eli]) < 30.]
if sanity_check:
psi.sel(pfull=500).plot.contourf(levels=np.arange(-300., 301., 50.))
plt.plot(edge_loc_lat)
plt.figure(2)
plt.plot(psi_min)
plt.plot(psi_min_summer, 'r')
plt.show()
low_lat = (edge_ll_summer <= 9.) & (edge_ll_summer >= 0.)
high_lat = (edge_ll_summer > 9.)
edge_low_lat = edge_ll_summer[low_lat]
psi_min_low_lat = psi_min_summer[low_lat]
edge_high_lat = edge_ll_summer[high_lat]
psi_min_high_lat = psi_min_summer[high_lat]
A = np.array([edge_low_lat**(1. / 5.), np.ones(edge_low_lat.shape)])
consts_low_lat = np.linalg.lstsq(
A.T, psi_min_low_lat)[0] # obtaining the parameters
print consts_low_lat
A = np.array([edge_high_lat**(3. / 4.), np.ones(edge_high_lat.shape)])
consts_high_lat = np.linalg.lstsq(
A.T, psi_min_high_lat)[0] # obtaining the parameters
print consts_high_lat
line_low = consts_low_lat[0] * np.arange(
1., 9.)**(1. / 5.) + consts_low_lat[1] # regression line
line_high = consts_high_lat[0] * np.arange(
10., 31.)**(3. / 4.) + consts_high_lat[1] # regression line
plt.plot(np.arange(1., 9.), line_low, 'r-')
plt.plot(np.arange(10., 31.), line_high, 'b-')
plt.plot(edge_ll_summer, psi_min_summer, 'kx')
ax = plt.gca()
plt.xlabel('Northernmost latitude of SH Hadley cell')
plt.ylabel('Peak strength of SH Hadley circulation')
plt.xscale('log', subsx=[2, 3, 4, 5])
plt.yscale('log', subsx=[2, 3, 4, 5])
plt.xlim(1, 30)
plt.ylim(50, 500)
plt.xticks([1, 5, 10, 20, 30])
ax.get_xaxis().set_major_formatter(tk.ScalarFormatter())
plt.tight_layout()
figname = 'psi_latvsstrength_ylog.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
#plt.show()
return edge_ll_summer, psi_min_summer
def ampsi(run,
before_after,
filename='plev_pentad',
timeav='pentad',
period_fac=1.,
lonin=[-1., 361.],
pres_mask=False):
rcParams['figure.figsize'] = 10, 15
rcParams['font.size'] = 25
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
omega = 7.2921150e-5
a = 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi / 180)
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, lons=lons, dp_in=50.)
psi /= 1.e9
ang_mom = (data.ucomp.sel(lon=lons).mean('lon') +
omega * a * coslat) * a * coslat
ang_mom = ang_mom / a**2 / omega * 15. #scale following SB08 to give neater units for plotting
am_before = ang_mom[before_after[0]:before_after[0] +
4, :, :].mean('xofyear')
am_after = ang_mom[before_after[1]:before_after[1] +
4, :, :].mean('xofyear')
psi_before = psi[::-1, before_after[0]:before_after[0] +
4, ::-1].mean('xofyear').transpose()
psi_after = psi[::-1, before_after[1]:before_after[1] +
4, ::-1].mean('xofyear').transpose()
if pres_mask:
pmask = xr.DataArray(np.ones(am_before.shape), am_before.coords)
p_max = (data.ps.sel(lon=lons).min('lon')).mean('xofyear')
for i in range(0, 64):
a = [
j for j in range(len(data.pfull))
if data.pfull[j] >= p_max[i] / 100.
]
pmask[a, i] = float('NaN')
am_before = am_before * pmask
am_after = am_after * pmask
psi_before = psi_before * pmask
psi_after = psi_after * pmask
# Two subplots
f, (ax1, ax2) = plt.subplots(2, sharex=True)
plt.set_cmap('RdBu_r')
#First plot
f1 = ax1.contour(data.lat,
data.pfull,
am_before,
colors='0.75',
extend='both',
levels=np.arange(0., 17.))
ax1.contour(data.lat,
data.pfull,
psi_before,
colors='k',
levels=np.arange(-300., 301., 50.))
ax1.invert_yaxis()
ax1.set_ylabel('Pressure, hPa')
ax1.set_xlim(-60, 60)
ax1.grid(True, linestyle=':')
#Second plot
f2 = ax2.contour(data.lat,
data.pfull,
am_after,
colors='0.75',
extend='both',
levels=np.arange(0., 17.))
ax2.contour(data.lat,
data.pfull,
psi_after,
colors='k',
levels=np.arange(-300., 301., 50.))
ax2.invert_yaxis()
ax2.grid(True, linestyle=':')
ax2.set_xlim(-60, 60)
ax2.set_ylabel('Pressure, hPa')
ax2.set_xlabel('Latitude')
plt.subplots_adjust(right=0.95, top=0.95, bottom=0.05, hspace=0.1)
if lonin == [-1., 361.]:
figname = 'ampsi.pdf'
else:
figname = 'ampsi_' + str(int(lonin[0])) + '_' + str(int(
lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
pd_dic = pentad_dic(1)
u = xr.DataArray(np.zeros((72, 40, 64, len(years))), [('pentad', range(1, 73)),
('pfull', rundata.pfull),
('lat', rundata.lat),
('year', years)])
psi = xr.DataArray(np.zeros((64, 72, 40, len(years))),
[('lat', rundata.lat), ('pentad', range(1, 73)),
('pfull', rundata.pfull), ('year', years)])
i = 0
for year in years:
print year
rundata = load_year_xr(run_fol, year)
rundata.coords['pentad'] = (rundata.time // 5) - 71
psi_do = mass_streamfunction(rundata, a=6376.0e3)
u[:, :, :, i] = rundata.ucomp.groupby('pentad').mean(('time', 'lon'))
psi[:, :, :, i] = psi_do.groupby('pentad').mean(('time'))
i = i + 1
psi_av = psi.mean('year') / 1e9
u_av = u.mean('year')
#data.ucomp.mean(('time','lon')).plot.contourf(x='lat',y='pfull',levels=range(-20,61,5))
cs = psi_av[:, 36:41, :].mean('pentad').plot.contour(x='lat',
y='pfull',
yincrease=False,
colors='k',
add_colorbar=False,
levels=np.arange(
-300., 301., 50.))
return ds
data_ap2 = u_psi('ap_2', [30, 39])
v_era = xr.open_dataset('/scratch/rg419/obs_and_reanalysis/era_v_alllevs.nc')
# Take pentad means
v_era.coords['xofyear'] = v_era.day_of_yr // 5 + 1.
v_era = v_era.groupby('xofyear').mean(('day_of_yr'))
lons = [
v_era.lon[i] for i in range(len(v_era.lon))
if v_era.lon[i] >= 60. and v_era.lon[i] < 150.
]
psi = mass_streamfunction(v_era, a=6376.0e3, lons=lons, dp_in=50.)
psi /= 1.e9
bef_aft = [18, 39]
era_psi_before = psi[:, bef_aft[0]:bef_aft[0] +
4, ::-1].mean('xofyear').transpose()
era_psi_after = psi[:, bef_aft[1]:bef_aft[1] +
4, ::-1].mean('xofyear').transpose()
lwid = 2
# Four subplots
f, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, sharex='col', sharey='row')
plt.set_cmap('RdBu_r')
#First plot
f1 = ax1.contourf(data_ap2.lat,
data_ap2.pfull,
def get_edge_psi_time(data,
lons,
sanity_check=False,
lev=500.,
thresh=0.,
intdown=True):
"""Get cell edge for climatological data"""
# 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
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
# 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]))
#print times_defd
# 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))
# Up to this point code is same as psi_edge_loc.py. Now take psi_red and look for min instead of taking abs and finding max
#psi_max = np.abs(psi_red).max('lat')
#psi_max_loc = psi_red.lat.values[np.abs(psi_red).argmax('lat').values]
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'])
# _nh Create list of which times the min of psi_rad 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.)
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[ispos] = -1.*psi_max
psi_max = -1. * psi_max
# 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, ignore point, as this is from NH cell
edge_loc[
i] = np.nan #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.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.show()
return edge_loc, psi_max
def hm_vars(run,
months,
filename='atmos_pentad',
timeav='pentad',
period_fac=1.):
data = time_means(run,
months,
filename=filename,
timeav=timeav,
period_fac=period_fac)
psi = mass_streamfunction(data, a=6376.0e3) / 1e9
mse = (cp * data.temp + L * data.sphum + g * data.height) / 1000.
mse_max_loc = data.lat[np.argmax(mse[:, 38, :, :].mean('lon').values,
axis=1)]
omega_max_loc = data.lat[np.argmin(data.omega[:,
27, :, :].mean('lon').values,
axis=1)]
tsurf_max_loc = data.lat[np.argmax(data.t_surf.mean('lon').values, axis=1)]
oei = (np.diff([1 if i > 0 else 0
for i in omega_max_loc]).tolist()).index(1) + 1
mei = (np.diff([1 if i > 0 else 0
for i in mse_max_loc]).tolist()).index(1) + 1
tei = (np.diff([1 if i > 0 else 0
for i in tsurf_max_loc]).tolist()).index(1) + 1
data.t_surf.mean('lon').plot.contourf(x='xofyear',
y='lat',
add_labels=False,
extend='both',
levels=np.arange(250., 311., 5.))
plt.grid()
plt.plot(data.xofyear, tsurf_max_loc)
plt.plot(
data.xofyear,
np.amax(tsurf_max_loc) * np.sin(
(data.xofyear - tei) * 5. * np.pi / 180. / period_fac))
plt.ylim(-45, 45)
plt.xlabel('Pentad')
plt.ylabel('Latitude')
plt.title('SST, K')
plt.savefig('/scratch/rg419/plots/seasons_and_rotation/' + run +
'_sst_hm.png')
plt.close()
mse[:, 38, :, :].mean('lon').plot.contourf(x='xofyear',
y='lat',
add_labels=False,
extend='both',
levels=np.arange(255, 361, 5))
plt.grid()
plt.plot(data.xofyear, mse_max_loc)
plt.plot(
data.xofyear,
np.amax(mse_max_loc) * np.sin(
(data.xofyear - mei) * 5. * np.pi / 180. / period_fac))
plt.ylim(-45, 45)
plt.xlabel('Pentad')
plt.ylabel('Latitude')
plt.title('MSE, kJ/kg (921 hPa)')
plt.savefig('/scratch/rg419/plots/seasons_and_rotation/' + run +
'_mse_hm.png')
plt.close()
data.omega[:, 27, :, :].mean('lon').plot.contourf(x='xofyear',
y='lat',
add_labels=False,
extend='both',
levels=np.arange(
-0.1, 0.11, 0.01))
plt.grid()
plt.plot(data.xofyear, omega_max_loc, 'r')
plt.plot(
data.xofyear,
np.amax(omega_max_loc) * np.sin(
(data.xofyear - oei) * 5. * np.pi / 180. / period_fac), 'g')
plt.ylim(-45, 45)
plt.xlabel('Pentad')
plt.ylabel('Latitude')
plt.title('w, Pa/s (501 hPa)')
plt.savefig('/scratch/rg419/plots/seasons_and_rotation/' + run +
'_w_hm.png')
plt.close()
psi[:, :, 27].plot.contourf(x='xofyear',
y='lat',
add_labels=False,
extend='both',
levels=np.arange(-450., 451., 50.))
plt.grid()
plt.ylim(-45, 45)
plt.xlabel('Pentad')
plt.ylabel('Latitude')
plt.title('Mass streamfunction, 10^9 kg/s (501 hPa)')
plt.savefig('/scratch/rg419/plots/seasons_and_rotation/' + run +
'_psi_hm.png')
plt.close()
print 'done'
def upsi(run, months, before_after, filename='plev_pentad', timeav='pentad', period_fac=1.,lonin=[-1.,361.]):
rcParams['figure.figsize'] = 6, 9
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/clean_diags/'+run+'/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
#data = time_means(run, months, filename=filename, timeav=timeav,period_fac=period_fac)
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/'+run+'.nc')
a= 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi/180)
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, lons=lons, dp_in=50.)
psi /= 1.e9
psi_before = psi[::-1,before_after[0]:before_after[0]+5,::-1].mean('xofyear').transpose()
psi_after = psi[::-1,before_after[1]:before_after[1]+5,::-1].mean('xofyear').transpose()
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat *np.pi/180)
abs_vort_before = (data.vor[before_after[0]:before_after[0]+5,:,:,:].mean('xofyear') + f)*86400.
abs_vort_after = (data.vor[before_after[1]:before_after[1]+5,:,:,:].mean('xofyear') + f)*86400.
# Two subplots
f, (ax1, ax2) = plt.subplots(2, sharex=True)
plt.set_cmap('RdBu_r')
#First plot
f1 = ax1.contourf(data.lat, data.pfull, abs_vort_before.sel(lon=lons).mean('lon'), levels=np.arange(-14,15,2), extend = 'both')
ax1.contour(data.lat, data.pfull, psi_before, colors='k', levels=np.arange(-300.,301.,60.), linewidths=2)
ax1.invert_yaxis()
ax1.set_ylabel('Pressure, hPa')
ax1.set_xlim(-60,60)
ax1.grid(True,linestyle=':')
#Second plot
f2 = ax2.contourf(data.lat, data.pfull, abs_vort_after.sel(lon=lons).mean('lon'), levels=np.arange(-14,15,2), extend = 'both')
ax2.contour(data.lat, data.pfull, psi_after, colors='k', levels=np.arange(-300.,301.,60.), linewidths=2)
ax2.grid(True,linestyle=':')
ax2.set_xlim(-60,60)
ax2.invert_yaxis()
ax2.set_ylabel('Pressure, hPa')
ax2.set_xlabel('Latitude')
plt.subplots_adjust(left=0.15, right=0.95, top=0.95, bottom=0., hspace=0.2)
#Colorbar
cb1=f.colorbar(f2, ax=[ax1,ax2], use_gridspec=True, orientation = 'horizontal',fraction=0.15, pad=0.1, aspect=30)
cb1.set_label('Absolute vorticity, day$^{-1}$')
if lonin == [-1.,361.]:
figname = 'upsi_vor.pdf'
else:
figname = 'upsi_vor_' + str(int(lonin[0]))+ '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(plot_dir+figname, format='pdf')
plt.close()
import xarray as xr
import numpy as np
from finite_difference import cfd
a = 6376.0e3
omega = 7.2921150e-5
data = time_means('topo_10m', [121, 145],
filename='atmos_daily',
timeav='pentad')
u = data.ucomp.mean('lon')
print 'u loaded'
a_cos = a * np.cos(data.lat * np.pi / 180.)
m = (u + omega * a_cos) * a_cos
psi = abs(mass_streamfunction(data, a=6376.0e3) / 1e9)
dpsidy = xr.DataArray(cfd(psi.values, psi.lat * np.pi / 180, 0),
[('lat', psi.lat), ('xofyear', psi.xofyear),
('pfull', psi.pfull)])
dpsidp = xr.DataArray(cfd(psi.values, psi.pfull * 100, 2),
[('lat', psi.lat), ('xofyear', psi.xofyear),
('pfull', psi.pfull)])
dmdy = xr.DataArray(cfd(m.values, m.lat * np.pi / 180, 2),
[('xofyear', psi.xofyear), ('pfull', m.pfull),
('lat', m.lat)])
dmdp = xr.DataArray(cfd(m.values, m.pfull * 100, 1), [('xofyear', psi.xofyear),
('pfull', m.pfull),
('lat', m.lat)])
