Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Type aliases for Fields #473

Merged
merged 37 commits into from
Jul 29, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
37 commits
Select commit Hold shift + click to select a range
0472a13
introduction of fields type aliases and partial implementation
nfarabullini May 30, 2024
7c3ee6c
fields type aliases partial implementation
nfarabullini May 30, 2024
80f8dbf
Field[[EdgeDim, KDim], wpfloat] -> EKwpField
nfarabullini May 31, 2024
33bbafb
Field[[KDim], int32] -> KintField
nfarabullini May 31, 2024
9381bdc
Field[[EdgeDim, KDim], bool] -> EKboolField
nfarabullini May 31, 2024
d46cfc2
Field[[CellDim], bool] -> CboolField
nfarabullini May 31, 2024
4dc27fa
Field[[EdgeDim, KDim], int32] -> EKintField
nfarabullini May 31, 2024
85dab66
Field[[EdgeDim, KDim], int32] -> EKintField
nfarabullini May 31, 2024
88b8fb2
Field[[CellDim], int32] -> CintField
nfarabullini May 31, 2024
01df8d8
Field[[CellDim, KDim], vpfloat] -> CKvpField
nfarabullini May 31, 2024
8140c47
left-overs
nfarabullini May 31, 2024
c4f181e
Field[[KDim], wpfloat] -> KwpField
nfarabullini May 31, 2024
675f1da
small edit
nfarabullini May 31, 2024
4011992
Field[[CellDim, KDim], float] -> CKfloatField and file rename
nfarabullini May 31, 2024
17030ff
Field[[EdgeDim, KDim], float] -> EKfloatField and file rename
nfarabullini May 31, 2024
e9696ef
Field[[CellDim], float] -> CfloatField
nfarabullini May 31, 2024
5a5d2a6
Field[[EdgeDim], float] -> EfloatField and Field[[KDim], float] -> Kf…
nfarabullini May 31, 2024
db641cd
Field[[EdgeDim, KDim], vpfloat] -> EKvpField and Field[[VertexDim, KD…
nfarabullini May 31, 2024
efeb00b
Field[[EdgeDim, KDim], vpfloat] -> EKvpField and Field[[VertexDim, KD…
nfarabullini May 31, 2024
bb56208
Field[[CellDim, KDim], bool] -> CKboolField and Field[[CellDim, KDim]…
nfarabullini May 31, 2024
f2abd94
moved field_type_aliases location
nfarabullini May 31, 2024
c284ddb
left over path edit
nfarabullini May 31, 2024
131aafe
left over path edit
nfarabullini May 31, 2024
8313ca3
fixed path edits
nfarabullini May 31, 2024
b3d6262
edits following review
nfarabullini Jun 10, 2024
597f8df
Merge branch 'main' into introduce_type_alias_fields
nfarabullini Jun 25, 2024
03312c4
Merge branch 'main' of https://github.com/C2SM/icon4py into introduce…
nfarabullini Jun 27, 2024
5abd0f8
ran pre-commit
nfarabullini Jun 27, 2024
21b3003
Merge branch 'main' into introduce_type_alias_fields
nfarabullini Jul 3, 2024
27695c9
ran pre-commit
nfarabullini Jul 3, 2024
d1dad2f
update with upstream
nfarabullini Jul 10, 2024
912e10f
additional field type aliases edits
nfarabullini Jul 10, 2024
7d48ee8
Merge branch 'main' into introduce_type_alias_fields
nfarabullini Jul 24, 2024
ae4a7b5
Merge branch 'main' into introduce_type_alias_fields
nfarabullini Jul 24, 2024
1fa22b3
Update model/common/src/icon4py/model/common/states/prognostic_state.py
nfarabullini Jul 29, 2024
8803e4d
Update model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/co…
nfarabullini Jul 29, 2024
13ad369
update with upstream
nfarabullini Jul 29, 2024
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
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
Loading