forked from 3rdparty/wrf-python
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
222 lines
6.9 KiB
222 lines
6.9 KiB
from __future__ import (absolute_import, division, print_function, |
|
unicode_literals) |
|
|
|
import numpy as np |
|
import numpy.ma as ma |
|
|
|
from .constants import Constants |
|
from .extension import (_interpz3d, _interp2dxy, _interp1d, _slp, _tk, _td, |
|
_rh, _uvmet, _smooth2d, _cape, _cloudfrac, _ctt, _dbz, |
|
_srhel, _udhel, _avo, _pvo, _eth, _wetbulb, _tv, |
|
_omega, _pw) |
|
from .util import from_var |
|
from .metadecorators import (set_alg_metadata, set_uvmet_alg_metadata, |
|
set_interp_metadata, set_cape_alg_metadata, |
|
set_cloudfrac_alg_metadata) |
|
from .interputils import get_xy |
|
|
|
@set_interp_metadata("xy") |
|
def xy(field, pivot_point=None, angle=None, start_point=None, end_point=None, |
|
meta=True): |
|
return get_xy(field, pivot_point, angle, start_point, end_point) |
|
|
|
|
|
@set_interp_metadata("1d") |
|
def interp1d(field, z_in, z_out, missingval=Constants.DEFAULT_FILL, |
|
meta=True): |
|
return _interp1d(field, z_in, z_out, missingval) |
|
|
|
|
|
@set_interp_metadata("2dxy") |
|
def interp2dxy(field3d, xy, meta=True): |
|
return _interp2dxy(field3d, xy) |
|
|
|
|
|
@set_interp_metadata("horiz") |
|
def interpz3d(field3d, z, desiredloc, missingval=Constants.DEFAULT_FILL, |
|
meta=True): |
|
return _interpz3d(field3d, z, desiredloc, missingval) |
|
|
|
|
|
@set_alg_metadata(2, "pres", refvarndims=3, units="hPa", |
|
description="sea level pressure") |
|
def slp(height, tkel, pres, qv, meta=True): |
|
return _slp(height, tkel, pres, qv) |
|
|
|
|
|
@set_alg_metadata(3, "pres", units="K", |
|
description="temperature") |
|
def tk(pres, theta, meta=True): |
|
return _tk(pres, theta) |
|
|
|
|
|
@set_alg_metadata(3, "pres", units="degC", |
|
description="dew point temperature") |
|
def td(pres, qv, meta=True): |
|
return _td(pres, qv) |
|
|
|
|
|
@set_alg_metadata(3, "pres", |
|
description="relative humidity", units=None) |
|
def rh(qv, pres, tkel, meta=True): |
|
return _rh(qv, pres, tkel) |
|
|
|
|
|
@set_uvmet_alg_metadata(latarg="lat", windarg="u") |
|
def uvmet(u, v, lat, lon, cen_long, cone, meta=True): |
|
return _uvmet(u, v, lat, lon, cen_long, cone) |
|
|
|
|
|
@set_alg_metadata(2, "field", |
|
description=from_var("field", "description"), |
|
units=from_var("field", "units")) |
|
def smooth2d(field, passes, meta=True): |
|
return _smooth2d(field, passes) |
|
|
|
|
|
@set_cape_alg_metadata(is2d=True, copyarg="pres_hpa") |
|
def cape_2d(pres_hpa, tkel, qvapor, height, terrain, psfc_hpa, ter_follow, |
|
missing=Constants.DEFAULT_FILL, meta=True): |
|
|
|
if isinstance(ter_follow, bool): |
|
ter_follow = 1 if ter_follow else 0 |
|
|
|
i3dflag = 0 |
|
cape_cin = _cape(pres_hpa, tkel, qvapor, height, terrain, psfc_hpa, |
|
missing, i3dflag, ter_follow) |
|
|
|
left_dims = cape_cin.shape[1:-3] |
|
right_dims = cape_cin.shape[-2:] |
|
|
|
resdim = (4,) + left_dims + right_dims |
|
|
|
# Make a new output array for the result |
|
result = np.zeros(resdim, cape_cin.dtype) |
|
|
|
# Cape 2D output is not flipped in the vertical, so index from the |
|
# end |
|
result[0,...,:,:] = cape_cin[0,...,-1,:,:] |
|
result[1,...,:,:] = cape_cin[1,...,-1,:,:] |
|
result[2,...,:,:] = cape_cin[1,...,-2,:,:] |
|
result[3,...,:,:] = cape_cin[1,...,-3,:,:] |
|
|
|
return ma.masked_values(result, missing) |
|
|
|
|
|
@set_cape_alg_metadata(is2d=False, copyarg="pres_hpa") |
|
def cape_3d(pres_hpa, tkel, qvapor, height, terrain, psfc_hpa, ter_follow, |
|
missing=Constants.DEFAULT_FILL, meta=True): |
|
|
|
if isinstance(ter_follow, bool): |
|
ter_follow = 1 if ter_follow else 0 |
|
|
|
i3dflag = 1 |
|
cape_cin = _cape(pres_hpa, tkel, qvapor, height, terrain, psfc_hpa, |
|
missing, i3dflag, ter_follow) |
|
|
|
return ma.masked_values(cape_cin, missing) |
|
|
|
|
|
@set_cloudfrac_alg_metadata(copyarg="pres") |
|
def cloudfrac(pres, relh, meta=True): |
|
return _cloudfrac(pres, relh) |
|
|
|
|
|
@set_alg_metadata(2, "pres_hpa", refvarndims=3, units="degC", |
|
description="cloud top temperature") |
|
def ctt(pres_hpa, tkel, qv, qcld, height, terrain, qice=None, meta=True): |
|
|
|
# Qice and QCLD need to be in g/kg |
|
if qice is None: |
|
qice = np.zeros(qv.shape, qv.dtype) |
|
haveqci = 0 |
|
else: |
|
haveqci = 1 if qice.any() else 0 |
|
|
|
return _ctt(pres_hpa, tkel, qice, qcld, qv, height, terrain, haveqci) |
|
|
|
|
|
|
|
@set_alg_metadata(3, "pres", units="dBZ", |
|
description="radar reflectivity") |
|
def dbz(pres, tkel, qv, qr, qs=None, qg=None, use_varint=False, |
|
use_liqskin=False, meta=True): |
|
|
|
if qs is None: |
|
qs = np.zeros(qv.shape, qv.dtype) |
|
|
|
if qg is None: |
|
qg = np.zeros(qv.shape, qv.dtype) |
|
|
|
sn0 = 1 if qs.any() else 0 |
|
ivarint = 1 if use_varint else 0 |
|
iliqskin = 1 if use_liqskin else 0 |
|
|
|
return _dbz(pres, tkel, qv, qr, qs, qg, sn0, ivarint, iliqskin) |
|
|
|
|
|
@set_alg_metadata(2, "terrain", units="m2 s-2", |
|
description="storm relative helicity") |
|
def srhel(u, v, z, terrain, top=3000.0, meta=True): |
|
# u, v get swapped in vertical |
|
_u = np.ascontiguousarray(u[...,::-1,:,:]) |
|
_v = np.ascontiguousarray(v[...,::-1,:,:]) |
|
_z = np.ascontiguousarray(z[...,::-1,:,:]) |
|
|
|
return _srhel(_u, _v, _z, terrain, top) |
|
|
|
|
|
@set_alg_metadata(2, "u", refvarndims=3, units="m2 s-2", |
|
description="updraft helicity") |
|
def udhel(zstag, mapfct, u, v, wstag, dx, dy, bottom=2000.0, top=5000.0, |
|
meta=True): |
|
return _udhel(zstag, mapfct, u, v, wstag, dx, dy, bottom, top) |
|
|
|
|
|
# Requires both u an v for dimnames |
|
@set_alg_metadata(3, "ustag", units="10-5 s-1", |
|
stagdim=-1, stagsubvar="vstag", |
|
description="absolute vorticity") |
|
def avo(ustag, vstag, msfu, msfv, msfm, cor, dx, dy, meta=True): |
|
return _avo(ustag, vstag, msfu, msfv, msfm, cor, dx, dy) |
|
|
|
|
|
@set_alg_metadata(3, "tkel", units="PVU", |
|
description="potential vorticity") |
|
def pvo(ustag, vstag, tkel, pres, msfu, msfv, msfm, cor, dx, dy, meta=True): |
|
return _pvo(ustag, vstag, tkel, pres, msfu, msfv, msfm, cor, dx, dy) |
|
|
|
|
|
@set_alg_metadata(3, "qv", units="K", |
|
description="equivalent potential temperature") |
|
def eth(qv, tkel, pres, meta=True): |
|
return _eth(qv, tkel, pres) |
|
|
|
|
|
@set_alg_metadata(3, "pres", units="K", |
|
description="wetbulb temperature") |
|
def wetbulb(pres, tkel, qv, meta=True): |
|
return _wetbulb(pres, tkel, qv) |
|
|
|
|
|
@set_alg_metadata(3, "tkel", units="K", |
|
description="virtual temperature") |
|
def tvirtual(tkel, qv, meta=True): |
|
return _tv(tkel, qv) |
|
|
|
|
|
@set_alg_metadata(3, "qv", units="Pa s-1", |
|
description="omega") |
|
def omega(qv, tkel, w, pres, meta=True): |
|
return _omega(qv, tkel, w, pres) |
|
|
|
|
|
@set_alg_metadata(2, "pres", refvarndims=3, units="kg m-2", |
|
description="precipitable water") |
|
def pw(pres, tkel, qv, height, meta=True): |
|
tv = _tv(tkel, qv) |
|
|
|
return _pw(pres, tv, qv, height) |
|
|
|
|
|
|
|
|