def get_geoht(daskArray, height=True, msl=True, omp_threads=1):
varname = wrapped_either(daskArray, ("PH", "GHT"))
if varname == "PH":
ph = fetch_variable(daskArray, "PH")
phb = fetch_variable(daskArray, "PHB")
hgt = fetch_variable(daskArray, "HGT")
dtype = ph.dtype
geopt = map_blocks(wrapped_add, ph, phb, dtype=dtype)
geopt_f = wrapped_destagger(geopt, -3)
else:
geopt = fetch_variable(daskArray, "GHT")
hgt = fetch_variable(daskArray, "HGT_M")
dtype = geopt.dtype
geopt_f = map_blocks(wrapped_mul, geopt, Constants.G, dtype=dtype)
if height:
if msl:
mslh = map_blocks(wrapped_div, geopt_f, Constants.G, dtype=dtype)
return mslh
else:
# Due to broadcasting with multifile/multitime, the 2D terrain
# array needs to be reshaped to a 3D array so the right dims
# line up
new_dims = list(hgt.shape)
new_dims.insert(-2, 1)
hgt = hgt.reshape(new_dims)
mslh = map_blocks(wrapped_div, geopt_f, Constants.G, dtype=dtype)
mslh_f = map_blocks(wrapped_sub, mslh, hgt, dtype=dtype)
return mslh_f
else:
return geopt_f
def get_winds_at_level(daskArray,
vertical_field=None,
requested_top=0.,
omp_threads=1):
varname = wrapped_either(daskArray, ("U", "UU"))
uS = fetch_variable(daskArray, varname)
u = wrapped_destagger(uS, -1)
varname = wrapped_either(daskArray, ("V", "VV"))
vS = fetch_variable(daskArray, varname)
v = wrapped_destagger(vS, -2)
del (varname)
del (uS)
del (vS)
if (requested_top == 0.):
return u[0], v[0]
else:
uLev = wrapped_interplevel(u,
vertical_field,
requested_top,
omp_threads=omp_threads)
vLev = wrapped_interplevel(v,
vertical_field,
requested_top,
omp_threads=omp_threads)
return uLev, vLev
def get_omega(daskArray, omp_threads=1):
t = fetch_variable(daskArray, "T")
p = fetch_variable(daskArray, "P")
w = fetch_variable(daskArray, "W")
pb = fetch_variable(daskArray, "PB")
qv = fetch_variable(daskArray, "QVAPOR")
dtype = t.dtype
wa = wrapped_destagger(w, -3)
full_t = map_blocks(wrapped_add, t, Constants.T_BASE, dtype=dtype)
full_p = map_blocks(wrapped_add, p, pb, dtype=dtype)
tk = map_blocks(tk_wrap, full_p, full_t, omp_threads, dtype=dtype)
del (t)
del (p)
del (pb)
del (full_t)
omega = map_blocks(omega_wrap,
qv,
tk,
wa,
full_p,
omp_threads,
dtype=dtype)
return omega
def get_full_p(daskArray, omp_threads=1):
p = fetch_variable(daskArray, "P")
pb = fetch_variable(daskArray, "PB")
total_p = map_blocks(wrapped_add, p, pb, omp_threads, dtype=p.dtype)
full_p = map_blocks(wrapped_div, total_p, 100, omp_threads, dtype=p.dtype)
return full_p
def get_cape3d(daskArray, omp_threads=1, num_workers=1):
missing = default_fill(np.float64)
t = fetch_variable(daskArray, "T")
p = fetch_variable(daskArray, "P")
pb = fetch_variable(daskArray, "PB")
qv = fetch_variable(daskArray, "QVAPOR")
ph = fetch_variable(daskArray, "PH")
phb = fetch_variable(daskArray, "PHB")
ter = fetch_variable(daskArray, "HGT")
psfc = fetch_variable(daskArray, "PSFC")
dtype = p.dtype
full_t = map_blocks(wrapped_add, t, Constants.T_BASE, dtype=dtype)
full_p = map_blocks(wrapped_add, p, pb, dtype=dtype)
tk = map_blocks(tk_wrap, full_p, full_t, omp_threads, dtype=dtype)
del (full_t)
del (t)
del (p)
geopt = map_blocks(wrapped_add, ph, phb, dtype=dtype)
geopt_unstag = wrapped_destagger(geopt, -3)
z = map_blocks(wrapped_div, geopt_unstag, Constants.G, dtype=dtype)
del (ph)
del (phb)
del (geopt)
del (geopt_unstag)
p_hpa = map_blocks(wrapped_mul,
full_p,
ConversionFactors.PA_TO_HPA,
dtype=dtype)
psfc_hpa = map_blocks(wrapped_mul,
psfc,
ConversionFactors.PA_TO_HPA,
dtype=dtype)
del (full_p)
del (psfc)
i3dflag = 1
ter_follow = 1
cape_cin = map_blocks(cape_wrap,
p_hpa,
tk,
qv,
z,
ter,
psfc_hpa,
missing,
i3dflag,
ter_follow,
omp_threads,
dtype=dtype)
comp = cape_cin.compute(num_workers=num_workers)
return comp
def get_tk(daskArray, omp_threads):
t = fetch_variable(daskArray, "T")
p = fetch_variable(daskArray, "P")
pb = fetch_variable(daskArray, "PB")
dtype = t.dtype
full_t = map_blocks(wrapped_add, t, Constants.T_BASE, dtype=dtype)
full_p = map_blocks(wrapped_add, p, pb, dtype=dtype)
del (t)
del (p)
del (pb)
tk = map_blocks(tk_wrap, full_p, full_t, omp_threads, dtype=dtype)
return tk
def get_dewpoint(daskArray, omp_threads=1):
p = fetch_variable(daskArray, "P")
pb = fetch_variable(daskArray, "PB")
qvapor = fetch_variable(daskArray, "QVAPOR", include_meta=True)
dtype = p.dtype
full_p = map_blocks(wrapped_add, p, pb, dtype=dtype)
full_p_div = map_blocks(wrapped_mul, full_p, 0.01, dtype=dtype)
del (p)
del (pb)
del (full_p)
qvapor = qvapor.where(qvapor >= 0, 0)
td = map_blocks(td_wrap, full_p_div, qvapor.data, omp_threads, dtype=dtype)
return td
def get_tv(daskArray, omp_threads=1, num_workers=1):
t = fetch_variable(daskArray, "T")
p = fetch_variable(daskArray, "P")
pb = fetch_variable(daskArray, "PB")
qv = fetch_variable(daskArray, "QVAPOR")
dtype = t.dtype
full_t = map_blocks(wrapped_add, t, Constants.T_BASE, dtype=dtype)
full_p = map_blocks(wrapped_add, p, pb, dtype=dtype)
del (t)
del (p)
del (pb)
tk = map_blocks(tk_wrap, full_p, full_t, omp_threads, dtype=dtype)
tv = map_blocks(tv_wrap, tk, qv, omp_threads, dtype=dtype)
return tv.compute(num_workers=num_workers)
def get_tw(daskArray, omp_threads=1):
t = fetch_variable(daskArray, "T")
p = fetch_variable(daskArray, "P")
pb = fetch_variable(daskArray, "PB")
qv = fetch_variable(daskArray, "QVAPOR")
dtype = t.dtype
full_t = map_blocks(wrapped_add, t, Constants.T_BASE, dtype=dtype)
full_p = map_blocks(wrapped_add, p, pb, dtype=dtype)
del (t)
del (p)
del (pb)
tk = map_blocks(tk_wrap, full_p, full_t, omp_threads, dtype=dtype)
tw = map_blocks(wetbulb_wrap, tk, qv, omp_threads, dtype=dtype)
return tw
def get_theta(daskArray, omp_threads=1):
t = fetch_variable(daskArray, "T")
full_t = map_blocks(wrapped_add,
t,
Constants.T_BASE,
omp_threads,
dtype=t.dtype)
return full_t
def get_rh(daskArray, omp_threads=1):
t = fetch_variable(daskArray, "T")
p = fetch_variable(daskArray, "P")
pb = fetch_variable(daskArray, "PB")
qvapor = fetch_variable(daskArray, "QVAPOR", include_meta=True)
dtype = t.dtype
full_t = map_blocks(wrapped_add, t, Constants.T_BASE, dtype=dtype)
full_p = map_blocks(wrapped_add, p, pb, dtype=dtype)
del (t)
del (p)
del (pb)
qvapor = qvapor.where(qvapor >= 0, 0)
tk = map_blocks(tk_wrap, full_p, full_t, omp_threads, dtype=dtype)
del (full_t)
rh = map_blocks(rh_wrap, qvapor.data, full_p, tk, omp_threads, dtype=dtype)
return rh
def get_pw(daskArray, omp_threads=1):
t = fetch_variable(daskArray, "T")
p = fetch_variable(daskArray, "P")
pb = fetch_variable(daskArray, "PB")
ph = fetch_variable(daskArray, "PH")
phb = fetch_variable(daskArray, "PHB")
qv = fetch_variable(daskArray, "QVAPOR")
dtype = t.dtype
full_p = map_blocks(wrapped_add, p, pb, dtype=dtype)
full_ph = map_blocks(wrapped_add, ph, phb, dtype=dtype)
ht = map_blocks(wrapped_div, full_ph, Constants.G, dtype=dtype)
full_t = map_blocks(wrapped_add, t, Constants.T_BASE, dtype=dtype)
del (p)
del (pb)
del (ph)
tk = map_blocks(tk_wrap, full_p, full_t, omp_threads, dtype=dtype)
tv = map_blocks(tv_wrap, tk, qv, omp_threads, dtype=dtype)
del (full_t)
del (tk)
pw = map_blocks(pw_wrap, full_p, tv, qv, ht, omp_threads, dtype=dtype)
return pw
def get_rvor(daskArray, omp_threads=1):
u = fetch_variable(daskArray, "U")
v = fetch_variable(daskArray, "V")
msfu = fetch_variable(daskArray, "MAPFAC_U")
msfv = fetch_variable(daskArray, "MAPFAC_V")
msfm = fetch_variable(daskArray, "MAPFAC_M")
cor = fetch_variable(daskArray, "F")
dx = daskArray.DX
dy = daskArray.DY
dtype = u.dtype
avo = map_blocks(avo_wrap,
u,
v,
msfu,
msfv,
msfm,
cor,
dx,
dy,
omp_threads,
dtype=dtype)
rvor = map_blocks(wrapped_sub, avo, cor, dtype=dtype)
return rvor
def get_avo(daskArray, omp_threads=1, num_workers=1):
u = fetch_variable(daskArray, "U")
v = fetch_variable(daskArray, "V")
msfu = fetch_variable(daskArray, "MAPFAC_U")
msfv = fetch_variable(daskArray, "MAPFAC_V")
msfm = fetch_variable(daskArray, "MAPFAC_M")
cor = fetch_variable(daskArray, "F")
dx = daskArray.DX
dy = daskArray.DY
dtype = u.dtype
avo = map_blocks(avo_wrap,
u,
v,
msfu,
msfv,
msfm,
cor,
dx,
dy,
omp_threads,
dtype=dtype)
return avo.compute(num_workers=num_workers)
def get_dbz(daskArray,
use_varint=False,
use_liqskin=False,
omp_threads=1,
num_workers=1):
t = fetch_variable(daskArray, "T")
p = fetch_variable(daskArray, "P")
pb = fetch_variable(daskArray, "PB")
qv = fetch_variable(daskArray, "QVAPOR")
qr = fetch_variable(daskArray, "QRAIN")
dtype = t.dtype
try:
qs = fetch_variable(daskArray, "QSNOW")
except KeyError:
qs = da.zeros(qv.shape, qv.dtype)
try:
qgraup = fetch_variable(daskArray, "QGRAUP")
except KeyError:
qgraup = da.zeros(qv.shape, qv.dtype)
full_t = map_blocks(wrapped_add, t, Constants.T_BASE, dtype=t.dtype)
full_p = map_blocks(wrapped_add, p, pb, dtype=p.dtype)
tk = map_blocks(tk_wrap, full_p, full_t, omp_threads, dtype=p.dtype)
sn0 = 1 if qs.any() else 0
ivarint = 1 if use_varint else 0
iliqskin = 1 if use_liqskin else 0
del (t)
del (p)
del (pb)
dbz = map_blocks(dbz_wrap,
full_p,
tk,
qv,
qr,
qs,
qgraup,
sn0,
ivarint,
iliqskin,
omp_threads,
dtype=dtype)
return dbz.compute(num_workers=num_workers)
def get_pvo(daskArray, omp_threads=1):
u = fetch_variable(daskArray, "U")
v = fetch_variable(daskArray, "V")
t = fetch_variable(daskArray, "T")
p = fetch_variable(daskArray, "P")
pb = fetch_variable(daskArray, "PB")
msfu = fetch_variable(daskArray, "MAPFAC_U")
msfv = fetch_variable(daskArray, "MAPFAC_V")
msfm = fetch_variable(daskArray, "MAPFAC_M")
cor = fetch_variable(daskArray, "F")
dx = daskArray.DX
dy = daskArray.DY
dtype = u.dtype
full_t = map_blocks(wrapped_add, t, 300, dtype=dtype)
full_p = map_blocks(wrapped_add, p, pb, dtype=dtype)
del (t)
del (p)
del (pb)
pvo = map_blocks(pvo_wrap,
u,
v,
full_t,
full_p,
msfu,
msfv,
msfm,
cor,
dx,
dy,
omp_threads,
dtype=dtype)
return pvo
def get_udhel(daskArray,
bottom=2000.0,
top=5000.0,
omp_threads=1,
num_workers=1):
wstag = fetch_variable(daskArray, "W")
ph = fetch_variable(daskArray, "PH")
phb = fetch_variable(daskArray, "PHB")
dtype = ph.dtype
mapfct = fetch_variable(daskArray, "MAPFAC_M")
dx = daskArray.DX
dy = daskArray.DY
varname = wrapped_either(daskArray, ("U", "UU"))
uS = fetch_variable(daskArray, varname)
u = wrapped_destagger(uS, -1)
varname = wrapped_either(daskArray, ("V", "VV"))
vS = fetch_variable(daskArray, varname)
v = wrapped_destagger(vS, -2)
del (uS)
del (vS)
geopt = map_blocks(wrapped_add, ph, phb, dtype=dtype)
zp = map_blocks(wrapped_div, geopt, Constants.G, dtype=dtype)
del (ph)
del (phb)
del (geopt)
udhel = map_blocks(udhel_wrap,
zp,
mapfct,
u,
v,
wstag,
dx,
dy,
bottom,
top,
omp_threads,
dtype=dtype)
return udhel.compute(num_workers=num_workers)
def get_srh(daskArray, top=3000.0, omp_threads=1):
lat_VN = wrapped_lat_varname(daskArray, stagger=None)
lats = fetch_variable(daskArray, lat_VN)
hgt = fetch_variable(daskArray, "HGT")
ph = fetch_variable(daskArray, "PH")
phb = fetch_variable(daskArray, "PHB")
dtype = ph.dtype
varname = wrapped_either(daskArray, ("U", "UU"))
uS = fetch_variable(daskArray, varname)
u = wrapped_destagger(uS, -1)
varname = wrapped_either(daskArray, ("V", "VV"))
vS = fetch_variable(daskArray, varname)
v = wrapped_destagger(vS, -2)
geopt = map_blocks(wrapped_add, ph, phb, dtype=dtype)
geopt_f = wrapped_destagger(geopt, -3)
z = map_blocks(wrapped_div, geopt_f, Constants.G, dtype=dtype)
del (ph)
del (phb)
del (geopt)
del (geopt_f)
u1 = np.ascontiguousarray(u[..., ::-1, :, :])
v1 = np.ascontiguousarray(v[..., ::-1, :, :])
z1 = np.ascontiguousarray(z[..., ::-1, :, :])
del (u)
del (v)
del (z)
srh = map_blocks(srh_wrap,
u1,
v1,
z1,
hgt,
lats,
top,
omp_threads,
dtype=dtype)
return srh
def get_slp(daskArray, omp_threads=1):
t = fetch_variable(daskArray, "T")
p = fetch_variable(daskArray, "P")
pb = fetch_variable(daskArray, "PB")
qvapor = fetch_variable(daskArray, "QVAPOR", include_meta=True)
ph = fetch_variable(daskArray, "PH")
phb = fetch_variable(daskArray, "PHB")
dtype = p.dtype
full_t = map_blocks(wrapped_add, t, Constants.T_BASE, dtype=dtype)
full_p = map_blocks(wrapped_add, p, pb, dtype=dtype)
qvapor = qvapor.where(qvapor >= 0, 0)
del (t)
del (p)
del (pb)
pre_full_ph = map_blocks(wrapped_add, ph, phb, dtype=dtype)
full_ph = map_blocks(wrapped_div, pre_full_ph, Constants.G, dtype=dtype)
destag_ph = wrapped_destagger(full_ph, -3)
del (full_ph)
del (ph)
del (phb)
tk = map_blocks(tk_wrap, full_p, full_t, omp_threads, dtype=dtype)
slp = map_blocks(slp_wrap,
destag_ph,
tk,
full_p,
qvapor.data,
omp_threads,
dtype=dtype)
slp_calc = slp
return slp_calc
def get_accum_precip(daskArray, omp_threads=1):
rainc = fetch_variable(daskArray, "RAINC")
rainnc = fetch_variable(daskArray, "RAINNC")
rainsum = map_blocks(wrapped_add, rainc, rainnc, dtype=rainc.dtype)
return rainsum
def get_cape2d(daskArray, omp_threads=1, num_workers=1):
missing = default_fill(np.float64)
t = fetch_variable(daskArray, "T")
p = fetch_variable(daskArray, "P")
pb = fetch_variable(daskArray, "PB")
qv = fetch_variable(daskArray, "QVAPOR")
ph = fetch_variable(daskArray, "PH")
phb = fetch_variable(daskArray, "PHB")
ter = fetch_variable(daskArray, "HGT")
psfc = fetch_variable(daskArray, "PSFC")
dtype = p.dtype
full_t = map_blocks(wrapped_add, t, Constants.T_BASE, dtype=dtype)
full_p = map_blocks(wrapped_add, p, pb, dtype=dtype)
tk = map_blocks(tk_wrap, full_p, full_t, omp_threads, dtype=dtype)
del (full_t)
del (t)
del (p)
geopt = map_blocks(wrapped_add, ph, phb, dtype=dtype)
geopt_unstag = wrapped_destagger(geopt, -3)
z = map_blocks(wrapped_div, geopt_unstag, Constants.G, dtype=dtype)
del (ph)
del (phb)
del (geopt)
del (geopt_unstag)
p_hpa = map_blocks(wrapped_mul,
full_p,
ConversionFactors.PA_TO_HPA,
dtype=dtype)
psfc_hpa = map_blocks(wrapped_mul,
psfc,
ConversionFactors.PA_TO_HPA,
dtype=dtype)
del (full_p)
del (psfc)
i3dflag = 0
ter_follow = 1
cape_cin = map_blocks(cape_wrap,
p_hpa,
tk,
qv,
z,
ter,
psfc_hpa,
missing,
i3dflag,
ter_follow,
omp_threads,
dtype=dtype)
calc_cape = cape_cin.compute(num_workers=num_workers)
left_dims = calc_cape.shape[1:-3]
right_dims = calc_cape.shape[-2:]
resdim = (4, ) + left_dims + right_dims
# Make a new output array for the result
result = np.zeros(resdim, calc_cape.dtype)
# Cape 2D output is not flipped in the vertical, so index from the
# end
result[0, ..., :, :] = calc_cape[0, ..., -1, :, :]
result[1, ..., :, :] = calc_cape[1, ..., -1, :, :]
result[2, ..., :, :] = calc_cape[1, ..., -2, :, :]
result[3, ..., :, :] = calc_cape[1, ..., -3, :, :]
return npma.masked_values(result, missing)
def get_full_p(daskArray, omp_threads=1, num_workers=1):
p = fetch_variable(daskArray, "P")
pb = fetch_variable(daskArray, "PB")
full_p = map_blocks(wrapped_add_then_div, p, pb, 100, dtype=p.dtype)
return full_p.compute(num_workers=num_workers)
