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',
