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_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_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_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_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)
