Skip to content

Commit

Permalink
Type aliases for Fields (#473)
Browse files Browse the repository at this point in the history
Introduction of Type Aliases for Fields of global dimensions

CellField -> Field[[CellDim], dtype]
EdgeField -> Field[[EdgeDim], dtype]
VertexField -> Field[[VertexDim], dtype]
KField -> Field[[KDim], dtype]

CellKField -> Field[[CellDim, KDim], dtype]
EdgeKField -> Field[[EdgeDim, KDim], dtype]
VertexKField -> Field[[VertexDim, KDim], dtype]
  • Loading branch information
nfarabullini authored Jul 29, 2024
1 parent b4e9967 commit 3d1e928
Show file tree
Hide file tree
Showing 184 changed files with 3,227 additions and 3,056 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -12,26 +12,27 @@
# SPDX-License-Identifier: GPL-3.0-or-later
from gt4py.next import GridType
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import Field, broadcast, where
from gt4py.next.ffront.fbuiltins import broadcast, where

from icon4py.model.common import field_type_aliases as fa
from icon4py.model.common.dimension import EdgeDim, KDim


@field_operator
def _btraj_dreg_stencil_01(
lcounterclock: bool,
p_vn: Field[[EdgeDim, KDim], float],
tangent_orientation: Field[[EdgeDim], float],
) -> Field[[EdgeDim, KDim], bool]:
p_vn: fa.EdgeKField[float],
tangent_orientation: fa.EdgeField[float],
) -> fa.EdgeKField[bool]:
tangent_orientation = broadcast(tangent_orientation, (EdgeDim, KDim))
return where(p_vn * tangent_orientation >= 0.0, lcounterclock, False)


@program(grid_type=GridType.UNSTRUCTURED)
def btraj_dreg_stencil_01(
lcounterclock: bool,
p_vn: Field[[EdgeDim, KDim], float],
tangent_orientation: Field[[EdgeDim], float],
lvn_sys_pos: Field[[EdgeDim, KDim], bool],
p_vn: fa.EdgeKField[float],
tangent_orientation: fa.EdgeField[float],
lvn_sys_pos: fa.EdgeKField[bool],
):
_btraj_dreg_stencil_01(lcounterclock, p_vn, tangent_orientation, out=lvn_sys_pos)
Original file line number Diff line number Diff line change
Expand Up @@ -14,16 +14,17 @@
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, sqrt, where

from icon4py.model.common import field_type_aliases as fa
from icon4py.model.common.dimension import E2EC, ECDim, EdgeDim, KDim


@field_operator
def _btraj_dreg_stencil_02(
p_vn: Field[[EdgeDim, KDim], float],
p_vt: Field[[EdgeDim, KDim], float],
p_vn: fa.EdgeKField[float],
p_vt: fa.EdgeKField[float],
edge_cell_length: Field[[ECDim], float],
p_dt: float,
) -> Field[[EdgeDim, KDim], int32]:
) -> fa.EdgeKField[int32]:
lvn_pos = where(p_vn >= 0.0, True, False)
traj_length = sqrt(p_vn * p_vn + p_vt * p_vt) * p_dt
e2c_length = where(lvn_pos, edge_cell_length(E2EC[0]), edge_cell_length(E2EC[1]))
Expand All @@ -34,10 +35,10 @@ def _btraj_dreg_stencil_02(

@program(grid_type=GridType.UNSTRUCTURED)
def btraj_dreg_stencil_02(
p_vn: Field[[EdgeDim, KDim], float],
p_vt: Field[[EdgeDim, KDim], float],
p_vn: fa.EdgeKField[float],
p_vt: fa.EdgeKField[float],
edge_cell_length: Field[[ECDim], float],
p_dt: float,
opt_famask_dsl: Field[[EdgeDim, KDim], int32],
opt_famask_dsl: fa.EdgeKField[int32],
):
_btraj_dreg_stencil_02(p_vn, p_vt, edge_cell_length, p_dt, out=opt_famask_dsl)
Original file line number Diff line number Diff line change
Expand Up @@ -14,41 +14,42 @@
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import Field, int32, where

from icon4py.model.common.dimension import E2EC, ECDim, EdgeDim, KDim
from icon4py.model.common import field_type_aliases as fa
from icon4py.model.common.dimension import E2EC, ECDim


@field_operator
def _btraj_dreg_stencil_03(
p_vn: Field[[EdgeDim, KDim], float],
p_vt: Field[[EdgeDim, KDim], float],
p_vn: fa.EdgeKField[float],
p_vt: fa.EdgeKField[float],
cell_idx: Field[[ECDim], int32],
cell_blk: Field[[ECDim], int32],
edge_verts_1_x: Field[[EdgeDim], float],
edge_verts_2_x: Field[[EdgeDim], float],
edge_verts_1_y: Field[[EdgeDim], float],
edge_verts_2_y: Field[[EdgeDim], float],
pos_on_tplane_e_1_x: Field[[EdgeDim], float],
pos_on_tplane_e_2_x: Field[[EdgeDim], float],
pos_on_tplane_e_1_y: Field[[EdgeDim], float],
pos_on_tplane_e_2_y: Field[[EdgeDim], float],
edge_verts_1_x: fa.EdgeField[float],
edge_verts_2_x: fa.EdgeField[float],
edge_verts_1_y: fa.EdgeField[float],
edge_verts_2_y: fa.EdgeField[float],
pos_on_tplane_e_1_x: fa.EdgeField[float],
pos_on_tplane_e_2_x: fa.EdgeField[float],
pos_on_tplane_e_1_y: fa.EdgeField[float],
pos_on_tplane_e_2_y: fa.EdgeField[float],
primal_normal_cell_x: Field[[ECDim], float],
primal_normal_cell_y: Field[[ECDim], float],
dual_normal_cell_x: Field[[ECDim], float],
dual_normal_cell_y: Field[[ECDim], float],
lvn_sys_pos: Field[[EdgeDim, KDim], bool],
lvn_sys_pos: fa.EdgeKField[bool],
p_dt: float,
) -> tuple[
Field[[EdgeDim, KDim], int32],
Field[[EdgeDim, KDim], int32],
Field[[EdgeDim, KDim], int32],
Field[[EdgeDim, KDim], float],
Field[[EdgeDim, KDim], float],
Field[[EdgeDim, KDim], float],
Field[[EdgeDim, KDim], float],
Field[[EdgeDim, KDim], float],
Field[[EdgeDim, KDim], float],
Field[[EdgeDim, KDim], float],
Field[[EdgeDim, KDim], float],
fa.EdgeKField[int32],
fa.EdgeKField[int32],
fa.EdgeKField[int32],
fa.EdgeKField[float],
fa.EdgeKField[float],
fa.EdgeKField[float],
fa.EdgeKField[float],
fa.EdgeKField[float],
fa.EdgeKField[float],
fa.EdgeKField[float],
fa.EdgeKField[float],
]:
# logical switch for MERGE operations: True for p_vn >= 0
lvn_pos = where(p_vn >= 0.0, True, False)
Expand Down Expand Up @@ -130,35 +131,35 @@ def _btraj_dreg_stencil_03(

@program(grid_type=GridType.UNSTRUCTURED)
def btraj_dreg_stencil_03(
p_vn: Field[[EdgeDim, KDim], float],
p_vt: Field[[EdgeDim, KDim], float],
p_vn: fa.EdgeKField[float],
p_vt: fa.EdgeKField[float],
cell_idx: Field[[ECDim], int32],
cell_blk: Field[[ECDim], int32],
edge_verts_1_x: Field[[EdgeDim], float],
edge_verts_2_x: Field[[EdgeDim], float],
edge_verts_1_y: Field[[EdgeDim], float],
edge_verts_2_y: Field[[EdgeDim], float],
pos_on_tplane_e_1_x: Field[[EdgeDim], float],
pos_on_tplane_e_2_x: Field[[EdgeDim], float],
pos_on_tplane_e_1_y: Field[[EdgeDim], float],
pos_on_tplane_e_2_y: Field[[EdgeDim], float],
edge_verts_1_x: fa.EdgeField[float],
edge_verts_2_x: fa.EdgeField[float],
edge_verts_1_y: fa.EdgeField[float],
edge_verts_2_y: fa.EdgeField[float],
pos_on_tplane_e_1_x: fa.EdgeField[float],
pos_on_tplane_e_2_x: fa.EdgeField[float],
pos_on_tplane_e_1_y: fa.EdgeField[float],
pos_on_tplane_e_2_y: fa.EdgeField[float],
primal_normal_cell_x: Field[[ECDim], float],
primal_normal_cell_y: Field[[ECDim], float],
dual_normal_cell_x: Field[[ECDim], float],
dual_normal_cell_y: Field[[ECDim], float],
lvn_sys_pos: Field[[EdgeDim, KDim], bool],
lvn_sys_pos: fa.EdgeKField[bool],
p_dt: float,
p_cell_idx: Field[[EdgeDim, KDim], int32],
p_cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32],
p_cell_blk: Field[[EdgeDim, KDim], int32],
p_coords_dreg_v_1_lon_dsl: Field[[EdgeDim, KDim], float],
p_coords_dreg_v_2_lon_dsl: Field[[EdgeDim, KDim], float],
p_coords_dreg_v_3_lon_dsl: Field[[EdgeDim, KDim], float],
p_coords_dreg_v_4_lon_dsl: Field[[EdgeDim, KDim], float],
p_coords_dreg_v_1_lat_dsl: Field[[EdgeDim, KDim], float],
p_coords_dreg_v_2_lat_dsl: Field[[EdgeDim, KDim], float],
p_coords_dreg_v_3_lat_dsl: Field[[EdgeDim, KDim], float],
p_coords_dreg_v_4_lat_dsl: Field[[EdgeDim, KDim], float],
p_cell_idx: fa.EdgeKField[int32],
p_cell_rel_idx_dsl: fa.EdgeKField[int32],
p_cell_blk: fa.EdgeKField[int32],
p_coords_dreg_v_1_lon_dsl: fa.EdgeKField[float],
p_coords_dreg_v_2_lon_dsl: fa.EdgeKField[float],
p_coords_dreg_v_3_lon_dsl: fa.EdgeKField[float],
p_coords_dreg_v_4_lon_dsl: fa.EdgeKField[float],
p_coords_dreg_v_1_lat_dsl: fa.EdgeKField[float],
p_coords_dreg_v_2_lat_dsl: fa.EdgeKField[float],
p_coords_dreg_v_3_lat_dsl: fa.EdgeKField[float],
p_coords_dreg_v_4_lat_dsl: fa.EdgeKField[float],
):
_btraj_dreg_stencil_03(
p_vn,
Expand Down
Loading

0 comments on commit 3d1e928

Please sign in to comment.