def absvrt(uwnd, vwnd, lon, lat, xdim, ydim, cyclic=True, sphere=True):
"""
Calculate absolute vorticity.
:param uwnd: ndarray, u-component wind.
:param vwnd: ndarray, v-component wind.
:param lon: array_like, longitude.
:param lat: array_like, latitude.
:param xdim: the longitude dimension index
:param ydim: the latitude dimension index
:param cyclic: east-west boundary is cyclic
:param sphere: sphere coordinate
:return:
"""
u, v = np.ma.getdata(uwnd), np.ma.getdata(vwnd)
mask = np.ma.getmask(uwnd) | np.ma.getmask(vwnd)
ndim = u.ndim
vor = rot(u, v, lon, lat, xdim, ydim, cyclic=cyclic, sphere=sphere)
f = arr.expand(constants.earth_f(lat), ndim, axis=ydim)
out = f + vor
out = np.ma.array(out, mask=mask)
out = arr.mrollaxis(out, ydim, 0)
out[0, ...] = np.ma.masked
out[-1, ...] = np.ma.masked
out = arr.mrollaxis(out, 0, ydim + 1)
return out
def ertelpv(uwnd, vwnd, temp, lon, lat, lev, xdim, ydim, zdim,
cyclic=True, punit=100., sphere=True):
"""
Calculate Ertel potential vorticity.
Hoskins, B.J., M.E. McIntyre and A.E. Robertson, 1985:
On the use and significance of isentropic potential
vorticity maps, `QJRMS`, 111, 877-946,
<http://onlinelibrary.wiley.com/doi/10.1002/qj.49711147002/abstract>
:param uwnd: ndarray, u component wind [m/s].
:param vwnd: ndarray, v component wind [m/s].
:param temp: ndarray, temperature [K].
:param lon: array_like, longitude [degrees].
:param lat: array_like, latitude [degrees].
:param lev: array_like, pressure level [punit*Pa].
:param xdim: west-east axis
:param ydim: south-north axis
:param zdim: vertical axis
:param cyclic: west-east cyclic boundary
:param punit: pressure level unit
:param sphere: sphere coordinates.
:return:
"""
u, v, t = np.ma.getdata(uwnd), np.ma.getdata(vwnd), np.ma.getdata(temp)
mask = np.ma.getmask(uwnd) | np.ma.getmask(vwnd) | np.ma.getmask(temp)
ndim = u.ndim
# potential temperature
theta = pottemp(t, lev, zdim, punit=punit)
# partial derivation
dthdp = dvardp(theta, lev, zdim, punit=punit)
dudp = dvardp(u, lev, zdim, punit=punit)
dvdp = dvardp(v, lev, zdim, punit=punit)
dthdx = dvardx(theta, lon, lat, xdim, ydim, cyclic=cyclic, sphere=sphere)
dthdy = dvardy(theta, lat, ydim, sphere=sphere)
# absolute vorticity
vor = rot(u, v, lon, lat, xdim, ydim, cyclic=cyclic, sphere=sphere)
f = arr.expand(constants.earth_f(lat), ndim, axis=ydim)
avor = f + vor
out = -g * (avor*dthdp - (dthdx*dvdp-dthdy*dudp))
out = np.ma.array(out, mask=mask)
out = arr.mrollaxis(out, ydim, 0)
out[0, ...] = np.ma.masked
out[-1, ...] = np.ma.masked
out = arr.mrollaxis(out, 0, ydim+1)
return out
def vmean(var, bottom, top, lev, zdim):
"""
Calculate vertical mean.
:param var: array_like.
:param bottom: bottom boundary of integration.
:param top: top boundary of integration.
:param lev: isobaric levels.
:param zdim: vertical dimension.
:return: array_like.
"""
var = np.ma.asarray(var)
lev = np.asarray(lev)
ndim = var.ndim
lev = lev[(lev <= bottom) & (lev >= top)]
lev_m = np.r_[bottom, (lev[1:] + lev[:-1]) / 2., top]
dp = lev_m[:-1] - lev_m[1:]
# roll lat dim axis to last
var = arr.mrollaxis(var, zdim, ndim)
out = var[..., (lev <= bottom) & (lev >= top)] * dp
out = out.sum(axis=-1) / (dp.sum())
return out
