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