def wqs2_stencil_w(ta: sd, den: sd, wqsat: sd, dqdt: sd):
with computation(PARALLEL), interval(...):
wqsat, dqdt = wqs2_fn_w(ta, den)
def func(in_field: gtscript.Field[np.float_],
out_field: gtscript.Field[np.float_]):
from gt4py.__externals__ import offset
with computation(PARALLEL), interval(...):
out_field[0, 0, offset] = in_field
def func(in_field: gtscript.Field[np.float_],
out_field: gtscript.Field[np.float_]):
with computation(PARALLEL), interval(...):
tmp[0, 0, 0] = 2 * in_field
out_field = tmp
def definition(inout_field: gtscript.Field[float]):
with computation(PARALLEL), interval(...):
inout_field = assert_in_func(inout_field)
def definition(
field_in: gtscript.Field[gtscript.IJK, np.float_],
field_out: gtscript.Field[gtscript.K, np.float_],
):
with computation(PARALLEL), interval(...):
field_out = field_in[0, 0, 0]
def definition_func(phi: gtscript.Field[np.float64], dx: float):
from __externals__ import lap
with computation(PARALLEL), interval(...):
phi = lap(lap(phi, dx), dx)
def definition(inout_field: gtscript.Field[float]):
from __externals__ import EXTERNAL
with computation(PARALLEL), interval(...):
compile_assert(EXTERNAL < 1)
inout_field = inout_field[0, 0, 0] + EXTERNAL
def func(in_field: gtscript.Field[np.float_]):
with computation(PARALLEL), interval(...):
in_field += min(abs(sin(add_external_const(in_field))), -0.5)
def func(in_field: gtscript.Field[np.float_]):
with computation(PARALLEL), interval(...):
in_field += sinus(in_field)
def func(in_field: gtscript.Field[np.float_]):
with computation(PARALLEL), interval(...):
in_field += 2.0
in_field -= 0.5
in_field /= 0.5
in_field *= 4.0
def definition(in_field: gtscript.Field[np.float_],
out_field: gtscript.Field[np.float_]):
with computation(PARALLEL):
with interval(...):
in_field = out_field
def satadjust_part3_laststep_qa(
qa: sd,
area: sd,
qpz: sd,
hs: sd,
tin: sd,
q_cond: sd,
q_sol: sd,
den: sd,
):
with computation(PARALLEL), interval(...):
it, ap1 = ap1_and_index(tin)
wqs1 = wqs1_fn_w(it, ap1, tin, den)
iqs1 = wqs1_fn_2(it, ap1, tin, den)
# Determine saturated specific humidity
if tin < T_WFR:
# ice phase
qstar = iqs1
elif tin >= TICE:
qstar = wqs1
else:
# qsw = wqs1
if q_cond > 1e-6:
rqi = q_sol / q_cond
else:
rqi = (TICE - tin) / (TICE - T_WFR)
qstar = rqi * iqs1 + (1.0 - rqi) * wqs1
# higher than 10 m is considered "land" and will have higher subgrid
# variability
mindw = min(1.0, abs(hs) / (10.0 * constants.GRAV))
dw = (
spec.namelist.dw_ocean
+ (spec.namelist.dw_land - spec.namelist.dw_ocean) * mindw
)
# "scale - aware" subgrid variability: 100 - km as the base
dbl_sqrt_area = dw * (area ** 0.5 / 100.0e3) ** 0.5
maxtmp = 0.01 if 0.01 > dbl_sqrt_area else dbl_sqrt_area
hvar = min(0.2, maxtmp)
# partial cloudiness by pdf:
# assuming subgrid linear distribution in horizontal; this is
# effectively a smoother for the binary cloud scheme;
# qa = 0.5 if qstar == qpz
rh = qpz / qstar
# icloud_f = 0: bug - fixed
# icloud_f = 1: old fvgfs gfdl) mp implementation
# icloud_f = 2: binary cloud scheme (0 / 1)
if rh > 0.75 and qpz > 1.0e-8:
dq = hvar * qpz
q_plus = qpz + dq
q_minus = qpz - dq
if spec.namelist.icloud_f == 2: # TODO untested
if qpz > qstar:
qa = 1.0
elif (qstar < q_plus) and (q_cond > 1.0e-8):
qa = min(1.0, ((q_plus - qstar) / dq) ** 2)
else:
qa = 0.0
else:
if qstar < q_minus:
qa = 1.0
else:
if qstar < q_plus:
if spec.namelist.icloud_f == 0:
qa = (q_plus - qstar) / (dq + dq)
else:
qa = (q_plus - qstar) / (2.0 * dq * (1.0 - q_cond))
else:
qa = 0.0
# impose minimum cloudiness if substantial q_cond exist
if q_cond > 1.0e-8:
qa = max(spec.namelist.cld_min, qa)
qa = min(1, qa)
else:
qa = 0.0
def satadjust_part2(
wqsat: sd,
dq2dt: sd,
pt1: sd,
pt: sd,
cvm: sd,
mc_air: sd,
tcp3: sd,
lhl: sd,
lhi: sd,
lcp2: sd,
icp2: sd,
qv: sd,
ql: sd,
q_liq: sd,
qi: sd,
q_sol: sd,
den: sd,
qr: sd,
qg: sd,
qs: sd,
cappa: sd,
dp: sd,
tin: sd,
te0: sd,
q_cond: sd,
q_con: sd,
sdt: float,
adj_fac: float,
zvir: float,
fac_i2s: float,
c_air: float,
consv_te: bool,
hydrostatic: bool,
do_qa: bool,
fac_v2l: float,
fac_l2v: float,
lv00: float,
d0_vap: float,
c_vap: float,
mdt: float,
fac_r2g: float,
fac_smlt: float,
fac_l2r: float,
last_step: bool,
rad_snow: bool,
rad_rain: bool,
rad_graupel: bool,
tintqs: bool,
):
with computation(PARALLEL), interval(...):
dq0 = 0.0
if last_step:
dq0 = compute_dq0(qv, wqsat, dq2dt, tcp3)
if dq0 > 0:
src = dq0
else:
# TODO: We'd like to use this abstraction rather than duplicate
# code, but inside the if conditional complains 'not
# implemented'.
# factor, src = ql_evaporation(wqsat, qv, ql, dq0,fac_l2v)
factor = -1.0 * min(1, fac_l2v * 10.0 * (1.0 - qv / wqsat))
src = -1.0 * min(ql, factor * dq0)
# TODO Causes a visit_if error 'NoneType' object has no attribute 'inputs'
# qv, ql, q_liq, cvm, pt1 = wqsat_correct(src, pt1, lhl, qv, ql,
# q_liq, q_sol, mc_air, c_vap)
# lhl, lhi, lcp2, icp2 = update_latent_heat_coefficient(pt1, cvm,
# lv00, d0_vap)
qv = qv - src
ql = ql + src
q_liq = q_liq + src
cvm = compute_cvm(mc_air, qv, c_vap, q_liq, q_sol)
pt1 = add_src_pt1(pt1, src, lhl, cvm) # pt1 + src * lhl / cvm
# TODO: Revisit when gt4py updated, causes an Assertion error
# lhl, lhi, lcp2, icp2 = update_latent_heat_coefficient(pt1, cvm,
# lv00, d0_vap)
lhl = lv00 + d0_vap * pt1
lhi = LI00 + DC_ICE * pt1
lcp2 = lhl / cvm
icp2 = lhi / cvm
# homogeneous freezing of cloud water to cloud ice
ql, qi, q_liq, q_sol, cvm, pt1 = homogenous_freezing(
qv, ql, qi, q_liq, q_sol, pt1, cvm, icp2, mc_air, lhi, c_vap
)
# update some of the latent heat coefficients
lhi, icp2 = update_latent_heat_coefficient_i(pt1, cvm)
exptc = exp(0.66 * (TICE0 - pt1))
# bigg mechanism (heterogeneous freezing of cloud water to cloud ice)
ql, qi, q_liq, q_sol, cvm, pt1 = heterogeneous_freezing(
exptc,
pt1,
cvm,
ql,
qi,
q_liq,
q_sol,
den,
icp2,
mdt,
mc_air,
lhi,
qv,
c_vap,
)
lhi, icp2 = update_latent_heat_coefficient_i(pt1, cvm)
# freezing of rain to graupel
qr, qg, q_liq, q_sol, cvm, pt1 = make_graupel(
pt1, cvm, fac_r2g, qr, qg, q_liq, q_sol, lhi, icp2, mc_air, qv, c_vap
)
lhi, icp2 = update_latent_heat_coefficient_i(pt1, cvm)
# melting of snow to rain or cloud water
qs, ql, qr, q_liq, q_sol, cvm, pt1 = melt_snow(
pt1,
cvm,
fac_smlt,
qs,
ql,
qr,
q_liq,
q_sol,
lhi,
icp2,
mc_air,
qv,
c_vap,
spec.namelist.qs_mlt,
)
# autoconversion from cloud water to rain
ql, qr = autoconversion_cloud_to_rain(ql, qr, fac_l2r, spec.namelist.ql0_max)
iqs2, dqsdt = wqs2_fn_2(pt1, den)
expsubl = exp(0.875 * log(qi * den))
lhl, lhi, lcp2, icp2 = update_latent_heat_coefficient(pt1, cvm, lv00, d0_vap)
tcp2 = lcp2 + icp2
qv, qi, q_sol, cvm, pt1 = sublimation(
pt1,
cvm,
expsubl,
qv,
qi,
q_liq,
q_sol,
iqs2,
tcp2,
den,
dqsdt,
sdt,
adj_fac,
mc_air,
c_vap,
lhl,
lhi,
spec.namelist.t_sub,
spec.namelist.qi_gen,
spec.namelist.qi_lim,
)
# virtual temp updated
q_con = q_liq + q_sol
tmp = 1.0 + zvir * qv
pt = pt1 * tmp * (1.0 - q_con)
tmp = constants.RDGAS * tmp
cappa = tmp / (tmp + cvm)
# fix negative graupel with available cloud ice
maxtmp = 0.0
if qg < 0:
maxtmp = 0.0 if 0.0 > qi else qi
tmp = -qg if -qg < maxtmp else maxtmp
qg = qg + tmp
qi = qi - tmp
else:
qg = qg
# autoconversion from cloud ice to snow
qim = spec.namelist.qi0_max / den
sink = 0.0
if qi > qim:
sink = fac_i2s * (qi - qim)
qi = qi - sink
qs = qs + sink
# fix energy conservation
if consv_te:
if hydrostatic:
te0 = dp * (te0 + c_air * pt1)
else:
te0 = dp * (te0 + cvm * pt1)
# update latent heat coefficient
cvm = mc_air + (qv + q_liq + q_sol) * c_vap
lhl, lhi, lcp2, icp2 = update_latent_heat_coefficient(pt1, cvm, lv00, d0_vap)
# compute cloud fraction
tin = 0.0
if do_qa and last_step:
# combine water species
if rad_snow:
if rad_graupel:
q_sol = qi + qs + qg
else:
q_sol = qi + qs
else:
q_sol = qi
if rad_rain:
q_liq = ql + qr
else:
q_liq = ql
q_cond = q_sol + q_liq
# use the "liquid - frozen water temperature" (tin) to compute
# saturated specific humidity
if tintqs:
tin = pt1
else:
tin = pt1 - (lcp2 * q_cond + icp2 * q_sol)
def satadjust_part1(
wqsat: sd,
dq2dt: sd,
dpln: sd,
den: sd,
pt1: sd,
cvm: sd,
mc_air: sd,
peln: sd,
qv: sd,
ql: sd,
q_liq: sd,
qi: sd,
qr: sd,
qs: sd,
q_sol: sd,
qg: sd,
pt: sd,
dp: sd,
delz: sd,
te0: sd,
qpz: sd,
lhl: sd,
lhi: sd,
lcp2: sd,
icp2: sd,
tcp3: sd,
zvir: float,
hydrostatic: bool,
consv_te: bool,
c_air: float,
c_vap: float,
fac_imlt: float,
d0_vap: float,
lv00: float,
fac_v2l: float,
fac_l2v: float,
):
with computation(FORWARD), interval(1, None):
if hydrostatic:
delz = delz[0, 0, -1]
with computation(PARALLEL), interval(...):
dpln = peln[0, 0, 1] - peln
q_liq = ql + qr
q_sol = qi + qs + qg
qpz = q_liq + q_sol
pt1 = pt / ((1.0 + zvir * qv) * (1.0 - qpz))
t0 = pt1 # true temperature
qpz = qpz + qv # total_wat conserved in this routine
# define air density based on hydrostatical property
den = (
dp / (dpln * constants.RDGAS * pt)
if hydrostatic
else -dp / (constants.GRAV * delz)
)
# define heat capacity and latend heat coefficient
mc_air = (1.0 - qpz) * c_air
cvm = compute_cvm(mc_air, qv, c_vap, q_liq, q_sol)
lhi, icp2 = update_latent_heat_coefficient_i(pt1, cvm)
# fix energy conservation
if consv_te:
if hydrostatic:
te0 = -c_air * t0
else:
te0 = -cvm * t0
# fix negative cloud ice with snow
if qi < 0.0:
qs = qs + qi
qi = 0.0
# melting of cloud ice to cloud water and rain
qi, ql, q_liq, q_sol, cvm, pt1 = melt_cloud_ice(
qv, qi, ql, q_liq, q_sol, pt1, icp2, fac_imlt, mc_air, c_vap, lhi, cvm
)
# update latend heat coefficient
lhi, icp2 = update_latent_heat_coefficient_i(pt1, cvm)
# fix negative snow with graupel or graupel with available snow
qs, qg = fix_negative_snow(qs, qg)
# after this point cloud ice & snow are positive definite
# fix negative cloud water with rain or rain with available cloud water
ql, qr = fix_negative_cloud_water(ql, qr)
# enforce complete freezing of cloud water to cloud ice below - 48 c
ql, qi, q_liq, q_sol, cvm, pt1 = complete_freezing(
qv, ql, qi, q_liq, q_sol, pt1, cvm, icp2, mc_air, lhi, c_vap
)
wqsat, dq2dt = wqs2_fn_w(pt1, den)
# update latent heat coefficient
lhl, lhi, lcp2, icp2 = update_latent_heat_coefficient(pt1, cvm, lv00, d0_vap)
diff_ice = dim(TICE, pt1) / 48.0
dimmin = min(1.0, diff_ice)
tcp3 = lcp2 + icp2 * dimmin
dq0 = (qv - wqsat) / (
1.0 + tcp3 * dq2dt
) # compute_dq0(qv, wqsat, dq2dt, tcp3) #(qv - wqsat) / (1.0 + tcp3 * dq2dt)
# TODO Might be able to get rid of these temporary allocations when not used?
if dq0 > 0: # whole grid - box saturated
src = min(
spec.namelist.sat_adj0 * dq0,
max(spec.namelist.ql_gen - ql, fac_v2l * dq0),
)
else:
# TODO We'd like to use this abstraction rather than duplicate
# code, but inside the if conditional complains 'not implemented'
# factor, src = ql_evaporation(wqsat, qv, ql, dq0,fac_l2v)
factor = -1.0 * min(1, fac_l2v * 10.0 * (1.0 - qv / wqsat))
src = -1.0 * min(ql, factor * dq0)
qv, ql, q_liq, cvm, pt1 = wqsat_correct(
src, pt1, lhl, qv, ql, q_liq, q_sol, mc_air, c_vap
)
# update latent heat coefficient
lhl, lhi, lcp2, icp2 = update_latent_heat_coefficient(pt1, cvm, lv00, d0_vap)
# TODO: Remove duplicate
diff_ice = dim(TICE, pt1) / 48.0
dimmin = min(1.0, diff_ice)
tcp3 = lcp2 + icp2 * dimmin
def definition_func(field: gtscript.Field[float]):
from __gtscript__ import K
with computation(PARALLEL), interval(K[2], -1):
field = 0
def func(in_field: gtscript.Field[np.float_]):
with computation(PARALLEL), interval(...):
in_field = not isfinite(in_field)
def definition_func(field: gtscript.Field[float]):
with computation(PARALLEL), interval(-1, 1):
field = 0
def func(in_field: gtscript.Field[np.float_]):
with computation(PARALLEL), interval(...):
in_field = asin(in_field) + 1 if 1 < in_field else sin(
in_field)
def definition_func(phi: gtscript.Field[np.float64]):
from __externals__ import test
with computation(PARALLEL), interval(...):
phi = test(phi)
def definition_func(inout_field: gtscript.Field[float]):
from __externals__ import some_call
with computation(PARALLEL), interval(...):
inout_field = func() + some_call()
def definition(inout_field: gtscript.Field[float]):
with computation(PARALLEL), interval(...):
compile_assert(GLOBAL_VERY_NESTED_CONSTANTS.nested.A > 1)
inout_field = inout_field[
0, 0, 0] + GLOBAL_VERY_NESTED_CONSTANTS.nested.A
def definition_func(inout_field: gtscript.Field[float]):
from gt4py.__gtscript__ import PARALLEL, computation, interval
with computation(PARALLEL), interval(...):
inout_field = func()
def definition(inout_field: gtscript.Field[float]):
with computation(PARALLEL), interval(...):
compile_assert(inout_field[0, 0, 0] < 0)
def definition_func(inout_field: gtscript.Field[float]):
from gt4py.__externals__ import NESTED_CONSTANTS
from gt4py.__gtscript__ import PARALLEL, computation, interval
with computation(PARALLEL), interval(...):
inout_field = inout_field[0, 0, 0] + NESTED_CONSTANTS.missing
def definition_func(in_field: gtscript.Field["dtype"],
out_field: gtscript.Field["dtype"],
param: "dtype"):
with computation(PARALLEL), interval(...):
out_field = in_field + param
def definition_func(inout_field: gtscript.Field[float]):
from gt4py.__externals__ import WRONG_VALUE_CONSTANT
with computation(PARALLEL), interval(...):
inout_field = inout_field[0, 0, 0] + WRONG_VALUE_CONSTANT
def func(in_field: gtscript.Field[np.float_],
out_field: gtscript.Field[np.float_]):
with computation(PARALLEL), interval(...):
out_field["a_key"] = in_field
def definition_func(field: gtscript.Field[float]):
from __externals__ import kstart
with computation(PARALLEL), interval(kstart, -1):
field = 0
def set_gz(zs: FloatFieldIJ, delz: FloatField, gz: FloatField):
with computation(BACKWARD):
with interval(-1, None):
gz[0, 0, 0] = zs
with interval(0, -1):
gz[0, 0, 0] = gz[0, 0, 1] - delz
def ap1_stencil(ta: sd, ap1: sd):
with computation(PARALLEL), interval(...):
ap1 = ap1_for_wqs2(ta)
