def calc_Koc(ds, detrend=False, plot=False): if detrend: ds['s'].data = signal.detrend(ds['s'].data, axis=0) + ds['s'].mean(dim='time').data s_mean = ds['s'].mean(dim='time') s_anom = ds['s'] - s_mean w = VectorWind(ds['u'], ds['v']) grad_s_mean = w.gradient(s_mean) grad_s_mean = VectorWind(grad_s_mean[0], grad_s_mean[1]) abs_grad_s_mean = grad_s_mean.magnitude() sq_abs_grad_s_mean = abs_grad_s_mean * abs_grad_s_mean grad_s_anom = w.gradient(s_anom) grad_s_anom = VectorWind(grad_s_anom[0], grad_s_anom[1]) abs_grad_s_anom = grad_s_anom.magnitude() sq_abs_grad_s_anom = abs_grad_s_anom * abs_grad_s_anom sq_abs_grad_s_anom = sq_abs_grad_s_anom.mean(dim='time') sq_abs_grad_s_mean = sq_abs_grad_s_mean.mean(dim='longitude') sq_abs_grad_s_anom = sq_abs_grad_s_anom.mean(dim='longitude') Koc = sq_abs_grad_s_anom / sq_abs_grad_s_mean if plot: plt.ion() fig1, ax1 = plt.subplots() ax1.plot(lats, sq_abs_grad_s_mean.data, label='mean') ax1.plot(lats, sq_abs_grad_s_anom.data, label='anom') ax11 = ax1.twinx() ax11.plot(lats, ds['q'].mean(dim=['time', 'longitude']).data) fig2, ax2 = plt.subplots() ax2.plot(lats, Koc, label='Koc') ax21 = ax2.twinx() ax21.plot(lats, ds['q'].mean(dim=['time', 'longitude']).data) plt.show() return Koc
def potential_vorticity_baroclinic(uwnd, vwnd, theta, coord, **kwargs): ''' Calculate potential vorticity on isobaric levels. Requires input uwnd, vwnd and tmp arrays to have (lat, lon, ...) format for Windspharm. Input ----- uwnd : zonal winds, array-like vwnd : meridional winds, array-like tmp : temperature, array-like theta : potential temperature, array-like coord : dimension name for pressure axis (eg. 'pfull') omega : planetary rotation rate, optional g : planetary gravitational acceleration, optional rsphere : planetary radius, in metres, optional ''' omega = kwargs.pop('omega', 7.08822e-05) g = kwargs.pop('g', 3.72076) rsphere = kwargs.pop('rsphere', 3.3962e6) w = VectorWind(uwnd.fillna(0), vwnd.fillna(0), rsphere=rsphere) relvort = w.vorticity() relvort = relvort.where(relvort != 0, other=np.nan) planvort = w.planetaryvorticity(omega=omega) absvort = relvort + planvort dthtady, dthtadx = w.gradient(theta.fillna(0)) dthtadx = dthtadx.where(dthtadx != 0, other=np.nan) dthtady = dthtady.where(dthtady != 0, other=np.nan) dthtadp = wrapped_gradient(theta, coord) dudp = wrapped_gradient(uwnd, coord) dvdp = wrapped_gradient(vwnd, coord) s = -dthtadp f = dvdp * dthtadx - dudp * dthtady ret = g * (absvort + f / s) * s return ret
# Read zonal and meridional wind components from file using the xarray module. # The components are in separate files. ds = xr.open_mfdataset( [example_data_path(f) for f in ('uwnd_mean.nc', 'vwnd_mean.nc')]) uwnd = ds['uwnd'] vwnd = ds['vwnd'] # Create a VectorWind instance to handle the computations. w = VectorWind(uwnd, vwnd) # Compute components of rossby wave source: absolute vorticity, divergence, # irrotational (divergent) wind components, gradients of absolute vorticity. eta = w.absolutevorticity() div = w.divergence() uchi, vchi = w.irrotationalcomponent() etax, etay = w.gradient(eta) etax.attrs['units'] = 'm**-1 s**-1' etay.attrs['units'] = 'm**-1 s**-1' # Combine the components to form the Rossby wave source term. S = eta * -1. * div - (uchi * etax + vchi * etay) # Pick out the field for December at 200 hPa. S_dec = S[S['time.month'] == 12] # Plot Rossby wave source. clevs = [-30, -25, -20, -15, -10, -5, 0, 5, 10, 15, 20, 25, 30] ax = plt.subplot(111, projection=ccrs.PlateCarree(central_longitude=180)) S_dec *= 1e11 fill = S_dec[0].plot.contourf(ax=ax, levels=clevs,
# Read zonal and meridional wind components from file using the xarray module. # The components are in separate files. ds = xr.open_mfdataset([example_data_path(f) for f in ('uwnd_mean.nc', 'vwnd_mean.nc')]) uwnd = ds['uwnd'] vwnd = ds['vwnd'] # Create a VectorWind instance to handle the computations. w = VectorWind(uwnd, vwnd) # Compute components of rossby wave source: absolute vorticity, divergence, # irrotational (divergent) wind components, gradients of absolute vorticity. eta = w.absolutevorticity() div = w.divergence() uchi, vchi = w.irrotationalcomponent() etax, etay = w.gradient(eta) etax.attrs['units'] = 'm**-1 s**-1' etay.attrs['units'] = 'm**-1 s**-1' # Combine the components to form the Rossby wave source term. S = eta * -1. * div - (uchi * etax + vchi * etay) # Pick out the field for December at 200 hPa. S_dec = S[S['time.month'] == 12] # Plot Rossby wave source. clevs = [-30, -25, -20, -15, -10, -5, 0, 5, 10, 15, 20, 25, 30] ax = plt.subplot(111, projection=ccrs.PlateCarree(central_longitude=180)) S_dec *= 1e11 fill = S_dec[0].plot.contourf(ax=ax, levels=clevs, cmap=plt.cm.RdBu_r, transform=ccrs.PlateCarree(), extend='both',