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