def plot_clim(inp_fol):
plot_dir = '/scratch/rg419/plots/climatology/' + inp_fol + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
sn_dic = season_dic(1)
data = time_means(inp_fol, [121, 181],
filename='atmos_pentad',
timeav='season')
rain = (data.convection_rain + data.condensation_rain) * 86400.
del data
data_p = time_means(inp_fol, [121, 181],
filename='atmos_pentad',
timeav='season')
ucomp = data_p.ucomp
vcomp = data_p.vcomp
del data_p
land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/exp/summer_holiday/input/land.nc'
land = xr.open_dataset(land_file)
land_plot = xr.DataArray(land.land_mask.values, [('lat', ucomp.lat),
('lon', ucomp.lon)])
for i in range(0, 4):
print i
rain[i, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(0., 31., 3.),
add_label=False)
#plt.ylim(-30,60)
#plt.xlim(60,180)
land_plot.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
colors='k',
add_colorbar=False,
add_labels=False)
plt.quiver(ucomp.lon[::3],
ucomp.lat[::1],
ucomp[i, 36, ::1, ::3],
vcomp[i, 36, ::1, ::3],
scale=500.,
headwidth=5)
plt.savefig(plot_dir + 'wind_and_rain_' + sn_dic[i] + '.png')
plt.close()
def wind_comp_uv(run, months, filename='atmos_pentad', timeav='pentad', period_fac=1.):
data = time_means(run, months, filename=filename, timeav=timeav,period_fac=period_fac)
omega = 7.2921150e-5
a= 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi/180)
f = 2 * omega * np.sin(data.lat *np.pi/180)
dudy = xr.DataArray( cfd((data.ucomp*coslat).values,data.lon*np.pi/180.,3), [('xofyear', data.xofyear ), ('pfull', data.pfull ), ('lat', data.lat), ('lon', data.lon )])
dvdy = xr.DataArray( cfd((data.vcomp*coslat).values,data.lon*np.pi/180.,3), [('xofyear', data.xofyear ), ('pfull', data.pfull ), ('lat', data.lat), ('lon', data.lon )])
dlat=xr.DataArray(data.latb.diff('latb').values*np.pi/180, coords=[('lat', data.lat)])
v_div = np.cumsum(dvdy * dlat, axis=2)/coslat
u_vort = np.cumsum(dudy * dlat, axis=2)/coslat
u_div = data.ucomp - u_vort
v_vort = data.vcomp - v_div
plt.figure(1)
v_vort[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-10,11,1))
plt.figure(2)
v_div[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-10,11,1))
plt.figure(3)
data.vcomp[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-10,11,1))
plt.figure(4)
u_vort[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-80,81,10))
plt.figure(5)
u_div[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-80,81,10))
plt.figure(6)
data.ucomp[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-80,81,10))
plt.show()
def hugo_plot(run,
months,
filename='atmos_pentad',
timeav='pentad',
period_fac=1.):
data = time_means(run,
months,
filename=filename,
timeav=timeav,
period_fac=period_fac)
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)]
plt.plot(tsurf_max_loc, omega_max_loc, 'x')
plt.plot([-50, 50], [-50, 50], 'k--')
plt.xlim(-50, 50)
plt.ylim(-50, 50)
plt.xlabel('Latitude of max SST')
plt.ylabel('Latitude of max w')
plt.savefig('/scratch/rg419/plots/seasons_and_rotation/hugo_' + run +
'.png')
plt.show()
def w_max_lat(run, months,filename='atmos_pentad', period_fac=1.):
data = time_means(run, months, filename=filename, timeav='pentad',period_fac=period_fac)
omega_max_loc = data.lat[np.argmin(data.omega[36:71,27,:,:].mean('lon').values, axis=1)]
omega_max = np.amin(data.omega[36:71,27,:,:].mean('lon').values, axis=1)
print 'peak ascent located'
plt.plot(omega_max_loc[0:11],omega_max[0:11],'x')
plt.plot(omega_max_loc[12:23],omega_max[12:23],'xr')
plt.plot(omega_max_loc[24:35],omega_max[24:35],'xg')
plt.show()
def p_and_t(run,
months,
filename='atmos_pentad',
timeav='pentad',
period_fac=1.,
lonin=[-1., 361.]):
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)
totp = (data.convection_rain + data.condensation_rain) * 86400.
lons = pick_lons(data, lonin)
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
# ax=totp[:,:,lons].mean('lon').plot.contourf(x='xofyear', y='lat',levels=np.arange(6.,31.,3.), add_label = False, add_colorbar=False)
ax = totp[:, :,
lons].mean('lon').plot.contourf(x='xofyear',
y='lat',
levels=np.arange(2., 15., 2.),
add_label=False,
add_colorbar=False,
extend='both')
totp[:, :,
lons].mean('lon').plot.contour(x='xofyear',
y='lat',
levels=np.arange(-92., 109., 100.),
add_label=False,
add_colorbar=False,
colors='k')
plt.grid(True, linestyle=':')
cb1 = plt.colorbar(ax)
cb1.set_label('Precip, mm/day')
cs = data.t_surf[:, :,
lons].mean('lon').plot.contour(x='xofyear',
y='lat',
levels=range(250, 321, 10),
add_label=False,
colors='w',
add_colorbar=False)
plt.clabel(cs, fontsize=15, inline_spacing=-1, fmt='%1.0f')
plt.ylim((-45, 45))
plt.xlim((1, 73))
plt.xlabel('')
plt.xticks(tickspace, labels, rotation=25)
plt.ylabel('Latitude')
plt.tight_layout()
plt.savefig(plot_dir + 'p_and_t.png')
plt.close()
def cross_prod(run,
months,
filename='plev_pentad',
timeav='month',
period_fac=1.,
land_mask=False,
level=9):
plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/cross_prod/' + str(
level) + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
mn_dic = month_dic(1)
data = time_means(run,
months,
filename=filename,
timeav=timeav,
period_fac=period_fac)
uwnd = data.ucomp[:, level, :, :]
vwnd = data.vcomp[:, level, :, :]
w = VectorWind(uwnd, vwnd)
uchi, vchi, upsi, vpsi = w.helmholtz()
cross_prod = (upsi * vchi - vpsi * uchi)
for i in range(0, 12):
ax = cross_prod[i, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(
-200., 201., 10.),
extend='both',
add_colorbar=False,
add_label=False)
cb1 = plt.colorbar(ax)
cb1.set_label('Vpsi x Vchi')
if land_mask:
land = xr.open_dataset(
'/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc')
land_plot = xr.DataArray(land.land_mask.values,
[('lat', data.lat), ('lon', data.lon)])
land_plot.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
colors='0.75',
add_colorbar=False,
add_labels=False)
plt.ylabel('Latitude')
plt.xlabel('Longitude')
plt.ylim(-10, 45)
plt.xlim(25, 150)
plt.tight_layout()
plt.savefig(plot_dir + str(i + 1) + '_' + str(mn_dic[i + 1]) + '.png')
plt.close()
def surf_ts(run,
months,
filename='atmos_pentad',
timeav='pentad',
period_fac=1.):
data = time_means(run,
months,
filename=filename,
timeav=timeav,
period_fac=period_fac)
plt.plot(data.xofyear, data.flux_t[:, 37, :].mean('lon'))
plt.xlabel('Pentad')
plt.ylabel('Sensible heat flux, W/m2 (15N)')
plt.savefig('/scratch/rg419/plots/seasons_and_rotation/' + run +
'_flux_t_ts.png')
plt.close()
plt.plot(data.xofyear, data.flux_lhe[:, 37, :].mean('lon'))
plt.xlabel('Pentad')
plt.ylabel('Latent heat flux, W/m2 (15N)')
plt.savefig('/scratch/rg419/plots/seasons_and_rotation/' + run +
'_flux_lhe_ts.png')
plt.close()
plt.plot(data.xofyear, data.flux_sw[:, 37, :].mean('lon'))
plt.xlabel('Pentad')
plt.ylabel('SW heat flux, W/m2 (15N)')
plt.savefig('/scratch/rg419/plots/seasons_and_rotation/' + run +
'_flux_sw_ts.png')
plt.close()
bsurf = stefan * np.power(data.t_surf, 4)
plt.plot(data.xofyear, (bsurf - data.flux_lw)[:, 37, :].mean('lon'))
plt.xlabel('Pentad')
plt.ylabel('LW heat flux, W/m2 (15N)')
plt.savefig('/scratch/rg419/plots/seasons_and_rotation/' + run +
'_flux_lw_ts.png')
plt.close()
plt.plot(data.xofyear,
(bsurf - data.flux_lw + data.flux_lhe)[:, 37, :].mean('lon'))
plt.xlabel('Pentad')
plt.ylabel('LW heat flux, W/m2 (15N)')
plt.savefig('/scratch/rg419/plots/seasons_and_rotation/' + run +
'_flux_lwlhe_ts.png')
plt.close()
plt.plot(data.xofyear, data.t_surf[:, 37, :].mean('lon'))
plt.xlabel('Pentad')
plt.ylabel('SST, K (15N)')
plt.savefig('/scratch/rg419/plots/seasons_and_rotation/' + run +
'_SST15N_ts.png')
plt.close()
print 'done'
def rain_and_sst(run,
months,
filename='plev_pentad',
period_fac=1.,
plot_land=True):
rcParams['figure.figsize'] = 20, 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)
data['rain'] = (('xofyear', 'lat', 'lon'),
(data.convection_rain + data.condensation_rain) * 86400.)
land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc'
land = xr.open_dataset(land_file)
data['land'] = (('lat', 'lon'), land.land_mask)
for i in range(0, 4):
print i
ax = data.rain[i, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(0., 31., 3.),
add_labels=False,
extend='max',
add_colorbar=False)
if plot_land:
data.land.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
colors='k',
add_colorbar=False,
add_labels=False)
cs = data.t_surf[i, :, :].plot.contour(x='lon',
y='lat',
levels=np.arange(
200., 350., 10.),
colors='w',
add_labels=False,
add_colorbar=False)
plt.clabel(cs, fontsize=15, inline_spacing=-1, fmt='%1.0f')
cb1 = plt.colorbar(ax)
cb1.set_label('Precipitation, mm/day')
plt.ylabel('Latitude')
plt.xlabel('Longitude')
plt.tight_layout()
plt.savefig(plot_dir + 'rain_and_sst_' + sn_dic[i] + '.png')
plt.close()
def ke_partition(run,
months,
filename='atmos_pentad',
timeav='pentad',
period_fac=1.):
data = time_means(run,
months,
filename=filename,
timeav=timeav,
period_fac=period_fac)
totp = (data.convection_rain + data.condensation_rain) * 86400.
cell_ar = cell_area(42, '/scratch/rg419/GFDL_model/GFDLmoistModel/')
cell_ar = xr.DataArray(cell_ar, [('lat', data.lat), ('lon', data.lon)])
uwnd = data.ucomp
vwnd = data.vcomp
w = VectorWind(uwnd, vwnd)
uchi, vchi, upsi, vpsi = w.helmholtz()
lats = [
i for i in range(len(data.lat))
if data.lat[i] >= 5. and data.lat[i] < 30.
]
lons = [
i for i in range(len(data.lon))
if data.lon[i] >= 60. and data.lon[i] < 150.
]
ke_chi = 0.5 * (uchi * uchi + vchi * vchi) * cell_ar
ke_chi_av = ke_chi[:, :, lats, lons].sum('lat').sum('lon') / cell_ar[
lats, lons].sum('lat').sum('lon')
ke_psi = 0.5 * (upsi * upsi + vpsi * vpsi) * cell_ar
ke_psi_av = ke_psi[:, :, lats, lons].sum('lat').sum('lon') / cell_ar[
lats, lons].sum('lat').sum('lon')
totp_av = (totp[:, lats, lons] * cell_ar[lats, lons]).sum('lat').sum(
'lon') / cell_ar[lats, lons].sum('lat').sum('lon')
ke_chi_av[:, 36].plot()
ke_psi_av[:, 36].plot()
plt.legend(['KE_chi', 'KE_psi'])
plt.xlabel('Pentad')
plt.ylabel('Kinetic Energy, m2/s2')
plt.savefig(plot_dir + 'KE_' + run + '.png')
plt.close()
totp_av.plot()
plt.xlabel('Pentad')
plt.ylabel('Precipitation, mm/day')
plt.savefig(plot_dir + 'precip_' + run + '.png')
plt.close()
def oei_centred_data(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
return omega_max_loc[oei - 15:oei + 16]
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 diabatic_heating(run,
months,
filename='plev_pentad',
period_fac=1.,
plot_land=True):
rcParams['figure.figsize'] = 20, 10
stefan = 5.6734e-8
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)
data['heat_rate'] = (('xofyear', 'pfull', 'lat', 'lon'),
(data.tdt_rad + data.dt_tg_condensation +
data.dt_tg_convection + data.dt_tg_diffusion) *
86400.)
land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc'
land = xr.open_dataset(land_file)
data['land'] = (('lat', 'lon'), land.land_mask)
for i in range(0, 4):
print i
ax = data.heat_rate[i, 4, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(
-10., 10., 1.),
add_labels=False,
extend='both',
add_colorbar=False)
if plot_land:
data.land.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
colors='k',
add_colorbar=False,
add_labels=False)
cb1 = plt.colorbar(ax)
cb1.set_label('Diabatic heating, K/day')
plt.ylabel('Latitude')
plt.xlabel('Longitude')
plt.tight_layout()
plt.savefig(plot_dir + 'diabatic_heating_' + sn_dic[i] + '.png')
plt.close()
def surface_fluxes(run,
months,
filename='plev_pentad',
period_fac=1.,
plot_land=True):
rcParams['figure.figsize'] = 20, 10
stefan = 5.6734e-8
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)
data['surf_flux'] = (('xofyear', 'lat',
'lon'), stefan * data.t_surf**4. - data.flux_lw -
data.flux_lw + data.flux_t + data.flux_lhe)
land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc'
land = xr.open_dataset(land_file)
data['land'] = (('lat', 'lon'), land.land_mask)
for i in range(0, 4):
print i
ax = data.surf_flux[i, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(
-700., 700., 50.),
add_labels=False,
extend='both',
add_colorbar=False)
if plot_land:
data.land.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
colors='k',
add_colorbar=False,
add_labels=False)
cb1 = plt.colorbar(ax)
cb1.set_label('Surface flux, W/m2')
plt.ylabel('Latitude')
plt.xlabel('Longitude')
plt.tight_layout()
plt.savefig(plot_dir + 'surface_fluxes_' + sn_dic[i] + '.png')
plt.close()
def omega_max(run,
months,
filename='atmos_pentad',
timeav='pentad',
period_fac=1.):
data = time_means(run,
months,
filename=filename,
timeav=timeav,
period_fac=period_fac)
omega_max_loc = data.lat[np.argmin(data.omega[:,
27, :, :].mean('lon').values,
axis=1)]
oei = (np.diff([1 if i > 0 else 0
for i in omega_max_loc]).tolist()).index(1) + 1
return omega_max_loc, oei
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 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 load_mom_vars(run, output_file):
data = time_means(run, [121, 481], filename='plev_pentad', timeav='pentad')
print 'data passed'
mom_budg = xr.Dataset({
'ucomp': data.ucomp,
'vcomp': data.vcomp,
'omega': data.omega,
'ucomp_sq': data.ucomp_sq,
'ucomp_vcomp': data.ucomp_vcomp,
'ucomp_omega': data.ucomp_omega,
'height': data.height,
'dt_ug_diffusion': data.dt_ug_diffusion
})
print 'mom_budg defined'
mom_budg.to_netcdf(path=output_file, mode='w')
print 'mom_budg saved'
data.close()
mom_budg.close()
print 'data closed'
def wind_components(run, months, filename='atmos_pentad', timeav='pentad', period_fac=1.):
data = time_means(run, months, filename=filename, timeav=timeav,period_fac=period_fac)
omega = 7.2921150e-5
a= 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi/180)
f = 2 * omega * np.sin(data.lat *np.pi/180)
dudx = xr.DataArray( cfd(data.ucomp.values,data.lon*np.pi/180.,3), [('xofyear', data.xofyear ), ('pfull', data.pfull ), ('lat', data.lat), ('lon', data.lon )])
dudx = dudx #/coslat/a Don't bother with geometric factor as will just be multiplied by in integration
dvdx = xr.DataArray( cfd(data.vcomp.values,data.lon*np.pi/180.,3), [('xofyear', data.xofyear ), ('pfull', data.pfull ), ('lat', data.lat), ('lon', data.lon )])
dvdx = dvdx #/coslat/a
dvdy_div = data.div - dudx
dudy_vort = dvdx - data.vor
dlat=xr.DataArray(data.latb.diff('latb').values*np.pi/180, coords=[('lat', data.lat)])
v_div = np.cumsum(dvdy_div * dlat, axis=2)/coslat
u_vort = np.cumsum(dudy_vort * dlat, axis=2)/coslat
u_div = data.ucomp - u_vort
v_vort = data.vcomp - v_div
plt.figure(1)
v_vort[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-10,11,1))
plt.figure(2)
v_div[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-10,11,1))
plt.figure(3)
data.vcomp[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-10,11,1))
plt.figure(4)
u_vort[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-80,81,10))
plt.figure(5)
u_div[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-80,81,10))
plt.figure(6)
data.ucomp[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-80,81,10))
plt.show()
#Create ozone single time file.
import xarray as xr
import numpy as np
from data_handling import time_means
from finite_difference import cfd
import matplotlib.pyplot as plt
data = time_means('aquaplanet_10m_test', [121, 145],
filename='plev_pentad',
timeav='month')
def eddy_decomp(a, b, ab):
#Decompose fields into mean, stationary and transient eddies
a_zmean = a.mean(('lon'))
a_zed = a - a_zmean
b_zmean = b.mean(('lon'))
b_zed = b - b_zmean
ab_ms = a_zmean * b_zmean
ab_stat = a_zmean * b_zed + a_zed * b_zmean + a_zed * b_zed
ab_trans = ab - ab_stat - ab_ms
return ab_ms, ab_stat, ab_trans
def calc_mean_heat(vt, wt):
#Calculate momentum advection by time and zonal mean
a = 6376.0e3 #radius used in model
def mom_budg_hm(run,
months,
lev=17,
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)
if lonin[1] > lonin[0]:
lons = [
i for i in range(len(data.lon))
if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]
]
else:
lons = [
i for i in range(len(data.lon))
if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]
]
#advective terms
flux_conv(data, lons, lev=lev)
#Coriolis
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
fv = data.vcomp[:, lev, :, :] * f * 86400.
fv = fv[:, :, lons].mean('lon')
#Geopotential gradient
dphidx = ddx(data.height[:, lev, :, :])
dphidx = 9.8 * dphidx[:, :, lons].mean('lon')
mom_mean = data.u_dudx + data.v_dudy + data.w_dudp
mom_trans = data.uu_trans_dx + data.uv_trans_dy + data.uw_trans_dp
mom_stat = data.uu_stat_dx + data.uv_stat_dy + data.uw_stat_dp
mom_sum = fv + dphidx + mom_mean + mom_trans + mom_stat
levels = np.arange(-10, 10.1, 2.)
# Six subplots
f, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2,
3,
sharex='col',
sharey='row')
plt.set_cmap('RdBu_r')
#First plot
f1 = fv.plot.contourf(ax=ax1,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax1.set_ylabel('Latitude')
ax1.set_ylim(-60, 60)
ax1.grid(True, linestyle=':')
#Second plot
mom_mean.plot.contourf(ax=ax2,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax2.grid(True, linestyle=':')
ax2.set_ylim(-60, 60)
#Third plot
dphidx.plot.contourf(ax=ax3,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax3.grid(True, linestyle=':')
ax3.set_ylim(-60, 60)
#Fourth plot
mom_trans.plot.contourf(ax=ax4,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax4.grid(True, linestyle=':')
ax4.set_ylabel('Latitude')
ax4.set_xlabel('Pentad')
ax4.set_ylim(-60, 60)
#Fifth plot
mom_stat.plot.contourf(ax=ax5,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax5.grid(True, linestyle=':')
ax5.set_xlabel('Pentad')
ax5.set_ylim(-60, 60)
#Sixth plot
mom_sum.plot.contourf(ax=ax6,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax6.grid(True, linestyle=':')
ax6.set_xlabel('Pentad')
ax6.set_ylim(-60, 60)
plt.subplots_adjust(right=0.95, top=0.95, bottom=0., hspace=0.1)
#Colorbar
cb1 = f.colorbar(f1,
ax=[ax1, ax2, ax3, ax4, ax5, ax6],
use_gridspec=True,
orientation='horizontal',
fraction=0.15,
pad=0.1,
aspect=30)
cb1.set_label('$ms^{-1}day^{-1}$')
plt.savefig(plot_dir + 'zon_mom_budg_60_150.pdf', format='pdf')
plt.close()
plt.title(plttitle)
plt.savefig('/scratch/rg419/plots/seasons_and_rotation/' + plttitle +
'_fft.png')
plt.close()
plt.plot(var_max_loc)
plt.xlabel('Pentad')
plt.ylabel('Latitude')
plt.title(plttitle)
plt.savefig('/scratch/rg419/plots/seasons_and_rotation/' + plttitle +
'_mean.png')
plt.close()
data = time_means('sn_3.000', [121, 481],
filename='atmos_pentad',
timeav='pentad',
period_fac=3.)
fft_max(data.omega[:, 27, :, :],
nmax=False,
period_fac=3.,
plttitle='sn_3.000_omega')
mse = (cp * data.temp + L * data.sphum + g * data.height) / 1000.
fft_max(mse[:, 38, :, :], period_fac=3., plttitle='sn_3.000_mse')
data = time_means('aquaplanet_10m', [121, 481],
filename='atmos_pentad',
timeav='pentad')
fft_max(data.omega[:, 27, :, :], nmax=False, plttitle='aquaplanet_10m_omega')
mse = (cp * data.temp + L * data.sphum + g * data.height) / 1000.
fft_max(mse[:, 38, :, :], plttitle='aquaplanet_10m_mse')
def vort_av(run,
months,
filename='atmos_pentad',
timeav='pentad',
period_fac=1.):
data = time_means(run,
months,
filename=filename,
timeav=timeav,
period_fac=period_fac)
omega = 7.2921150e-5
a = 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi / 180)
f = 2 * omega * np.sin(data.lat * np.pi / 180)
abs_vort = data.vor + f
dvortdx = xr.DataArray(cfd(abs_vort.values, data.lon * np.pi / 180., 3),
[('xofyear', data.xofyear), ('pfull', data.pfull),
('lat', data.lat), ('lon', data.lon)])
dvortdx = dvortdx / coslat / a
dvortdy = xr.DataArray(cfd(abs_vort.values, data.lat * np.pi / 180., 2),
[('xofyear', data.xofyear), ('pfull', data.pfull),
('lat', data.lat), ('lon', data.lon)])
dvortdy = dvortdy / a
dvortdp = xr.DataArray(cfd(abs_vort.values, data.pfull * 100., 1),
[('xofyear', data.xofyear), ('pfull', data.pfull),
('lat', data.lat), ('lon', data.lon)])
vort_adv = data.ucomp * dvortdx + data.vcomp * dvortdy
vort_vert = data.omega * dvortdp
vort_div = abs_vort * data.div
precip = (data.convection_rain + data.condensation_rain) * 86400.
plt.figure(1)
vort_adv[:, 15, :, :].mean('lon').plot.contourf(x='xofyear',
y='lat',
levels=np.arange(
-4.5e-10, 4.6e-10,
0.5e-10))
plt.figure(2)
vort_div[:, 15, :, :].mean('lon').plot.contourf(x='xofyear',
y='lat',
levels=np.arange(
-4.5e-10, 4.6e-10,
0.5e-10))
plt.figure(3)
vort_vert[:, 15, :, :].mean('lon').plot.contourf(x='xofyear', y='lat')
plt.figure(4)
abs_vort[:, 15, :, :].mean('lon').plot.contourf(x='xofyear',
y='lat',
levels=np.arange(
-2.4e-4, 2.5e-4,
0.2e-4))
precip.mean('lon').plot.contour(x='xofyear',
y='lat',
levels=np.arange(6., 31., 6.),
add_label=False,
add_colorbar=False,
colors='k')
plt.show()
"""Check updated netcdfs load ok"""
import numpy as np
from data_handling import time_means
import xarray as xr
import matplotlib.pyplot as plt
test = time_means('ap_1_rd', [1, 13], filename='atmos_pentad', timeav='pentad')
test.ucomp[71, 37, :, :].plot.contourf()
plt.show()
def load_data(inp_fol):
data = time_means(inp_fol, [121, 481], 'pentad')
mombudg = mom_budg.mombudg_fn(data)
mombudg.to_netcdf(path='/scratch/rg419/plots/monsoon_analysis/' + inp_fol +
'/mom_stat_data.nc',
mode='w')
except:
raise NameError('Coord ' + timedir + ' not found')
field_dt = cfd( field.values, field.coords[timedir].values*secperunit, field.get_axis_num(timedir), cyclic=cyclic)/secperunit/2.
field_dt = xr.DataArray( field_dt, dims = field.dims, coords = field.coords )
return field_dt
if __name__ == '__main__':
"""Examples/sanity check"""
import matplotlib.pyplot as plt
from data_handling import time_means
data = time_means('full_qflux', [121,481], filename='plev_pentad', timeav='pentad')
dudx = ddx(data.ucomp[40,17,:,:])
dudy = ddy(data.ucomp[40,17,:,:])
dudp = ddp(data.ucomp[40,:,:,:].mean('lon'))
dudt = ddt(data.ucomp[:,17,32,0])
data.ucomp[40,17,:,:].plot.contourf(levels=np.arange(-50.,51.,10.))
plt.figure(2)
dudx.plot.contourf(levels=np.arange(-0.000016,0.0000161,0.000002))
plt.figure(3)
dudy.plot.contourf(levels=np.arange(-0.000016,0.0000161,0.000002))
plt.figure(4)
data.ucomp[40,:,:,:].mean('lon').plot.contourf(levels=np.arange(-50.,51.,10.), yincrease=False)
plt.figure(5)
# 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')
import xarray as xr
import matplotlib.pyplot as plt
import gc
plt.rc('text', usetex=True)
font = {
'family': 'sans-serif',
'sans-serif': ['Helvetica'],
'weight': 'bold',
'size': 18
}
plt.rc('font', **font)
data = time_means('aquaplanet_2m', [121, 481],
filename='atmos_daily',
timeav='pentad')
def plot_precip(data, lons=range(0, 128)):
print 'precip'
data['totp'] = (('xofyear', 'lat', 'lon'),
(data.convection_rain + data.condensation_rain) * 86400.)
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
ax = data.totp[:, :, lons].mean('lon').plot.contourf(x='xofyear',
y='lat',
levels=np.arange(
6., 25., 3.),
def advection_hm(run,
months,
lev=150,
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)
if lonin[1] > lonin[0]:
lons = [
data.lon[i] for i in range(len(data.lon))
if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]
]
else:
lons = [
data.lon[i] for i in range(len(data.lon))
if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]
]
#advective terms
partition_advection(data, lons, lev=150)
data['mom_mean'] = (('xofyear', 'lat'),
data.u_dudx_zav + data.v_dudy_zav + data.w_dudp_zav)
data['mom_trans'] = (('xofyear', 'lat'), data.uu_trans_dx +
data.uv_trans_dy + data.uw_trans_dp)
data['mom_stat'] = (('xofyear', 'lat'),
data.u_dudx_stat + data.v_dudy_stat + data.w_dudp_stat)
levels = np.arange(-10, 10.1, 2.)
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
fig_dicts = [{
'p1': 'mom_mean',
'p2': 'u_dudx_zav',
'p3': 'v_dudy_zav',
'p4': 'w_dudp_zav',
'namebase': 'meanstate_adv'
}, {
'p1': 'mom_stat',
'p2': 'u_dudx_stat',
'p3': 'v_dudy_stat',
'p4': 'w_dudp_stat',
'namebase': 'stat_adv'
}, {
'p1': 'mom_trans',
'p2': 'uu_trans_dx',
'p3': 'uv_trans_dy',
'p4': 'uw_trans_dp',
'namebase': 'trans_adv'
}]
for fig_dict in fig_dicts:
# Six subplots
fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2,
3,
sharex='col',
sharey='row')
plt.set_cmap('RdBu_r')
#First plot
f1 = data[fig_dict['p1']].plot.contourf(ax=ax1,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax1.set_ylabel('Latitude')
ax1.set_ylim(-60, 60)
ax1.grid(True, linestyle=':')
#Fourth plot
data[fig_dict['p2']].plot.contourf(ax=ax4,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax4.grid(True, linestyle=':')
ax4.set_ylabel('Latitude')
ax4.set_xticks(tickspace)
ax4.set_xticklabels(labels, rotation=25)
ax4.set_ylim(-60, 60)
#Fifth plot
data[fig_dict['p3']].plot.contourf(ax=ax5,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax5.grid(True, linestyle=':')
ax5.set_xticks(tickspace)
ax5.set_xticklabels(labels, rotation=25)
ax5.set_ylim(-60, 60)
#Sixth plot
data[fig_dict['p4']].plot.contourf(ax=ax6,
x='xofyear',
y='lat',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
ax6.grid(True, linestyle=':')
ax6.set_xticks(tickspace)
ax6.set_xticklabels(labels, rotation=25)
ax6.set_ylim(-60, 60)
plt.subplots_adjust(right=0.95, top=0.95, bottom=0., hspace=0.1)
#Colorbar
cb1 = fig.colorbar(f1,
ax=[ax1, ax2, ax3, ax4, ax5, ax6],
use_gridspec=True,
orientation='horizontal',
fraction=0.15,
pad=0.1,
aspect=30)
cb1.set_label('$ms^{-1}day^{-1}$')
if lonin == [-1., 361.]:
figname = fig_dict['namebase'] + '.pdf'
else:
figname = fig_dict['namebase'] + '_' + str(int(
lonin[0])) + '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()