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