def compute_al(q: FloatField, dya: FloatFieldIJ):
"""
Interpolate q at interface.
Inputs:
q: transported scalar centered along the y-direction
dya: dy on A-grid (?)
Returns:
q interpolated to y-interfaces
"""
from __externals__ import j_end, j_start, jord
compile_assert(jord < 8)
al = ppm.p1 * (q[0, -1, 0] + q) + ppm.p2 * (q[0, -2, 0] + q[0, 1, 0])
if __INLINED(jord < 0):
compile_assert(False)
al = max(al, 0.0)
with horizontal(region[:, j_start - 1], region[:, j_end]):
al = ppm.c1 * q[0, -2, 0] + ppm.c2 * q[0, -1, 0] + ppm.c3 * q
with horizontal(region[:, j_start], region[:, j_end + 1]):
al = 0.5 * (((2.0 * dya[0, -1] + dya[0, -2]) * q[0, -1, 0] -
dya[0, -1] * q[0, -2, 0]) / (dya[0, -2] + dya[0, -1]) +
((2.0 * dya[0, 0] + dya[0, 1]) * q[0, 0, 0] -
dya[0, 0] * q[0, 1, 0]) / (dya[0, 0] + dya[0, 1]))
with horizontal(region[:, j_start + 1], region[:, j_end + 2]):
al = ppm.c3 * q[0, -1, 0] + ppm.c2 * q[0, 0, 0] + ppm.c1 * q[0, 1, 0]
return al
def divergence_corner(
u: sd,
v: sd,
ua: sd,
va: sd,
dxc: sd,
dyc: sd,
sin_sg1: sd,
sin_sg2: sd,
sin_sg3: sd,
sin_sg4: sd,
cos_sg1: sd,
cos_sg2: sd,
cos_sg3: sd,
cos_sg4: sd,
rarea_c: sd,
divg_d: sd,
):
"""Calculate divg on d-grid.
Args:
u: x-velocity (input)
v: y-velocity (input)
ua: x-velocity on a (input)
va: y-velocity on a (input)
dxc: grid spacing in x-direction (input)
dyc: grid spacing in y-direction (input)
sin_sg1: grid sin(sg1) (input)
sin_sg2: grid sin(sg2) (input)
sin_sg3: grid sin(sg3) (input)
sin_sg4: grid sin(sg4) (input)
cos_sg1: grid cos(sg1) (input)
cos_sg2: grid cos(sg2) (input)
cos_sg3: grid cos(sg3) (input)
cos_sg4: grid cos(sg4) (input)
rarea_c: inverse cell areas on c-grid (input)
divg_d: divergence on d-grid (output)
"""
from __externals__ import i_end, i_start, j_end, j_start
with computation(PARALLEL), interval(...):
uf = ((u - 0.25 * (va[0, -1, 0] + va) *
(cos_sg4[0, -1, 0] + cos_sg2)) * dyc * 0.5 *
(sin_sg4[0, -1, 0] + sin_sg2))
with horizontal(region[:, j_start], region[:, j_end + 1]):
uf = u * dyc * 0.5 * (sin_sg4[0, -1, 0] + sin_sg2)
vf = ((v - 0.25 * (ua[-1, 0, 0] + ua) *
(cos_sg3[-1, 0, 0] + cos_sg1)) * dxc * 0.5 *
(sin_sg3[-1, 0, 0] + sin_sg1))
with horizontal(region[i_start, :], region[i_end + 1, :]):
vf = v * dxc * 0.5 * (sin_sg3[-1, 0, 0] + sin_sg1)
divg_d = vf[0, -1, 0] - vf + uf[-1, 0, 0] - uf
with horizontal(region[i_start, j_start], region[i_end + 1, j_start]):
divg_d -= vf[0, -1, 0]
with horizontal(region[i_end + 1, j_end + 1], region[i_start,
j_end + 1]):
divg_d += vf
divg_d *= rarea_c
def final_fluxes(
q_advected_y_x_advected_mean: FloatField,
q_x_advected_mean: FloatField,
q_advected_x_y_advected_mean: FloatField,
q_y_advected_mean: FloatField,
x_unit_flux: FloatField,
y_unit_flux: FloatField,
x_flux: FloatField,
y_flux: FloatField,
):
"""
Compute final x and y fluxes of q from different numerical representations.
Corresponds roughly to eq. 4.17 of FV3 documentation, except that the flux
is in units of q rather than in units of q per interface area per time.
This corresponds to eq 4.17 with both sides multiplied by
e.g. x_unit_flux / u^* (similarly for y/v).
Combining the advection operators in this way is done to cancel leading-order
numerical splitting error.
"""
with computation(PARALLEL), interval(...):
with horizontal(region[:, :-1]):
x_flux = (0.5 *
(q_advected_y_x_advected_mean + q_x_advected_mean) *
x_unit_flux)
with horizontal(region[:-1, :]):
y_flux = (0.5 *
(q_advected_x_y_advected_mean + q_y_advected_mean) *
y_unit_flux)
def fill_corners_3cells_mult_y(
q: FloatField,
q_corner: FloatField,
sw_mult: float,
se_mult: float,
nw_mult: float,
ne_mult: float,
):
"""
Fills cell quantity q using corners from q_corner and multipliers in y-dir.
"""
from __externals__ import i_end, i_start, j_end, j_start
q = fill_corners_2cells_mult_y(q, q_corner, sw_mult, se_mult, nw_mult,
ne_mult)
# Southwest
with horizontal(region[i_start - 1, j_start - 3]):
q = sw_mult * q_corner[3, 2, 0]
# Southeast
with horizontal(region[i_end + 1, j_start - 3]):
q = se_mult * q_corner[-3, 2, 0]
# Northwest
with horizontal(region[i_start - 1, j_end + 3]):
q = nw_mult * q_corner[3, -2, 0]
# Northeast
with horizontal(region[i_end + 1, j_end + 3]):
q = ne_mult * q_corner[-3, -2, 0]
return q
def vc_contra_y_edge(vc: FloatField, cosa_v: FloatFieldIJ,
uc_contra: FloatField, vc_contra: FloatField):
from __externals__ import i_end, i_start, j_end, j_start, local_je, local_js
# This works for 6 ranks, but not 54:
# with horizontal(region[i_start - 1: i_start + 1, j_start + 2:j_end], \
# region[i_end : i_end + 2, j_start+2:j_end]):
# vt = vc - 0.25 * cosa_v * (
# ut[0, -1, 0] + ut[1, -1, 0] + ut + ut[1, 0, 0]
# )
# original bounds with stencil calls
# j1 = grid().js + 2 if grid().south_edge else grid().js
# j2 = grid().je if grid().north_edge else grid().je + 2
# TODO: this is a hack, copying vt to vtmp to 'correct' the edges
# Can we *just* apply edge calculations in the correct regions without overcomputing
# rank 0, 1, 2: local_js + 2:local_je + 2
# rank 3, 4, 5: local_js:local_je + 2
# rank 6, 7, 8: local_js:local_je
with computation(PARALLEL), interval(...):
vtmp = vc_contra
with horizontal(
region[i_start - 1:i_start + 1, local_js:local_je + 2],
region[i_end:i_end + 2, local_js:local_je + 2],
):
u_contra = 0.25 * (uc_contra[0, -1, 0] + uc_contra[1, -1, 0] +
uc_contra + uc_contra[1, 0, 0])
vc_contra = contravariant(vc, u_contra, cosa_v, 1.0)
with horizontal(
region[i_start - 1:i_start + 1, j_start:j_start + 2],
region[i_end:i_end + 2, j_start:j_start + 2],
region[i_start - 1:i_start + 1, j_end:j_end + 2],
region[i_end:i_end + 2, j_end:j_end + 2],
):
vc_contra = vtmp
def update_u_and_v(
ut: FloatField,
vt: FloatField,
u: FloatField,
v: FloatField,
damp_vt: FloatFieldK,
):
"""
Updates u and v after calculation of heat source from vorticity damping.
Args:
ut (in)
vt (in)
u (in/out)
v (in/out)
damp_vt (in): column scalar for damping vorticity
"""
from __externals__ import local_ie, local_is, local_je, local_js
with computation(PARALLEL), interval(...):
if damp_vt > 1e-5:
with horizontal(region[local_is : local_ie + 1, local_js : local_je + 2]):
u += vt
with horizontal(region[local_is : local_ie + 2, local_js : local_je + 1]):
v -= ut
def main_uc_contra(
uc: FloatField,
vc: FloatField,
cosa_u: FloatFieldIJ,
rsin_u: FloatFieldIJ,
uc_contra: FloatField,
):
"""
Args:
uc: covariant c-grid x-wind (in)
vc: covariant c-grid y-wind (in)
cosa_u: ??? (in)
rsin_u: ??? (in)
uc_contra: contravariant c-grid x-wind (out)
"""
from __externals__ import j_end, j_start, local_ie, local_is
with computation(PARALLEL), interval(...):
utmp = uc_contra
with horizontal(region[local_is - 1:local_ie + 3, :]):
# for C-grid, v must be regridded to lie at the same point as u
v = 0.25 * (vc[-1, 0, 0] + vc + vc[-1, 1, 0] + vc[0, 1, 0])
uc_contra = contravariant(uc, v, cosa_u, rsin_u)
with horizontal(region[:, j_start - 1:j_start + 1],
region[:, j_end:j_end + 2]):
uc_contra = utmp
def stencil(in_f: gtscript.Field[np.float_]):
with computation(PARALLEL), interval(...):
in_f = in_f + 1.0
with horizontal(region[I[0], :]):
in_f = 1.0
with horizontal(region[:, J[-1]]):
in_f = 2.0
def ord4_transform(
u: FloatField,
v: FloatField,
a11: FloatField,
a12: FloatField,
a21: FloatField,
a22: FloatField,
dx: FloatField,
dy: FloatField,
ua: FloatField,
va: FloatField,
):
with computation(PARALLEL), interval(...):
from __externals__ import i_end, i_start, j_end, j_start
utmp = C2 * (u[0, -1, 0] + u[0, 2, 0]) + C1 * (u + u[0, 1, 0])
vtmp = C2 * (v[-1, 0, 0] + v[2, 0, 0]) + C1 * (v + v[1, 0, 0])
# south/north edge
with horizontal(region[:, j_start], region[:, j_end]):
vtmp = 2.0 * ((v * dy) +
(v[1, 0, 0] * dy[1, 0, 0])) / (dy + dy[1, 0, 0])
utmp = 2.0 * (u * dx + u[0, 1, 0] * dx[0, 1, 0]) / (dx +
dx[0, 1, 0])
# west/east edge
with horizontal(region[i_start, :], region[i_end, :]):
utmp = 2.0 * ((u * dx) +
(u[0, 1, 0] * dx[0, 1, 0])) / (dx + dx[0, 1, 0])
vtmp = 2.0 * ((v * dy) +
(v[1, 0, 0] * dy[1, 0, 0])) / (dy + dy[1, 0, 0])
# Transform local a-grid winds into latitude-longitude coordinates
ua = a11 * utmp + a12 * vtmp
va = a21 * utmp + a22 * vtmp
def a2b_interpolation(
tmp_qout_edges: FloatField,
qout: FloatField,
qx: FloatField,
qy: FloatField,
):
from __externals__ import i_end, i_start, j_end, j_start
with computation(PARALLEL), interval(...):
qxx = a2 * (qx[0, -2, 0] + qx[0, 1, 0]) + a1 * (qx[0, -1, 0] + qx)
qyy = a2 * (qy[-2, 0, 0] + qy[1, 0, 0]) + a1 * (qy[-1, 0, 0] + qy)
# TODO(rheag) use a function with an offset when that works consistently
with horizontal(region[:, j_start + 1]):
qxx_upper = a2 * (qx[0, -1, 0] + qx[0, 2, 0]) + a1 * (qx +
qx[0, 1, 0])
qxx = c1 * (qx[0, -1, 0] + qx) + c2 * (tmp_qout_edges[0, -1, 0] +
qxx_upper)
with horizontal(region[:, j_end]):
qxx_lower = a2 * (qx[0, -3, 0] + qx) + a1 * (qx[0, -2, 0] +
qx[0, -1, 0])
qxx = c1 * (qx[0, -1, 0] + qx) + c2 * (tmp_qout_edges[0, 1, 0] +
qxx_lower)
with horizontal(region[i_start + 1, :]):
qyy_right = a2 * (qy[-1, 0, 0] + qy[2, 0, 0]) + a1 * (qy +
qy[1, 0, 0])
qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * (tmp_qout_edges[-1, 0, 0] +
qyy_right)
with horizontal(region[i_end, :]):
qyy_left = a2 * (qy[-3, 0, 0] + qy) + a1 * (qy[-2, 0, 0] +
qy[-1, 0, 0])
qyy = c1 * (qy[-1, 0, 0] + qy) + c2 * (tmp_qout_edges[1, 0, 0] +
qyy_left)
qout = 0.5 * (qxx + qyy)
def compute_al(q: FloatField, dxa: FloatFieldIJ):
"""
Interpolate q at interface.
Inputs:
q: transported scalar centered along the x-direction
dxa: dx on A-grid (?)
Returns:
q interpolated to x-interfaces
"""
from __externals__ import i_end, i_start, iord
compile_assert(iord < 8)
al = ppm.p1 * (q[-1, 0, 0] + q) + ppm.p2 * (q[-2, 0, 0] + q[1, 0, 0])
if __INLINED(iord < 0):
compile_assert(False)
al = max(al, 0.0)
with horizontal(region[i_start - 1, :], region[i_end, :]):
al = ppm.c1 * q[-2, 0, 0] + ppm.c2 * q[-1, 0, 0] + ppm.c3 * q
with horizontal(region[i_start, :], region[i_end + 1, :]):
al = 0.5 * (((2.0 * dxa[-1, 0] + dxa[-2, 0]) * q[-1, 0, 0] -
dxa[-1, 0] * q[-2, 0, 0]) / (dxa[-2, 0] + dxa[-1, 0]) +
((2.0 * dxa[0, 0] + dxa[1, 0]) * q[0, 0, 0] -
dxa[0, 0] * q[1, 0, 0]) / (dxa[0, 0] + dxa[1, 0]))
with horizontal(region[i_start + 1, :], region[i_end + 2, :]):
al = ppm.c3 * q[-1, 0, 0] + ppm.c2 * q[0, 0, 0] + ppm.c1 * q[1, 0, 0]
return al
def edge_pe(pe: FloatField, delp: FloatField, ptop: float):
"""
This corresponds to the pe_halo routine in FV3core
Updading the interface pressure from the pressure differences
Arguments:
pe: The pressure on the interfaces of the cell
delp: The pressure difference between vertical grid cells
ptop: The pressure level at the top of the grid
"""
from __externals__ import local_ie, local_is, local_je, local_js
with computation(FORWARD):
with interval(0, 1):
with horizontal(
region[local_is - 1, local_js:local_je + 1],
region[local_ie + 1, local_js:local_je + 1],
region[local_is - 1:local_ie + 2, local_js - 1],
region[local_is - 1:local_ie + 2, local_je + 1],
):
pe[0, 0, 0] = ptop
with interval(1, None):
with horizontal(
region[local_is - 1, local_js:local_je + 1],
region[local_ie + 1, local_js:local_je + 1],
region[local_is - 1:local_ie + 2, local_js - 1],
region[local_is - 1:local_ie + 2, local_je + 1],
):
pe[0, 0, 0] = pe[0, 0, -1] + delp[0, 0, -1]
def definition(field_in, field_out):
with computation(PARALLEL), interval(...):
field_out = ( # noqa: F841 # local variable 'field_out' is assigned to but never used
field_in)
with horizontal(region[I[0], :], region[I[-1], :]):
field_out = ( # noqa: F841 # local variable 'field_out' is assigned to but never used
field_in + 1.0)
with horizontal(region[:, J[0]], region[:, J[-1]]):
field_out = ( # noqa: F841 # local variable 'field_out' is assigned to but never used
field_in - 1.0)
def uc_contra_corners(
cosa_u: FloatFieldIJ,
cosa_v: FloatFieldIJ,
uc: FloatField,
vc: FloatField,
uc_contra: FloatField,
uc_contra_copy: FloatField,
vc_contra: FloatField,
):
"""
The following code (and vt_corners) solves a 2x2 system to
get the interior parallel-to-edge uc,vc values near the corners
(ex: for the sw corner ut(2,1) and vt(1,2) are solved for simultaneously).
It then computes the halo uc, vc values so as to be consistent with the
computations on the facing panel.
The system solved is:
ut(2,1) = uc(2,1) - avg(vt)*cosa_u(2,1)
vt(1,2) = vc(1,2) - avg(ut)*cosa_v(1,2)
in which avg(vt) includes vt(1,2) and avg(ut) includes ut(2,1)
"""
from __externals__ import i_end, i_start, j_end, j_start
with computation(PARALLEL), interval(...):
damp = 1.0 / (1.0 - 0.0625 * cosa_u * cosa_v[-1, 0])
with horizontal(region[i_start + 1, j_start - 1], region[i_start + 1,
j_end]):
uc_contra = (uc - 0.25 * cosa_u *
(vc_contra[-1, 1, 0] + vc_contra[0, 1, 0] +
vc_contra + vc[-1, 0, 0] - 0.25 * cosa_v[-1, 0] *
(uc_contra_copy[-1, 0, 0] + uc_contra_copy[-1, -1, 0]
+ uc_contra_copy[0, -1, 0]))) * damp
damp = 1.0 / (1.0 - 0.0625 * cosa_u * cosa_v[-1, 1])
with horizontal(region[i_start + 1, j_start], region[i_start + 1,
j_end + 1]):
damp = 1.0 / (1.0 - 0.0625 * cosa_u * cosa_v[-1, 1])
uc_contra = (uc - 0.25 * cosa_u *
(vc_contra[-1, 0, 0] + vc_contra + vc_contra[0, 1, 0]
+ vc[-1, 1, 0] - 0.25 * cosa_v[-1, 1] *
(uc_contra_copy[-1, 0, 0] + uc_contra_copy[-1, 1, 0]
+ uc_contra_copy[0, 1, 0]))) * damp
damp = 1.0 / (1.0 - 0.0625 * cosa_u * cosa_v)
with horizontal(region[i_end, j_start - 1], region[i_end, j_end]):
uc_contra = (uc - 0.25 * cosa_u *
(vc_contra[0, 1, 0] + vc_contra[-1, 1, 0] +
vc_contra[-1, 0, 0] + vc - 0.25 * cosa_v *
(uc_contra_copy[1, 0, 0] + uc_contra_copy[1, -1, 0] +
uc_contra_copy[0, -1, 0]))) * damp
damp = 1.0 / (1.0 - 0.0625 * cosa_u * cosa_v[0, 1])
with horizontal(region[i_end, j_start], region[i_end, j_end + 1]):
uc_contra = (uc - 0.25 * cosa_u *
(vc_contra + vc_contra[-1, 0, 0] + vc_contra[-1, 1, 0]
+ vc[0, 1, 0] - 0.25 * cosa_v[0, 1] *
(uc_contra_copy[1, 0, 0] + uc_contra_copy[1, 1, 0] +
uc_contra_copy[0, 1, 0]))) * damp
def vbke(vc: sd, uc: sd, cosa: sd, rsina: sd, vt: sd, vb: sd, dt4: float,
dt5: float):
from __externals__ import i_end, i_start, j_end, j_start
with computation(PARALLEL), interval(...):
vb = dt5 * (vc[-1, 0, 0] + vc - (uc[0, -1, 0] + uc) * cosa) * rsina
if __INLINED(spec.namelist.grid_type < 3):
with horizontal(region[i_start, :], region[i_end + 1, :]):
vb = dt4 * (-vt[-2, 0, 0] + 3.0 *
(vt[-1, 0, 0] + vt) - vt[1, 0, 0])
with horizontal(region[:, j_start], region[:, j_end + 1]):
vb = dt5 * (vt[-1, 0, 0] + vt)
def fill_4corners_y(q: sd):
from __externals__ import i_end, i_start, j_end, j_start
# copy field
q_out = q
# Southwest
with horizontal(region[i_start - 1, j_start - 1]):
q_out = q[1, 0, 0]
with horizontal(region[i_start - 1, j_start - 2]):
q_out = q[2, 1, 0]
# Southeast
with horizontal(region[i_end + 1, j_start - 1]):
q_out = q[-1, 0, 0]
with horizontal(region[i_end + 1, j_start - 2]):
q_out = q[-2, 1, 0]
# Northwest
with horizontal(region[i_start - 1, j_end + 1]):
q_out = q[1, 0, 0]
with horizontal(region[i_start - 1, j_end + 2]):
q_out = q[2, -1, 0]
# Northeast
with horizontal(region[i_end + 1, j_end + 1]):
q_out = q[-1, 0, 0]
with horizontal(region[i_end + 1, j_end + 2]):
q_out = q[-2, -1, 0]
return q_out
def all_corners_ke(ke, u, v, ut, vt, dt):
from __externals__ import i_end, i_start, j_end, j_start
# Assumption: not __INLINED(grid.nested)
with horizontal(region[i_start, j_start]):
ke = corner_ke(u, v, ut, vt, dt, 0, 0, -1, 1)
with horizontal(region[i_end + 1, j_start]):
ke = corner_ke(u, v, ut, vt, dt, -1, 0, 0, -1)
with horizontal(region[i_end + 1, j_end + 1]):
ke = corner_ke(u, v, ut, vt, dt, -1, -1, 0, 1)
with horizontal(region[i_start, j_end + 1]):
ke = corner_ke(u, v, ut, vt, dt, 0, -1, -1, -1)
return ke
def east_west_edges(
u: FloatField,
ua: FloatField,
uc: FloatField,
utc: FloatField,
utmp: FloatField,
v: FloatField,
sin_sg1: FloatFieldIJ,
sin_sg3: FloatFieldIJ,
cosa_u: FloatFieldIJ,
rsin_u: FloatFieldIJ,
dxa: FloatFieldIJ,
):
from __externals__ import i_end, i_start, local_je, local_js
with computation(PARALLEL), interval(...):
# West
with horizontal(region[i_start - 1, local_js - 1:local_je + 2]):
uc = vol_conserv_cubic_interp_func_x(utmp)
with horizontal(region[i_start, local_js - 1:local_je + 2]):
utc = edge_interpolate4_x(ua, dxa)
uc = utc * sin_sg3[-1, 0] if utc > 0 else utc * sin_sg1
with horizontal(region[i_start + 1, local_js - 1:local_je + 2]):
uc = vol_conserv_cubic_interp_func_x_rev(utmp)
with horizontal(region[i_start - 1, local_js - 1:local_je + 2]):
utc = contravariant(uc, v, cosa_u, rsin_u)
with horizontal(region[i_start + 1, local_js - 1:local_je + 2]):
utc = contravariant(uc, v, cosa_u, rsin_u)
# East
with horizontal(region[i_end, local_js - 1:local_je + 2]):
uc = vol_conserv_cubic_interp_func_x(utmp)
with horizontal(region[i_end + 1, local_js - 1:local_je + 2]):
utc = edge_interpolate4_x(ua, dxa)
uc = utc * sin_sg3[-1, 0] if utc > 0 else utc * sin_sg1
with horizontal(region[i_end + 2, local_js - 1:local_je + 2]):
uc = vol_conserv_cubic_interp_func_x_rev(utmp)
with horizontal(region[i_end, local_js - 1:local_je + 2]):
utc = contravariant(uc, v, cosa_u, rsin_u)
with horizontal(region[i_end + 2, local_js - 1:local_je + 2]):
utc = contravariant(uc, v, cosa_u, rsin_u)
def stencil(in_f: gtscript.Field[np.float_]):
from gt4py.__externals__ import i0 # forget to add 'ia'
with computation(PARALLEL), interval(...):
in_f = in_f + 1.0
with horizontal(region[i0 : 1 + ia, :]):
in_f = 1.0
def add_1_in_region_stencil(q_in: FloatField, q_out: FloatField):
from __externals__ import i_start
with computation(PARALLEL), interval(...):
q_out = q_in
with horizontal(region[i_start, :]):
q_out = q_in + 1.0
def redo_divg_d(
uc: FloatField,
vc: FloatField,
divg_d: FloatField,
adjustment_factor: FloatFieldIJ,
):
from __externals__ import do_adjustment, i_end, i_start, j_end, j_start
with computation(PARALLEL), interval(...):
divg_d = uc[0, -1, 0] - uc + vc[-1, 0, 0] - vc
with horizontal(region[i_start, j_start], region[i_end + 1, j_start]):
divg_d = vc[-1, 0, 0] - vc - uc
with horizontal(region[i_start, j_end + 1], region[i_end + 1, j_end + 1]):
divg_d = uc[0, -1, 0] + vc[-1, 0, 0] - vc
if __INLINED(do_adjustment):
divg_d = divg_d * adjustment_factor
def ppm_volume_mean_x(
qin: FloatField,
qx: FloatField,
dxa: FloatFieldIJ,
):
from __externals__ import i_end, i_start
with computation(PARALLEL), interval(...):
qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) + b1 * (qin[-1, 0, 0] + qin)
with horizontal(region[i_start, :]):
qx = qx_edge_west(qin, dxa)
with horizontal(region[i_start + 1, :]):
qx = qx_edge_west2(qin, dxa)
with horizontal(region[i_end + 1, :]):
qx = qx_edge_east(qin, dxa)
with horizontal(region[i_end, :]):
qx = qx_edge_east2(qin, dxa)
def main_vc_contra(
uc: FloatField,
vc: FloatField,
cosa_v: FloatFieldIJ,
rsin_v: FloatFieldIJ,
vc_contra: FloatField,
):
from __externals__ import j_end, j_start, local_je, local_js
with computation(PARALLEL), interval(...):
vtmp = vc_contra
with horizontal(region[:, local_js - 1:local_je + 3]):
# for C-grid, u must be regridded to lie at same point as v
u = 0.25 * (uc[0, -1, 0] + uc[1, -1, 0] + uc + uc[1, 0, 0])
vc_contra = contravariant(vc, u, cosa_v, rsin_v)
with horizontal(region[:, j_start], region[:, j_end + 1]):
vc_contra = vtmp
def region_func():
from gt4py.__externals__ import ie
field = 0.0
with horizontal(region[ie, :]):
field = 1.0
return field
def flux_x(cx, dxa, dy, sin_sg3, sin_sg1, xfx):
from __externals__ import local_ie, local_is, local_je, local_js
with horizontal(region[local_is : local_ie + 2, local_js - 3 : local_je + 4]):
xfx = (
cx * dxa[-1, 0] * dy * sin_sg3[-1, 0] if cx > 0 else cx * dxa * dy * sin_sg1
)
return xfx
def vort_differencing(
vort: FloatField,
vort_x_delta: FloatField,
vort_y_delta: FloatField,
dcon: FloatFieldK,
):
from __externals__ import local_ie, local_is, local_je, local_js
with computation(PARALLEL), interval(...):
if dcon[0] > dcon_threshold:
# Creating a gtscript function for the ub/vb computation
# results in an "NotImplementedError" error for Jenkins
# Inlining the ub/vb computation in this stencil resolves the Jenkins error
with horizontal(region[local_is : local_ie + 1, local_js : local_je + 2]):
vort_x_delta = vort - vort[1, 0, 0]
with horizontal(region[local_is : local_ie + 2, local_js : local_je + 1]):
vort_y_delta = vort - vort[0, 1, 0]
def flux_y(cy, dya, dx, sin_sg4, sin_sg2, yfx):
from __externals__ import local_ie, local_is, local_je, local_js
with horizontal(region[local_is - 3 : local_ie + 4, local_js : local_je + 2]):
yfx = (
cy * dya[0, -1] * dx * sin_sg4[0, -1] if cy > 0 else cy * dya * dx * sin_sg2
)
return yfx
def ppm_volume_mean_y(
qin: FloatField,
qy: FloatField,
dya: FloatFieldIJ,
):
from __externals__ import j_end, j_start
with computation(PARALLEL), interval(...):
qy = b2 * (qin[0, -2, 0] + qin[0, 1, 0]) + b1 * (qin[0, -1, 0] + qin)
with horizontal(region[:, j_start]):
qy = qy_edge_south(qin, dya)
with horizontal(region[:, j_start + 1]):
qy = qy_edge_south2(qin, dya)
with horizontal(region[:, j_end + 1]):
qy = qy_edge_north(qin, dya)
with horizontal(region[:, j_end]):
qy = qy_edge_north2(qin, dya)
def copy_stencil_interval(q_in: FloatField, q_out: FloatField):
from __externals__ import (
local_ie,
local_is,
local_je,
local_js,
nord0,
nord1,
nord2,
nord3,
)
with computation(PARALLEL), interval(0, 1):
if __INLINED(nord0 == 0):
with horizontal(
region[local_is - 1 : local_ie + 2, local_js - 1 : local_je + 2]
):
q_out = q_in
else:
q_out = q_in
with computation(PARALLEL), interval(1, 2):
if __INLINED(nord1 == 0):
with horizontal(
region[local_is - 1 : local_ie + 2, local_js - 1 : local_je + 2]
):
q_out = q_in
else:
q_out = q_in
with computation(PARALLEL), interval(2, 3):
if __INLINED(nord2 == 0):
with horizontal(
region[local_is - 1 : local_ie + 2, local_js - 1 : local_je + 2]
):
q_out = q_in
else:
q_out = q_in
with computation(PARALLEL), interval(3, None):
if __INLINED(nord3 == 0):
with horizontal(
region[local_is - 1 : local_ie + 2, local_js - 1 : local_je + 2]
):
q_out = q_in
else:
q_out = q_in
def d2_damp_interval(q: FloatField, d2: FloatField, damp: FloatFieldK):
from __externals__ import (
local_ie,
local_is,
local_je,
local_js,
nord0,
nord1,
nord2,
nord3,
)
with computation(PARALLEL), interval(0, 1):
if __INLINED(nord0 == 0):
with horizontal(
region[local_is - 1 : local_ie + 2, local_js - 1 : local_je + 2]
):
d2 = damp * q
else:
d2 = damp * q
with computation(PARALLEL), interval(1, 2):
if __INLINED(nord1 == 0):
with horizontal(
region[local_is - 1 : local_ie + 2, local_js - 1 : local_je + 2]
):
d2 = damp * q
else:
d2 = damp * q
with computation(PARALLEL), interval(2, 3):
if __INLINED(nord2 == 0):
with horizontal(
region[local_is - 1 : local_ie + 2, local_js - 1 : local_je + 2]
):
d2 = damp * q
else:
d2 = damp * q
with computation(PARALLEL), interval(3, None):
if __INLINED(nord3 == 0):
with horizontal(
region[local_is - 1 : local_ie + 2, local_js - 1 : local_je + 2]
):
d2 = damp * q
else:
d2 = damp * q
