Skip to content

Commit

Permalink
edits following review
Browse files Browse the repository at this point in the history
  • Loading branch information
@nfarabullini
nfarabullini committed Sep 6, 2024
1 parent fcc293f commit 85bdff4
Showing 1 changed file with 99 additions and 65 deletions.
164 changes: 99 additions & 65 deletions model/common/src/icon4py/model/common/metrics/metrics_factory.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
)
from icon4py.model.common import constants, dimension as dims
from icon4py.model.common.decomposition import definitions as decomposition
from icon4py.model.common.grid import horizontal
from icon4py.model.common.grid import horizontal as h_grid, vertical as v_grid
from icon4py.model.common.interpolation.stencils import cell_2_edge_interpolation
from icon4py.model.common.io import cf_utils
from icon4py.model.common.metrics import compute_vwind_impl_wgt, metric_fields as mf
Expand Down Expand Up @@ -49,9 +49,9 @@
grid_id = dt_utils.get_grid_id_for_experiment(dt_utils.REGIONAL_EXPERIMENT)
grid_savepoint = data_provider.from_savepoint_grid(grid_id, root, level)
nlev = grid_savepoint.num(dims.KDim)
cell_domain = horizontal.domain(dims.CellDim)
edge_domain = horizontal.domain(dims.EdgeDim)
vertex_domain = horizontal.domain(dims.VertexDim)
cell_domain = h_grid.domain(dims.CellDim)
edge_domain = h_grid.domain(dims.EdgeDim)
vertex_domain = h_grid.domain(dims.VertexDim)
#######

# start build up factory:
Expand Down Expand Up @@ -116,10 +116,7 @@
height_provider = factory.ProgramFieldProvider(
func=mf.compute_z_mc,
domain={
dims.CellDim: (
horizontal._local(dims.CellDim),
horizontal._end(dims.CellDim),
),
dims.CellDim: (cell_domain(h_grid.Zone.LOCAL), cell_domain(h_grid.Zone.END)),
dims.KDim: (0, nlev),
},
fields={"z_mc": "height"},
Expand All @@ -136,10 +133,13 @@
},
domain={
dims.CellDim: (
icon_grid.start_index(cell_domain(horizontal.Zone.LOCAL)),
cell_domain(horizontal.Zone.LOCAL),
cell_domain(h_grid.Zone.LOCAL),
cell_domain(h_grid.Zone.LOCAL),
),
dims.KDim: (
v_grid.domain(dims.KDim)(v_grid.Zone.TOP),
v_grid.domain(dims.KDim)(v_grid.Zone.BOTTOM),
),
dims.KDim: (0, nlev + 1),
},
fields={"ddqz_z_half": "ddqz_z_half"},
params={"nlev": nlev},
Expand All @@ -153,10 +153,13 @@
},
domain={
dims.CellDim: (
icon_grid.start_index(cell_domain(horizontal.Zone.LOCAL)),
icon_grid.end_index(cell_domain(horizontal.Zone.LOCAL)),
cell_domain(h_grid.Zone.LOCAL),
cell_domain(h_grid.Zone.LOCAL),
),
dims.KDim: (
v_grid.domain(dims.KDim)(v_grid.Zone.TOP),
v_grid.domain(dims.KDim)(v_grid.Zone.BOTTOM),
),
dims.KDim: (0, nlev),
},
fields={"ddqz_z_full": "ddqz_z_full", "inv_ddqz_z_full": "inv_ddqz_z_full"},
)
Expand All @@ -172,7 +175,10 @@
"vct_a": "vct_a",
},
domain={
dims.KDim: (0, nlev),
dims.KDim: (
v_grid.domain(dims.KDim)(v_grid.Zone.TOP),
v_grid.domain(dims.KDim)(v_grid.Zone.BOTTOM),
)
},
fields={"scalfac_dd3d": "scalfac_dd3d"},
params={
Expand Down Expand Up @@ -217,8 +223,8 @@
},
domain={
dims.CellDim: (
icon_grid.start_index(cell_domain(horizontal.Zone.LOCAL)),
icon_grid.end_index(cell_domain(horizontal.Zone.LOCAL)),
cell_domain(h_grid.Zone.LOCAL),
cell_domain(h_grid.Zone.LOCAL),
),
dims.KDim: (1, nlev),
},
Expand All @@ -236,10 +242,13 @@
},
domain={
dims.CellDim: (
icon_grid.start_index(cell_domain(horizontal.Zone.LOCAL)),
icon_grid.end_index(cell_domain(horizontal.Zone.LOCAL)),
cell_domain(h_grid.Zone.LOCAL),
cell_domain(h_grid.Zone.LOCAL),
),
dims.KDim: (
v_grid.domain(dims.KDim)(v_grid.Zone.TOP),
v_grid.domain(dims.KDim)(v_grid.Zone.BOTTOM),
),
dims.KDim: (0, nlev),
},
fields={"d2dexdz2_fac1_mc": "d2dexdz2_fac1_mc", "d2dexdz2_fac2_mc": "d2dexdz2_fac2_mc"},
params={
Expand All @@ -261,10 +270,13 @@
},
domain={
dims.CellDim: (
icon_grid.start_index(vertex_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_2)),
icon_grid.end_index(vertex_domain(horizontal.Zone.INTERIOR)),
vertex_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2),
vertex_domain(h_grid.Zone.INTERIOR),
),
dims.KDim: (
v_grid.domain(dims.KDim)(v_grid.Zone.TOP),
v_grid.domain(dims.KDim)(v_grid.Zone.BOTTOM),
),
dims.KDim: (0, nlev + 1),
},
fields={"z_ifv": "z_ifv"},
)
Expand All @@ -279,8 +291,8 @@
},
domain={
dims.CellDim: (
icon_grid.start_index(edge_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_3)),
icon_grid.end_index(edge_domain(horizontal.Zone.INTERIOR)),
edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3),
edge_domain(h_grid.Zone.INTERIOR),
),
dims.KDim: (nlev, nlev + 1),
},
Expand Down Expand Up @@ -308,8 +320,8 @@
},
domain={
dims.EdgeDim: (
icon_grid.start_index(edge_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_2)),
icon_grid.end_index(edge_domain(horizontal.Zone.INTERIOR)),
edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2),
edge_domain(h_grid.Zone.INTERIOR),
),
dims.KDim: (nlev, nlev + 1),
},
Expand All @@ -322,10 +334,13 @@
func=compute_vwind_impl_wgt.compute_vwind_impl_wgt,
domain={
dims.CellDim: (
icon_grid.start_index(edge_domain(horizontal.Zone.LOCAL)),
icon_grid.end_index(edge_domain(horizontal.Zone.LOCAL)),
edge_domain(h_grid.Zone.LOCAL),
edge_domain(h_grid.Zone.LOCAL),
),
dims.KDim: (
v_grid.domain(dims.KDim)(v_grid.Zone.TOP),
v_grid.domain(dims.KDim)(v_grid.Zone.BOTTOM),
),
dims.KDim: (0, nlev),
},
fields=["vwind_impl_wgt"],
deps={
Expand Down Expand Up @@ -355,8 +370,8 @@
},
domain={
dims.CellDim: (
icon_grid.start_index(cell_domain(horizontal.Zone.LOCAL)),
icon_grid.end_index(cell_domain(horizontal.Zone.LOCAL)),
cell_domain(h_grid.Zone.LOCAL),
cell_domain(h_grid.Zone.LOCAL),
),
},
fields={"vwind_expl_wgt": "vwind_expl_wgt"},
Expand All @@ -370,10 +385,13 @@
},
domain={
dims.CellDim: (
icon_grid.start_index(cell_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_2)),
icon_grid.end_index(cell_domain(horizontal.Zone.LOCAL)),
cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2),
cell_domain(h_grid.Zone.LOCAL),
),
dims.KDim: (
v_grid.domain(dims.KDim)(v_grid.Zone.TOP),
v_grid.domain(dims.KDim)(v_grid.Zone.BOTTOM),
),
dims.KDim: (0, nlev),
},
fields={"exner_exfac": "exner_exfac"},
params={"exner_expol": "exner_expol"},
Expand All @@ -388,10 +406,13 @@
},
domain={
dims.CellDim: (
icon_grid.start_index(edge_domain(horizontal.Zone.LOCAL)),
icon_grid.end_index(edge_domain(horizontal.Zone.LOCAL)),
edge_domain(h_grid.Zone.LOCAL),
edge_domain(h_grid.Zone.LOCAL),
),
dims.KDim: (
v_grid.domain(dims.KDim)(v_grid.Zone.TOP),
v_grid.domain(dims.KDim)(v_grid.Zone.BOTTOM),
),
dims.KDim: (0, nlev + 1),
},
fields={"wgtfac_e": "wgtfac_e"},
)
Expand All @@ -402,10 +423,8 @@
deps={"z_ifc_sliced": "z_ifc_sliced"},
domain={
dims.EdgeDim: (
icon_grid.end_index(
edge_domain(horizontal.Zone.NUDGING)
), # TODO: check if this is really end (also in mf)
icon_grid.end_index(edge_domain(horizontal.Zone.LOCAL)),
edge_domain(h_grid.Zone.NUDGING), # TODO: check if this is really end (also in mf)
edge_domain(h_grid.Zone.LOCAL),
)
},
fields={"z_aux2": "z_aux2"},
Expand All @@ -416,10 +435,13 @@
deps={"in_field": "height", "coeff": "c_lin_e"},
domain={
dims.EdgeDim: (
icon_grid.start_index(edge_domain(horizontal.Zone.LOCAL)),
icon_grid.end_index(edge_domain(horizontal.Zone.LOCAL)),
edge_domain(h_grid.Zone.LOCAL),
edge_domain(h_grid.Zone.LOCAL),
),
dims.KDim: (
v_grid.domain(dims.KDim)(v_grid.Zone.TOP),
v_grid.domain(dims.KDim)(v_grid.Zone.BOTTOM),
),
dims.KDim: (0, nlev),
},
fields={"z_me": "z_me"},
)
Expand All @@ -433,10 +455,13 @@
},
domain={
dims.EdgeDim: (
icon_grid.start_index(edge_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_3)),
icon_grid.end_index(edge_domain(horizontal.Zone.LOCAL)),
edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3),
edge_domain(h_grid.Zone.LOCAL),
),
dims.KDim: (
v_grid.domain(dims.KDim)(v_grid.Zone.TOP),
v_grid.domain(dims.KDim)(v_grid.Zone.BOTTOM),
),
dims.KDim: (0, nlev),
},
fields={"flat_idx": "flat_idx"},
)
Expand All @@ -458,10 +483,13 @@
},
domain={
dims.EdgeDim: (
icon_grid.start_index(edge_domain(horizontal.Zone.NUDGING)),
icon_grid.end_index(edge_domain(horizontal.Zone.LOCAL)),
edge_domain(h_grid.Zone.NUDGING),
edge_domain(h_grid.Zone.LOCAL),
),
dims.KDim: (
v_grid.domain(dims.KDim)(v_grid.Zone.TOP),
v_grid.domain(dims.KDim)(v_grid.Zone.BOTTOM),
),
dims.KDim: (0, nlev),
},
fields={"pg_edgeidx": "pg_edgeidx", "pg_vertidx": "pg_vertidx"},
)
Expand All @@ -471,10 +499,13 @@
deps={"pg_edgeidx": "pg_edgeidx", "pg_vertidx": "pg_vertidx"},
domain={
dims.EdgeDim: (
icon_grid.start_index(edge_domain(horizontal.Zone.LOCAL)),
icon_grid.end_index(edge_domain(horizontal.Zone.LOCAL)),
edge_domain(h_grid.Zone.LOCAL),
edge_domain(h_grid.Zone.LOCAL),
),
dims.KDim: (
v_grid.domain(dims.KDim)(v_grid.Zone.TOP),
v_grid.domain(dims.KDim)(v_grid.Zone.BOTTOM),
),
dims.KDim: (0, nlev),
},
fields={"pg_edgeidx_dsl": "pg_edgeidx_dsl"},
)
Expand All @@ -492,10 +523,13 @@
},
domain={
dims.CellDim: (
icon_grid.start_index(edge_domain(horizontal.Zone.NUDGING)),
icon_grid.end_index(edge_domain(horizontal.Zone.LOCAL)),
edge_domain(h_grid.Zone.NUDGING),
edge_domain(h_grid.Zone.LOCAL),
),
dims.KDim: (
v_grid.domain(dims.KDim)(v_grid.Zone.TOP),
v_grid.domain(dims.KDim)(v_grid.Zone.BOTTOM),
),
dims.KDim: (0, nlev),
},
fields={"pg_exdist_dsl": "pg_exdist_dsl"},
)
Expand All @@ -509,8 +543,8 @@
},
domain={
dims.CellDim: (
icon_grid.start_index(cell_domain(horizontal.Zone.HALO)),
icon_grid.end_index(cell_domain(horizontal.Zone.LOCAL)),
cell_domain(h_grid.Zone.HALO),
cell_domain(h_grid.Zone.LOCAL),
),
},
fields={"mask_prog_halo_c": "mask_prog_halo_c"},
Expand All @@ -525,8 +559,8 @@
},
domain={
dims.CellDim: (
icon_grid.start_index(cell_domain(horizontal.Zone.HALO)),
icon_grid.end_index(cell_domain(horizontal.Zone.LOCAL)),
cell_domain(h_grid.Zone.HALO),
cell_domain(h_grid.Zone.LOCAL),
),
},
fields={"bdy_halo_c": "bdy_halo_c"},
Expand All @@ -541,14 +575,14 @@
},
domain={
dims.CellDim: (
icon_grid.start_index(cell_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_2)),
icon_grid.end_index(cell_domain(horizontal.Zone.LOCAL)),
cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2),
cell_domain(h_grid.Zone.LOCAL),
)
},
fields={"hmask_dd3d": "hmask_dd3d"},
params={
"grf_nudge_start_e": gtx.int32(horizontal._GRF_NUDGEZONE_START_EDGES),
"grf_nudgezone_width": gtx.int32(horizontal._GRF_NUDGEZONE_WIDTH),
"grf_nudge_start_e": gtx.int32(h_grid._GRF_NUDGEZONE_START_EDGES),
"grf_nudgezone_width": gtx.int32(h_grid._GRF_NUDGEZONE_WIDTH),
},
)
fields_factory.register_provider(compute_hmask_dd3d_provider)

0 comments on commit 85bdff4

Please sign in to comment.