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 parallel(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 parallel(region[:, j_start], region[:, j_end + 1]):
vb = dt5 * (vt[-1, 0, 0] + vt)
def fill_4corners_x(q: sd):
from __externals__ import i_end, i_start, j_end, j_start
# copy field
q_out = q
# Southwest
with parallel(region[i_start - 2, j_start - 1]):
q_out = q[1, 2, 0]
with parallel(region[i_start - 1, j_start - 1]):
q_out = q[0, 1, 0]
# Southeast
with parallel(region[i_end + 2, j_start - 1]):
q_out = q[-1, 2, 0]
with parallel(region[i_end + 1, j_start - 1]):
q_out = q[0, 1, 0]
# Northwest
with parallel(region[i_start - 1, j_end + 1]):
q_out = q[0, -1, 0]
with parallel(region[i_start - 2, j_end + 1]):
q_out = q[1, -2, 0]
# Northeast
with parallel(region[i_end + 1, j_end + 1]):
q_out = q[0, -1, 0]
with parallel(region[i_end + 2, j_end + 1]):
q_out = q[-1, -2, 0]
return q_out
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 parallel(region[:, j_start - 1], region[:, j_end]):
vort = vort * sin_sg4 + u[0, 1, 0] * cos_sg4 if va <= 0.0 else vort
with parallel(region[:, j_start], region[:, j_end + 1]):
vort = vort * sin_sg2 + u * cos_sg2 if va > 0.0 else vort
with parallel(region[i_end, :], region[i_start - 1, :]):
ke = ke * sin_sg3 + v[1, 0, 0] * cos_sg3 if ua <= 0.0 else ke
with parallel(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 circulation_cgrid(uc: sd, vc: sd, dxc: sd, dyc: sd, vort_c: sd):
"""Update vort_c.
Args:
uc: x-velocity on C-grid (input)
vc: y-velocity on C-grid (input)
dxc: grid spacing in x-dir (input)
dyc: grid spacing in y-dir (input)
vort_c: C-grid vorticity (output)
"""
from __externals__ import i_end, i_start, j_end, j_start
with computation(PARALLEL), interval(...):
fx = dxc * uc
fy = dyc * vc
vort_c = fx[0, -1, 0] - fx - fy[-1, 0, 0] + fy
with parallel(region[i_start, j_start], region[i_start, j_end + 1]):
vort_c += fy[-1, 0, 0]
with parallel(region[i_end + 1, j_start], region[i_end + 1, j_end + 1]):
vort_c -= fy[0, 0, 0]
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 parallel(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 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 parallel(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 parallel(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 parallel(region[i_start, j_start], region[i_end + 1, j_start]):
divg_d -= vf[0, -1, 0]
with parallel(region[i_end + 1, j_end + 1], region[i_start, j_end + 1]):
divg_d += vf
divg_d *= rarea_c
