Thanks to visit codestin.com
Credit goes to github.com

Skip to content
This repository was archived by the owner on Dec 22, 2022. It is now read-only.

Merge/fxadv #525

Draft
oelbert wants to merge 5 commits into
base: master
Choose a base branch
from
Draft

Merge/fxadv #525

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
005686e
initial commit
Jun 4, 2021
70498d8
trying to merge fxadv
Jun 4, 2021
a7c4554
working fxadv
Jun 4, 2021
b783540
removing commented code
Jun 4, 2021
0e14383
more clarifying code
Jun 4, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
111 changes: 50 additions & 61 deletions fv3core/stencils/c_sw.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ def nonhydro_y_fluxes(delp: FloatField, pt: FloatField, w: FloatField, vtc: Floa
return fy, fy1, fy2


def transportdelp(
def transportdelp_update_vorticity_and_kineticenergy(
delp: FloatField,
pt: FloatField,
utc: FloatField,
Expand All @@ -79,8 +79,25 @@ def transportdelp(
delpc: FloatField,
ptc: FloatField,
wc: FloatField,
ke: FloatField,
vort: FloatField,
ua: FloatField,
va: FloatField,
uc: FloatField,
vc: FloatField,
u: FloatField,
v: FloatField,
sin_sg1: FloatFieldIJ,
cos_sg1: FloatFieldIJ,
sin_sg2: FloatFieldIJ,
cos_sg2: FloatFieldIJ,
sin_sg3: FloatFieldIJ,
cos_sg3: FloatFieldIJ,
sin_sg4: FloatFieldIJ,
cos_sg4: FloatFieldIJ,
dt2: float,
):
"""Transport delp.
"""Transport delp then update vorticity and kinetic energy

Args:
delp: What is transported (input)
Expand All @@ -92,11 +109,18 @@ def transportdelp(
delpc: Updated delp (output)
ptc: Updated pt (output)
wc: Updated w (output)
ke: kinetic energy (inout)
vort: vorticity (inout)
ua/uc/u: u wind on the a/c/d grid (in)
va/vc/v: v wind on the a/c/d grid (in)
sin_sg/cos_sg 1/2/3/4: variables that specify grid geometry grid (in)
dt2: length of half a timestep (in)
"""

from __externals__ import grid_type
from __externals__ import grid_type, i_end, i_start, j_end, j_start

with computation(PARALLEL), interval(...):
# transport delP
compile_assert(grid_type < 3)
# additional assumption (not grid.nested)

Expand All @@ -116,6 +140,25 @@ def transportdelp(
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

with computation(PARALLEL), interval(...):
# update vorticity and kinetic energy
compile_assert(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 divergence_corner(
u: FloatField,
Expand Down Expand Up @@ -214,46 +257,6 @@ def circulation_cgrid(
vort_c -= fy[0, 0, 0]


def update_vorticity_and_kinetic_energy(
ke: FloatField,
vort: FloatField,
ua: FloatField,
va: FloatField,
uc: FloatField,
vc: FloatField,
u: FloatField,
v: FloatField,
sin_sg1: FloatFieldIJ,
cos_sg1: FloatFieldIJ,
sin_sg2: FloatFieldIJ,
cos_sg2: FloatFieldIJ,
sin_sg3: FloatFieldIJ,
cos_sg3: FloatFieldIJ,
sin_sg4: FloatFieldIJ,
cos_sg4: FloatFieldIJ,
dt2: float,
):
from __externals__ import grid_type, i_end, i_start, j_end, j_start

with computation(PARALLEL), interval(...):
compile_assert(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 update_x_velocity(
vorticity: FloatField,
ke: FloatField,
Expand Down Expand Up @@ -369,27 +372,15 @@ def __init__(self, grid, namelist):
ax_offsets_transportdelp = axis_offsets(
self.grid, geo_origin, domain_transportdelp
)
self._transportdelp = FrozenStencil(
func=transportdelp,
self._transportdelp_updatevorticity_and_ke = FrozenStencil(
func=transportdelp_update_vorticity_and_kineticenergy,
externals={
"grid_type": grid_type,
**ax_offsets_transportdelp,
},
origin=geo_origin,
domain=(self.grid.nic + 2, self.grid.njc + 2, self.grid.npz),
)
origin_vort_ke = (self.grid.is_ - 1, self.grid.js - 1, 0)
domain_vort_ke = (self.grid.nic + 2, self.grid.njc + 2, self.grid.npz)
ax_offsets_vort_ke = axis_offsets(self.grid, origin_vort_ke, domain_vort_ke)
self._update_vorticity_and_kinetic_energy = FrozenStencil(
func=update_vorticity_and_kinetic_energy,
externals={
"grid_type": grid_type,
**ax_offsets_vort_ke,
},
origin=origin_vort_ke,
domain=domain_vort_ke,
)
self._circulation_cgrid = FrozenStencil(
func=circulation_cgrid,
externals={
Expand All @@ -400,7 +391,7 @@ def __init__(self, grid, namelist):
)
self._absolute_vorticity = FrozenStencil(
func=absolute_vorticity,
origin=self.grid.compute_origin(),
origin=origin,
domain=(self.grid.nic + 1, self.grid.njc + 1, self.grid.npz),
)

Expand Down Expand Up @@ -546,7 +537,7 @@ def __call__(
self.grid.sin_sg2,
dt2,
)
self._transportdelp(
self._transportdelp_updatevorticity_and_ke(
delp,
pt,
ut,
Expand All @@ -556,8 +547,6 @@ def __call__(
self.delpc,
self.ptc,
omga,
)
self._update_vorticity_and_kinetic_energy(
self._ke,
self._vort,
ua,
Expand Down
Loading