def update_zonal_velocity(
vorticity: sd,
ke: sd,
velocity: sd,
velocity_c: sd,
cosa: sd,
sina: sd,
rdxc: sd,
dt2: float,
):
from __externals__ import i_end, i_start, namelist
with computation(PARALLEL), interval(...):
assert __INLINED(namelist.grid_type < 3)
# additional assumption: not __INLINED(spec.grid.nested)
tmp_flux = dt2 * (velocity - velocity_c * cosa) / sina
with parallel(region[i_start, :], region[i_end + 1, :]):
tmp_flux = dt2 * velocity
flux = vorticity[0, 0, 0] if tmp_flux > 0.0 else vorticity[0, 1, 0]
velocity_c = velocity_c + tmp_flux * flux + rdxc * (ke[-1, 0, 0] - ke)
def transportdelp(delp: sd, pt: sd, utc: sd, vtc: sd, w: sd, rarea: sd,
delpc: sd, ptc: sd, wc: sd):
"""Transport delp.
Args:
delp: What is transported (input)
pt: Pressure (input)
utc: x-velocity on C-grid (input)
vtc: y-velocity on C-grid (input)
w: z-velocity on C-grid (input)
rarea: Inverse areas (input) -- IJ field
delpc: Updated delp (output)
ptc: Updated pt (output)
wc: Updated w (output)
"""
from __externals__ import namelist
with computation(PARALLEL), interval(...):
assert __INLINED(namelist.grid_type < 3)
# additional assumption (not grid.nested)
delp = fill_4corners_x(delp)
pt = fill_4corners_x(pt)
w = fill_4corners_x(w)
fx, fx1, fx2 = nonhydro_x_fluxes(delp, pt, w, utc)
delp = fill_4corners_y(delp)
pt = fill_4corners_y(pt)
w = fill_4corners_y(w)
fy, fy1, fy2 = nonhydro_y_fluxes(delp, pt, w, vtc)
delpc = delp + (fx1 - fx1[1, 0, 0] + fy1 - fy1[0, 1, 0]) * rarea
ptc = (pt * delp +
(fx - fx[1, 0, 0] + fy - fy[0, 1, 0]) * rarea) / delpc
wc = (w * delp +
(fx2 - fx2[1, 0, 0] + fy2 - fy2[0, 1, 0]) * rarea) / delpc
def update_meridional_velocity(
vorticity: sd,
ke: sd,
velocity: sd,
velocity_c: sd,
cosa: sd,
sina: sd,
rdyc: sd,
dt2: float,
):
from __externals__ import j_end, j_start, namelist
with computation(PARALLEL), interval(...):
assert __INLINED(namelist.grid_type < 3)
# additional assumption: not __INLINED(spec.grid.nested)
tmp_flux = dt2 * (velocity - velocity_c * cosa) / sina
with horizontal(region[:, j_start], region[:, j_end + 1]):
tmp_flux = dt2 * velocity
flux = vorticity[0, 0, 0] if tmp_flux > 0.0 else vorticity[1, 0, 0]
velocity_c = velocity_c - tmp_flux * flux + rdyc * (ke[0, -1, 0] - ke)
def update_vorticity_and_kinetic_energy(
ke: sd,
vort: sd,
ua: sd,
va: sd,
uc: sd,
vc: sd,
u: sd,
v: sd,
sin_sg1: sd,
cos_sg1: sd,
sin_sg2: sd,
cos_sg2: sd,
sin_sg3: sd,
cos_sg3: sd,
sin_sg4: sd,
cos_sg4: sd,
dt2: float,
):
from __externals__ import i_end, i_start, j_end, j_start, namelist
with computation(PARALLEL), interval(...):
assert __INLINED(namelist.grid_type < 3)
ke = uc if ua > 0.0 else uc[1, 0, 0]
vort = vc if va > 0.0 else vc[0, 1, 0]
with horizontal(region[:, j_start - 1], region[:, j_end]):
vort = vort * sin_sg4 + u[0, 1, 0] * cos_sg4 if va <= 0.0 else vort
with horizontal(region[:, j_start], region[:, j_end + 1]):
vort = vort * sin_sg2 + u * cos_sg2 if va > 0.0 else vort
with horizontal(region[i_end, :], region[i_start - 1, :]):
ke = ke * sin_sg3 + v[1, 0, 0] * cos_sg3 if ua <= 0.0 else ke
with horizontal(region[i_end + 1, :], region[i_start, :]):
ke = ke * sin_sg1 + v * cos_sg1 if ua > 0.0 else ke
ke = 0.5 * dt2 * (ua * ke + va * vort)
def definition2(f_in: gtscript.Field[np.float_]):
from __externals__ import flag
with computation(PARALLEL), interval(...):
if __INLINED(flag):
B = f_in
def copy_corners_x_nord(q_in: FloatField, q_out: FloatField):
from __externals__ import i_end, i_start, j_end, j_start, nord0, nord1, nord2, nord3
with computation(PARALLEL), interval(0, 1):
if __INLINED(nord0 > 0):
with horizontal(
region[i_start - 3, j_start - 3], region[i_end + 3, j_start - 3]
):
q_out = q_in[0, 5, 0]
with horizontal(
region[i_start - 2, j_start - 3], region[i_end + 3, j_start - 2]
):
q_out = q_in[-1, 4, 0]
with horizontal(
region[i_start - 1, j_start - 3], region[i_end + 3, j_start - 1]
):
q_out = q_in[-2, 3, 0]
with horizontal(
region[i_start - 3, j_start - 2], region[i_end + 2, j_start - 3]
):
q_out = q_in[1, 4, 0]
with horizontal(
region[i_start - 2, j_start - 2], region[i_end + 2, j_start - 2]
):
q_out = q_in[0, 3, 0]
with horizontal(
region[i_start - 1, j_start - 2], region[i_end + 2, j_start - 1]
):
q_out = q_in[-1, 2, 0]
with horizontal(
region[i_start - 3, j_start - 1], region[i_end + 1, j_start - 3]
):
q_out = q_in[2, 3, 0]
with horizontal(
region[i_start - 2, j_start - 1], region[i_end + 1, j_start - 2]
):
q_out = q_in[1, 2, 0]
with horizontal(
region[i_start - 1, j_start - 1], region[i_end + 1, j_start - 1]
):
q_out = q_in[0, 1, 0]
with horizontal(
region[i_start - 3, j_end + 1], region[i_end + 1, j_end + 3]
):
q_out = q_in[2, -3, 0]
with horizontal(
region[i_start - 2, j_end + 1], region[i_end + 1, j_end + 2]
):
q_out = q_in[1, -2, 0]
with horizontal(
region[i_start - 1, j_end + 1], region[i_end + 1, j_end + 1]
):
q_out = q_in[0, -1, 0]
with horizontal(
region[i_start - 3, j_end + 2], region[i_end + 2, j_end + 3]
):
q_out = q_in[1, -4, 0]
with horizontal(
region[i_start - 2, j_end + 2], region[i_end + 2, j_end + 2]
):
q_out = q_in[0, -3, 0]
with horizontal(
region[i_start - 1, j_end + 2], region[i_end + 2, j_end + 1]
):
q_out = q_in[-1, -2, 0]
with horizontal(
region[i_start - 3, j_end + 3], region[i_end + 3, j_end + 3]
):
q_out = q_in[0, -5, 0]
with horizontal(
region[i_start - 2, j_end + 3], region[i_end + 3, j_end + 2]
):
q_out = q_in[-1, -4, 0]
with horizontal(
region[i_start - 1, j_end + 3], region[i_end + 3, j_end + 1]
):
q_out = q_in[-2, -3, 0]
with computation(PARALLEL), interval(1, 2):
if __INLINED(nord1 > 0):
with horizontal(
region[i_start - 3, j_start - 3], region[i_end + 3, j_start - 3]
):
q_out = q_in[0, 5, 0]
with horizontal(
region[i_start - 2, j_start - 3], region[i_end + 3, j_start - 2]
):
q_out = q_in[-1, 4, 0]
with horizontal(
region[i_start - 1, j_start - 3], region[i_end + 3, j_start - 1]
):
q_out = q_in[-2, 3, 0]
with horizontal(
region[i_start - 3, j_start - 2], region[i_end + 2, j_start - 3]
):
q_out = q_in[1, 4, 0]
with horizontal(
region[i_start - 2, j_start - 2], region[i_end + 2, j_start - 2]
):
q_out = q_in[0, 3, 0]
with horizontal(
region[i_start - 1, j_start - 2], region[i_end + 2, j_start - 1]
):
q_out = q_in[-1, 2, 0]
with horizontal(
region[i_start - 3, j_start - 1], region[i_end + 1, j_start - 3]
):
q_out = q_in[2, 3, 0]
with horizontal(
region[i_start - 2, j_start - 1], region[i_end + 1, j_start - 2]
):
q_out = q_in[1, 2, 0]
with horizontal(
region[i_start - 1, j_start - 1], region[i_end + 1, j_start - 1]
):
q_out = q_in[0, 1, 0]
with horizontal(
region[i_start - 3, j_end + 1], region[i_end + 1, j_end + 3]
):
q_out = q_in[2, -3, 0]
with horizontal(
region[i_start - 2, j_end + 1], region[i_end + 1, j_end + 2]
):
q_out = q_in[1, -2, 0]
with horizontal(
region[i_start - 1, j_end + 1], region[i_end + 1, j_end + 1]
):
q_out = q_in[0, -1, 0]
with horizontal(
region[i_start - 3, j_end + 2], region[i_end + 2, j_end + 3]
):
q_out = q_in[1, -4, 0]
with horizontal(
region[i_start - 2, j_end + 2], region[i_end + 2, j_end + 2]
):
q_out = q_in[0, -3, 0]
with horizontal(
region[i_start - 1, j_end + 2], region[i_end + 2, j_end + 1]
):
q_out = q_in[-1, -2, 0]
with horizontal(
region[i_start - 3, j_end + 3], region[i_end + 3, j_end + 3]
):
q_out = q_in[0, -5, 0]
with horizontal(
region[i_start - 2, j_end + 3], region[i_end + 3, j_end + 2]
):
q_out = q_in[-1, -4, 0]
with horizontal(
region[i_start - 1, j_end + 3], region[i_end + 3, j_end + 1]
):
q_out = q_in[-2, -3, 0]
with computation(PARALLEL), interval(2, 3):
if __INLINED(nord2 > 0):
with horizontal(
region[i_start - 3, j_start - 3], region[i_end + 3, j_start - 3]
):
q_out = q_in[0, 5, 0]
with horizontal(
region[i_start - 2, j_start - 3], region[i_end + 3, j_start - 2]
):
q_out = q_in[-1, 4, 0]
with horizontal(
region[i_start - 1, j_start - 3], region[i_end + 3, j_start - 1]
):
q_out = q_in[-2, 3, 0]
with horizontal(
region[i_start - 3, j_start - 2], region[i_end + 2, j_start - 3]
):
q_out = q_in[1, 4, 0]
with horizontal(
region[i_start - 2, j_start - 2], region[i_end + 2, j_start - 2]
):
q_out = q_in[0, 3, 0]
with horizontal(
region[i_start - 1, j_start - 2], region[i_end + 2, j_start - 1]
):
q_out = q_in[-1, 2, 0]
with horizontal(
region[i_start - 3, j_start - 1], region[i_end + 1, j_start - 3]
):
q_out = q_in[2, 3, 0]
with horizontal(
region[i_start - 2, j_start - 1], region[i_end + 1, j_start - 2]
):
q_out = q_in[1, 2, 0]
with horizontal(
region[i_start - 1, j_start - 1], region[i_end + 1, j_start - 1]
):
q_out = q_in[0, 1, 0]
with horizontal(
region[i_start - 3, j_end + 1], region[i_end + 1, j_end + 3]
):
q_out = q_in[2, -3, 0]
with horizontal(
region[i_start - 2, j_end + 1], region[i_end + 1, j_end + 2]
):
q_out = q_in[1, -2, 0]
with horizontal(
region[i_start - 1, j_end + 1], region[i_end + 1, j_end + 1]
):
q_out = q_in[0, -1, 0]
with horizontal(
region[i_start - 3, j_end + 2], region[i_end + 2, j_end + 3]
):
q_out = q_in[1, -4, 0]
with horizontal(
region[i_start - 2, j_end + 2], region[i_end + 2, j_end + 2]
):
q_out = q_in[0, -3, 0]
with horizontal(
region[i_start - 1, j_end + 2], region[i_end + 2, j_end + 1]
):
q_out = q_in[-1, -2, 0]
with horizontal(
region[i_start - 3, j_end + 3], region[i_end + 3, j_end + 3]
):
q_out = q_in[0, -5, 0]
with horizontal(
region[i_start - 2, j_end + 3], region[i_end + 3, j_end + 2]
):
q_out = q_in[-1, -4, 0]
with horizontal(
region[i_start - 1, j_end + 3], region[i_end + 3, j_end + 1]
):
q_out = q_in[-2, -3, 0]
with computation(PARALLEL), interval(3, None):
if __INLINED(nord3 > 0):
with horizontal(
region[i_start - 3, j_start - 3], region[i_end + 3, j_start - 3]
):
q_out = q_in[0, 5, 0]
with horizontal(
region[i_start - 2, j_start - 3], region[i_end + 3, j_start - 2]
):
q_out = q_in[-1, 4, 0]
with horizontal(
region[i_start - 1, j_start - 3], region[i_end + 3, j_start - 1]
):
q_out = q_in[-2, 3, 0]
with horizontal(
region[i_start - 3, j_start - 2], region[i_end + 2, j_start - 3]
):
q_out = q_in[1, 4, 0]
with horizontal(
region[i_start - 2, j_start - 2], region[i_end + 2, j_start - 2]
):
q_out = q_in[0, 3, 0]
with horizontal(
region[i_start - 1, j_start - 2], region[i_end + 2, j_start - 1]
):
q_out = q_in[-1, 2, 0]
with horizontal(
region[i_start - 3, j_start - 1], region[i_end + 1, j_start - 3]
):
q_out = q_in[2, 3, 0]
with horizontal(
region[i_start - 2, j_start - 1], region[i_end + 1, j_start - 2]
):
q_out = q_in[1, 2, 0]
with horizontal(
region[i_start - 1, j_start - 1], region[i_end + 1, j_start - 1]
):
q_out = q_in[0, 1, 0]
with horizontal(
region[i_start - 3, j_end + 1], region[i_end + 1, j_end + 3]
):
q_out = q_in[2, -3, 0]
with horizontal(
region[i_start - 2, j_end + 1], region[i_end + 1, j_end + 2]
):
q_out = q_in[1, -2, 0]
with horizontal(
region[i_start - 1, j_end + 1], region[i_end + 1, j_end + 1]
):
q_out = q_in[0, -1, 0]
with horizontal(
region[i_start - 3, j_end + 2], region[i_end + 2, j_end + 3]
):
q_out = q_in[1, -4, 0]
with horizontal(
region[i_start - 2, j_end + 2], region[i_end + 2, j_end + 2]
):
q_out = q_in[0, -3, 0]
with horizontal(
region[i_start - 1, j_end + 2], region[i_end + 2, j_end + 1]
):
q_out = q_in[-1, -2, 0]
with horizontal(
region[i_start - 3, j_end + 3], region[i_end + 3, j_end + 3]
):
q_out = q_in[0, -5, 0]
with horizontal(
region[i_start - 2, j_end + 3], region[i_end + 3, j_end + 2]
):
q_out = q_in[-1, -4, 0]
with horizontal(
region[i_start - 1, j_end + 3], region[i_end + 3, j_end + 1]
):
q_out = q_in[-2, -3, 0]
def finalize(
u0: FloatField,
v0: FloatField,
w0: FloatField,
t0: FloatField,
ua: FloatField,
va: FloatField,
ta: FloatField,
w: FloatField,
u_dt: FloatField,
v_dt: FloatField,
q0_vapor: FloatField,
q0_liquid: FloatField,
q0_rain: FloatField,
q0_ice: FloatField,
q0_snow: FloatField,
q0_graupel: FloatField,
q0_o3mr: FloatField,
q0_sgs_tke: FloatField,
q0_cld: FloatField,
qvapor: FloatField,
qliquid: FloatField,
qrain: FloatField,
qice: FloatField,
qsnow: FloatField,
qgraupel: FloatField,
qo3mr: FloatField,
qsgs_tke: FloatField,
qcld: FloatField,
timestep: float,
):
from __externals__ import fv_sg_adj, hydrostatic
with computation(PARALLEL), interval(...):
fra = timestep / fv_sg_adj
if fra < 1.0:
t0 = readjust_by_frac(t0, ta, fra)
u0 = readjust_by_frac(u0, ua, fra)
v0 = readjust_by_frac(v0, va, fra)
if __INLINED(not hydrostatic):
w0 = readjust_by_frac(w0, w, fra)
q0_vapor = readjust_by_frac(q0_vapor, qvapor, fra)
q0_liquid = readjust_by_frac(q0_liquid, qliquid, fra)
q0_rain = readjust_by_frac(q0_rain, qrain, fra)
q0_ice = readjust_by_frac(q0_ice, qice, fra)
q0_snow = readjust_by_frac(q0_snow, qsnow, fra)
q0_graupel = readjust_by_frac(q0_graupel, qgraupel, fra)
q0_o3mr = readjust_by_frac(q0_o3mr, qo3mr, fra)
q0_sgs_tke = readjust_by_frac(q0_sgs_tke, qsgs_tke, fra)
q0_cld = readjust_by_frac(q0_cld, qcld, fra)
rdt = 1.0 / timestep
u_dt = rdt * (u0 - ua)
v_dt = rdt * (v0 - va)
ta = t0
ua = u0
va = v0
w = w0
qvapor = q0_vapor
qliquid = q0_liquid
qrain = q0_rain
qice = q0_ice
qsnow = q0_snow
qgraupel = q0_graupel
qo3mr = q0_o3mr
qsgs_tke = q0_sgs_tke
qcld = q0_cld
def heat_source_from_vorticity_damping(
vort_x_delta: FloatField,
vort_y_delta: FloatField,
ut: FloatField,
vt: FloatField,
u: FloatField,
v: FloatField,
delp: FloatField,
rsin2: FloatFieldIJ,
cosa_s: FloatFieldIJ,
rdx: FloatFieldIJ,
rdy: FloatFieldIJ,
heat_source: FloatField,
heat_source_total: FloatField,
dissipation_estimate: FloatField,
kinetic_energy_fraction_to_damp: FloatFieldK,
):
"""
Calculates heat source from vorticity damping implied by energy conservation.
Args:
vort_x_delta (in)
vort_y_delta (in)
ut (in)
vt (in)
u (in)
v (in)
delp (in)
rsin2 (in)
cosa_s (in)
rdx (in): 1 / dx
rdy (in): 1 / dy
heat_source (out): heat source from vorticity damping
implied by energy conservation
heat_source_total: (out) accumulated heat source
dissipation_estimate (out): dissipation estimate, only calculated if
calculate_dissipation_estimate is 1
kinetic_energy_fraction_to_damp (in): according to its comment in fv_arrays,
the fraction of kinetic energy to explicitly damp and convert into heat.
TODO: confirm this description is accurate, why is it multiplied
by 0.25 below?
"""
from __externals__ import d_con, do_skeb, local_ie, local_is, local_je, local_js
with computation(PARALLEL), interval(...):
ubt = (vort_x_delta + vt) * rdx
fy = u * rdx
gy = fy * ubt
vbt = (vort_y_delta - ut) * rdy
fx = v * rdy
gx = fx * vbt
if (kinetic_energy_fraction_to_damp > dcon_threshold) or do_skeb:
u2 = fy + fy[0, 1, 0]
du2 = ubt + ubt[0, 1, 0]
v2 = fx + fx[1, 0, 0]
dv2 = vbt + vbt[1, 0, 0]
dampterm = heat_damping_term(
ubt, vbt, gx, gy, rsin2, cosa_s, u2, v2, du2, dv2
)
heat_source = delp * (
heat_source - 0.25 * kinetic_energy_fraction_to_damp * dampterm
)
if __INLINED((d_con > dcon_threshold) or do_skeb):
with horizontal(region[local_is : local_ie + 1, local_js : local_je + 1]):
heat_source_total = heat_source_total + heat_source
# TODO: do_skeb could be renamed to calculate_dissipation_estimate
if __INLINED(do_skeb):
dissipation_estimate -= dampterm
def ray_fast_wind_compute(
u: FloatField,
v: FloatField,
w: FloatField,
dp: FloatFieldK,
pfull: FloatFieldK,
dt: float,
ptop: float,
rf_cutoff_nudge: float,
ks: int,
):
from __externals__ import hydrostatic, local_ie, local_je, rf_cutoff, tau
# dm_stencil
with computation(PARALLEL), interval(...):
# TODO -- in the fortran model rf is only computed once, repeating
# the computation every time ray_fast is run is inefficient
if pfull < rf_cutoff:
rf = compute_rff_vals(pfull, dt, rf_cutoff, tau * SDAY, ptop)
with computation(FORWARD):
with interval(0, 1):
if pfull < rf_cutoff_nudge: # TODO and kaxes(k) < ks:
dm = dp
with interval(1, None):
dm = dm[0, 0, -1]
if pfull < rf_cutoff_nudge: # TODO and kaxes(k) < ks:
dm += dp
with computation(BACKWARD), interval(0, -1):
if pfull < rf_cutoff_nudge:
dm = dm[0, 0, 1]
# ray_fast_wind(u)
with computation(FORWARD):
with interval(0, 1):
with horizontal(region[:local_ie + 1, :]):
if pfull < rf_cutoff:
dmdir = dm_layer(rf, dp, u)
u *= rf
else:
dm = 0
with interval(1, None):
with horizontal(region[:local_ie + 1, :]):
dmdir = dmdir[0, 0, -1]
if pfull < rf_cutoff:
dmdir += dm_layer(rf, dp, u)
u *= rf
with computation(BACKWARD), interval(0, -1):
if pfull < rf_cutoff:
dmdir = dmdir[0, 0, 1]
with computation(PARALLEL), interval(...):
with horizontal(region[:local_ie + 1, :]):
if pfull < rf_cutoff_nudge: # TODO and axes(k) < ks:
u += dmdir / dm
# ray_fast_wind(v)
with computation(FORWARD):
with interval(0, 1):
with horizontal(region[:, :local_je + 1]):
if pfull < rf_cutoff:
dmdir = dm_layer(rf, dp, v)
v *= rf
else:
dm = 0
with interval(1, None):
with horizontal(region[:, :local_je + 1]):
dmdir = dmdir[0, 0, -1]
if pfull < rf_cutoff:
dmdir += dm_layer(rf, dp, v)
v *= rf
with computation(BACKWARD), interval(0, -1):
if pfull < rf_cutoff:
dmdir = dmdir[0, 0, 1]
with computation(PARALLEL), interval(...):
with horizontal(region[:, :local_je + 1]):
if pfull < rf_cutoff_nudge: # TODO and axes(k) < ks:
v += dmdir / dm
# ray_fast_w
with computation(PARALLEL), interval(...):
with horizontal(region[:local_ie + 1, :local_je + 1]):
if __INLINED(not hydrostatic):
if pfull < rf_cutoff:
w *= rf
def assert_in_func(field):
assert __INLINED(GLOBAL_CONSTANT < 2), "An error occurred"
return field[0, 0, 0] + GLOBAL_CONSTANT
def set_initial_vals(
gam: FloatField,
q: FloatField,
delp: FloatField,
a4_1: FloatField,
a4_2: FloatField,
a4_3: FloatField,
a4_4: FloatField,
q_bot: FloatField,
qs: FloatFieldIJ,
):
from __externals__ import iv, kord
with computation(FORWARD):
with interval(0, 1):
# set top
if __INLINED(iv == -2):
# gam = 0.5
q = 1.5 * a4_1
else:
grid_ratio = delp[0, 0, 1] / delp
bet = grid_ratio * (grid_ratio + 0.5)
q = (
(grid_ratio + grid_ratio) * (grid_ratio + 1.0) * a4_1
+ a4_1[0, 0, 1]
) / bet
gam = (1.0 + grid_ratio * (grid_ratio + 1.5)) / bet
with interval(1, 2):
if __INLINED(iv == -2):
gam = 0.5
grid_ratio = delp[0, 0, -1] / delp
bet = 2.0 + grid_ratio + grid_ratio - gam
q = (3.0 * (a4_1[0, 0, -1] + a4_1) - q[0, 0, -1]) / bet
with computation(FORWARD), interval(1, -1):
if __INLINED(iv != -2):
# set middle
d4 = delp[0, 0, -1] / delp
bet = 2.0 + d4 + d4 - gam[0, 0, -1]
q = (3.0 * (a4_1[0, 0, -1] + d4 * a4_1) - q[0, 0, -1]) / bet
gam = d4 / bet
with computation(FORWARD):
with interval(2, -1):
if __INLINED(iv == -2):
old_grid_ratio = delp[0, 0, -2] / delp[0, 0, -1]
old_bet = 2.0 + old_grid_ratio + old_grid_ratio - gam[0, 0, -1]
gam = old_grid_ratio / old_bet
grid_ratio = delp[0, 0, -1] / delp
with computation(FORWARD):
with interval(2, -2):
if __INLINED(iv == -2):
# set middle
bet = 2.0 + grid_ratio + grid_ratio - gam
q = (3.0 * (a4_1[0, 0, -1] + a4_1) - q[0, 0, -1]) / bet
# gam[0, 0, 1] = grid_ratio / bet
with interval(-2, -1):
if __INLINED(iv == -2):
# set bottom
q = (3.0 * (a4_1[0, 0, -1] + a4_1) - grid_ratio * qs - q[0, 0, -1]) / (
2.0 + grid_ratio + grid_ratio - gam
)
q_bot = qs
with interval(-1, None):
if __INLINED(iv == -2):
q_bot = qs
q = qs
else:
# set bottom
d4 = delp[0, 0, -2] / delp[0, 0, -1]
a_bot = 1.0 + d4 * (d4 + 1.5)
q = (
2.0 * d4 * (d4 + 1.0) * a4_1[0, 0, -1]
+ a4_1[0, 0, -2]
- a_bot * q[0, 0, -1]
) / (d4 * (d4 + 0.5) - a_bot * gam[0, 0, -1])
with computation(BACKWARD), interval(0, -1):
if __INLINED(iv != -2):
q = q - gam * q[0, 0, 1]
with computation(BACKWARD), interval(0, -2):
if __INLINED(iv == -2):
q = q - gam[0, 0, 1] * q[0, 0, 1]
# set_avals
with computation(PARALLEL), interval(...):
if __INLINED(kord > 16):
a4_2 = q
a4_3 = q[0, 0, 1]
a4_4 = 3.0 * (2.0 * a4_1 - (a4_2 + a4_3))
with computation(PARALLEL), interval(-1, None):
if __INLINED(kord > 16):
a4_3 = q_bot
def calc_t1g(dt, gama):
if __INLINED(MOIST_CAPPA):
t1g = 2. * dt * dt
else:
t1g = gama * 2. * dt * dt
return t1g
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 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 definition(inout_field: gtscript.Field[float]):
from __externals__ import EXTERNAL
with computation(PARALLEL), interval(...):
assert __INLINED(EXTERNAL < 1), "An error occurred"
inout_field = inout_field[0, 0, 0] + EXTERNAL
def apply_constraints(
q: FloatField,
gam: FloatField,
a4_1: FloatField,
a4_2: FloatField,
a4_3: FloatField,
a4_4: FloatField,
ext5: BoolField,
ext6: BoolField,
extm: BoolField,
):
from __externals__ import iv, kord
# apply constraints
with computation(PARALLEL), interval(1, None):
a4_1_0 = a4_1[0, 0, -1]
tmp = a4_1_0 if a4_1_0 > a4_1 else a4_1
tmp2 = a4_1_0 if a4_1_0 < a4_1 else a4_1
gam = a4_1 - a4_1_0
with computation(PARALLEL), interval(1, 2):
# do top
if q >= tmp:
q = tmp
if q <= tmp2:
q = tmp2
with computation(FORWARD):
with interval(2, -1):
# do middle
if (gam[0, 0, -1] * gam[0, 0, 1]) > 0:
if q >= tmp:
q = tmp
if q <= tmp2:
q = tmp2
elif gam[0, 0, -1] > 0:
# there's a local maximum
if q <= tmp2:
q = tmp2
else:
# there's a local minimum
if q >= tmp:
q = tmp
if __INLINED(iv == 0):
if q < 0.0:
q = 0.0
# q = constrain_interior(q, gam, a4_1)
with interval(-1, None):
# do bottom
if q >= tmp:
q = tmp
if q <= tmp2:
q = tmp2
with computation(PARALLEL), interval(...):
# re-set a4_2 and a4_3
a4_2 = q
a4_3 = q[0, 0, 1]
# set_extm
with computation(PARALLEL):
with interval(0, 1):
extm = (a4_2 - a4_1) * (a4_3 - a4_1) > 0.0
with interval(1, -1):
extm = gam * gam[0, 0, 1] < 0.0
with interval(-1, None):
extm = (a4_2 - a4_1) * (a4_3 - a4_1) > 0.0
# set_exts
with computation(PARALLEL), interval(...):
if __INLINED(kord > 9):
x0 = 2.0 * a4_1 - (a4_2 + a4_3)
x1 = abs(a4_2 - a4_3)
a4_4 = 3.0 * x0
ext5 = abs(x0) > x1
ext6 = abs(a4_4) > x1
def definition(inout_field: gtscript.Field[float]):
with computation(PARALLEL), interval(...):
assert __INLINED(GLOBAL_VERY_NESTED_CONSTANTS.nested.A > 1), "An error occurred"
inout_field = inout_field[0, 0, 0] + GLOBAL_VERY_NESTED_CONSTANTS.nested.A
def set_interpolation_coefficients(
q: FloatField,
gam: FloatField,
a4_1: FloatField,
a4_2: FloatField,
a4_3: FloatField,
a4_4: FloatField,
ext5: BoolField,
ext6: BoolField,
extm: BoolField,
qmin: float,
):
from __externals__ import iv, kord
# set_top_as_iv0
with computation(PARALLEL), interval(0, 1):
if __INLINED(iv == 0):
if a4_2 < 0.0:
a4_2 = 0.0
with computation(PARALLEL), interval(0, 2):
if __INLINED(iv == 0):
a4_4 = 3.0 * (2.0 * a4_1 - (a4_2 + a4_3))
# set_top_as_iv1
with computation(PARALLEL), interval(0, 1):
if __INLINED(iv == -1):
if a4_2 * a4_1 <= 0.0:
a4_2 = 0.0
with computation(PARALLEL), interval(0, 2):
if __INLINED(iv == -1):
a4_4 = 3.0 * (2.0 * a4_1 - (a4_2 + a4_3))
# set_top_as_iv2
with computation(PARALLEL):
with interval(0, 1):
if __INLINED(iv == 2):
a4_2 = a4_1
a4_3 = a4_1
a4_4 = 0.0
with interval(1, 2):
if __INLINED(iv == 2):
a4_4 = 3.0 * (2.0 * a4_1 - (a4_2 + a4_3))
# set_top_as_else
with computation(PARALLEL), interval(0, 2):
if __INLINED(iv < -1 or iv == 1 or iv > 2):
a4_4 = 3.0 * (2.0 * a4_1 - (a4_2 + a4_3))
with computation(PARALLEL):
with interval(0, 1):
if __INLINED(iv != 2):
a4_1, a4_2, a4_3, a4_4 = posdef_constraint_iv1(a4_1, a4_2, a4_3, a4_4)
with interval(1, 2):
a4_1, a4_2, a4_3, a4_4 = remap_constraint(a4_1, a4_2, a4_3, a4_4, extm)
with computation(PARALLEL), interval(2, -2):
# set_inner_as_kordsmall
if __INLINED(kord < 9):
# left edges?
pmp_1 = a4_1 - gam[0, 0, 1]
lac_1 = pmp_1 + 1.5 * gam[0, 0, 2]
tmp_min = (
a4_1
if (a4_1 < pmp_1) and (a4_1 < lac_1)
else pmp_1
if pmp_1 < lac_1
else lac_1
)
tmp_max0 = a4_2 if a4_2 > tmp_min else tmp_min
tmp_max = (
a4_1
if (a4_1 > pmp_1) and (a4_1 > lac_1)
else pmp_1
if pmp_1 > lac_1
else lac_1
)
a4_2 = tmp_max0 if tmp_max0 < tmp_max else tmp_max
# right edges?
pmp_2 = a4_1 + 2.0 * gam[0, 0, 1]
lac_2 = pmp_2 - 1.5 * gam[0, 0, -1]
tmp_min = (
a4_1
if (a4_1 < pmp_2) and (a4_1 < lac_2)
else pmp_2
if pmp_2 < lac_2
else lac_2
)
tmp_max0 = a4_3 if a4_3 > tmp_min else tmp_min
tmp_max = (
a4_1
if (a4_1 > pmp_2) and (a4_1 > lac_2)
else pmp_2
if pmp_2 > lac_2
else lac_2
)
a4_3 = tmp_max0 if tmp_max0 < tmp_max else tmp_max
a4_4 = 3.0 * (2.0 * a4_1 - (a4_2 + a4_3))
# set_inner_as_kord9
if __INLINED(kord == 9):
pmp_1 = a4_1 - 2.0 * gam[0, 0, 1]
lac_1 = pmp_1 + 1.5 * gam[0, 0, 2]
pmp_2 = a4_1 + 2.0 * gam
lac_2 = pmp_2 - 1.5 * gam[0, 0, -1]
tmp_min = a4_1
tmp_max = a4_2
tmp_max0 = a4_1
if (
(extm and extm[0, 0, -1])
or (extm and extm[0, 0, 1])
or (extm and (qmin > 0.0 and a4_1 < qmin))
):
a4_2 = a4_1
a4_3 = a4_1
a4_4 = 0.0
else:
a4_4 = 6.0 * a4_1 - 3.0 * (a4_2 + a4_3)
if abs(a4_4) > abs(a4_2 - a4_3):
tmp_min = (
a4_1
if (a4_1 < pmp_1) and (a4_1 < lac_1)
else pmp_1
if pmp_1 < lac_1
else lac_1
)
tmp_max0 = a4_2 if a4_2 > tmp_min else tmp_min
tmp_max = (
a4_1
if (a4_1 > pmp_1) and (a4_1 > lac_1)
else pmp_1
if pmp_1 > lac_1
else lac_1
)
a4_2 = tmp_max0 if tmp_max0 < tmp_max else tmp_max
tmp_min = (
a4_1
if (a4_1 < pmp_2) and (a4_1 < lac_2)
else pmp_2
if pmp_2 < lac_2
else lac_2
)
tmp_max0 = a4_3 if a4_3 > tmp_min else tmp_min
tmp_max = (
a4_1
if (a4_1 > pmp_2) and (a4_1 > lac_2)
else pmp_2
if pmp_2 > lac_2
else lac_2
)
a4_3 = tmp_max0 if tmp_max0 < tmp_max else tmp_max
a4_4 = 6.0 * a4_1 - 3.0 * (a4_2 + a4_3)
else:
a4_2 = a4_2
# set_inner_as_kord10
if __INLINED(kord == 10):
pmp_1 = a4_1 - 2.0 * gam[0, 0, 1]
lac_1 = pmp_1 + 1.5 * gam[0, 0, 2]
pmp_2 = a4_1 + 2.0 * gam
lac_2 = pmp_2 - 1.5 * gam[0, 0, -1]
tmp_min2 = (
a4_1
if (a4_1 < pmp_1) and (a4_1 < lac_1)
else pmp_1
if pmp_1 < lac_1
else lac_1
)
tmp_max2 = (
a4_1
if (a4_1 > pmp_1) and (a4_1 > lac_1)
else pmp_1
if pmp_1 > lac_1
else lac_1
)
tmp2 = a4_2 if a4_2 > tmp_min2 else tmp_min2
tmp_min3 = a4_1 if a4_1 < pmp_2 else pmp_2
tmp_min3 = lac_2 if lac_2 < tmp_min3 else tmp_min3
tmp_max3 = a4_1 if a4_1 > pmp_2 else pmp_2
tmp_max3 = lac_2 if lac_2 > tmp_max3 else tmp_max3
tmp3 = a4_3 if a4_3 > tmp_min3 else tmp_min3
if ext5:
if ext5[0, 0, -1] or ext5[0, 0, 1]:
a4_2 = a4_1
a4_3 = a4_1
elif ext6[0, 0, -1] or ext6[0, 0, 1]:
a4_2 = tmp2 if tmp2 < tmp_max2 else tmp_max2
a4_3 = tmp3 if tmp3 < tmp_max3 else tmp_max3
else:
a4_2 = a4_2
elif ext6:
if ext5[0, 0, -1] or ext5[0, 0, 1]:
a4_2 = tmp2 if tmp2 < tmp_max2 else tmp_max2
a4_3 = tmp3 if tmp3 < tmp_max3 else tmp_max3
else:
a4_2 = a4_2
else:
a4_2 = a4_2
a4_4 = 3.0 * (2.0 * a4_1 - (a4_2 + a4_3))
# remap_constraint
if __INLINED(iv == 0):
a4_1, a4_2, a4_3, a4_4 = posdef_constraint_iv0(a4_1, a4_2, a4_3, a4_4)
# set_bottom_as_iv0, set_bottom_as_iv1
# TODO(rheag) temporary workaround to gtc:gt:gpu bug
# this computation can get out of order with the one that follows
with computation(FORWARD), interval(-1, None):
if __INLINED(iv == 0):
if a4_3 < 0.0:
a4_3 = 0.0
if __INLINED(iv == -1):
if a4_3 * a4_1 <= 0.0:
a4_3 = 0.0
with computation(PARALLEL), interval(-2, None):
# set_bottom_as_iv0, set_bottom_as_iv1, set_bottom_as_else
a4_4 = 3.0 * (2.0 * a4_1 - (a4_2 + a4_3))
with computation(FORWARD):
with interval(-2, -1):
a4_1, a4_2, a4_3, a4_4 = remap_constraint(a4_1, a4_2, a4_3, a4_4, extm)
with interval(-1, None):
a4_1, a4_2, a4_3, a4_4 = posdef_constraint_iv1(a4_1, a4_2, a4_3, a4_4)
def definition(inout_field: gtscript.Field[float]):
with computation(PARALLEL), interval(...):
assert __INLINED(inout_field[0, 0, 0] < 0), "An error occurred"
def moist_cv_pt_pressure(
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,
pe: FloatField,
pe2: FloatField,
ak: FloatFieldK,
bk: FloatFieldK,
dp2: FloatField,
ps: FloatFieldIJ,
pn2: FloatField,
peln: FloatField,
r_vir: float,
):
from __externals__ import hydrostatic, kord_tm
# moist_cv.moist_pt
with computation(PARALLEL), interval(0, -1):
if __INLINED(kord_tm < 0):
cvm, gz, q_con, cappa, pt = moist_pt_func(
qvapor,
qliquid,
qrain,
qsnow,
qice,
qgraupel,
q_con,
gz,
cvm,
pt,
cappa,
delp,
delz,
r_vir,
)
# delz_adjust
if __INLINED(not hydrostatic):
delz = -delz / delp
# pressure_updates
with computation(FORWARD):
with interval(-1, None):
ps = pe
with computation(PARALLEL):
with interval(0, 1):
pn2 = peln
with interval(1, -1):
pe2 = ak + bk * ps
with interval(-1, None):
pn2 = peln
with computation(BACKWARD), interval(0, -1):
dp2 = pe2[0, 0, 1] - pe2
# copy_stencil
with computation(PARALLEL), interval(0, -1):
delp = dp2
