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 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 interpolate_uc_vc_to_cell_corners(
uc_cov, vc_cov, cosa, rsina, uc_contra, vc_contra
):
"""
Convert covariant C-grid winds to contravariant B-grid (cell-corner) winds.
"""
from __externals__ import i_end, i_start, j_end, j_start
# In the original Fortran, this routine was given dt4 (0.25 * dt)
# and dt5 (0.5 * dt), and its outputs were wind times timestep. This has
# been refactored so the timestep is later explicitly multiplied, when
# the wind is integrated forward in time.
# TODO: ask Lucas why we interpolate then convert to contravariant in tile center,
# but convert to contravariant and then interpolate on tile edges.
ub_cov = 0.5 * (uc_cov[0, -1, 0] + uc_cov)
vb_cov = 0.5 * (vc_cov[-1, 0, 0] + vc_cov)
ub_contra = contravariant(ub_cov, vb_cov, cosa, rsina)
vb_contra = contravariant(vb_cov, ub_cov, cosa, rsina)
# ASSUME : if __INLINED(namelist.grid_type < 3):
with horizontal(region[:, j_start], region[:, j_end + 1]):
ub_contra = 0.25 * (
-uc_contra[0, -2, 0]
+ 3.0 * (uc_contra[0, -1, 0] + uc_contra)
- uc_contra[0, 1, 0]
)
with horizontal(region[i_start, :], region[i_end + 1, :]):
ub_contra = 0.5 * (uc_contra[0, -1, 0] + uc_contra)
with horizontal(region[i_start, :], region[i_end + 1, :]):
vb_contra = 0.25 * (
-vc_contra[-2, 0, 0]
+ 3.0 * (vc_contra[-1, 0, 0] + vc_contra)
- vc_contra[1, 0, 0]
)
with horizontal(region[:, j_start], region[:, j_end + 1]):
vb_contra = 0.5 * (vc_contra[-1, 0, 0] + vc_contra)
return ub_contra, vb_contra
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 uc_contra_x_edge(uc: FloatField, cosa_u: FloatFieldIJ,
vc_contra: FloatField, uc_contra: FloatField):
from __externals__ import i_end, i_start, j_end, j_start, local_ie, local_is
with computation(PARALLEL), interval(...):
# TODO: parallel to what done for the vt_y_edge section
utmp = uc_contra
with horizontal(
region[local_is:local_ie + 2, j_start - 1:j_start + 1],
region[local_is:local_ie + 2, j_end:j_end + 2],
):
v_contra = 0.25 * (vc_contra[-1, 0, 0] + vc_contra +
vc_contra[-1, 1, 0] + vc_contra[0, 1, 0])
uc_contra = contravariant(uc, v_contra, cosa_u, 1.0)
with horizontal(
region[i_start:i_start + 2, j_start - 1:j_start + 1],
region[i_start:i_start + 2, j_end:j_end + 2],
region[i_end:i_end + 2, j_start - 1:j_start + 1],
region[i_end:i_end + 2, j_end:j_end + 2],
):
uc_contra = utmp
def get_delpc(
u: FloatField,
v: FloatField,
ua: FloatField,
va: FloatField,
cosa_u: FloatFieldIJ,
sina_u: FloatFieldIJ,
dxc: FloatFieldIJ,
dyc: FloatFieldIJ,
uc: FloatField,
vc: FloatField,
sin_sg1: FloatFieldIJ,
sin_sg2: FloatFieldIJ,
sin_sg3: FloatFieldIJ,
sin_sg4: FloatFieldIJ,
cosa_v: FloatFieldIJ,
sina_v: FloatFieldIJ,
rarea_c: FloatFieldIJ,
delpc: FloatField,
u_contra_dyc: FloatField,
v_contra_dxc: FloatField,
):
from __externals__ import i_end, i_start, j_end, j_start
# in the Fortran, u_contra_dyc is called ke and v_contra_dxc is called vort
with computation(PARALLEL), interval(...):
# TODO: why does vc_from_va sometimes have different sign than vc?
vc_from_va = 0.5 * (va[0, -1, 0] + va)
# TODO: why do we use vc_from_va and not just vc?
u_contra = contravariant(u, vc_from_va, cosa_v, sina_v)
with horizontal(region[:, j_start], region[:, j_end + 1]):
u_contra = u * sin_sg4[0, -1] if vc > 0 else u * sin_sg2
u_contra_dyc = u_contra * dyc
# TODO: why does uc_from_ua sometimes have different sign than uc?
uc_from_ua = 0.5 * (ua[-1, 0, 0] + ua)
# TODO: why do we use uc_from_ua and not just uc?
v_contra = contravariant(v, uc_from_ua, cosa_u, sina_u)
with horizontal(region[i_start, :], region[i_end + 1, :]):
v_contra = v * sin_sg3[-1, 0] if uc > 0 else v * sin_sg1
v_contra_dxc = v_contra * dxc
with horizontal(
region[i_start, j_end + 1],
region[i_end + 1, j_end + 1],
region[i_start, j_start - 1],
region[i_end + 1, j_start - 1],
):
# TODO: seems odd that this adjustment is only needed for `v_contra_dxc`
# but is not needed for `u_contra_dyc`. Is this a bug? Add a comment
# describing what this adjustment is doing and why.
v_contra_dxc = 0.0
with computation(PARALLEL), interval(...):
delpc = (
v_contra_dxc[0, -1, 0]
- v_contra_dxc
+ u_contra_dyc[-1, 0, 0]
- u_contra_dyc
)
delpc = rarea_c * delpc