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()