def tv_wrap(temp_k, qvapor, omp_threads=1):
from wrf.extension import _tv, omp_set_num_threads
omp_set_num_threads(omp_threads)
result = _tv(temp_k, qvapor)
return result
def eth_wrap(qv, tk, full_p, omp_threads=1):
from wrf.extension import _eth, omp_set_num_threads
omp_set_num_threads(omp_threads)
result = _eth(qv, tk, full_p)
return result
def td_wrap(full_p, qvapor, omp_threads=1):
from wrf.extension import _td, omp_set_num_threads
omp_set_num_threads(omp_threads)
result = _td(full_p, qvapor)
return result
def slp_wrap(destag_ph, tk, full_p, qvapor, omp_threads=1):
from wrf.extension import _slp, omp_set_num_threads
omp_set_num_threads(omp_threads)
result = _slp(destag_ph, tk, full_p, qvapor)
return result
def tk_wrap(full_p, full_t, omp_threads=1):
from wrf.extension import _tk, omp_set_num_threads
omp_set_num_threads(omp_threads)
result = _tk(full_p, full_t)
return result
def rh_wrap(qvapor, full_p, tk, omp_threads=1):
from wrf.extension import _rh, omp_set_num_threads
omp_set_num_threads(omp_threads)
result = _rh(qvapor, full_p, tk)
return result
def avo_wrap(u, v, msfu, msfv, msfm, cor, dx, dy, omp_threads=1):
from wrf.extension import _avo, omp_set_num_threads
omp_set_num_threads(omp_threads)
result = _avo(u, v, msfu, msfv, msfm, cor, dx, dy)
return result
def omega_wrap(qv, tk, wa, full_p, omp_threads=1):
from wrf.extension import _omega, omp_set_num_threads
omp_set_num_threads(omp_threads)
result = _omega(qv, tk, wa, full_p)
return result
def pw_wrap(full_p, tv, qv, ht, omp_threads=1):
from wrf.extension import _pw, omp_set_num_threads
omp_set_num_threads(omp_threads)
result = _pw(full_p, tv, qv, ht)
return result
def udhel_wrap(zp, mapfct, u, v, wstag, dx, dy, bottom, top, omp_threads=1):
from wrf.extension import _udhel, omp_set_num_threads
omp_set_num_threads(omp_threads)
result = _udhel(zp, mapfct, u, v, wstag, dx, dy, bottom, top)
return result
def srh_wrap(u1, v1, z1, ter, lats, top, omp_threads=1):
from wrf.extension import _srhel, omp_set_num_threads
omp_set_num_threads(omp_threads)
result = _srhel(u1, v1, z1, ter, lats, top)
return result
def wetbulb_wrap(full_p, tk, qv, omp_threads=1):
# NOTE: _wetbulb is potentially not thread-safe, this may require a re-write at some point, but just to be aware.
from wrf.extension import _wetbulb, omp_set_num_threads
omp_set_num_threads(omp_threads)
result = _wetbulb(full_p, tk, qv)
return result
def wrapped_interpz3d_lev2d(field3d,
z,
lev2d,
missingval,
outview=None,
omp_threads=1):
from wrf.extension import _interpz3d_lev2d, omp_set_num_threads
omp_set_num_threads(omp_threads)
result = _interpz3d_lev2d(field3d, z, lev2d, missingval, outview)
return result
def dbz_wrap(full_p,
tk,
qv,
qr,
qs,
qg,
sn0,
ivarint,
iliqskin,
omp_threads=1):
from wrf.extension import _dbz, omp_set_num_threads
omp_set_num_threads(omp_threads)
result = _dbz(full_p, tk, qv, qr, qs, qg, sn0, ivarint, iliqskin)
return result
def cape_wrap(p_hpa,
tk,
qv,
z,
ter,
psfc_hpa,
missing,
i3dflag,
ter_follow,
omp_threads=1):
# NOTE: _cape is potentially not thread-safe, this may require a re-write at some point, but just to be aware.
from wrf.extension import _cape, omp_set_num_threads
omp_set_num_threads(omp_threads)
result = _cape(p_hpa, tk, qv, z, ter, psfc_hpa, missing, i3dflag,
ter_follow)
return result
def wrapped_interplevel(field3d,
vert,
desiredlev,
missing=default_fill(np.float64),
omp_threads=1,
num_workers=1):
from wrf.extension import omp_set_num_threads
import dask.array.ma as ma
from dask.array import map_blocks
omp_set_num_threads(omp_threads)
dtype = field3d.dtype
_desiredlev = da.asarray(desiredlev)
if _desiredlev.ndim == 0:
_desiredlev = np.array([desiredlev], np.float64)
levsare2d = False
else:
levsare2d = _desiredlev.ndim >= 2
desiredlev_da = da.from_array(_desiredlev)
if not levsare2d:
result = map_blocks(wrapped_interpz3d,
field3d,
vert,
desiredlev_da,
missing,
omp_threads=omp_threads,
dtype=dtype)
else:
result = map_blocks(wrapped_interpz3d_lev2d,
field3d,
vert,
desiredlev_da,
missing,
omp_threads=omp_threads,
dtype=dtype)
#TODO: Fill this
masked = ma.masked_values(result, missing)
return masked.compute(num_workers=num_workers)
