def fv_setup(
qvapor: FloatField,
qliquid: FloatField,
qrain: FloatField,
qsnow: FloatField,
qice: FloatField,
qgraupel: FloatField,
q_con: FloatField,
cvm: FloatField,
pkz: FloatField,
pt: FloatField,
cappa: FloatField,
delp: FloatField,
delz: FloatField,
dp1: FloatField,
):
with computation(PARALLEL), interval(...):
from __externals__ import moist_phys
if __INLINED(moist_phys):
cvm, q_con = moist_cv_nwat6_fn(
qvapor, qliquid, qrain, qsnow, qice, qgraupel
) # if (nwat == 6) else moist_cv_default_fn(cv_air)
dp1 = constants.ZVIR * qvapor
cappa = constants.RDGAS / (constants.RDGAS + cvm / (1.0 + dp1))
pkz = exp(
cappa
* log(constants.RDG * delp * pt * (1.0 + dp1) * (1.0 - q_con) / delz)
)
else:
dp1 = 0
pkz = exp(constants.KAPPA * log(constants.RDG * delp * pt / delz))
def fvsetup_stencil(
qvapor: FloatField,
qliquid: FloatField,
qrain: FloatField,
qsnow: FloatField,
qice: FloatField,
qgraupel: FloatField,
q_con: FloatField,
cvm: FloatField,
pkz: FloatField,
pt: FloatField,
cappa: FloatField,
delp: FloatField,
delz: FloatField,
dp1: FloatField,
zvir: float,
nwat: int,
moist_phys: bool,
):
with computation(PARALLEL), interval(...):
# TODO: The conditional with gtscript function triggers and undefined
# temporary variable, even though there are no new temporaries
# if moist_phys:
cvm, q_con = moist_cv_nwat6_fn(
qvapor, qliquid, qrain, qsnow, qice,
qgraupel) # if (nwat == 6) else moist_cv_default_fn(cv_air)
dp1 = zvir * qvapor
cappa = constants.RDGAS / (constants.RDGAS + cvm / (1.0 + dp1))
pkz = exp(cappa * log(constants.RDG * delp * pt * (1.0 + dp1) *
(1.0 - q_con) / delz))
def native_functions(field_a: Field3D, field_b: Field3D):
with computation(PARALLEL), interval(...):
abs_res = abs(field_a)
max_res = max(abs_res, field_b)
min_res = min(max_res, 42)
mod_res = mod(min_res, 37.5)
sin_res = sin(mod_res)
asin_res = asin(sin_res)
cos_res = cos(asin_res)
acos_res = acos(cos_res)
tan_res = tan(acos_res)
atan_res = atan(tan_res)
sinh_res = sinh(atan_res)
asinh_res = asinh(sinh_res)
cosh_res = cosh(asinh_res)
acosh_res = acosh(cosh_res)
tanh_res = tanh(acosh_res)
atanh_res = atanh(tanh_res)
sqrt_res = afunc(atanh_res)
exp_res = exp(sqrt_res)
log_res = log(exp_res)
gamma_res = gamma(log_res)
cbrt_res = cbrt(gamma_res)
floor_res = floor(cbrt_res)
ceil_res = ceil(floor_res)
trunc_res = trunc(ceil_res)
field_a = (trunc_res if isfinite(trunc_res) else field_a if
isinf(trunc_res) else field_b if isnan(trunc_res) else 0.0)
def precompute(
cp3: sd,
gz: sd,
dm: sd,
q_con: sd,
pem: sd,
peg: sd,
dz: sd,
gm: sd,
pef: sd,
pm: sd,
ptop: float,
):
with computation(FORWARD):
with interval(0, 1):
pem = ptop
peg = ptop
pef = ptop
with interval(1, None):
pem = pem[0, 0, -1] + dm[0, 0, -1]
peg = peg[0, 0, -1] + dm[0, 0, -1] * (1.0 - q_con[0, 0, -1])
pef = ptop
with computation(PARALLEL), interval(0, -1):
dz = gz[0, 0, 1] - gz
with computation(PARALLEL), interval(...):
gm = 1.0 / (1.0 - cp3)
dm = dm / constants.GRAV
with computation(PARALLEL), interval(0, -1):
pm = (peg[0, 0, 1] - peg) / log(peg[0, 0, 1] / peg)
def moist_pt(
qvapor: sd,
qliquid: sd,
qrain: sd,
qsnow: sd,
qice: sd,
qgraupel: sd,
q_con: sd,
gz: sd,
cvm: sd,
pt: sd,
cappa: sd,
delp: sd,
delz: sd,
r_vir: float,
nwat: int,
):
with computation(PARALLEL), interval(...):
cvm, gz = moist_cv_nwat6_fn(
qvapor, qliquid, qrain, qsnow, qice,
qgraupel) # if (nwat == 6) else moist_cv_default_fn(cv_air)
q_con[0, 0, 0] = gz
cappa = set_cappa(qvapor, cvm, r_vir)
pt = pt * exp(cappa /
(1.0 - cappa) * log(constants.RDG * delp / delz * pt))
def edge_pe(pe: sd, delp: sd, pk3: sd, ptop: float, akap: float):
with computation(FORWARD):
with interval(0, 1):
pe[0, 0, 0] = ptop
with interval(1, None):
pe[0, 0, 0] = pe[0, 0, -1] + delp[0, 0, -1]
with computation(PARALLEL), interval(1, None):
pk3 = exp(akap * log(pe))
def compute_pkz_tempadjust(
delp: sd, delz: sd, cappa: sd, heat_source: sd, delt: sd, pt: sd, pkz: sd
):
with computation(PARALLEL), interval(...):
pkz = exp(cappa / (1.0 - cappa) * log(constants.RDG * delp / delz * pt))
pkz = (constants.RDG * delp / delz * pt) ** (cappa / (1.0 - cappa))
dtmp = heat_source / (constants.CV_AIR * delp)
deltmin = sign(min(delt, abs(dtmp)), dtmp)
pt = pt + deltmin / pkz
def precompute(
delp: FloatField,
cappa: FloatField,
pe: FloatField,
pe_init: FloatField,
dm: FloatField,
zh: FloatField,
q_con: FloatField,
pem: FloatField,
peln: FloatField,
pk3: FloatField,
gm: FloatField,
dz: FloatField,
pm: FloatField,
ptop: float,
peln1: float,
ptk: float,
):
with computation(PARALLEL), interval(...):
dm = delp
pe_init = pe
with computation(FORWARD):
with interval(0, 1):
pem = ptop
peln = peln1
pk3 = ptk
peg = ptop
pelng = peln1
with interval(1, None):
# TODO consolidate with riem_solver_c, same functions, math functions
pem = pem[0, 0, -1] + dm[0, 0, -1]
peln = log(pem)
# Excluding contribution from condensates
# peln used during remap; pk3 used only for p_grad
peg = peg[0, 0, -1] + dm[0, 0, -1] * (1.0 - q_con[0, 0, -1])
pelng = log(peg)
# interface pk is using constant akap
pk3 = exp(constants.KAPPA * peln)
with computation(PARALLEL), interval(...):
gm = 1.0 / (1.0 - cappa)
dm = dm * constants.RGRAV
with computation(PARALLEL), interval(0, -1):
pm = (peg[0, 0, 1] - peg) / (pelng[0, 0, 1] - pelng)
dz = zh[0, 0, 1] - zh
def pn2_pk_delp(
dp2: FloatField,
delp: FloatField,
pe2: FloatField,
pn2: FloatField,
pk: FloatField,
akap: float,
):
with computation(PARALLEL), interval(...):
delp = dp2
pn2 = log(pe2)
pk = exp(akap * pn2)
def precompute(
cp3: sd,
dm: sd,
zh: sd,
q_con: sd,
pem: sd,
peln: sd,
pk3: sd,
peg: sd,
pelng: sd,
gm: sd,
dz: sd,
pm: sd,
ptop: float,
peln1: float,
ptk: float,
rgrav: float,
akap: float,
):
with computation(FORWARD):
with interval(0, 1):
pem = ptop
peln = peln1
pk3 = ptk
peg = ptop
pelng = peln1
with interval(1, None):
# TODO consolidate with riem_solver_c, same functions, math functions
pem = pem[0, 0, -1] + dm[0, 0, -1]
peln = log(pem)
peg = peg[0, 0, -1] + dm[0, 0, -1] * (1.0 - q_con[0, 0, -1])
pelng = log(peg)
pk3 = exp(akap * peln)
with computation(PARALLEL), interval(...):
gm = 1.0 / (1.0 - cp3)
dm = dm * rgrav
with computation(PARALLEL), interval(0, -1):
pm = (peg[0, 0, 1] - peg) / (pelng[0, 0, 1] - pelng)
dz = zh[0, 0, 1] - zh
def pn2_and_pk(pe2: FloatField, pn2: FloatField, pk: FloatField, akap: float):
from __externals__ import local_je
# copy_j_adjacent
with computation(PARALLEL), interval(1, None):
with horizontal(region[:, local_je + 1 : local_je + 2]):
# TODO: Fix silly hack due to pe2 being 2d, so pe[:, je+1, 1:npz] should be
# the same as it was for pe[:, je, 1:npz] (unchanged)
pe2 = pe2[0, -1, 0]
# pn2_and_pk
with computation(PARALLEL), interval(...):
pn2 = log(pe2)
pk = exp(akap * pn2)
def qs_table2_fn(i):
tem0 = tem_lower(i)
table2 = 0.0
if i < 1600:
# compute es over ice between - 160 deg c and 0 deg c.
table2 = table_ice_oneline(tem0)
else:
# compute es over water between 0 deg c and 102 deg c.
table2 = table_vapor_oneline(tem0)
table2_m1 = 0.0
table2_p1 = 0.0
# TODO is there way to express the code below with something closer to this?:
# if i == 1599:
# table2 = 0.25 * (table2(i-1) + 2.0 * qs_table_fn(i) + qs_table2_fn(i+1))
# if i == 1600:
# table2 = 0.25 * (table2(i-1) + 2.0 * qs_table_fn(i) + qs_table2_fn(i+1))
table = 0.0
if i == 1599:
# table(i)
table = table_ice_oneline(tem0)
tem0 = 253.16 + DELT * (i - 1400) # tem_upper(i - 1400)
# table_vapor_oneline(tem0)
table = (0.05 * (TICE - tem0)) * table + (0.05 * (tem0 - 253.16)) * (
E00
* exp(
(DC_VAP * log(tem0 / TICE) + (tem0 - TICE) / (tem0 * TICE) * LV0)
/ constants.RVGAS
)
)
# table2(i - 1)
tem0 = TMIN + DELT * 1598 # tem_lower(1598)
table2_m1 = table_ice_oneline(tem0)
# table2(i + 1)
tem0 = TMIN + DELT * 1600 # tem_lower(1600)
table2_p1 = table_vapor_oneline(tem0)
table2 = 0.25 * (table2_m1 + 2.0 * table + table2_p1)
if i == 1600:
# table(i)
tem0 = 253.16 + DELT * (i - 1400) # tem_upper(i - 1400)
table = table_vapor_oneline(tem0)
# table2(i - 1)
tem0 = TMIN + DELT * 1599 # tem_lower(1599)
table2_m1 = table_ice_oneline(tem0)
# table2(i + 1)
tem0 = TMIN + DELT * 1601 # tem_lower(1601)
table2_p1 = table_vapor_oneline(tem0)
table2 = 0.25 * (table2_m1 + 2.0 * table + table2_p1)
return table2
def compute_pkz_tempadjust(
delp: FloatField,
delz: FloatField,
cappa: FloatField,
heat_source: FloatField,
pt: FloatField,
pkz: FloatField,
delt_time_factor: float,
):
"""
Adjust air temperature from heating due to vorticity damping.
Heating is limited by deltmax times the length of a timestep, with the
highest levels limited further.
Args:
delp: Pressure thickness of atmosphere layers (in)
delz: Vertical thickness of atmosphere layers (in)
cappa: Power to raise pressure to (in)
heat_source: heat source from vorticity damping implied by
energy conservation (in)
pt: Air potential temperature (inout)
pkz: Layer mean pressure raised to the power of Kappa (in)
delta_time_factor: scaled time step (in)
"""
with computation(PARALLEL), interval(...):
pkz = exp(cappa / (1.0 - cappa) *
log(constants.RDG * delp / delz * pt))
pkz = (constants.RDG * delp / delz * pt)**(cappa / (1.0 - cappa))
dtmp = heat_source / (constants.CV_AIR * delp)
with computation(PARALLEL):
with interval(0, 1):
deltmin = sign(min(delt_time_factor * 0.1, abs(dtmp)), dtmp)
pt = pt + deltmin / pkz
with interval(1, 2):
deltmin = sign(min(delt_time_factor * 0.5, abs(dtmp)), dtmp)
pt = pt + deltmin / pkz
with interval(2, None):
deltmin = sign(min(delt_time_factor, abs(dtmp)), dtmp)
pt = pt + deltmin / pkz
def moist_pt_func(
qvapor: FloatField,
qliquid: FloatField,
qrain: FloatField,
qsnow: FloatField,
qice: FloatField,
qgraupel: FloatField,
q_con: FloatField,
gz: FloatField,
cvm: FloatField,
pt: FloatField,
cappa: FloatField,
delp: FloatField,
delz: FloatField,
r_vir: float,
):
cvm, gz = moist_cv_nwat6_fn(
qvapor, qliquid, qrain, qsnow, qice, qgraupel
) # if (nwat == 6) else moist_cv_default_fn(cv_air)
q_con = gz
cappa = set_cappa(qvapor, cvm, r_vir)
pt = pt * exp(cappa / (1.0 - cappa) * log(constants.RDG * delp / delz * pt))
return cvm, gz, q_con, cappa, pt
def compute_rf_vals(pfull, bdt, rf_cutoff, tau0, ptop):
return (bdt / tau0 * sin(0.5 * constants.PI * log(rf_cutoff / pfull) /
log(rf_cutoff / ptop))**2)
def table_ice_oneline(tem):
return E00 * exp(
(D2ICE * log(tem / TICE) + (tem - TICE) / (tem * TICE) * LI2) / constants.RVGAS
)
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 table_vapor_oneline(tem):
return E00 * exp(
(DC_VAP * log(tem / TICE) + (tem - TICE) / (tem * TICE) * LV0) / constants.RVGAS
)
def riem_solver_c(ms: int, dt: float, akap: float, cappa: Field[float, IJK],
cp: float, ptop: float, hs: Field[float,
IJ], w3: Field[float, IJK],
pt: Field[float, IJK], q_con: Field[float, IJK],
delp: Field[float, IJK], gz: Field[float, IJK],
pef: Field[float, IJK], ws: Field[float, IJK], p_fac: float,
scale_m: float, pe: Field[float, IJK]):
"""
C-grid Riemann solver.
Args:
ms: ...
dt: Timestep
akap: TBD
cappa: TBD
cp: TBD
ptop: TBD
hs: TBD
w3: TBD
pt: TBD
q_con: TBD
delp: TBD
gz: TBD [inout]
pef: TBD [out]
ws: TBD
p_fac: TBD
scale_m: TBD
"""
from __externals__ import A_IMP
with computation(PARALLEL), interval(...):
dm = delp
gama = 1. / (1. - akap)
with computation(FORWARD):
with interval(0, 1):
pef = ptop
pem = ptop
if __INLINED(USE_COND):
peg = ptop
with interval(2, None):
pem = pem[0, 0, -1] + dm[0, 0, -1]
if __INLINED(USE_COND):
peg = peg[0, 0, -1] + dm[0, 0, -1] * (1. - q_con[0, 0, -1])
with computation(PARALLEL), interval(0, -1):
dz2 = gz[0, 0, 1] - gz[0, 0, 0]
if __INLINED(USE_COND):
pm2 = (peg[0, 0, 1] - peg[0, 0, 0]) / log(
peg[0, 0, 1] / peg[0, 0, 0])
if __INLINED(MOIST_CAPPA):
cp2 = cappa
gm2 = 1. / (1. - cp2)
else:
pm2 = dm / log(pem[0, 0, 1] / pem[0, 0, 0])
dm = dm * 1. / GRAV
w2 = w3
# SIM1_solver {
# Line 232
with computation(PARALLEL), interval(0, -1):
if __INLINED(A_IMP > 0.5):
if __INLINED(MOIST_CAPPA):
fac = gm2
else:
fac = gama
pe = exp(fac * log(-dm / dz2 * RDGAS * pt)) - pm2
w1 = w2
# Line 245
with computation(PARALLEL), interval(1, -2):
if __INLINED(A_IMP > 0.5):
g_rat = dm[0, 0, 0] / dm[0, 0, 1]
bb = 2. * (1. + g_rat)
dd = 3. * (pe + g_rat) * pe[0, 0, 1]
# with interval(-2, -1):
# bb = 2.0
# dd = 3.0 * pe
# Line 255
with computation(FORWARD):
# Set special conditions on pp, bb, dd
with interval(0, 1):
if __INLINED(A_IMP > 0.5):
pp = 0
with interval(1, 2):
if __INLINED(A_IMP > 0.5):
pp = dd[0, 0, -1] / bb[0, 0, -1]
with interval(-2, -1):
if __INLINED(A_IMP > 0.5):
bb = 2.
dd = 3 * pe
with computation(FORWARD), interval(2, 3):
if __INLINED(A_IMP > 0.5):
bet = bb[0, 0, -1]
# Line 265
with computation(PARALLEL), interval(2, -1):
if __INLINED(A_IMP > 0.5):
gam = g_rat[0, 0, -1] / bet[0, 0, 0]
bet = bb - gam
with computation(PARALLEL), interval(3, None):
if __INLINED(A_IMP > 0.5):
pp = (dd[0, 0, 0] - pp[0, 0, 0]) / bet[0, 0, -1]
# Line 275
with computation(BACKWARD), interval(1, -1): # this may need to be -2
if __INLINED(A_IMP > 0.5):
pp = pp - gam * pp[0, 0, 1]
aa = (pem[0, 0, 0] + pp[0, 0, 0]) / (dz2[0, 0, -1] + dz2[0, 0, 0])
t1g = calc_t1g(dt, gama)
if __INLINED(MOIST_CAPPA):
aa *= t1g * 0.5 * (gm2[0, 0, -1] + gm2[0, 0, 0])
else:
aa *= t1g
# Line 295
with computation(PARALLEL), interval(1, 2):
if __INLINED(A_IMP > 0.5):
bet = dm[0, 0, 0] - aa[0, 0, 1]
w2 = (dm[0, 0, 0] * w1[0, 0, 0] + dt * pp[0, 0, 1]) / bet
# Line 302
with computation(FORWARD), interval(2, -2):
if __INLINED(A_IMP > 0.5):
gam = aa / bet
bet = dm - (aa * (gam + 1) + aa[0, 0, 1])
w2 = (dm * w1 + dt *
(pp[0, 0, 1] - pp[0, 0, 0]) - aa * w2[0, 0, -1]) / bet
# Line 312
with computation(FORWARD), interval(-2, -1):
if __INLINED(A_IMP > 0.5):
t1g = calc_t1g(dt, gam)
if __INLINED(MOIST_CAPPA):
p1 = t1g * gm2 / dz2 * (pem[0, 0, 1] + pp[0, 0, 1])
else:
p1 = t1g / dz2 * (pem[0, 0, 1] + pp[0, 0, 1])
gam = aa / bet
bet = dm - (aa * (1 + gam) + p1)
w2 = (dm * w1 + dt * (pp[0, 0, 1] - pp[0, 0, 0]) - p1 * ws -
aa * w2[0, 0, -1]) / bet
# Line 325
with computation(BACKWARD), interval(0, -2):
if __INLINED(A_IMP > 0.5):
w2 -= gam[0, 0, 1] * w2[0, 0, 1]
# Line 332
with computation(FORWARD):
with interval(0, 2):
if __INLINED(A_IMP > 0.5):
pe = 0.
with interval(2, -1):
if __INLINED(A_IMP > 0.5):
rdt = 1. / dt
pe = pe[0, 0, -1] + dm * (w2 - w1) * rdt
# Line 346
with computation(BACKWARD):
with interval(-2, -1):
if __INLINED(A_IMP > 0.5):
r3 = 1. / 3.
p1 = (pe + 2. * pe[0, 0, 1]) * r3
dz2 = calc_dz2(dm, pt, cp2, akap, p_fac, pm2, p1)
with interval(0, -2):
if __INLINED(A_IMP > 0.5):
r3 = 1. / 3.
p1 = (pe + bb * pe[0, 0, 1] +
g_rat * pe[0, 0, 2]) * r3 - g_rat * p1
dz2 = calc_dz2(dm, pt, cp2, akap, p_fac, pm2, p1)
# } SIM1_solver
# Line 177
with computation(PARALLEL), interval(2, None):
pef = pe + pem
with computation(BACKWARD), interval(1, -2):
gz = gz[0, 0, 1] - dz2 * GRAV
def satadjust(
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,
cappa: sd,
q_sol: sd,
qg: sd,
pt: sd,
dp: sd,
tin: sd,
delz: sd,
te0: sd,
q_cond: sd,
q_con: sd,
qa: sd,
area: sd,
qpz: sd,
hs: sd,
pkz: sd,
lhl: sd,
lhi: sd,
lcp2: sd,
icp2: sd,
tcp3: sd,
sdt: float,
zvir: float,
fac_i2s: float,
do_qa: bool,
hydrostatic: bool,
consv_te: bool,
c_air: float,
c_vap: float,
mdt: float,
fac_r2g: float,
fac_smlt: float,
fac_l2r: float,
fac_imlt: float,
d0_vap: float,
lv00: float,
fac_v2l: float,
fac_l2v: float,
last_step: bool,
rad_snow: bool,
rad_rain: bool,
rad_graupel: bool,
tintqs: bool,
):
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
dq0 = 0.0
if last_step:
wqsat, dq2dt = wqs2_fn_w(pt1, den)
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
if last_step:
adj_fac = 1.0
else:
adj_fac = spec.namelist.sat_adj0
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)
# CK : Additions from satadjust_part3_laststep_qa
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
if not hydrostatic:
pkz = compute_pkz_func(dp, delz, pt, cappa)
def init_ph_columns(ak: sd, bk: sd, pfull: sd, ph1: sd, ph2: sd, p_ref: float):
with computation(PARALLEL), interval(...):
ph1 = ak + bk * p_ref
ph2 = ak[0, 0, 1] + bk[0, 0, 1] * p_ref
pfull = (ph2 - ph1) / log(ph2 / ph1)
def compute_pkz_func(delp, delz, pt, cappa):
# TODO use the exponential form for closer answer matching
return exp(cappa * log(constants.RDG * delp / delz * pt))
def calc_dz2(dm2, pt, cp2, akap, p_fac, pm2, p1):
if __INLINED(MOIST_CAPPA):
fac = cp2 - 1.
else:
fac = akap - 1
return -dm2 * RDGAS * pt * exp(fac * log(max(p_fac * pm2, p1 + pm2)))
def sim1_solver(
w: sd,
dm: sd,
gm: sd,
dz: sd,
ptr: sd,
pm: sd,
pe: sd,
pem: sd,
wsr: sd,
cp3: sd,
dt: float,
t1g: float,
rdt: float,
p_fac: float,
):
# TODO: We only want to bottom level of wsr, so this could be removed once
# wsr_top is a 2d field.
with computation(FORWARD):
with interval(0, 1):
wsr_top = wsr
with interval(1, None):
wsr_top = wsr_top[0, 0, -1]
with computation(PARALLEL), interval(0, -1):
pe = exp(gm * log(-dm / dz * constants.RDGAS * ptr)) - pm
w1 = w
with computation(FORWARD):
with interval(0, -2):
g_rat = dm / dm[0, 0, 1]
bb = 2.0 * (1.0 + g_rat)
dd = 3.0 * (pe + g_rat * pe[0, 0, 1])
with interval(-2, -1):
bb = 2.0
dd = 3.0 * pe
# bet[i,j,k] = bb[i,j,0]
with computation(FORWARD):
with interval(0, 1):
bet = bb
with interval(1, -1):
bet = bet[0, 0, -1]
# stencils: w_solver
# {
with computation(PARALLEL):
with interval(0, 1):
pp = 0.0
with interval(1, 2):
pp = dd[0, 0, -1] / bet
with computation(FORWARD), interval(1, -1):
gam = g_rat[0, 0, -1] / bet[0, 0, -1]
bet = bb - gam
with computation(FORWARD), interval(2, None):
pp = (dd[0, 0, -1] - pp[0, 0, -1]) / bet[0, 0, -1]
with computation(BACKWARD), interval(1, -1):
pp = pp - gam * pp[0, 0, 1]
# w solver
aa = t1g * 0.5 * (gm[0, 0, -1] + gm) / (dz[0, 0, -1] + dz) * (pem + pp)
# }
# updates on bet:
with computation(FORWARD):
with interval(0, 1):
bet = dm[0, 0, 0] - aa[0, 0, 1]
with interval(1, None):
bet = bet[0, 0, -1]
# w_pe_dz_compute
# {
with computation(FORWARD):
with interval(0, 1):
w = (dm * w1 + dt * pp[0, 0, 1]) / bet
with interval(1, -2):
gam = aa / bet[0, 0, -1]
bet = dm - (aa + aa[0, 0, 1] + aa * gam)
w = (dm * w1 + dt * (pp[0, 0, 1] - pp) - aa * w[0, 0, -1]) / bet
with interval(-2, -1):
p1 = t1g * gm / dz * (pem[0, 0, 1] + pp[0, 0, 1])
gam = aa / bet[0, 0, -1]
bet = dm - (aa + p1 + aa * gam)
w = (dm * w1 + dt *
(pp[0, 0, 1] - pp) - p1 * wsr_top - aa * w[0, 0, -1]) / bet
with computation(BACKWARD), interval(0, -2):
w = w - gam[0, 0, 1] * w[0, 0, 1]
with computation(FORWARD):
with interval(0, 1):
pe = 0.0
with interval(1, None):
pe = pe[0, 0,
-1] + dm[0, 0, -1] * (w[0, 0, -1] - w1[0, 0, -1]) * rdt
with computation(BACKWARD):
with interval(-2, -1):
p1 = (pe + 2.0 * pe[0, 0, 1]) * 1.0 / 3.0
with interval(0, -2):
p1 = (pe + bb * pe[0, 0, 1] +
g_rat * pe[0, 0, 2]) * 1.0 / 3.0 - g_rat * p1[0, 0, 1]
with computation(PARALLEL), interval(0, -1):
maxp = p_fac * pm if p_fac * dm > p1 + pm else p1 + pm
dz = -dm * constants.RDGAS * ptr * exp((cp3 - 1.0) * log(maxp))
def pn2_and_pk(pe2: sd, pn2: sd, pk: sd, akap: float):
with computation(PARALLEL), interval(...):
pn2 = log(pe2)
pk = exp(akap * pn2)
