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