From af437713913d46e0b2122517442c4f30b4378667 Mon Sep 17 00:00:00 2001 From: Andreas Date: Fri, 22 Mar 2024 18:09:36 +0100 Subject: [PATCH 01/65] further interpolation coefficients --- .../interpolation/interpolation_fields3.py | 105 ++++++++++++++++++ .../test_interpolation_fields3.py | 98 ++++++++++++++++ 2 files changed, 203 insertions(+) create mode 100644 model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py create mode 100644 model/common/tests/interpolation_tests/test_interpolation_fields3.py diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py new file mode 100644 index 000000000..0e9e07889 --- /dev/null +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py @@ -0,0 +1,105 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +from gt4py.next import where +from gt4py.next.ffront.decorator import field_operator +from gt4py.next.ffront.fbuiltins import Field + +from icon4py.model.common.dimension import C2E, V2E, C2EDim, CellDim, EdgeDim, V2EDim, VertexDim + +def grad_fd_norm( + psi_c: np.array, + inv_dual_edge_length: np.array, + e2c: np.array, + second_boundary_layer_start_index: np.int32, + second_boundary_layer_end_index: np.int32, + nlev, +) -> np.array: + llb = second_boundary_layer_start_index + grad_norm_psi_e = np.zeros([second_boundary_layer_end_index, nlev]) + for i in range(nlev): + grad_norm_psi_e[llb:, i] = (psi_c[e2c[llb:, 1], i] - psi_c[e2c[llb:, 0], i]) * inv_dual_edge_length[llb:] + return grad_norm_psi_e + +def grad_fd_tang( + psi_c: np.array, + inv_primal_edge_length: np.array, + e2v: np.array, + third_boundary_layer_start_index: np.int32, + second_boundary_layer_end_index: np.int32, + nlev, +) -> np.array: + llb = third_boundary_layer_start_index + grad_tang_psi_e = np.zeros([second_boundary_layer_end_index, nlev]) + for i in range(nlev): + grad_norm_psi_e[llb:, i] = (psi_v[e2v[llb:, 1], i] - psi_v[e2v[llb:, 0], i]) * inv_primal_edge_length[llb:] + return grad_tang_psi_e + +def compute_ddxn_z_half_e( + z_ifc: np.array, + inv_dual_edge_length: np.array, + e2c: np.array, + second_boundary_layer_start_index: np.int32, + second_boundary_layer_end_index: np.int32, +) -> np.array: + nlev = z_ifc.shape[1] + ddxn_z_half_e = grad_fd_norm(z_ifc, inv_dual_edge_length, e2c, second_boundary_layer_start_index, second_boundary_layer_end_index, nlev) + return ddxn_z_half_e + +def compute_ddxt_z_half_e( + z_ifv: np.array, + inv_primal_edge_length: np.array, + e2v: np.array, + third_boundary_layer_start_index: np.int32, + second_boundary_layer_end_index: np.int32, +) -> np.array: + nlev = z_ifv.shape[1] + return ddxt_z_half_e + +def compute_ddxnt_z_full( + z_ddxnt_z_half_e: np.array, +) -> np.array: + ddxnt_z_full = 0.5 * (z_ddxnt_z_half_e[:, :z_ddxnt_z_half_e.shape[1]-1] + z_ddxnt_z_half_e[:, 1:]) + return ddxnt_z_full + +def cells2verts_scalar( +) -> np.array: +# p_vert_out(jv,jk,jb) = & +# c_int(jv,1,jb) * p_cell_in(iidx(jv,jb,1),jk,iblk(jv,jb,1)) + & +# c_int(jv,2,jb) * p_cell_in(iidx(jv,jb,2),jk,iblk(jv,jb,2)) + & +# c_int(jv,3,jb) * p_cell_in(iidx(jv,jb,3),jk,iblk(jv,jb,3)) + & +# c_int(jv,4,jb) * p_cell_in(iidx(jv,jb,4),jk,iblk(jv,jb,4)) + & +# c_int(jv,5,jb) * p_cell_in(iidx(jv,jb,5),jk,iblk(jv,jb,5)) + & +# c_int(jv,6,jb) * p_cell_in(iidx(jv,jb,6),jk,iblk(jv,jb,6)) + return p_vert_out + +def compute_cells_aw_verts( + dual_area: np.array, + edge_vert_length: np.array, + edge_cell_length: np.array, +) -> np.array: + cells_aw_verts = cells_aw_verts + 0.5 / dual_area * edge_vert_length * edge_cell_length + return cells_aw_verts diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields3.py b/model/common/tests/interpolation_tests/test_interpolation_fields3.py new file mode 100644 index 000000000..9281fa444 --- /dev/null +++ b/model/common/tests/interpolation_tests/test_interpolation_fields3.py @@ -0,0 +1,98 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.iterator.builtins import int32 + +from icon4py.model.common.dimension import ( + C2E2CDim, + C2EDim, + CellDim, + E2C2EDim, + E2CDim, + E2VDim, + EdgeDim, + V2CDim, + V2EDim, + VertexDim, +) +from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex +from icon4py.model.common.interpolation.interpolation_fields3 import ( + compute_ddxn_z_half_e, + compute_ddxnt_z_full, +) +from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package + data_provider, + datapath, + download_ser_data, + experiment, + processor_props, + ranked_data_path, +) +from icon4py.model.common.test_utils.helpers import zero_field + + +@pytest.mark.datatest +def test_compute_ddxn_z_half_e(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): # fixture + z_ifc = metrics_savepoint.z_ifc().asnumpy() + inv_dual_edge_length = grid_savepoint.inv_dual_edge_length().asnumpy() + e2c = icon_grid.connectivities[E2CDim] +# edge_cell_length = grid_savepoint.edge_cell_length() +# owner_mask = grid_savepoint.e_owner_mask() + ddxn_z_full_ref = metrics_savepoint.ddxn_z_full().asnumpy() + lateral_boundary = np.arange(2) + lateral_boundary[0] = icon_grid.get_start_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, + ) + lateral_boundary[1] = icon_grid.get_end_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, + ) + ddxn_z_half_e = compute_ddxn_z_half_e( + z_ifc, + inv_dual_edge_length, + e2c, + lateral_boundary[0], + lateral_boundary[1], + ) + ddxn_z_full = compute_ddxnt_z_full( + ddxn_z_half_e, + ) + + assert np.allclose(ddxn_z_full, ddxn_z_full_ref) + + +@pytest.mark.datatest +def test_compute_ddxt_z_half_e(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): +# dual_area = grid_savepoint.dual_area().asnumpy() +# edge_vert_length = grid_savepoint.edge_vert_length().asnumpy() + edge_cell_length = grid_savepoint.edge_cell_length().asnumpy() + cells_aw_verts = compute_cells_aw_verts( + dual_area, + edge_vert_length, + edge_cell_length, + ) From b1a7e5d268f711fd2c9fc0c2ac3a705c617bdd9d Mon Sep 17 00:00:00 2001 From: Andreas Date: Tue, 26 Mar 2024 15:47:38 +0100 Subject: [PATCH 02/65] progress --- .../interpolation/interpolation_fields3.py | 25 +++++++++++++++++-- .../common/test_utils/serialbox_utils.py | 6 +++++ .../test_interpolation_fields3.py | 18 +++++++++++-- 3 files changed, 45 insertions(+), 4 deletions(-) diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py index 0e9e07889..00c56e69a 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py @@ -86,8 +86,15 @@ def compute_ddxnt_z_full( return ddxnt_z_full def cells2verts_scalar( + c_int: np.array, + p_cell_in: np.array, + v2c: np.array, ) -> np.array: -# p_vert_out(jv,jk,jb) = & +# iz_ifc: np.array, +# cells_aw_verts: np.array, + p_vert_out = np.zeros() + for i in range(6): + p_vert_out[:, i] = p_vert_out[:, i] + p_cell_in[v2c[:, 0]] # c_int(jv,1,jb) * p_cell_in(iidx(jv,jb,1),jk,iblk(jv,jb,1)) + & # c_int(jv,2,jb) * p_cell_in(iidx(jv,jb,2),jk,iblk(jv,jb,2)) + & # c_int(jv,3,jb) * p_cell_in(iidx(jv,jb,3),jk,iblk(jv,jb,3)) + & @@ -100,6 +107,20 @@ def compute_cells_aw_verts( dual_area: np.array, edge_vert_length: np.array, edge_cell_length: np.array, + e2c: np.array, + v2c: np.array, + v2e: np.array, + e2v: np.array, + second_boundary_layer_end_index: np.int32, ) -> np.array: - cells_aw_verts = cells_aw_verts + 0.5 / dual_area * edge_vert_length * edge_cell_length + cells_aw_verts = np.zeros([second_boundary_layer_end_index, 6]) + index = np.zeros([second_boundary_layer_end_index, 6]) + for j in range (6): + index[:, j] = np.arange(second_boundary_layer_end_index) + idx_ve = np.where(index == e2v[v2e, 0], 0, 1) + + for i in range(2): + for j in range (6): + for k in range (6): + cells_aw_verts[:, k] = np.where(e2c[v2e[:, j], i] == v2c[:, k], cells_aw_verts[:, k] + 0.5 / dual_area[:] * edge_vert_length[v2e[:, j], idx_ve[:, j]] * edge_cell_length[v2e[:, j], i], cells_aw_verts[:, k]) return cells_aw_verts diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py index ade294653..9b0ddd9ea 100644 --- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -134,6 +134,12 @@ def _read(self, name: str, offset=0, dtype=int): class IconGridSavepoint(IconSavepoint): + def v_dual_area(self): + return self._get_field("v_dual_area", VertexDim) + + def edge_vert_length(self): + return self._get_field("edge_vert_length", EdgeDim, E2C2VDim) + def vct_a(self): return self._get_field("vct_a", KDim) diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields3.py b/model/common/tests/interpolation_tests/test_interpolation_fields3.py index 9281fa444..cf7dfff52 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_fields3.py +++ b/model/common/tests/interpolation_tests/test_interpolation_fields3.py @@ -43,6 +43,7 @@ from icon4py.model.common.interpolation.interpolation_fields3 import ( compute_ddxn_z_half_e, compute_ddxnt_z_full, + compute_cells_aw_verts, ) from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package data_provider, @@ -88,11 +89,24 @@ def test_compute_ddxn_z_half_e(grid_savepoint, interpolation_savepoint, icon_gri @pytest.mark.datatest def test_compute_ddxt_z_half_e(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): -# dual_area = grid_savepoint.dual_area().asnumpy() -# edge_vert_length = grid_savepoint.edge_vert_length().asnumpy() + dual_area = grid_savepoint.v_dual_area().asnumpy() + edge_vert_length = grid_savepoint.edge_vert_length().asnumpy() edge_cell_length = grid_savepoint.edge_cell_length().asnumpy() + e2c = icon_grid.connectivities[E2CDim] + v2c = icon_grid.connectivities[V2CDim] + v2e = icon_grid.connectivities[V2EDim] + e2v = icon_grid.connectivities[E2VDim] + second_boundary_layer_end_index = icon_grid.get_end_index( + VertexDim, + HorizontalMarkerIndex.lateral_boundary(VertexDim) - 1, + ) cells_aw_verts = compute_cells_aw_verts( dual_area, edge_vert_length, edge_cell_length, + e2c, + v2c, + v2e, + e2v, + second_boundary_layer_end_index, ) From c2d201c43649d4460198e94fddae32912f8903e2 Mon Sep 17 00:00:00 2001 From: Andreas Date: Thu, 28 Mar 2024 12:02:27 +0100 Subject: [PATCH 03/65] progress --- .../interpolation/interpolation_fields3.py | 27 ++++++------- .../test_interpolation_fields3.py | 40 ++++++++++++++++--- 2 files changed, 47 insertions(+), 20 deletions(-) diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py index 00c56e69a..b2eea0460 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py @@ -35,27 +35,25 @@ def grad_fd_norm( inv_dual_edge_length: np.array, e2c: np.array, second_boundary_layer_start_index: np.int32, - second_boundary_layer_end_index: np.int32, nlev, ) -> np.array: llb = second_boundary_layer_start_index - grad_norm_psi_e = np.zeros([second_boundary_layer_end_index, nlev]) + grad_norm_psi_e = np.zeros([len(e2c[:, 0]), nlev]) for i in range(nlev): grad_norm_psi_e[llb:, i] = (psi_c[e2c[llb:, 1], i] - psi_c[e2c[llb:, 0], i]) * inv_dual_edge_length[llb:] return grad_norm_psi_e def grad_fd_tang( - psi_c: np.array, + psi_v: np.array, inv_primal_edge_length: np.array, e2v: np.array, third_boundary_layer_start_index: np.int32, - second_boundary_layer_end_index: np.int32, nlev, ) -> np.array: llb = third_boundary_layer_start_index - grad_tang_psi_e = np.zeros([second_boundary_layer_end_index, nlev]) + grad_tang_psi_e = np.zeros([len(e2v[:, 0]), nlev]) for i in range(nlev): - grad_norm_psi_e[llb:, i] = (psi_v[e2v[llb:, 1], i] - psi_v[e2v[llb:, 0], i]) * inv_primal_edge_length[llb:] + grad_tang_psi_e[llb:, i] = (psi_v[e2v[llb:, 1], i] - psi_v[e2v[llb:, 0], i]) * inv_primal_edge_length[llb:] return grad_tang_psi_e def compute_ddxn_z_half_e( @@ -63,10 +61,9 @@ def compute_ddxn_z_half_e( inv_dual_edge_length: np.array, e2c: np.array, second_boundary_layer_start_index: np.int32, - second_boundary_layer_end_index: np.int32, ) -> np.array: nlev = z_ifc.shape[1] - ddxn_z_half_e = grad_fd_norm(z_ifc, inv_dual_edge_length, e2c, second_boundary_layer_start_index, second_boundary_layer_end_index, nlev) + ddxn_z_half_e = grad_fd_norm(z_ifc, inv_dual_edge_length, e2c, second_boundary_layer_start_index, nlev) return ddxn_z_half_e def compute_ddxt_z_half_e( @@ -74,9 +71,9 @@ def compute_ddxt_z_half_e( inv_primal_edge_length: np.array, e2v: np.array, third_boundary_layer_start_index: np.int32, - second_boundary_layer_end_index: np.int32, ) -> np.array: nlev = z_ifv.shape[1] + ddxt_z_half_e = grad_fd_tang(z_ifv, inv_primal_edge_length, e2v, third_boundary_layer_start_index, nlev) return ddxt_z_half_e def compute_ddxnt_z_full( @@ -86,15 +83,15 @@ def compute_ddxnt_z_full( return ddxnt_z_full def cells2verts_scalar( - c_int: np.array, p_cell_in: np.array, + c_int: np.array, v2c: np.array, ) -> np.array: -# iz_ifc: np.array, -# cells_aw_verts: np.array, - p_vert_out = np.zeros() - for i in range(6): - p_vert_out[:, i] = p_vert_out[:, i] + p_cell_in[v2c[:, 0]] + kdim = len(p_cell_in[0, :]) + p_vert_out = np.zeros([len(c_int[:, 0]), kdim]) + for j in range(kdim): + for i in range(6): + p_vert_out[:, j] = p_vert_out[:, j] + c_int[:, i] * p_cell_in[v2c[:, i], j] # c_int(jv,1,jb) * p_cell_in(iidx(jv,jb,1),jk,iblk(jv,jb,1)) + & # c_int(jv,2,jb) * p_cell_in(iidx(jv,jb,2),jk,iblk(jv,jb,2)) + & # c_int(jv,3,jb) * p_cell_in(iidx(jv,jb,3),jk,iblk(jv,jb,3)) + & diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields3.py b/model/common/tests/interpolation_tests/test_interpolation_fields3.py index cf7dfff52..aff90ab00 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_fields3.py +++ b/model/common/tests/interpolation_tests/test_interpolation_fields3.py @@ -44,6 +44,8 @@ compute_ddxn_z_half_e, compute_ddxnt_z_full, compute_cells_aw_verts, + cells2verts_scalar, + compute_ddxt_z_half_e, ) from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package data_provider, @@ -57,7 +59,7 @@ @pytest.mark.datatest -def test_compute_ddxn_z_half_e(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): # fixture +def test_compute_ddxn_z_full_e(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): # fixture z_ifc = metrics_savepoint.z_ifc().asnumpy() inv_dual_edge_length = grid_savepoint.inv_dual_edge_length().asnumpy() e2c = icon_grid.connectivities[E2CDim] @@ -78,7 +80,6 @@ def test_compute_ddxn_z_half_e(grid_savepoint, interpolation_savepoint, icon_gri inv_dual_edge_length, e2c, lateral_boundary[0], - lateral_boundary[1], ) ddxn_z_full = compute_ddxnt_z_full( ddxn_z_half_e, @@ -88,18 +89,33 @@ def test_compute_ddxn_z_half_e(grid_savepoint, interpolation_savepoint, icon_gri @pytest.mark.datatest -def test_compute_ddxt_z_half_e(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): +def test_compute_ddxt_z_full_e(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): + z_ifc = metrics_savepoint.z_ifc().asnumpy() dual_area = grid_savepoint.v_dual_area().asnumpy() edge_vert_length = grid_savepoint.edge_vert_length().asnumpy() edge_cell_length = grid_savepoint.edge_cell_length().asnumpy() + inv_dual_edge_length = grid_savepoint.inv_dual_edge_length().asnumpy() + ddxt_z_full_ref = metrics_savepoint.ddxt_z_full().asnumpy() e2c = icon_grid.connectivities[E2CDim] v2c = icon_grid.connectivities[V2CDim] v2e = icon_grid.connectivities[V2EDim] e2v = icon_grid.connectivities[E2VDim] - second_boundary_layer_end_index = icon_grid.get_end_index( + second_boundary_layer_start_index_vertex = icon_grid.get_start_index( + VertexDim, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, + ) + second_boundary_layer_end_index_vertex = icon_grid.get_end_index( VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) - 1, ) + second_boundary_layer_start_index_cell = icon_grid.get_start_index( + CellDim, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, + ) + second_boundary_layer_end_index_cell = icon_grid.get_end_index( + CellDim, + HorizontalMarkerIndex.lateral_boundary(CellDim) - 1, + ) cells_aw_verts = compute_cells_aw_verts( dual_area, edge_vert_length, @@ -108,5 +124,19 @@ def test_compute_ddxt_z_half_e(grid_savepoint, interpolation_savepoint, icon_gri v2c, v2e, e2v, - second_boundary_layer_end_index, + second_boundary_layer_end_index_vertex, ) + z_ifv = cells2verts_scalar(z_ifc, cells_aw_verts, v2c) + ddxt_z_half_e = compute_ddxt_z_half_e( + z_ifv, + inv_dual_edge_length, + e2v, + second_boundary_layer_start_index_vertex, + ) + ddxt_z_full = compute_ddxnt_z_full( + ddxt_z_half_e, + ) + + print(ddxt_z_full_ref) + print(ddxt_z_full) + assert np.allclose(ddxt_z_full, ddxt_z_full_ref) From 89949b4203ec08ff6aa4d7c541623f1a2a8ffc0a Mon Sep 17 00:00:00 2001 From: Andreas Date: Tue, 2 Apr 2024 14:48:15 +0200 Subject: [PATCH 04/65] ddxt_z_full --- .../interpolation/interpolation_fields3.py | 29 ++++++++++--------- .../common/test_utils/serialbox_utils.py | 3 ++ .../test_interpolation_fields3.py | 24 ++++++++++----- 3 files changed, 35 insertions(+), 21 deletions(-) diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py index b2eea0460..45f23f54f 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py @@ -46,14 +46,15 @@ def grad_fd_norm( def grad_fd_tang( psi_v: np.array, inv_primal_edge_length: np.array, + tangent_orientation: np.array, e2v: np.array, third_boundary_layer_start_index: np.int32, nlev, ) -> np.array: llb = third_boundary_layer_start_index - grad_tang_psi_e = np.zeros([len(e2v[:, 0]), nlev]) + grad_tang_psi_e = np.zeros([e2v.shape[0], nlev]) for i in range(nlev): - grad_tang_psi_e[llb:, i] = (psi_v[e2v[llb:, 1], i] - psi_v[e2v[llb:, 0], i]) * inv_primal_edge_length[llb:] + grad_tang_psi_e[llb:, i] = tangent_orientation[llb:] * (psi_v[e2v[llb:, 1], i] - psi_v[e2v[llb:, 0], i]) * inv_primal_edge_length[llb:] return grad_tang_psi_e def compute_ddxn_z_half_e( @@ -69,11 +70,12 @@ def compute_ddxn_z_half_e( def compute_ddxt_z_half_e( z_ifv: np.array, inv_primal_edge_length: np.array, + tangent_orientation: np.array, e2v: np.array, third_boundary_layer_start_index: np.int32, ) -> np.array: nlev = z_ifv.shape[1] - ddxt_z_half_e = grad_fd_tang(z_ifv, inv_primal_edge_length, e2v, third_boundary_layer_start_index, nlev) + ddxt_z_half_e = grad_fd_tang(z_ifv, inv_primal_edge_length, tangent_orientation, e2v, third_boundary_layer_start_index, nlev) return ddxt_z_half_e def compute_ddxnt_z_full( @@ -82,34 +84,33 @@ def compute_ddxnt_z_full( ddxnt_z_full = 0.5 * (z_ddxnt_z_half_e[:, :z_ddxnt_z_half_e.shape[1]-1] + z_ddxnt_z_half_e[:, 1:]) return ddxnt_z_full -def cells2verts_scalar( +def compute_cells2verts_scalar( p_cell_in: np.array, c_int: np.array, v2c: np.array, + second_boundary_layer_start_index: np.int32, ) -> np.array: - kdim = len(p_cell_in[0, :]) - p_vert_out = np.zeros([len(c_int[:, 0]), kdim]) + llb = second_boundary_layer_start_index + kdim = p_cell_in.shape[1] + p_vert_out = np.zeros([c_int.shape[0], kdim]) for j in range(kdim): for i in range(6): - p_vert_out[:, j] = p_vert_out[:, j] + c_int[:, i] * p_cell_in[v2c[:, i], j] -# c_int(jv,1,jb) * p_cell_in(iidx(jv,jb,1),jk,iblk(jv,jb,1)) + & -# c_int(jv,2,jb) * p_cell_in(iidx(jv,jb,2),jk,iblk(jv,jb,2)) + & -# c_int(jv,3,jb) * p_cell_in(iidx(jv,jb,3),jk,iblk(jv,jb,3)) + & -# c_int(jv,4,jb) * p_cell_in(iidx(jv,jb,4),jk,iblk(jv,jb,4)) + & -# c_int(jv,5,jb) * p_cell_in(iidx(jv,jb,5),jk,iblk(jv,jb,5)) + & -# c_int(jv,6,jb) * p_cell_in(iidx(jv,jb,6),jk,iblk(jv,jb,6)) + p_vert_out[llb:, j] = p_vert_out[llb:, j] + c_int[llb:, i] * p_cell_in[v2c[llb:, i], j] return p_vert_out def compute_cells_aw_verts( dual_area: np.array, edge_vert_length: np.array, edge_cell_length: np.array, + owner_mask: np.array, e2c: np.array, v2c: np.array, v2e: np.array, e2v: np.array, + second_boundary_layer_start_index: np.int32, second_boundary_layer_end_index: np.int32, ) -> np.array: + llb = second_boundary_layer_start_index cells_aw_verts = np.zeros([second_boundary_layer_end_index, 6]) index = np.zeros([second_boundary_layer_end_index, 6]) for j in range (6): @@ -119,5 +120,5 @@ def compute_cells_aw_verts( for i in range(2): for j in range (6): for k in range (6): - cells_aw_verts[:, k] = np.where(e2c[v2e[:, j], i] == v2c[:, k], cells_aw_verts[:, k] + 0.5 / dual_area[:] * edge_vert_length[v2e[:, j], idx_ve[:, j]] * edge_cell_length[v2e[:, j], i], cells_aw_verts[:, k]) + cells_aw_verts[llb:, k] = np.where(np.logical_and(owner_mask[:], e2c[v2e[llb:, j], i] == v2c[llb:, k]), cells_aw_verts[llb:, k] + 0.5 / dual_area[llb:] * edge_vert_length[v2e[llb:, j], idx_ve[llb:, j]] * edge_cell_length[v2e[llb:, j], i], cells_aw_verts[llb:, k]) return cells_aw_verts diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py index 9b0ddd9ea..5100ca68a 100644 --- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -217,6 +217,9 @@ def edge_end_index(self): # one off accounts for being exclusive [from:to) return self.serializer.read("e_end_index", self.savepoint) + def v_owner_mask(self): + return self._get_field("v_owner_mask", VertexDim, dtype=bool) + def c_owner_mask(self): return self._get_field("c_owner_mask", CellDim, dtype=bool) diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields3.py b/model/common/tests/interpolation_tests/test_interpolation_fields3.py index aff90ab00..61549c7d6 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_fields3.py +++ b/model/common/tests/interpolation_tests/test_interpolation_fields3.py @@ -44,7 +44,7 @@ compute_ddxn_z_half_e, compute_ddxnt_z_full, compute_cells_aw_verts, - cells2verts_scalar, + compute_cells2verts_scalar, compute_ddxt_z_half_e, ) from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package @@ -94,7 +94,9 @@ def test_compute_ddxt_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri dual_area = grid_savepoint.v_dual_area().asnumpy() edge_vert_length = grid_savepoint.edge_vert_length().asnumpy() edge_cell_length = grid_savepoint.edge_cell_length().asnumpy() - inv_dual_edge_length = grid_savepoint.inv_dual_edge_length().asnumpy() + tangent_orientation = grid_savepoint.tangent_orientation().asnumpy() + inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths().asnumpy() + owner_mask = grid_savepoint.v_owner_mask() ddxt_z_full_ref = metrics_savepoint.ddxt_z_full().asnumpy() e2c = icon_grid.connectivities[E2CDim] v2c = icon_grid.connectivities[V2CDim] @@ -116,27 +118,35 @@ def test_compute_ddxt_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) - 1, ) + third_boundary_layer_start_index_edge = icon_grid.get_start_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, + ) cells_aw_verts = compute_cells_aw_verts( dual_area, edge_vert_length, edge_cell_length, + owner_mask, e2c, v2c, v2e, e2v, + second_boundary_layer_start_index_vertex, second_boundary_layer_end_index_vertex, ) - z_ifv = cells2verts_scalar(z_ifc, cells_aw_verts, v2c) + cells_aw_verts_ref = interpolation_savepoint.c_intp().asnumpy() + assert np.allclose(cells_aw_verts, cells_aw_verts_ref) + + z_ifv = compute_cells2verts_scalar(z_ifc, cells_aw_verts, v2c, second_boundary_layer_start_index_vertex) ddxt_z_half_e = compute_ddxt_z_half_e( z_ifv, - inv_dual_edge_length, + inv_primal_edge_length, + tangent_orientation, e2v, - second_boundary_layer_start_index_vertex, + third_boundary_layer_start_index_edge, ) ddxt_z_full = compute_ddxnt_z_full( ddxt_z_half_e, ) - print(ddxt_z_full_ref) - print(ddxt_z_full) assert np.allclose(ddxt_z_full, ddxt_z_full_ref) From b05bbee7f6f3b507ead7e1004703ee2e9fb80ee6 Mon Sep 17 00:00:00 2001 From: Andreas Date: Tue, 9 Apr 2024 16:47:48 +0200 Subject: [PATCH 05/65] model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py --- .../test_interpolation_fields3.py | 29 +++++++++++++------ 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields3.py b/model/common/tests/interpolation_tests/test_interpolation_fields3.py index 61549c7d6..346626cfb 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_fields3.py +++ b/model/common/tests/interpolation_tests/test_interpolation_fields3.py @@ -38,6 +38,7 @@ V2CDim, V2EDim, VertexDim, + KDim, ) from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex from icon4py.model.common.interpolation.interpolation_fields3 import ( @@ -60,31 +61,41 @@ @pytest.mark.datatest def test_compute_ddxn_z_full_e(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): # fixture - z_ifc = metrics_savepoint.z_ifc().asnumpy() - inv_dual_edge_length = grid_savepoint.inv_dual_edge_length().asnumpy() + z_ifc = metrics_savepoint.z_ifc() + inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() e2c = icon_grid.connectivities[E2CDim] # edge_cell_length = grid_savepoint.edge_cell_length() # owner_mask = grid_savepoint.e_owner_mask() ddxn_z_full_ref = metrics_savepoint.ddxn_z_full().asnumpy() - lateral_boundary = np.arange(2) - lateral_boundary[0] = icon_grid.get_start_index( + horizontal_start = icon_grid.get_start_index( EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, ) - lateral_boundary[1] = icon_grid.get_end_index( + horizontal_end = icon_grid.get_end_index( EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, ) - ddxn_z_half_e = compute_ddxn_z_half_e( + vertical_start = 0 + vertical_end = 66 + ddxn_z_half_e = zero_field(icon_grid, EdgeDim, KDim) + ddxn_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) + compute_ddxn_z_half_e( z_ifc, inv_dual_edge_length, - e2c, - lateral_boundary[0], + out=ddxn_z_half_e, + offset_provider={"E2C" : icon_grid.get_offset_provider("E2C") }, + domain={ + EdgeDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) + print(ddxn_z_half_e.asnumpy().shape) ddxn_z_full = compute_ddxnt_z_full( - ddxn_z_half_e, + ddxn_z_half_e.asnumpy(), ) + print(ddxn_z_full.shape) + print(ddxn_z_full_ref.shape) assert np.allclose(ddxn_z_full, ddxn_z_full_ref) From 8aac5e262ca5602966dc509d319e1d51aac0c055 Mon Sep 17 00:00:00 2001 From: Andreas Date: Tue, 9 Apr 2024 16:49:34 +0200 Subject: [PATCH 06/65] progress gt4py --- .../interpolation/interpolation_fields3.py | 29 +++++++------------ 1 file changed, 11 insertions(+), 18 deletions(-) diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py index 45f23f54f..70399f6cb 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py @@ -28,19 +28,14 @@ from gt4py.next.ffront.decorator import field_operator from gt4py.next.ffront.fbuiltins import Field -from icon4py.model.common.dimension import C2E, V2E, C2EDim, CellDim, EdgeDim, V2EDim, VertexDim +from icon4py.model.common.dimension import C2E, E2C, V2E, C2EDim, CellDim, EdgeDim, V2EDim, VertexDim, KDim +@field_operator def grad_fd_norm( - psi_c: np.array, - inv_dual_edge_length: np.array, - e2c: np.array, - second_boundary_layer_start_index: np.int32, - nlev, -) -> np.array: - llb = second_boundary_layer_start_index - grad_norm_psi_e = np.zeros([len(e2c[:, 0]), nlev]) - for i in range(nlev): - grad_norm_psi_e[llb:, i] = (psi_c[e2c[llb:, 1], i] - psi_c[e2c[llb:, 0], i]) * inv_dual_edge_length[llb:] + psi_c: Field[[CellDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], +) -> Field[[EdgeDim, KDim], float]: + grad_norm_psi_e = (psi_c(E2C[1]) - psi_c(E2C[0])) * inv_dual_edge_length return grad_norm_psi_e def grad_fd_tang( @@ -57,14 +52,12 @@ def grad_fd_tang( grad_tang_psi_e[llb:, i] = tangent_orientation[llb:] * (psi_v[e2v[llb:, 1], i] - psi_v[e2v[llb:, 0], i]) * inv_primal_edge_length[llb:] return grad_tang_psi_e +@field_operator def compute_ddxn_z_half_e( - z_ifc: np.array, - inv_dual_edge_length: np.array, - e2c: np.array, - second_boundary_layer_start_index: np.int32, -) -> np.array: - nlev = z_ifc.shape[1] - ddxn_z_half_e = grad_fd_norm(z_ifc, inv_dual_edge_length, e2c, second_boundary_layer_start_index, nlev) + z_ifc: Field[[CellDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], +) -> Field[[EdgeDim, KDim], float]: + ddxn_z_half_e = grad_fd_norm(z_ifc, inv_dual_edge_length) return ddxn_z_half_e def compute_ddxt_z_half_e( From 9ecd7e345c09b5b01248c8751bf2d8ae664111c8 Mon Sep 17 00:00:00 2001 From: Andreas Date: Wed, 10 Apr 2024 11:18:50 +0200 Subject: [PATCH 07/65] progress gt4py --- .../interpolation/interpolation_fields3.py | 42 ++++++++--------- .../test_interpolation_fields3.py | 46 ++++++++++++------- 2 files changed, 48 insertions(+), 40 deletions(-) diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py index 70399f6cb..22c3fbf75 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py @@ -1,4 +1,4 @@ -# ICON4Py - ICON inspired code in Python and GT4Py +# icon4pY - icon INSPIRED CODE IN pYTHON AND gt4Py # # Copyright (c) 2022, ETH Zurich and MeteoSwiss # All rights reserved. @@ -28,7 +28,7 @@ from gt4py.next.ffront.decorator import field_operator from gt4py.next.ffront.fbuiltins import Field -from icon4py.model.common.dimension import C2E, E2C, V2E, C2EDim, CellDim, EdgeDim, V2EDim, VertexDim, KDim +from icon4py.model.common.dimension import C2E, E2C, V2E, E2V, C2EDim, CellDim, EdgeDim, V2EDim, VertexDim, KDim, Koff @field_operator def grad_fd_norm( @@ -38,18 +38,13 @@ def grad_fd_norm( grad_norm_psi_e = (psi_c(E2C[1]) - psi_c(E2C[0])) * inv_dual_edge_length return grad_norm_psi_e +@field_operator def grad_fd_tang( - psi_v: np.array, - inv_primal_edge_length: np.array, - tangent_orientation: np.array, - e2v: np.array, - third_boundary_layer_start_index: np.int32, - nlev, -) -> np.array: - llb = third_boundary_layer_start_index - grad_tang_psi_e = np.zeros([e2v.shape[0], nlev]) - for i in range(nlev): - grad_tang_psi_e[llb:, i] = tangent_orientation[llb:] * (psi_v[e2v[llb:, 1], i] - psi_v[e2v[llb:, 0], i]) * inv_primal_edge_length[llb:] + psi_v: Field[[VertexDim, KDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + tangent_orientation: Field[[EdgeDim], float], +) -> Field[[EdgeDim, KDim], float]: + grad_tang_psi_e = tangent_orientation * (psi_v(E2V[1]) - psi_v(E2V[0])) * inv_primal_edge_length return grad_tang_psi_e @field_operator @@ -60,21 +55,20 @@ def compute_ddxn_z_half_e( ddxn_z_half_e = grad_fd_norm(z_ifc, inv_dual_edge_length) return ddxn_z_half_e +@field_operator def compute_ddxt_z_half_e( - z_ifv: np.array, - inv_primal_edge_length: np.array, - tangent_orientation: np.array, - e2v: np.array, - third_boundary_layer_start_index: np.int32, -) -> np.array: - nlev = z_ifv.shape[1] - ddxt_z_half_e = grad_fd_tang(z_ifv, inv_primal_edge_length, tangent_orientation, e2v, third_boundary_layer_start_index, nlev) + z_ifv: Field[[VertexDim, KDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + tangent_orientation: Field[[EdgeDim], float], +) -> Field[[EdgeDim, KDim], float]: + ddxt_z_half_e = grad_fd_tang(z_ifv, inv_primal_edge_length, tangent_orientation) return ddxt_z_half_e +@field_operator def compute_ddxnt_z_full( - z_ddxnt_z_half_e: np.array, -) -> np.array: - ddxnt_z_full = 0.5 * (z_ddxnt_z_half_e[:, :z_ddxnt_z_half_e.shape[1]-1] + z_ddxnt_z_half_e[:, 1:]) + z_ddxnt_z_half_e: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + ddxnt_z_full = 0.5 * (z_ddxnt_z_half_e + z_ddxnt_z_half_e(Koff[1])) return ddxnt_z_full def compute_cells2verts_scalar( diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields3.py b/model/common/tests/interpolation_tests/test_interpolation_fields3.py index 346626cfb..1efdfa13f 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_fields3.py +++ b/model/common/tests/interpolation_tests/test_interpolation_fields3.py @@ -26,6 +26,7 @@ import numpy as np import pytest from gt4py.next.iterator.builtins import int32 +from gt4py.next import as_field from icon4py.model.common.dimension import ( C2E2CDim, @@ -77,7 +78,6 @@ def test_compute_ddxn_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri ) vertical_start = 0 vertical_end = 66 - ddxn_z_half_e = zero_field(icon_grid, EdgeDim, KDim) ddxn_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) compute_ddxn_z_half_e( z_ifc, @@ -89,14 +89,14 @@ def test_compute_ddxn_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri KDim: (vertical_start, vertical_end), }, ) - print(ddxn_z_half_e.asnumpy().shape) - ddxn_z_full = compute_ddxnt_z_full( - ddxn_z_half_e.asnumpy(), + ddxn_z_full = zero_field(icon_grid, EdgeDim, KDim) + compute_ddxnt_z_full( + ddxn_z_half_e, + out=ddxn_z_full, + offset_provider={"Koff" : icon_grid.get_offset_provider("Koff")}, ) - print(ddxn_z_full.shape) - print(ddxn_z_full_ref.shape) - assert np.allclose(ddxn_z_full, ddxn_z_full_ref) + assert np.allclose(ddxn_z_full.asnumpy(), ddxn_z_full_ref) @pytest.mark.datatest @@ -105,8 +105,8 @@ def test_compute_ddxt_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri dual_area = grid_savepoint.v_dual_area().asnumpy() edge_vert_length = grid_savepoint.edge_vert_length().asnumpy() edge_cell_length = grid_savepoint.edge_cell_length().asnumpy() - tangent_orientation = grid_savepoint.tangent_orientation().asnumpy() - inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths().asnumpy() + tangent_orientation = grid_savepoint.tangent_orientation() + inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() owner_mask = grid_savepoint.v_owner_mask() ddxt_z_full_ref = metrics_savepoint.ddxt_z_full().asnumpy() e2c = icon_grid.connectivities[E2CDim] @@ -129,10 +129,14 @@ def test_compute_ddxt_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) - 1, ) - third_boundary_layer_start_index_edge = icon_grid.get_start_index( + horizontal_start_edge = icon_grid.get_start_index( EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, ) + horizontal_end_edge = icon_grid.get_end_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, + ) cells_aw_verts = compute_cells_aw_verts( dual_area, edge_vert_length, @@ -149,15 +153,25 @@ def test_compute_ddxt_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri assert np.allclose(cells_aw_verts, cells_aw_verts_ref) z_ifv = compute_cells2verts_scalar(z_ifc, cells_aw_verts, v2c, second_boundary_layer_start_index_vertex) - ddxt_z_half_e = compute_ddxt_z_half_e( - z_ifv, + vertical_start = 0 + vertical_end = 66 + ddxt_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) + compute_ddxt_z_half_e( + as_field((VertexDim, KDim), z_ifv), inv_primal_edge_length, tangent_orientation, - e2v, - third_boundary_layer_start_index_edge, + out=ddxt_z_half_e, + offset_provider={"E2V" : icon_grid.get_offset_provider("E2V") }, + domain={ + EdgeDim: (horizontal_start_edge, horizontal_end_edge), + KDim: (vertical_start, vertical_end), + }, ) - ddxt_z_full = compute_ddxnt_z_full( + ddxt_z_full = zero_field(icon_grid, EdgeDim, KDim) + compute_ddxnt_z_full( ddxt_z_half_e, + out=ddxt_z_full, + offset_provider={"Koff" : icon_grid.get_offset_provider("Koff")}, ) - assert np.allclose(ddxt_z_full, ddxt_z_full_ref) + assert np.allclose(ddxt_z_full.asnumpy(), ddxt_z_full_ref) From bac95e4a7c78b81cf4e932fa2e6e45283573d32b Mon Sep 17 00:00:00 2001 From: Andreas Date: Wed, 10 Apr 2024 13:38:15 +0200 Subject: [PATCH 08/65] progress numpy to gt4py --- .../interpolation/interpolation_fields3.py | 20 ++++-------- .../test_interpolation_fields3.py | 32 ++++++++++++------- 2 files changed, 28 insertions(+), 24 deletions(-) diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py index 22c3fbf75..0aa4d38e9 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py @@ -24,11 +24,11 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -from gt4py.next import where +from gt4py.next import where, neighbor_sum from gt4py.next.ffront.decorator import field_operator from gt4py.next.ffront.fbuiltins import Field -from icon4py.model.common.dimension import C2E, E2C, V2E, E2V, C2EDim, CellDim, EdgeDim, V2EDim, VertexDim, KDim, Koff +from icon4py.model.common.dimension import C2E, E2C, V2E, E2V, V2C, C2EDim, CellDim, EdgeDim, V2EDim, V2CDim, VertexDim, KDim, Koff @field_operator def grad_fd_norm( @@ -71,18 +71,12 @@ def compute_ddxnt_z_full( ddxnt_z_full = 0.5 * (z_ddxnt_z_half_e + z_ddxnt_z_half_e(Koff[1])) return ddxnt_z_full +@field_operator def compute_cells2verts_scalar( - p_cell_in: np.array, - c_int: np.array, - v2c: np.array, - second_boundary_layer_start_index: np.int32, -) -> np.array: - llb = second_boundary_layer_start_index - kdim = p_cell_in.shape[1] - p_vert_out = np.zeros([c_int.shape[0], kdim]) - for j in range(kdim): - for i in range(6): - p_vert_out[llb:, j] = p_vert_out[llb:, j] + c_int[llb:, i] * p_cell_in[v2c[llb:, i], j] + p_cell_in: Field[[CellDim, KDim], float], + c_int: Field[[VertexDim, V2CDim], float], +) -> Field[[VertexDim, KDim], float]: + p_vert_out = neighbor_sum(c_int * p_cell_in(V2C), axis=V2CDim) return p_vert_out def compute_cells_aw_verts( diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields3.py b/model/common/tests/interpolation_tests/test_interpolation_fields3.py index 1efdfa13f..2b6d99c34 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_fields3.py +++ b/model/common/tests/interpolation_tests/test_interpolation_fields3.py @@ -101,7 +101,7 @@ def test_compute_ddxn_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri @pytest.mark.datatest def test_compute_ddxt_z_full_e(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): - z_ifc = metrics_savepoint.z_ifc().asnumpy() + z_ifc = metrics_savepoint.z_ifc() dual_area = grid_savepoint.v_dual_area().asnumpy() edge_vert_length = grid_savepoint.edge_vert_length().asnumpy() edge_cell_length = grid_savepoint.edge_cell_length().asnumpy() @@ -113,19 +113,19 @@ def test_compute_ddxt_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri v2c = icon_grid.connectivities[V2CDim] v2e = icon_grid.connectivities[V2EDim] e2v = icon_grid.connectivities[E2VDim] - second_boundary_layer_start_index_vertex = icon_grid.get_start_index( + horizontal_start_vertex = icon_grid.get_start_index( VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, ) - second_boundary_layer_end_index_vertex = icon_grid.get_end_index( + horizontal_end_vertex = icon_grid.get_end_index( VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) - 1, ) - second_boundary_layer_start_index_cell = icon_grid.get_start_index( + horizontal_start_cell = icon_grid.get_start_index( CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, ) - second_boundary_layer_end_index_cell = icon_grid.get_end_index( + horizontal_end_cell = icon_grid.get_end_index( CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) - 1, ) @@ -137,6 +137,8 @@ def test_compute_ddxt_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, ) + vertical_start = 0 + vertical_end = 66 cells_aw_verts = compute_cells_aw_verts( dual_area, edge_vert_length, @@ -146,18 +148,26 @@ def test_compute_ddxt_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri v2c, v2e, e2v, - second_boundary_layer_start_index_vertex, - second_boundary_layer_end_index_vertex, + horizontal_start_vertex, + horizontal_end_vertex, ) cells_aw_verts_ref = interpolation_savepoint.c_intp().asnumpy() assert np.allclose(cells_aw_verts, cells_aw_verts_ref) - z_ifv = compute_cells2verts_scalar(z_ifc, cells_aw_verts, v2c, second_boundary_layer_start_index_vertex) - vertical_start = 0 - vertical_end = 66 + z_ifv = zero_field(icon_grid, VertexDim, KDim, extend={KDim: 1}) + compute_cells2verts_scalar( + z_ifc, + as_field((VertexDim, V2CDim), cells_aw_verts), + out=z_ifv, + offset_provider={"V2C" : icon_grid.get_offset_provider("V2C") }, + domain={ + VertexDim: (horizontal_start_vertex, horizontal_end_vertex), + KDim: (vertical_start, vertical_end), + }, + ) ddxt_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) compute_ddxt_z_half_e( - as_field((VertexDim, KDim), z_ifv), + z_ifv, inv_primal_edge_length, tangent_orientation, out=ddxt_z_half_e, From 596ec48423eb5f084372f9fc0ad8420254097ba2 Mon Sep 17 00:00:00 2001 From: Andreas Date: Fri, 19 Apr 2024 09:19:36 +0200 Subject: [PATCH 09/65] progress ddqz_z_full_e --- .../interpolation/interpolation_fields3.py | 19 ++++++- .../test_interpolation_fields3.py | 51 +++++++++++++++++++ 2 files changed, 69 insertions(+), 1 deletion(-) diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py index 0aa4d38e9..ad9ec0bbd 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py @@ -28,7 +28,7 @@ from gt4py.next.ffront.decorator import field_operator from gt4py.next.ffront.fbuiltins import Field -from icon4py.model.common.dimension import C2E, E2C, V2E, E2V, V2C, C2EDim, CellDim, EdgeDim, V2EDim, V2CDim, VertexDim, KDim, Koff +from icon4py.model.common.dimension import C2E, E2C, V2E, E2V, V2C, C2EDim, CellDim, E2CDim, EdgeDim, V2EDim, V2CDim, VertexDim, KDim, Koff @field_operator def grad_fd_norm( @@ -103,3 +103,20 @@ def compute_cells_aw_verts( for k in range (6): cells_aw_verts[llb:, k] = np.where(np.logical_and(owner_mask[:], e2c[v2e[llb:, j], i] == v2c[llb:, k]), cells_aw_verts[llb:, k] + 0.5 / dual_area[llb:] * edge_vert_length[v2e[llb:, j], idx_ve[llb:, j]] * edge_cell_length[v2e[llb:, j], i], cells_aw_verts[llb:, k]) return cells_aw_verts + +#@field_operator +#def compute_cells_aw_verts_new( +# dual_area: Field[VertexDim, float], +# edge_vert_length: Field[[EdgeDim, V2CDim], float], +# edge_cell_length: Field[[EdgeDim, E2CDim], float], +#) -> Field[[VertexDim, V2CDim], float]: +# cells_aw_verts = 0.5 * neighbor_sum(where(E2C(V2E) == V2C, neighbor_sum(edge_vert_length(V2E) * edge_cell_length(V2E) / dual_area, axis=V2EDim), 0.0), axis=E2CDim) +# return cells_aw_verts + +@field_operator +def compute_cells2edges_scalar( + inv_p_cell_in: Field[[CellDim, KDim], float], + c_int: Field[[EdgeDim, E2CDim], float], +) -> Field[[EdgeDim, KDim], float]: + p_vert_out = neighbor_sum(c_int / inv_p_cell_in(E2C), axis=E2CDim) + return p_vert_out diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields3.py b/model/common/tests/interpolation_tests/test_interpolation_fields3.py index 2b6d99c34..4ea59e2f9 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_fields3.py +++ b/model/common/tests/interpolation_tests/test_interpolation_fields3.py @@ -48,6 +48,7 @@ compute_cells_aw_verts, compute_cells2verts_scalar, compute_ddxt_z_half_e, + compute_cells2edges_scalar, ) from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package data_provider, @@ -185,3 +186,53 @@ def test_compute_ddxt_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri ) assert np.allclose(ddxt_z_full.asnumpy(), ddxt_z_full_ref) + +@pytest.mark.datatest +def test_compute_ddqz_z_full_e(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): + inv_ddqz_z_full = metrics_savepoint.inv_ddqz_z_full() + c_lin_e = interpolation_savepoint.c_lin_e() + e2c = icon_grid.connectivities[E2CDim] + v2c = icon_grid.connectivities[V2CDim] + v2e = icon_grid.connectivities[V2EDim] + e2v = icon_grid.connectivities[E2VDim] + horizontal_start_vertex = icon_grid.get_start_index( + VertexDim, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, + ) + horizontal_end_vertex = icon_grid.get_end_index( + VertexDim, + HorizontalMarkerIndex.lateral_boundary(VertexDim) - 1, + ) + horizontal_start_cell = icon_grid.get_start_index( + CellDim, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, + ) + horizontal_end_cell = icon_grid.get_end_index( + CellDim, + HorizontalMarkerIndex.lateral_boundary(CellDim) - 1, + ) + horizontal_start_edge = icon_grid.get_start_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, + ) + horizontal_end_edge = icon_grid.get_end_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, + ) + vertical_start = 0 + vertical_end = 65 + ddqz_z_full_e = zero_field(icon_grid, EdgeDim, KDim) + compute_cells2edges_scalar( + inv_ddqz_z_full, + c_lin_e, + out=ddqz_z_full_e, + offset_provider={"E2C" : icon_grid.get_offset_provider("E2C") }, + domain={ + EdgeDim: (horizontal_start_edge, horizontal_end_edge), + KDim: (vertical_start, vertical_end), + }, + ) + ddqz_z_full_e_ref = metrics_savepoint.ddqz_z_full_e().asnumpy() + print(ddqz_z_full_e_ref) + print(ddqz_z_full_e.asnumpy()) + assert np.allclose(ddqz_z_full_e.asnumpy(), ddqz_z_full_e_ref) From cb68e6cff94164d8f2c3d399f8761ada948ac70b Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 19 Apr 2024 12:18:10 +0200 Subject: [PATCH 10/65] partial implementation of vwind_impl_wgt and vwind_expl_wgt --- .../model/common/metrics/metric_fields.py | 96 ++++++++++++++++++- .../tests/metric_tests/test_metric_fields.py | 47 +++++++++ 2 files changed, 141 insertions(+), 2 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index f89b9f6dd..eaeaae211 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -11,9 +11,9 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next import Field, GridType, field_operator, int32, program, where +from gt4py.next import Field, GridType, abs, field_operator, int32, maximum, minimum, program, where -from icon4py.model.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import C2E, CellDim, EdgeDim, KDim, Koff from icon4py.model.common.math.helpers import ( average_k_level_up, difference_k_level_down, @@ -151,3 +151,95 @@ def compute_ddqz_z_full( out=(ddqz_z_full, inv_ddqz_z_full), domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, ) + + +@field_operator +def _compute_vwind_impl_wgt( + z_ddxn_z_half_e: Field[[EdgeDim], wpfloat], + z_ddxt_z_half_e: Field[[EdgeDim], wpfloat], + dual_edge_length: Field[[EdgeDim], wpfloat], + vwind_offctr: wpfloat, +) -> Field[[CellDim], wpfloat]: + z_ddx_1 = maximum(abs(z_ddxn_z_half_e(C2E[1])), abs(z_ddxt_z_half_e(C2E[1]))) + z_ddx_2 = maximum(abs(z_ddxn_z_half_e(C2E[2])), abs(z_ddxt_z_half_e(C2E[2]))) + z_ddx_3 = maximum(abs(z_ddxn_z_half_e(C2E[3])), abs(z_ddxt_z_half_e(C2E[3]))) + z_ddx_1_2 = maximum(z_ddx_1, z_ddx_2) + z_maxslope = maximum(z_ddx_1_2, z_ddx_3) + + z_diff_1_2 = maximum( + abs(z_ddxn_z_half_e(C2E[1]) * dual_edge_length(C2E[1])), + abs(z_ddxn_z_half_e(C2E[2]) * dual_edge_length(C2E[2])), + ) + z_diff = maximum(z_diff_1_2, abs(z_ddxn_z_half_e(C2E[3]) * dual_edge_length(C2E[3]))) + z_offctr_1 = maximum(vwind_offctr, 0.425 * z_maxslope**0.75) + z_offctr = maximum(z_offctr_1, minimum(0.25, 2.5e-4 * (z_diff - 250.0))) + z_offctr = minimum(maximum(vwind_offctr, 0.75), z_offctr) + vwind_impl_wgt = 0.5 + z_offctr + return vwind_impl_wgt + + +@program(grid_type=GridType.UNSTRUCTURED) +def compute_vwind_impl_wgt( + z_ddxn_z_half_e: Field[[EdgeDim], wpfloat], + z_ddxt_z_half_e: Field[[EdgeDim], wpfloat], + dual_edge_length: Field[[EdgeDim], wpfloat], + vwind_impl_wgt: Field[[CellDim], wpfloat], + vwind_offctr: wpfloat, + horizontal_start: int32, + horizontal_end: int32, +): + """ + Compute vwind_impl_wgt. + + See mo_vertical_grid.f90 + + Args: + z_ddxn_z_half_e: intermediate storage for field + z_ddxt_z_half_e: intermediate storage for field + dual_edge_length: dual_edge_length + vwind_impl_wgt: (output) offcentering in vertical mass flux + vwind_offctr: off-centering in vertical wind solver + horizontal_start: horizontal start index + horizontal_end: horizontal end index + + """ + _compute_vwind_impl_wgt( + z_ddxn_z_half_e=z_ddxn_z_half_e, + z_ddxt_z_half_e=z_ddxt_z_half_e, + dual_edge_length=dual_edge_length, + vwind_offctr=vwind_offctr, + out=vwind_impl_wgt, + domain={CellDim: (horizontal_start, horizontal_end)}, + ) + + +@field_operator +def _compute_vwind_expl_wgt(vwind_impl_wgt: Field[[CellDim], wpfloat]) -> Field[[CellDim], wpfloat]: + return 1.0 - vwind_impl_wgt + + +@program(grid_type=GridType.UNSTRUCTURED) +def compute_vwind_expl_wgt( + vwind_impl_wgt: Field[[CellDim], wpfloat], + vwind_expl_wgt: Field[[CellDim], wpfloat], + horizontal_start: int32, + horizontal_end: int32, +): + """ + Compute vwind_expl_wgt. + + See mo_vertical_grid.f90 + + Args: + vwind_impl_wgt: offcentering in vertical mass flux + vwind_expl_wgt: (output) 1 - of vwind_impl_wgt + horizontal_start: horizontal start index + horizontal_end: horizontal end index + + """ + + _compute_vwind_expl_wgt( + vwind_impl_wgt=vwind_impl_wgt, + out=vwind_expl_wgt, + domain={CellDim: (horizontal_start, horizontal_end)}, + ) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index f4ba7dad2..f76750fcd 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -20,6 +20,7 @@ from icon4py.model.common.metrics.metric_fields import ( compute_ddqz_z_full, compute_ddqz_z_half, + compute_vwind_expl_wgt, compute_z_mc, ) from icon4py.model.common.test_utils.helpers import ( @@ -120,3 +121,49 @@ def test_compute_ddqz_z_full(icon_grid, metrics_savepoint, backend): ) assert dallclose(inv_ddqz_z_full.asnumpy(), inv_ddqz_full_ref.asnumpy()) + + +# @pytest.mark.datatest +# def test_compute_vwind_impl_wgt(icon_grid, metrics_savepoint, backend): +# if is_roundtrip(backend): +# pytest.skip("skipping: slow backend") +# +# z_ddxn_z_half_e +# z_ddxt_z_half_e +# dual_edge_length +# vwind_offctr +# vwind_impl_wgt_full = zero_field(icon_grid, CellDim, KDim) +# vwind_impl_wgt_ref = metrics_savepoint.vwind_impl_wgt() +# +# compute_vwind_impl_wgt.with_backend(backend)( +# z_ddxn_z_half_e=z_ddxn_z_half_e, +# z_ddxt_z_half_e=z_ddxt_z_half_e, +# dual_edge_length=dual_edge_length, +# vwind_impl_wgt=vwind_impl_wgt_full, +# vwind_offctr=vwind_offctr, +# horizontal_start=int32(0), +# horizontal_end=icon_grid.num_cells, +# offset_provider = {"C2E": icon_grid.get_offset_provider("C2E")}, +# ) +# +# assert dallclose(vwind_impl_wgt_full.asnumpy(), vwind_impl_wgt_ref.asnumpy()) + + +@pytest.mark.datatest +def test_compute_vwind_expl_wgt(icon_grid, metrics_savepoint, backend): + if is_roundtrip(backend): + pytest.skip("skipping: slow backend") + + vwind_expl_wgt_full = zero_field(icon_grid, CellDim) + vwind_expl_wgt_ref = metrics_savepoint.vwind_expl_wgt() + vwind_impl_wgt = metrics_savepoint.vwind_impl_wgt() + + compute_vwind_expl_wgt.with_backend(backend)( + vwind_impl_wgt=vwind_impl_wgt, + vwind_expl_wgt=vwind_expl_wgt_full, + horizontal_start=int32(0), + horizontal_end=icon_grid.num_cells, + offset_provider={"C2E": icon_grid.get_offset_provider("C2E")}, + ) + + assert dallclose(vwind_expl_wgt_full.asnumpy(), vwind_expl_wgt_ref.asnumpy()) From 3d715780063b7072688c55d52803140c3e963399 Mon Sep 17 00:00:00 2001 From: Andreas Date: Tue, 23 Apr 2024 13:31:04 +0200 Subject: [PATCH 11/65] small progress --- .../interpolation_tests/test_interpolation_fields3.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields3.py b/model/common/tests/interpolation_tests/test_interpolation_fields3.py index 4ea59e2f9..b67d4f72b 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_fields3.py +++ b/model/common/tests/interpolation_tests/test_interpolation_fields3.py @@ -59,7 +59,10 @@ ranked_data_path, ) from icon4py.model.common.test_utils.helpers import zero_field - +from icon4py.model.common.test_utils.datatest_utils import ( + GLOBAL_EXPERIMENT, + REGIONAL_EXPERIMENT, +) @pytest.mark.datatest def test_compute_ddxn_z_full_e(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): # fixture @@ -78,7 +81,7 @@ def test_compute_ddxn_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, ) vertical_start = 0 - vertical_end = 66 + vertical_end = icon_grid.num_levels + 1 ddxn_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) compute_ddxn_z_half_e( z_ifc, @@ -188,6 +191,7 @@ def test_compute_ddxt_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri assert np.allclose(ddxt_z_full.asnumpy(), ddxt_z_full_ref) @pytest.mark.datatest +@pytest.mark.parametrize("experiment", (REGIONAL_EXPERIMENT,GLOBAL_EXPERIMENT)) def test_compute_ddqz_z_full_e(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): inv_ddqz_z_full = metrics_savepoint.inv_ddqz_z_full() c_lin_e = interpolation_savepoint.c_lin_e() @@ -220,8 +224,9 @@ def test_compute_ddqz_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, ) vertical_start = 0 - vertical_end = 65 + vertical_end = icon_grid.num_levels ddqz_z_full_e = zero_field(icon_grid, EdgeDim, KDim) +# if icon_grid.limited_area compute_cells2edges_scalar( inv_ddqz_z_full, c_lin_e, From 61ac8524cd5afd834f04a7d81db5351441f79872 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Wed, 24 Apr 2024 10:19:50 +0200 Subject: [PATCH 12/65] Update model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py Co-authored-by: Magdalena --- .../common/interpolation/interpolation_fields3.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py index ad9ec0bbd..f0176ea3a 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py @@ -47,7 +47,18 @@ def grad_fd_tang( grad_tang_psi_e = tangent_orientation * (psi_v(E2V[1]) - psi_v(E2V[0])) * inv_primal_edge_length return grad_tang_psi_e -@field_operator +@program +def compute_ddxn_z_half_e( + z_ifc: Field[[CellDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + ddxn_z_half_e: Field[[EdgeDim, KDim], float], + horizontal_lower: int32, + horizontal_upper: int32, + vertical_lower: int32, + vertical_upper: int32, +) : + grad_fd_norm(z_ifc, inv_dual_edge_length, out=ddxn_z_half_e, domain={EdgeDim: (horizontal_lower, horizontal_upper), KDim: (vertical_lower, vertical_upper)}) + def compute_ddxn_z_half_e( z_ifc: Field[[CellDim, KDim], float], inv_dual_edge_length: Field[[EdgeDim], float], From c489f8a61d82a42e8b53b247fb20f5be4fb5e023 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Wed, 24 Apr 2024 11:14:27 +0200 Subject: [PATCH 13/65] applied changes after PR review --- .../icon4py/model/common/grid/horizontal.py | 52 ++++++- .../interpolation/interpolation_fields3.py | 122 ++++++++-------- .../src/icon4py/model/common/math/helpers.py | 22 ++- .../test_interpolation_fields3.py | 133 +++++++----------- 4 files changed, 185 insertions(+), 144 deletions(-) diff --git a/model/common/src/icon4py/model/common/grid/horizontal.py b/model/common/src/icon4py/model/common/grid/horizontal.py index 22030c5d7..678670711 100644 --- a/model/common/src/icon4py/model/common/grid/horizontal.py +++ b/model/common/src/icon4py/model/common/grid/horizontal.py @@ -17,7 +17,18 @@ from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.common import dimension -from icon4py.model.common.dimension import E2C, CellDim, E2CDim, ECDim, ECVDim, EdgeDim, KDim +from icon4py.model.common.dimension import ( + E2C, + V2C, + CellDim, + E2CDim, + ECDim, + ECVDim, + EdgeDim, + KDim, + V2CDim, + VertexDim, +) from icon4py.model.common.type_alias import wpfloat @@ -343,3 +354,42 @@ def boundary_nudging_start(cls, dim: Dimension) -> int: raise ValueError( f"nudging start level only exists for {CellDim} and {EdgeDim}" ) from err + + +@field_operator +def _compute_cells2edges_scalar( + inv_p_cell_in: Field[[CellDim, KDim], float], + c_int: Field[[EdgeDim, E2CDim], float], +) -> Field[[EdgeDim, KDim], float]: + p_vert_out = neighbor_sum(c_int / inv_p_cell_in(E2C), axis=E2CDim) + return p_vert_out + + +@program +def compute_cells2edges_scalar( + inv_p_cell_in: Field[[CellDim, KDim], float], + c_int: Field[[EdgeDim, E2CDim], float], + p_vert_out: Field[[EdgeDim, KDim], float], + horizontal_start_edge: int32, + horizontal_end_edge: int32, + vertical_start: int32, + vertical_end: int32, +): + _compute_cells2edges_scalar( + inv_p_cell_in, + c_int, + out=p_vert_out, + domain={ + EdgeDim: (horizontal_start_edge, horizontal_end_edge), + KDim: (vertical_start, vertical_end), + }, + ) + + +@field_operator +def _compute_cells2verts_scalar( + p_cell_in: Field[[CellDim, KDim], float], + c_int: Field[[VertexDim, V2CDim], float], +) -> Field[[VertexDim, KDim], float]: + p_vert_out = neighbor_sum(c_int * p_cell_in(V2C), axis=V2CDim) + return p_vert_out diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py index f0176ea3a..c4e20b50f 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py @@ -1,4 +1,6 @@ # icon4pY - icon INSPIRED CODE IN pYTHON AND gt4Py +# TODO: This license is not consistent with license used in the project. +# Delete the inconsistent license and above line and rerun pre-commit to insert a good license. # # Copyright (c) 2022, ETH Zurich and MeteoSwiss # All rights reserved. @@ -24,22 +26,21 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -from gt4py.next import where, neighbor_sum -from gt4py.next.ffront.decorator import field_operator -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next import Field, field_operator, int32, program -from icon4py.model.common.dimension import C2E, E2C, V2E, E2V, V2C, C2EDim, CellDim, E2CDim, EdgeDim, V2EDim, V2CDim, VertexDim, KDim, Koff +from icon4py.model.common.dimension import ( + E2V, + CellDim, + EdgeDim, + KDim, + Koff, + VertexDim, +) +from icon4py.model.common.math.helpers import grad_fd_norm -@field_operator -def grad_fd_norm( - psi_c: Field[[CellDim, KDim], float], - inv_dual_edge_length: Field[[EdgeDim], float], -) -> Field[[EdgeDim, KDim], float]: - grad_norm_psi_e = (psi_c(E2C[1]) - psi_c(E2C[0])) * inv_dual_edge_length - return grad_norm_psi_e @field_operator -def grad_fd_tang( +def _grad_fd_tang( psi_v: Field[[VertexDim, KDim], float], inv_primal_edge_length: Field[[EdgeDim], float], tangent_orientation: Field[[EdgeDim], float], @@ -47,6 +48,7 @@ def grad_fd_tang( grad_tang_psi_e = tangent_orientation * (psi_v(E2V[1]) - psi_v(E2V[0])) * inv_primal_edge_length return grad_tang_psi_e + @program def compute_ddxn_z_half_e( z_ifc: Field[[CellDim, KDim], float], @@ -54,41 +56,50 @@ def compute_ddxn_z_half_e( ddxn_z_half_e: Field[[EdgeDim, KDim], float], horizontal_lower: int32, horizontal_upper: int32, - vertical_lower: int32, + vertical_lower: int32, vertical_upper: int32, -) : - grad_fd_norm(z_ifc, inv_dual_edge_length, out=ddxn_z_half_e, domain={EdgeDim: (horizontal_lower, horizontal_upper), KDim: (vertical_lower, vertical_upper)}) - -def compute_ddxn_z_half_e( - z_ifc: Field[[CellDim, KDim], float], - inv_dual_edge_length: Field[[EdgeDim], float], -) -> Field[[EdgeDim, KDim], float]: - ddxn_z_half_e = grad_fd_norm(z_ifc, inv_dual_edge_length) - return ddxn_z_half_e +): + grad_fd_norm( + z_ifc, + inv_dual_edge_length, + out=ddxn_z_half_e, + domain={ + EdgeDim: (horizontal_lower, horizontal_upper), + KDim: (vertical_lower, vertical_upper), + }, + ) -@field_operator + +@program def compute_ddxt_z_half_e( z_ifv: Field[[VertexDim, KDim], float], inv_primal_edge_length: Field[[EdgeDim], float], tangent_orientation: Field[[EdgeDim], float], -) -> Field[[EdgeDim, KDim], float]: - ddxt_z_half_e = grad_fd_tang(z_ifv, inv_primal_edge_length, tangent_orientation) - return ddxt_z_half_e + ddxt_z_half_e: Field[[EdgeDim, KDim], float], + horizontal_lower: int32, + horizontal_upper: int32, + vertical_lower: int32, + vertical_upper: int32, +): + _grad_fd_tang( + z_ifv, + inv_primal_edge_length, + tangent_orientation, + out=ddxt_z_half_e, + domain={ + EdgeDim: (horizontal_lower, horizontal_upper), + KDim: (vertical_lower, vertical_upper), + }, + ) + @field_operator -def compute_ddxnt_z_full( +def _compute_ddxnt_z_full( z_ddxnt_z_half_e: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: ddxnt_z_full = 0.5 * (z_ddxnt_z_half_e + z_ddxnt_z_half_e(Koff[1])) return ddxnt_z_full -@field_operator -def compute_cells2verts_scalar( - p_cell_in: Field[[CellDim, KDim], float], - c_int: Field[[VertexDim, V2CDim], float], -) -> Field[[VertexDim, KDim], float]: - p_vert_out = neighbor_sum(c_int * p_cell_in(V2C), axis=V2CDim) - return p_vert_out def compute_cells_aw_verts( dual_area: np.array, @@ -99,35 +110,24 @@ def compute_cells_aw_verts( v2c: np.array, v2e: np.array, e2v: np.array, - second_boundary_layer_start_index: np.int32, - second_boundary_layer_end_index: np.int32, + horizontal_start: np.int32, + horizontal_end: np.int32, ) -> np.array: - llb = second_boundary_layer_start_index - cells_aw_verts = np.zeros([second_boundary_layer_end_index, 6]) - index = np.zeros([second_boundary_layer_end_index, 6]) - for j in range (6): - index[:, j] = np.arange(second_boundary_layer_end_index) + llb = horizontal_start + cells_aw_verts = np.zeros([horizontal_end, 6]) + index = np.repeat(np.arange(horizontal_end, dtype=float), 6).reshape(horizontal_end, 6) idx_ve = np.where(index == e2v[v2e, 0], 0, 1) for i in range(2): - for j in range (6): - for k in range (6): - cells_aw_verts[llb:, k] = np.where(np.logical_and(owner_mask[:], e2c[v2e[llb:, j], i] == v2c[llb:, k]), cells_aw_verts[llb:, k] + 0.5 / dual_area[llb:] * edge_vert_length[v2e[llb:, j], idx_ve[llb:, j]] * edge_cell_length[v2e[llb:, j], i], cells_aw_verts[llb:, k]) + for j in range(6): + for k in range(6): + cells_aw_verts[llb:, k] = np.where( + np.logical_and(owner_mask[:], e2c[v2e[llb:, j], i] == v2c[llb:, k]), + cells_aw_verts[llb:, k] + + 0.5 + / dual_area[llb:] + * edge_vert_length[v2e[llb:, j], idx_ve[llb:, j]] + * edge_cell_length[v2e[llb:, j], i], + cells_aw_verts[llb:, k], + ) return cells_aw_verts - -#@field_operator -#def compute_cells_aw_verts_new( -# dual_area: Field[VertexDim, float], -# edge_vert_length: Field[[EdgeDim, V2CDim], float], -# edge_cell_length: Field[[EdgeDim, E2CDim], float], -#) -> Field[[VertexDim, V2CDim], float]: -# cells_aw_verts = 0.5 * neighbor_sum(where(E2C(V2E) == V2C, neighbor_sum(edge_vert_length(V2E) * edge_cell_length(V2E) / dual_area, axis=V2EDim), 0.0), axis=E2CDim) -# return cells_aw_verts - -@field_operator -def compute_cells2edges_scalar( - inv_p_cell_in: Field[[CellDim, KDim], float], - c_int: Field[[EdgeDim, E2CDim], float], -) -> Field[[EdgeDim, KDim], float]: - p_vert_out = neighbor_sum(c_int / inv_p_cell_in(E2C), axis=E2CDim) - return p_vert_out diff --git a/model/common/src/icon4py/model/common/math/helpers.py b/model/common/src/icon4py/model/common/math/helpers.py index fec02c984..2f960e589 100644 --- a/model/common/src/icon4py/model/common/math/helpers.py +++ b/model/common/src/icon4py/model/common/math/helpers.py @@ -13,7 +13,7 @@ from gt4py.next import Field, field_operator -from icon4py.model.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim, Koff from icon4py.model.common.type_alias import wpfloat @@ -69,3 +69,23 @@ def difference_k_level_up( """ return half_level_field - half_level_field(Koff[1]) + + +@field_operator +def grad_fd_norm( + psi_c: Field[[CellDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], +) -> Field[[EdgeDim, KDim], float]: + """ + Calculate the gradient value of adjacent interface levels. + + Computes the difference of two offseted values multiplied by a field of the offseted dimension + Args: + psi_c: Field[[CellDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + + Returns: Field[[EdgeDim, KDim], float] + + """ + grad_norm_psi_e = (psi_c(E2C[1]) - psi_c(E2C[0])) * inv_dual_edge_length + return grad_norm_psi_e diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields3.py b/model/common/tests/interpolation_tests/test_interpolation_fields3.py index b67d4f72b..5d399c562 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_fields3.py +++ b/model/common/tests/interpolation_tests/test_interpolation_fields3.py @@ -25,30 +25,27 @@ import numpy as np import pytest -from gt4py.next.iterator.builtins import int32 from gt4py.next import as_field from icon4py.model.common.dimension import ( - C2E2CDim, - C2EDim, - CellDim, - E2C2EDim, E2CDim, E2VDim, EdgeDim, + KDim, V2CDim, V2EDim, VertexDim, - KDim, ) -from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex +from icon4py.model.common.grid.horizontal import ( + HorizontalMarkerIndex, + _compute_cells2verts_scalar, + compute_cells2edges_scalar, +) from icon4py.model.common.interpolation.interpolation_fields3 import ( - compute_ddxn_z_half_e, - compute_ddxnt_z_full, + _compute_ddxnt_z_full, compute_cells_aw_verts, - compute_cells2verts_scalar, + compute_ddxn_z_half_e, compute_ddxt_z_half_e, - compute_cells2edges_scalar, ) from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package data_provider, @@ -58,19 +55,19 @@ processor_props, ranked_data_path, ) -from icon4py.model.common.test_utils.helpers import zero_field from icon4py.model.common.test_utils.datatest_utils import ( GLOBAL_EXPERIMENT, REGIONAL_EXPERIMENT, ) +from icon4py.model.common.test_utils.helpers import zero_field + @pytest.mark.datatest -def test_compute_ddxn_z_full_e(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): # fixture +def test_compute_ddxn_z_full_e( + grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint +): z_ifc = metrics_savepoint.z_ifc() inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() - e2c = icon_grid.connectivities[E2CDim] -# edge_cell_length = grid_savepoint.edge_cell_length() -# owner_mask = grid_savepoint.e_owner_mask() ddxn_z_full_ref = metrics_savepoint.ddxn_z_full().asnumpy() horizontal_start = icon_grid.get_start_index( EdgeDim, @@ -86,25 +83,27 @@ def test_compute_ddxn_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri compute_ddxn_z_half_e( z_ifc, inv_dual_edge_length, - out=ddxn_z_half_e, - offset_provider={"E2C" : icon_grid.get_offset_provider("E2C") }, - domain={ - EdgeDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, + ddxn_z_half_e, + horizontal_start, + horizontal_end, + vertical_start, + vertical_end, + offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, ) ddxn_z_full = zero_field(icon_grid, EdgeDim, KDim) - compute_ddxnt_z_full( + _compute_ddxnt_z_full( ddxn_z_half_e, out=ddxn_z_full, - offset_provider={"Koff" : icon_grid.get_offset_provider("Koff")}, + offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) assert np.allclose(ddxn_z_full.asnumpy(), ddxn_z_full_ref) @pytest.mark.datatest -def test_compute_ddxt_z_full_e(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): +def test_compute_ddxt_z_full_e( + grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint +): z_ifc = metrics_savepoint.z_ifc() dual_area = grid_savepoint.v_dual_area().asnumpy() edge_vert_length = grid_savepoint.edge_vert_length().asnumpy() @@ -125,21 +124,13 @@ def test_compute_ddxt_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) - 1, ) - horizontal_start_cell = icon_grid.get_start_index( - CellDim, - HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, - ) - horizontal_end_cell = icon_grid.get_end_index( - CellDim, - HorizontalMarkerIndex.lateral_boundary(CellDim) - 1, - ) horizontal_start_edge = icon_grid.get_start_index( EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, ) horizontal_end_edge = icon_grid.get_end_index( EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, ) vertical_start = 0 vertical_end = 66 @@ -159,11 +150,11 @@ def test_compute_ddxt_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri assert np.allclose(cells_aw_verts, cells_aw_verts_ref) z_ifv = zero_field(icon_grid, VertexDim, KDim, extend={KDim: 1}) - compute_cells2verts_scalar( + _compute_cells2verts_scalar( z_ifc, as_field((VertexDim, V2CDim), cells_aw_verts), out=z_ifv, - offset_provider={"V2C" : icon_grid.get_offset_provider("V2C") }, + offset_provider={"V2C": icon_grid.get_offset_provider("V2C")}, domain={ VertexDim: (horizontal_start_vertex, horizontal_end_vertex), KDim: (vertical_start, vertical_end), @@ -174,70 +165,50 @@ def test_compute_ddxt_z_full_e(grid_savepoint, interpolation_savepoint, icon_gri z_ifv, inv_primal_edge_length, tangent_orientation, - out=ddxt_z_half_e, - offset_provider={"E2V" : icon_grid.get_offset_provider("E2V") }, - domain={ - EdgeDim: (horizontal_start_edge, horizontal_end_edge), - KDim: (vertical_start, vertical_end), - }, + ddxt_z_half_e, + horizontal_start_edge, + horizontal_end_edge, + vertical_start, + vertical_end, + offset_provider={"E2V": icon_grid.get_offset_provider("E2V")}, ) ddxt_z_full = zero_field(icon_grid, EdgeDim, KDim) - compute_ddxnt_z_full( + _compute_ddxnt_z_full( ddxt_z_half_e, out=ddxt_z_full, - offset_provider={"Koff" : icon_grid.get_offset_provider("Koff")}, + offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) assert np.allclose(ddxt_z_full.asnumpy(), ddxt_z_full_ref) + @pytest.mark.datatest -@pytest.mark.parametrize("experiment", (REGIONAL_EXPERIMENT,GLOBAL_EXPERIMENT)) -def test_compute_ddqz_z_full_e(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): +@pytest.mark.parametrize("experiment", (REGIONAL_EXPERIMENT, GLOBAL_EXPERIMENT)) +def test_compute_ddqz_z_full_e( + grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint +): inv_ddqz_z_full = metrics_savepoint.inv_ddqz_z_full() c_lin_e = interpolation_savepoint.c_lin_e() - e2c = icon_grid.connectivities[E2CDim] - v2c = icon_grid.connectivities[V2CDim] - v2e = icon_grid.connectivities[V2EDim] - e2v = icon_grid.connectivities[E2VDim] - horizontal_start_vertex = icon_grid.get_start_index( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, - ) - horizontal_end_vertex = icon_grid.get_end_index( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) - 1, - ) - horizontal_start_cell = icon_grid.get_start_index( - CellDim, - HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, - ) - horizontal_end_cell = icon_grid.get_end_index( - CellDim, - HorizontalMarkerIndex.lateral_boundary(CellDim) - 1, - ) + ddqz_z_full_e_ref = metrics_savepoint.ddqz_z_full_e().asnumpy() horizontal_start_edge = icon_grid.get_start_index( EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, ) - horizontal_end_edge = icon_grid.get_end_index( + horizontal_end_edge = icon_grid.get_end_index( # noqa: F841 EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, - ) + HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, + ) # for the global experiment, this outputs 0 vertical_start = 0 vertical_end = icon_grid.num_levels ddqz_z_full_e = zero_field(icon_grid, EdgeDim, KDim) -# if icon_grid.limited_area compute_cells2edges_scalar( - inv_ddqz_z_full, - c_lin_e, - out=ddqz_z_full_e, - offset_provider={"E2C" : icon_grid.get_offset_provider("E2C") }, - domain={ - EdgeDim: (horizontal_start_edge, horizontal_end_edge), - KDim: (vertical_start, vertical_end), - }, + inv_p_cell_in=inv_ddqz_z_full, + c_int=c_lin_e, + p_vert_out=ddqz_z_full_e, + horizontal_start_edge=horizontal_start_edge, + horizontal_end_edge=ddqz_z_full_e.shape[0], # horizontal_end_edge, + vertical_start=vertical_start, + vertical_end=vertical_end, + offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, ) - ddqz_z_full_e_ref = metrics_savepoint.ddqz_z_full_e().asnumpy() - print(ddqz_z_full_e_ref) - print(ddqz_z_full_e.asnumpy()) assert np.allclose(ddqz_z_full_e.asnumpy(), ddqz_z_full_e_ref) From 249220556f1d07944d248d1edcc4c656950e9ff3 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Wed, 24 Apr 2024 11:17:55 +0200 Subject: [PATCH 14/65] one more change --- .../tests/interpolation_tests/test_interpolation_fields3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields3.py b/model/common/tests/interpolation_tests/test_interpolation_fields3.py index 5d399c562..7d6811aef 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_fields3.py +++ b/model/common/tests/interpolation_tests/test_interpolation_fields3.py @@ -133,7 +133,7 @@ def test_compute_ddxt_z_full_e( HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, ) vertical_start = 0 - vertical_end = 66 + vertical_end = icon_grid.num_levels + 1 cells_aw_verts = compute_cells_aw_verts( dual_area, edge_vert_length, From 7412635486b497752f6f9c77a129db2956e6ecd8 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Wed, 24 Apr 2024 11:21:37 +0200 Subject: [PATCH 15/65] fixed pre-commit --- .../common/interpolation/interpolation_fields3.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py index c4e20b50f..c121c2acc 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py @@ -1,17 +1,3 @@ -# icon4pY - icon INSPIRED CODE IN pYTHON AND gt4Py -# TODO: This license is not consistent with license used in the project. -# Delete the inconsistent license and above line and rerun pre-commit to insert a good license. -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later # ICON4Py - ICON inspired code in Python and GT4Py # # Copyright (c) 2022, ETH Zurich and MeteoSwiss From da8d98a146bb3e9c6a89407db0cf7be7a1ae6dbf Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Wed, 24 Apr 2024 14:26:13 +0200 Subject: [PATCH 16/65] edits following review --- .../icon4py/model/common/grid/horizontal.py | 16 +- .../interpolation/interpolation_fields.py | 70 ++++++ .../interpolation/interpolation_fields3.py | 119 ---------- .../src/icon4py/model/common/math/helpers.py | 32 ++- .../model/common/metrics/metric_fields.py | 28 ++- .../test_interpolation_fields.py | 63 +++++- .../test_interpolation_fields3.py | 214 ------------------ .../tests/metric_tests/test_metric_fields.py | 142 +++++++++++- 8 files changed, 332 insertions(+), 352 deletions(-) delete mode 100644 model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py delete mode 100644 model/common/tests/interpolation_tests/test_interpolation_fields3.py diff --git a/model/common/src/icon4py/model/common/grid/horizontal.py b/model/common/src/icon4py/model/common/grid/horizontal.py index 678670711..01d7178d9 100644 --- a/model/common/src/icon4py/model/common/grid/horizontal.py +++ b/model/common/src/icon4py/model/common/grid/horizontal.py @@ -357,17 +357,17 @@ def boundary_nudging_start(cls, dim: Dimension) -> int: @field_operator -def _compute_cells2edges_scalar( - inv_p_cell_in: Field[[CellDim, KDim], float], +def _compute_cells2edges( + p_cell_in: Field[[CellDim, KDim], float], c_int: Field[[EdgeDim, E2CDim], float], ) -> Field[[EdgeDim, KDim], float]: - p_vert_out = neighbor_sum(c_int / inv_p_cell_in(E2C), axis=E2CDim) + p_vert_out = neighbor_sum(c_int * p_cell_in(E2C), axis=E2CDim) return p_vert_out @program -def compute_cells2edges_scalar( - inv_p_cell_in: Field[[CellDim, KDim], float], +def compute_cells2edges( + p_cell_in: Field[[CellDim, KDim], float], c_int: Field[[EdgeDim, E2CDim], float], p_vert_out: Field[[EdgeDim, KDim], float], horizontal_start_edge: int32, @@ -375,8 +375,8 @@ def compute_cells2edges_scalar( vertical_start: int32, vertical_end: int32, ): - _compute_cells2edges_scalar( - inv_p_cell_in, + _compute_cells2edges( + p_cell_in, c_int, out=p_vert_out, domain={ @@ -387,7 +387,7 @@ def compute_cells2edges_scalar( @field_operator -def _compute_cells2verts_scalar( +def _compute_cells2verts( p_cell_in: Field[[CellDim, KDim], float], c_int: Field[[VertexDim, V2CDim], float], ) -> Field[[VertexDim, KDim], float]: diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py index cd7fa4ab7..2b2618eef 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py @@ -24,6 +24,14 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +from gt4py.next import Field, int32, program + +from icon4py.model.common.dimension import ( + EdgeDim, + KDim, + VertexDim, +) +from icon4py.model.common.math.helpers import _grad_fd_tang, average_ek_level_up def compute_c_lin_e( @@ -49,3 +57,65 @@ def compute_c_lin_e( c_lin_e[0:second_boundary_layer_start_index, :] = 0.0 mask = np.transpose(np.tile(owner_mask, (2, 1))) return np.where(mask, c_lin_e, 0.0) + + +def compute_cells_aw_verts( + dual_area: np.array, + edge_vert_length: np.array, + edge_cell_length: np.array, + owner_mask: np.array, + e2c: np.array, + v2c: np.array, + v2e: np.array, + e2v: np.array, + horizontal_start: np.int32, + horizontal_end: np.int32, +) -> np.array: + llb = horizontal_start + cells_aw_verts = np.zeros([horizontal_end, 6]) + index = np.repeat(np.arange(horizontal_end, dtype=float), 6).reshape(horizontal_end, 6) + idx_ve = np.where(index == e2v[v2e, 0], 0, 1) + + for i in range(2): + for j in range(6): + for k in range(6): + cells_aw_verts[llb:, k] = np.where( + np.logical_and(owner_mask[:], e2c[v2e[llb:, j], i] == v2c[llb:, k]), + cells_aw_verts[llb:, k] + + 0.5 + / dual_area[llb:] + * edge_vert_length[v2e[llb:, j], idx_ve[llb:, j]] + * edge_cell_length[v2e[llb:, j], i], + cells_aw_verts[llb:, k], + ) + return cells_aw_verts + + +@program +def compute_ddxt_z_half_e( + z_ifv: Field[[VertexDim, KDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + tangent_orientation: Field[[EdgeDim], float], + ddxt_z_half_e: Field[[EdgeDim, KDim], float], + horizontal_lower: int32, + horizontal_upper: int32, + vertical_lower: int32, + vertical_upper: int32, +): + _grad_fd_tang( + z_ifv, + inv_primal_edge_length, + tangent_orientation, + out=ddxt_z_half_e, + domain={ + EdgeDim: (horizontal_lower, horizontal_upper), + KDim: (vertical_lower, vertical_upper), + }, + ) + + +@program +def compute_ddxnt_z_full( + z_ddxnt_z_half_e: Field[[EdgeDim, KDim], float], ddxn_z_full: Field[[EdgeDim, KDim], float] +): + average_ek_level_up(z_ddxnt_z_half_e, out=ddxn_z_full) diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py deleted file mode 100644 index c121c2acc..000000000 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields3.py +++ /dev/null @@ -1,119 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -from gt4py.next import Field, field_operator, int32, program - -from icon4py.model.common.dimension import ( - E2V, - CellDim, - EdgeDim, - KDim, - Koff, - VertexDim, -) -from icon4py.model.common.math.helpers import grad_fd_norm - - -@field_operator -def _grad_fd_tang( - psi_v: Field[[VertexDim, KDim], float], - inv_primal_edge_length: Field[[EdgeDim], float], - tangent_orientation: Field[[EdgeDim], float], -) -> Field[[EdgeDim, KDim], float]: - grad_tang_psi_e = tangent_orientation * (psi_v(E2V[1]) - psi_v(E2V[0])) * inv_primal_edge_length - return grad_tang_psi_e - - -@program -def compute_ddxn_z_half_e( - z_ifc: Field[[CellDim, KDim], float], - inv_dual_edge_length: Field[[EdgeDim], float], - ddxn_z_half_e: Field[[EdgeDim, KDim], float], - horizontal_lower: int32, - horizontal_upper: int32, - vertical_lower: int32, - vertical_upper: int32, -): - grad_fd_norm( - z_ifc, - inv_dual_edge_length, - out=ddxn_z_half_e, - domain={ - EdgeDim: (horizontal_lower, horizontal_upper), - KDim: (vertical_lower, vertical_upper), - }, - ) - - -@program -def compute_ddxt_z_half_e( - z_ifv: Field[[VertexDim, KDim], float], - inv_primal_edge_length: Field[[EdgeDim], float], - tangent_orientation: Field[[EdgeDim], float], - ddxt_z_half_e: Field[[EdgeDim, KDim], float], - horizontal_lower: int32, - horizontal_upper: int32, - vertical_lower: int32, - vertical_upper: int32, -): - _grad_fd_tang( - z_ifv, - inv_primal_edge_length, - tangent_orientation, - out=ddxt_z_half_e, - domain={ - EdgeDim: (horizontal_lower, horizontal_upper), - KDim: (vertical_lower, vertical_upper), - }, - ) - - -@field_operator -def _compute_ddxnt_z_full( - z_ddxnt_z_half_e: Field[[EdgeDim, KDim], float], -) -> Field[[EdgeDim, KDim], float]: - ddxnt_z_full = 0.5 * (z_ddxnt_z_half_e + z_ddxnt_z_half_e(Koff[1])) - return ddxnt_z_full - - -def compute_cells_aw_verts( - dual_area: np.array, - edge_vert_length: np.array, - edge_cell_length: np.array, - owner_mask: np.array, - e2c: np.array, - v2c: np.array, - v2e: np.array, - e2v: np.array, - horizontal_start: np.int32, - horizontal_end: np.int32, -) -> np.array: - llb = horizontal_start - cells_aw_verts = np.zeros([horizontal_end, 6]) - index = np.repeat(np.arange(horizontal_end, dtype=float), 6).reshape(horizontal_end, 6) - idx_ve = np.where(index == e2v[v2e, 0], 0, 1) - - for i in range(2): - for j in range(6): - for k in range(6): - cells_aw_verts[llb:, k] = np.where( - np.logical_and(owner_mask[:], e2c[v2e[llb:, j], i] == v2c[llb:, k]), - cells_aw_verts[llb:, k] - + 0.5 - / dual_area[llb:] - * edge_vert_length[v2e[llb:, j], idx_ve[llb:, j]] - * edge_cell_length[v2e[llb:, j], i], - cells_aw_verts[llb:, k], - ) - return cells_aw_verts diff --git a/model/common/src/icon4py/model/common/math/helpers.py b/model/common/src/icon4py/model/common/math/helpers.py index 2f960e589..506fe9ed7 100644 --- a/model/common/src/icon4py/model/common/math/helpers.py +++ b/model/common/src/icon4py/model/common/math/helpers.py @@ -13,12 +13,12 @@ from gt4py.next import Field, field_operator -from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim, Koff +from icon4py.model.common.dimension import E2C, E2V, CellDim, EdgeDim, KDim, Koff, VertexDim from icon4py.model.common.type_alias import wpfloat @field_operator -def average_k_level_up( +def average_ck_level_up( half_level_field: Field[[CellDim, KDim], wpfloat], ) -> Field[[CellDim, KDim], wpfloat]: """ @@ -35,6 +35,24 @@ def average_k_level_up( return 0.5 * (half_level_field + half_level_field(Koff[1])) +@field_operator +def average_ek_level_up( + half_level_field: Field[[EdgeDim, KDim], wpfloat], +) -> Field[[EdgeDim, KDim], wpfloat]: + """ + Calculate the mean value of adjacent interface levels. + + Computes the average of two adjacent interface levels upwards over a cell field for storage + in the corresponding full levels. + Args: + half_level_field: Field[[CellDim, KDim], wpfloat] + + Returns: Field[[CellDim, KDim], wpfloat] full level field + + """ + return 0.5 * (half_level_field + half_level_field(Koff[1])) + + @field_operator def difference_k_level_down( half_level_field: Field[[CellDim, KDim], wpfloat], @@ -89,3 +107,13 @@ def grad_fd_norm( """ grad_norm_psi_e = (psi_c(E2C[1]) - psi_c(E2C[0])) * inv_dual_edge_length return grad_norm_psi_e + + +@field_operator +def _grad_fd_tang( + psi_v: Field[[VertexDim, KDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + tangent_orientation: Field[[EdgeDim], float], +) -> Field[[EdgeDim, KDim], float]: + grad_tang_psi_e = tangent_orientation * (psi_v(E2V[1]) - psi_v(E2V[0])) * inv_primal_edge_length + return grad_tang_psi_e diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 449f8c516..e59465abe 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -25,11 +25,12 @@ where, ) -from icon4py.model.common.dimension import CellDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, Koff from icon4py.model.common.math.helpers import ( - average_k_level_up, + average_ck_level_up, difference_k_level_down, difference_k_level_up, + grad_fd_norm, ) from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -63,7 +64,7 @@ def compute_z_mc( vertical_end:int32 end index of vertical domain """ - average_k_level_up( + average_ck_level_up( z_ifc, out=z_mc, domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, @@ -438,3 +439,24 @@ def compute_d2dexdz2_fac_mc( out=d2dexdz2_fac2_mc, domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, ) + + +@program +def compute_ddxn_z_half_e( + z_ifc: Field[[CellDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + ddxn_z_half_e: Field[[EdgeDim, KDim], float], + horizontal_lower: int32, + horizontal_upper: int32, + vertical_lower: int32, + vertical_upper: int32, +): + grad_fd_norm( + z_ifc, + inv_dual_edge_length, + out=ddxn_z_half_e, + domain={ + EdgeDim: (horizontal_lower, horizontal_upper), + KDim: (vertical_lower, vertical_upper), + }, + ) diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields.py b/model/common/tests/interpolation_tests/test_interpolation_fields.py index 335cd6de9..cb42c70fb 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_fields.py +++ b/model/common/tests/interpolation_tests/test_interpolation_fields.py @@ -26,9 +26,29 @@ import numpy as np import pytest -from icon4py.model.common.dimension import EdgeDim -from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex -from icon4py.model.common.interpolation.interpolation_fields import compute_c_lin_e +from icon4py.model.common.dimension import ( + E2CDim, + E2VDim, + EdgeDim, + V2CDim, + V2EDim, + VertexDim, +) +from icon4py.model.common.grid.horizontal import ( + HorizontalMarkerIndex, +) +from icon4py.model.common.interpolation.interpolation_fields import ( + compute_c_lin_e, + compute_cells_aw_verts, +) +from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package + data_provider, + datapath, + download_ser_data, + experiment, + processor_props, + ranked_data_path, +) @pytest.mark.datatest @@ -49,3 +69,40 @@ def test_compute_c_lin_e(grid_savepoint, interpolation_savepoint, icon_grid): ) assert np.allclose(c_lin_e, c_lin_e_ref.asnumpy()) + + +@pytest.mark.datatest +def test_compute_cells_aw_verts( + grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint +): + dual_area = grid_savepoint.v_dual_area().asnumpy() + edge_vert_length = grid_savepoint.edge_vert_length().asnumpy() + edge_cell_length = grid_savepoint.edge_cell_length().asnumpy() + owner_mask = grid_savepoint.v_owner_mask() + e2c = icon_grid.connectivities[E2CDim] + v2c = icon_grid.connectivities[V2CDim] + v2e = icon_grid.connectivities[V2EDim] + e2v = icon_grid.connectivities[E2VDim] + horizontal_start_vertex = icon_grid.get_start_index( + VertexDim, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, + ) + horizontal_end_vertex = icon_grid.get_end_index( + VertexDim, + HorizontalMarkerIndex.lateral_boundary(VertexDim) - 1, + ) + + cells_aw_verts = compute_cells_aw_verts( + dual_area, + edge_vert_length, + edge_cell_length, + owner_mask, + e2c, + v2c, + v2e, + e2v, + horizontal_start_vertex, + horizontal_end_vertex, + ) + cells_aw_verts_ref = interpolation_savepoint.c_intp().asnumpy() + assert np.allclose(cells_aw_verts, cells_aw_verts_ref) diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields3.py b/model/common/tests/interpolation_tests/test_interpolation_fields3.py deleted file mode 100644 index 7d6811aef..000000000 --- a/model/common/tests/interpolation_tests/test_interpolation_fields3.py +++ /dev/null @@ -1,214 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np -import pytest -from gt4py.next import as_field - -from icon4py.model.common.dimension import ( - E2CDim, - E2VDim, - EdgeDim, - KDim, - V2CDim, - V2EDim, - VertexDim, -) -from icon4py.model.common.grid.horizontal import ( - HorizontalMarkerIndex, - _compute_cells2verts_scalar, - compute_cells2edges_scalar, -) -from icon4py.model.common.interpolation.interpolation_fields3 import ( - _compute_ddxnt_z_full, - compute_cells_aw_verts, - compute_ddxn_z_half_e, - compute_ddxt_z_half_e, -) -from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package - data_provider, - datapath, - download_ser_data, - experiment, - processor_props, - ranked_data_path, -) -from icon4py.model.common.test_utils.datatest_utils import ( - GLOBAL_EXPERIMENT, - REGIONAL_EXPERIMENT, -) -from icon4py.model.common.test_utils.helpers import zero_field - - -@pytest.mark.datatest -def test_compute_ddxn_z_full_e( - grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint -): - z_ifc = metrics_savepoint.z_ifc() - inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() - ddxn_z_full_ref = metrics_savepoint.ddxn_z_full().asnumpy() - horizontal_start = icon_grid.get_start_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, - ) - horizontal_end = icon_grid.get_end_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, - ) - vertical_start = 0 - vertical_end = icon_grid.num_levels + 1 - ddxn_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) - compute_ddxn_z_half_e( - z_ifc, - inv_dual_edge_length, - ddxn_z_half_e, - horizontal_start, - horizontal_end, - vertical_start, - vertical_end, - offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, - ) - ddxn_z_full = zero_field(icon_grid, EdgeDim, KDim) - _compute_ddxnt_z_full( - ddxn_z_half_e, - out=ddxn_z_full, - offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, - ) - - assert np.allclose(ddxn_z_full.asnumpy(), ddxn_z_full_ref) - - -@pytest.mark.datatest -def test_compute_ddxt_z_full_e( - grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint -): - z_ifc = metrics_savepoint.z_ifc() - dual_area = grid_savepoint.v_dual_area().asnumpy() - edge_vert_length = grid_savepoint.edge_vert_length().asnumpy() - edge_cell_length = grid_savepoint.edge_cell_length().asnumpy() - tangent_orientation = grid_savepoint.tangent_orientation() - inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() - owner_mask = grid_savepoint.v_owner_mask() - ddxt_z_full_ref = metrics_savepoint.ddxt_z_full().asnumpy() - e2c = icon_grid.connectivities[E2CDim] - v2c = icon_grid.connectivities[V2CDim] - v2e = icon_grid.connectivities[V2EDim] - e2v = icon_grid.connectivities[E2VDim] - horizontal_start_vertex = icon_grid.get_start_index( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, - ) - horizontal_end_vertex = icon_grid.get_end_index( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) - 1, - ) - horizontal_start_edge = icon_grid.get_start_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, - ) - horizontal_end_edge = icon_grid.get_end_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, - ) - vertical_start = 0 - vertical_end = icon_grid.num_levels + 1 - cells_aw_verts = compute_cells_aw_verts( - dual_area, - edge_vert_length, - edge_cell_length, - owner_mask, - e2c, - v2c, - v2e, - e2v, - horizontal_start_vertex, - horizontal_end_vertex, - ) - cells_aw_verts_ref = interpolation_savepoint.c_intp().asnumpy() - assert np.allclose(cells_aw_verts, cells_aw_verts_ref) - - z_ifv = zero_field(icon_grid, VertexDim, KDim, extend={KDim: 1}) - _compute_cells2verts_scalar( - z_ifc, - as_field((VertexDim, V2CDim), cells_aw_verts), - out=z_ifv, - offset_provider={"V2C": icon_grid.get_offset_provider("V2C")}, - domain={ - VertexDim: (horizontal_start_vertex, horizontal_end_vertex), - KDim: (vertical_start, vertical_end), - }, - ) - ddxt_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) - compute_ddxt_z_half_e( - z_ifv, - inv_primal_edge_length, - tangent_orientation, - ddxt_z_half_e, - horizontal_start_edge, - horizontal_end_edge, - vertical_start, - vertical_end, - offset_provider={"E2V": icon_grid.get_offset_provider("E2V")}, - ) - ddxt_z_full = zero_field(icon_grid, EdgeDim, KDim) - _compute_ddxnt_z_full( - ddxt_z_half_e, - out=ddxt_z_full, - offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, - ) - - assert np.allclose(ddxt_z_full.asnumpy(), ddxt_z_full_ref) - - -@pytest.mark.datatest -@pytest.mark.parametrize("experiment", (REGIONAL_EXPERIMENT, GLOBAL_EXPERIMENT)) -def test_compute_ddqz_z_full_e( - grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint -): - inv_ddqz_z_full = metrics_savepoint.inv_ddqz_z_full() - c_lin_e = interpolation_savepoint.c_lin_e() - ddqz_z_full_e_ref = metrics_savepoint.ddqz_z_full_e().asnumpy() - horizontal_start_edge = icon_grid.get_start_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, - ) - horizontal_end_edge = icon_grid.get_end_index( # noqa: F841 - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, - ) # for the global experiment, this outputs 0 - vertical_start = 0 - vertical_end = icon_grid.num_levels - ddqz_z_full_e = zero_field(icon_grid, EdgeDim, KDim) - compute_cells2edges_scalar( - inv_p_cell_in=inv_ddqz_z_full, - c_int=c_lin_e, - p_vert_out=ddqz_z_full_e, - horizontal_start_edge=horizontal_start_edge, - horizontal_end_edge=ddqz_z_full_e.shape[0], # horizontal_end_edge, - vertical_start=vertical_start, - vertical_end=vertical_end, - offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, - ) - assert np.allclose(ddqz_z_full_e.asnumpy(), ddqz_z_full_e_ref) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 6ba890fef..145a08a0a 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -19,16 +19,30 @@ from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.common import constants -from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, V2CDim, VertexDim +from icon4py.model.common.grid.horizontal import ( + HorizontalMarkerIndex, + _compute_cells2verts, + compute_cells2edges, +) +from icon4py.model.common.interpolation.interpolation_fields import ( + compute_ddxnt_z_full, + compute_ddxt_z_half_e, +) from icon4py.model.common.metrics.metric_fields import ( compute_coeff_dwdz, compute_d2dexdz2_fac_mc, compute_ddqz_z_full, compute_ddqz_z_half, + compute_ddxn_z_half_e, compute_rayleigh_w, compute_scalfac_dd3d, compute_z_mc, ) +from icon4py.model.common.test_utils.datatest_utils import ( + GLOBAL_EXPERIMENT, + REGIONAL_EXPERIMENT, +) from icon4py.model.common.test_utils.helpers import ( StencilTest, dallclose, @@ -215,8 +229,8 @@ def test_compute_d2dexdz2_fac_mc(icon_grid, metrics_savepoint, grid_savepoint, b z_ifc = metrics_savepoint.z_ifc() z_mc = zero_field(icon_grid, CellDim, KDim) compute_z_mc.with_backend(backend)( - z_ifc, - z_mc, + z_ifc=z_ifc, + z_mc=z_mc, horizontal_start=int32(0), horizontal_end=icon_grid.num_cells, vertical_start=int32(0), @@ -255,3 +269,125 @@ def test_compute_d2dexdz2_fac_mc(icon_grid, metrics_savepoint, grid_savepoint, b assert dallclose(d2dexdz2_fac1_mc_full.asnumpy(), d2dexdz2_fac1_mc_ref.asnumpy()) assert dallclose(d2dexdz2_fac2_mc_full.asnumpy(), d2dexdz2_fac2_mc_ref.asnumpy()) + + +@pytest.mark.datatest +@pytest.mark.parametrize("experiment", (REGIONAL_EXPERIMENT, GLOBAL_EXPERIMENT)) +def test_compute_ddqz_z_full_e( + grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint +): + ddqz_z_full = as_field((CellDim, KDim), 1.0 / metrics_savepoint.inv_ddqz_z_full().asnumpy()) + c_lin_e = interpolation_savepoint.c_lin_e() + ddqz_z_full_e_ref = metrics_savepoint.ddqz_z_full_e().asnumpy() + vertical_start = 0 + vertical_end = icon_grid.num_levels + ddqz_z_full_e = zero_field(icon_grid, EdgeDim, KDim) + compute_cells2edges( + p_cell_in=ddqz_z_full, + c_int=c_lin_e, + p_vert_out=ddqz_z_full_e, + horizontal_start_edge=0, + horizontal_end_edge=ddqz_z_full_e.shape[0], + vertical_start=vertical_start, + vertical_end=vertical_end, + offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, + ) + assert np.allclose(ddqz_z_full_e.asnumpy(), ddqz_z_full_e_ref) + + +@pytest.mark.datatest +def test_compute_ddxn_z_full_e( + grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint +): + z_ifc = metrics_savepoint.z_ifc() + inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() + ddxn_z_full_ref = metrics_savepoint.ddxn_z_full().asnumpy() + horizontal_start = icon_grid.get_start_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, + ) + horizontal_end = icon_grid.get_end_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, + ) + vertical_start = 0 + vertical_end = icon_grid.num_levels + 1 + ddxn_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) + compute_ddxn_z_half_e( + z_ifc=z_ifc, + inv_dual_edge_length=inv_dual_edge_length, + ddxn_z_half_e=ddxn_z_half_e, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=vertical_start, + vertical_end=vertical_end, + offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, + ) + ddxn_z_full = zero_field(icon_grid, EdgeDim, KDim) + compute_ddxnt_z_full( + z_ddxnt_z_half_e=ddxn_z_half_e, + ddxn_z_full=ddxn_z_full, + offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, + ) + + assert np.allclose(ddxn_z_full.asnumpy(), ddxn_z_full_ref) + + +@pytest.mark.datatest +def test_compute_ddxt_z_full_e( + grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint +): + z_ifc = metrics_savepoint.z_ifc() + tangent_orientation = grid_savepoint.tangent_orientation() + inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() + ddxt_z_full_ref = metrics_savepoint.ddxt_z_full().asnumpy() + horizontal_start_vertex = icon_grid.get_start_index( + VertexDim, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, + ) + horizontal_end_vertex = icon_grid.get_end_index( + VertexDim, + HorizontalMarkerIndex.lateral_boundary(VertexDim) - 1, + ) + horizontal_start_edge = icon_grid.get_start_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, + ) + horizontal_end_edge = icon_grid.get_end_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, + ) + vertical_start = 0 + vertical_end = icon_grid.num_levels + 1 + cells_aw_verts = interpolation_savepoint.c_intp().asnumpy() + z_ifv = zero_field(icon_grid, VertexDim, KDim, extend={KDim: 1}) + _compute_cells2verts( + z_ifc, + as_field((VertexDim, V2CDim), cells_aw_verts), + out=z_ifv, + offset_provider={"V2C": icon_grid.get_offset_provider("V2C")}, + domain={ + VertexDim: (horizontal_start_vertex, horizontal_end_vertex), + KDim: (vertical_start, vertical_end), + }, + ) + ddxt_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) + compute_ddxt_z_half_e( + z_ifv=z_ifv, + inv_primal_edge_length=inv_primal_edge_length, + tangent_orientation=tangent_orientation, + ddxt_z_half_e=ddxt_z_half_e, + horizontal_lower=horizontal_start_edge, + horizontal_upper=horizontal_end_edge, + vertical_lower=vertical_start, + vertical_upper=vertical_end, + offset_provider={"E2V": icon_grid.get_offset_provider("E2V")}, + ) + ddxt_z_full = zero_field(icon_grid, EdgeDim, KDim) + compute_ddxnt_z_full( + z_ddxnt_z_half_e=ddxt_z_half_e, + ddxn_z_full=ddxt_z_full, + offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, + ) + + assert np.allclose(ddxt_z_full.asnumpy(), ddxt_z_full_ref) From 53c5ac60efcaf4fbac15f30abad1b7f49184035b Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Wed, 24 Apr 2024 14:30:14 +0200 Subject: [PATCH 17/65] docstring edit --- model/common/src/icon4py/model/common/math/helpers.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/model/common/src/icon4py/model/common/math/helpers.py b/model/common/src/icon4py/model/common/math/helpers.py index 506fe9ed7..1487023df 100644 --- a/model/common/src/icon4py/model/common/math/helpers.py +++ b/model/common/src/icon4py/model/common/math/helpers.py @@ -42,12 +42,12 @@ def average_ek_level_up( """ Calculate the mean value of adjacent interface levels. - Computes the average of two adjacent interface levels upwards over a cell field for storage + Computes the average of two adjacent interface levels upwards over an edge field for storage in the corresponding full levels. Args: - half_level_field: Field[[CellDim, KDim], wpfloat] + half_level_field: Field[[EdgeDim, KDim], wpfloat] - Returns: Field[[CellDim, KDim], wpfloat] full level field + Returns: Field[[EdgeDim, KDim], wpfloat] full level field """ return 0.5 * (half_level_field + half_level_field(Koff[1])) From ec468eb98086c84c715cb110b2ead584fb33b367 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 25 Apr 2024 14:04:37 +0200 Subject: [PATCH 18/65] small edits --- .../interpolation/interpolation_fields.py | 30 ++----------------- .../src/icon4py/model/common/math/helpers.py | 4 +-- .../model/common/metrics/metric_fields.py | 30 +++++++++++++++++-- .../tests/metric_tests/test_metric_fields.py | 6 ++-- 4 files changed, 34 insertions(+), 36 deletions(-) diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py index 2b2618eef..2a767f389 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py @@ -24,14 +24,13 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -from gt4py.next import Field, int32, program +from gt4py.next import Field, program from icon4py.model.common.dimension import ( EdgeDim, KDim, - VertexDim, ) -from icon4py.model.common.math.helpers import _grad_fd_tang, average_ek_level_up +from icon4py.model.common.math.helpers import edge_kdim def compute_c_lin_e( @@ -91,31 +90,8 @@ def compute_cells_aw_verts( return cells_aw_verts -@program -def compute_ddxt_z_half_e( - z_ifv: Field[[VertexDim, KDim], float], - inv_primal_edge_length: Field[[EdgeDim], float], - tangent_orientation: Field[[EdgeDim], float], - ddxt_z_half_e: Field[[EdgeDim, KDim], float], - horizontal_lower: int32, - horizontal_upper: int32, - vertical_lower: int32, - vertical_upper: int32, -): - _grad_fd_tang( - z_ifv, - inv_primal_edge_length, - tangent_orientation, - out=ddxt_z_half_e, - domain={ - EdgeDim: (horizontal_lower, horizontal_upper), - KDim: (vertical_lower, vertical_upper), - }, - ) - - @program def compute_ddxnt_z_full( z_ddxnt_z_half_e: Field[[EdgeDim, KDim], float], ddxn_z_full: Field[[EdgeDim, KDim], float] ): - average_ek_level_up(z_ddxnt_z_half_e, out=ddxn_z_full) + edge_kdim(z_ddxnt_z_half_e, out=ddxn_z_full) diff --git a/model/common/src/icon4py/model/common/math/helpers.py b/model/common/src/icon4py/model/common/math/helpers.py index 1487023df..8c0cd3b42 100644 --- a/model/common/src/icon4py/model/common/math/helpers.py +++ b/model/common/src/icon4py/model/common/math/helpers.py @@ -18,7 +18,7 @@ @field_operator -def average_ck_level_up( +def cell_kdim( half_level_field: Field[[CellDim, KDim], wpfloat], ) -> Field[[CellDim, KDim], wpfloat]: """ @@ -36,7 +36,7 @@ def average_ck_level_up( @field_operator -def average_ek_level_up( +def edge_kdim( half_level_field: Field[[EdgeDim, KDim], wpfloat], ) -> Field[[EdgeDim, KDim], wpfloat]: """ diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index e59465abe..3600c2bf9 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -25,9 +25,10 @@ where, ) -from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, Koff +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, Koff, VertexDim from icon4py.model.common.math.helpers import ( - average_ck_level_up, + _grad_fd_tang, + cell_kdim, difference_k_level_down, difference_k_level_up, grad_fd_norm, @@ -64,7 +65,7 @@ def compute_z_mc( vertical_end:int32 end index of vertical domain """ - average_ck_level_up( + cell_kdim( z_ifc, out=z_mc, domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, @@ -460,3 +461,26 @@ def compute_ddxn_z_half_e( KDim: (vertical_lower, vertical_upper), }, ) + + +@program +def compute_ddxt_z_half_e( + z_ifv: Field[[VertexDim, KDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + tangent_orientation: Field[[EdgeDim], float], + ddxt_z_half_e: Field[[EdgeDim, KDim], float], + horizontal_lower: int32, + horizontal_upper: int32, + vertical_lower: int32, + vertical_upper: int32, +): + _grad_fd_tang( + z_ifv, + inv_primal_edge_length, + tangent_orientation, + out=ddxt_z_half_e, + domain={ + EdgeDim: (horizontal_lower, horizontal_upper), + KDim: (vertical_lower, vertical_upper), + }, + ) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 145a08a0a..30633d836 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -25,16 +25,14 @@ _compute_cells2verts, compute_cells2edges, ) -from icon4py.model.common.interpolation.interpolation_fields import ( - compute_ddxnt_z_full, - compute_ddxt_z_half_e, -) +from icon4py.model.common.interpolation.interpolation_fields import compute_ddxnt_z_full from icon4py.model.common.metrics.metric_fields import ( compute_coeff_dwdz, compute_d2dexdz2_fac_mc, compute_ddqz_z_full, compute_ddqz_z_half, compute_ddxn_z_half_e, + compute_ddxt_z_half_e, compute_rayleigh_w, compute_scalfac_dd3d, compute_z_mc, From 54363af1794c2ce3a9a80ab290658ddb984f0c8a Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 26 Apr 2024 20:37:58 +0200 Subject: [PATCH 19/65] edits --- .../common/src/icon4py/model/common/config.py | 5 +- .../model/common/metrics/metric_fields.py | 30 +++++----- .../stencils/compute_zdiff_gradp_dsl.py | 53 +++++++++++++++--- .../tests/metric_tests/test_metric_fields.py | 55 ++++++++++++------- 4 files changed, 101 insertions(+), 42 deletions(-) diff --git a/model/common/src/icon4py/model/common/config.py b/model/common/src/icon4py/model/common/config.py index 8c9f0eb95..e30be12d9 100644 --- a/model/common/src/icon4py/model/common/config.py +++ b/model/common/src/icon4py/model/common/config.py @@ -18,12 +18,13 @@ from functools import cached_property import numpy as np + +from gt4py.next.program_processors.runners.double_roundtrip import backend as run_roundtrip from gt4py.next.program_processors.runners.gtfn import ( run_gtfn_cached, run_gtfn_gpu_cached, ) -from gt4py.next.program_processors.runners.roundtrip import backend as run_roundtrip - +#from gt4py.next.program_processors.runners.roundtrip import backend as run_roundtrip class Device(Enum): CPU = "CPU" diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 2be14a6f5..ad7b9d65c 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -88,15 +88,19 @@ def _compute_ddqz_z_half( k: Field[[KDim], int32], nlev: int32, ) -> Field[[CellDim, KDim], wpfloat]: + ddqz_z_half = 2.0 * (z_mc(Koff[-1]) - z_ifc) ddqz_z_half = where( (k > int32(0)) & (k < nlev), difference_k_level_down(z_mc), - where(k == 0, 2.0 * (z_ifc - z_mc), 2.0 * (z_mc(Koff[-1]) - z_ifc)), + ddqz_z_half ) + ddqz_z_half = where(k == 0, 2.0 * (z_ifc - z_mc), ddqz_z_half) + + #ddqz_z_half = where(k == nlev, 2.0 * (z_mc(Koff[-1]) - z_ifc), ddqz_z_half) return ddqz_z_half -@program(grid_type=GridType.UNSTRUCTURED, backend=None) +@program(grid_type=GridType.UNSTRUCTURED) def compute_ddqz_z_half( z_ifc: Field[[CellDim, KDim], wpfloat], z_mc: Field[[CellDim, KDim], wpfloat], @@ -456,18 +460,18 @@ def compute_ddxn_z_half_e( z_ifc: Field[[CellDim, KDim], float], inv_dual_edge_length: Field[[EdgeDim], float], ddxn_z_half_e: Field[[EdgeDim, KDim], float], - horizontal_lower: int32, - horizontal_upper: int32, - vertical_lower: int32, - vertical_upper: int32, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): grad_fd_norm( z_ifc, inv_dual_edge_length, out=ddxn_z_half_e, domain={ - EdgeDim: (horizontal_lower, horizontal_upper), - KDim: (vertical_lower, vertical_upper), + EdgeDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), }, ) @@ -502,17 +506,17 @@ def _compute_vwind_impl_wgt( dual_edge_length: Field[[EdgeDim], wpfloat], vwind_offctr: wpfloat, ) -> Field[[CellDim], wpfloat]: - z_ddx_1 = maximum(abs(z_ddxn_z_half_e(C2E[1])), abs(z_ddxt_z_half_e(C2E[1]))) - z_ddx_2 = maximum(abs(z_ddxn_z_half_e(C2E[2])), abs(z_ddxt_z_half_e(C2E[2]))) - z_ddx_3 = maximum(abs(z_ddxn_z_half_e(C2E[3])), abs(z_ddxt_z_half_e(C2E[3]))) + z_ddx_1 = maximum(abs(z_ddxn_z_half_e(C2E[0])), abs(z_ddxt_z_half_e(C2E[0]))) + z_ddx_2 = maximum(abs(z_ddxn_z_half_e(C2E[1])), abs(z_ddxt_z_half_e(C2E[1]))) + z_ddx_3 = maximum(abs(z_ddxn_z_half_e(C2E[2])), abs(z_ddxt_z_half_e(C2E[2]))) z_ddx_1_2 = maximum(z_ddx_1, z_ddx_2) z_maxslope = maximum(z_ddx_1_2, z_ddx_3) z_diff_1_2 = maximum( + abs(z_ddxn_z_half_e(C2E[0]) * dual_edge_length(C2E[0])), abs(z_ddxn_z_half_e(C2E[1]) * dual_edge_length(C2E[1])), - abs(z_ddxn_z_half_e(C2E[2]) * dual_edge_length(C2E[2])), ) - z_diff = maximum(z_diff_1_2, abs(z_ddxn_z_half_e(C2E[3]) * dual_edge_length(C2E[3]))) + z_diff = maximum(z_diff_1_2, abs(z_ddxn_z_half_e(C2E[2]) * dual_edge_length(C2E[2]))) z_offctr_1 = maximum(vwind_offctr, 0.425 * z_maxslope**0.75) z_offctr = maximum(z_offctr_1, minimum(0.25, 2.5e-4 * (z_diff - 250.0))) z_offctr = minimum(maximum(vwind_offctr, 0.75), z_offctr) diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py index 95de8fc68..7936ddb72 100644 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py @@ -27,18 +27,57 @@ def compute_zdiff_gradp_dsl( z_me = np.sum(z_mc[e2c] * c_lin_e, axis=1) zdiff_gradp[:, 0, :] = z_me - z_mc[e2c[:, 0]] zdiff_gradp[:, 1, :] = z_me - z_mc[e2c[:, 1]] + + for k in range(nlev): + for e in range(nedges): # TODO: boundary here is not nedges, rather p_patch(jg)%edges%start_block(2) + if z_me[e, k] <= z_ifc[e2c[e, 0], k] and \ + z_me[e, k] >= z_ifc[e2c[e, 0], k+1] and \ + z_me[e, k] <= z_ifc[e2c[e, 1], k] and \ + z_me[e, k] >= z_ifc[e2c[e, 1], k + 1]: + flat_idx[e] = k + + for je in range(nedges): + jk_start = flat_idx[je] + for jk in range(flat_idx[je] + 1, nlev): + for jk1 in range(jk_start, nlev): + if jk1 == nlev or z_me[je, jk] <= z_ifc[e2c[je, 0], jk1] and z_me[je, jk] >= z_ifc[e2c[je, 0], jk1+1]: + zdiff_gradp[je, 0, jk1] = z_me[je, jk] - z_mc[e2c[je, 0], jk1] + #zdiff_gradp[je, 1, jk1] = z_aux2[je] - z_mc[e2c[je, 1], jk1] + jk_start = jk1 + break + extrapol_dist = 5.0 - nlevp1 = nlev + 1 - z_aux1 = np.maximum(z_ifc[e2c[:, 0], nlevp1], z_ifc[e2c[:, 1], nlevp1]) + z_aux1 = np.maximum(z_ifc[e2c[:, 0], nlev], z_ifc[e2c[:, 1], nlev]) z_aux2 = z_aux1 - extrapol_dist - jk_start = flat_idx + + for je in range(nedges): + jk_start = flat_idx[je] + for jk in range(flat_idx[je] + 1, nlev): + for jk1 in range(jk_start, nlev): + if jk1 == nlev or z_me[je, jk] <= z_ifc[e2c[je, 1], jk1] and z_me[je, jk] >= z_ifc[e2c[je, 1], jk1+1]: + #zdiff_gradp[je, 0, jk1] = z_aux2[je] - z_mc[e2c[je, 0], jk1] + zdiff_gradp[je, 1, jk1] = z_me[je, jk] - z_mc[e2c[je, 1], jk1] + jk_start = jk1 + break + + + for je in range(nedges): + jk_start = flat_idx[je] + for jk in range(flat_idx[je] + 1, nlev): + if z_me[je, jk] < z_aux2[je]: + for jk1 in range(jk_start, nlev): + if jk1 == nlev or z_aux2[je] <= z_ifc[e2c[je, 0], jk1] and z_aux2[je] >= z_ifc[e2c[je, 0], jk1+1]: + zdiff_gradp[je, 0, jk1] = z_aux2[je] - z_mc[e2c[je, 0], jk1] + jk_start = jk1 + break + for je in range(nedges): + jk_start = flat_idx[je] for jk in range(flat_idx[je] + 1, nlev): if z_me[je, jk] < z_aux2[je]: - for jk1 in (jk_start, nlev): - if jk1 == nlev or z_aux2 <= z_ifc[e2c[je, 1], jk1] and z_aux2 >= z_ifc[e2c[je, 1], jk1+1]: - zdiff_gradp[je, 0, jk1] = z_aux2 - z_mc[e2c[je, 0], jk1] - zdiff_gradp[je, 1, jk1] = z_aux2 - z_mc[e2c[je, 1], jk1] + for jk1 in range(jk_start, nlev): + if jk1 == nlev or z_aux2[je] <= z_ifc[e2c[je, 1], jk1] and z_aux2[je] >= z_ifc[e2c[je, 1], jk1+1]: + zdiff_gradp[je, 1, jk1] = z_aux2[je] - z_mc[e2c[je, 1], jk1] jk_start = jk1 break diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index ade6c89a2..ad16f180d 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -99,13 +99,14 @@ def input_data(self, grid) -> dict: @pytest.mark.datatest def test_compute_ddq_z_half(icon_grid, metrics_savepoint, backend): - if is_python(backend): - pytest.skip("skipping: unsupported backend") + backend = None + # if is_python(backend): + # pytest.skip("skipping: unsupported backend") ddq_z_half_ref = metrics_savepoint.ddqz_z_half() z_ifc = metrics_savepoint.z_ifc() - z_mc = zero_field(icon_grid, CellDim, KDim) + z_mc = zero_field(icon_grid, CellDim, KDim, extend={KDim: 1}) nlevp1 = icon_grid.num_levels + 1 - k_index = as_field((KDim,), np.arange(nlevp1, dtype=int32)) + k_index = as_field((KDim,), np.arange(0, nlevp1, dtype=int32)) compute_z_mc.with_backend(backend)( z_ifc, z_mc, @@ -115,26 +116,30 @@ def test_compute_ddq_z_half(icon_grid, metrics_savepoint, backend): vertical_end=int32(icon_grid.num_levels), offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) - ddq_z_half = zero_field(icon_grid, CellDim, KDim, extend={KDim: 1}) + ddqz_z_half_zeros = zero_field(icon_grid, CellDim, KDim, extend={KDim: 1}) + # TODO: line below should be in program, however bounds do not allow to go from 0 to 66 for K + ddqz_z_half_np = np.where(k_index == 0, 2.0 * (z_ifc.asnumpy() - z_mc.asnumpy()), ddqz_z_half_zeros.asnumpy()) + ddqz_z_half = as_field((CellDim, KDim), ddqz_z_half_np) compute_ddqz_z_half.with_backend(backend=backend)( z_ifc=z_ifc, z_mc=z_mc, k=k_index, nlev=icon_grid.num_levels, - ddqz_z_half=ddq_z_half, + ddqz_z_half=ddqz_z_half, horizontal_start=0, horizontal_end=icon_grid.num_cells, - vertical_start=0, + vertical_start=1, vertical_end=nlevp1, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) - assert dallclose(ddq_z_half.asnumpy(), ddq_z_half_ref.asnumpy()) + assert dallclose(ddqz_z_half.asnumpy(), ddq_z_half_ref.asnumpy()) @pytest.mark.datatest def test_compute_ddqz_z_full(icon_grid, metrics_savepoint, backend): + backend = None if is_roundtrip(backend): pytest.skip("skipping: slow backend") z_ifc = metrics_savepoint.z_ifc() @@ -159,6 +164,7 @@ def test_compute_ddqz_z_full(icon_grid, metrics_savepoint, backend): # TODO: convert this to a stenciltest once it is possible to have only KDim in domain @pytest.mark.datatest def test_compute_scalfac_dd3d(icon_grid, metrics_savepoint, grid_savepoint, backend): + backend = None scalfac_dd3d_ref = metrics_savepoint.scalfac_dd3d() scalfac_dd3d_full = zero_field(icon_grid, KDim) divdamp_trans_start = 12500.0 @@ -182,6 +188,7 @@ def test_compute_scalfac_dd3d(icon_grid, metrics_savepoint, grid_savepoint, back # TODO: convert this to a stenciltest once it is possible to have only KDim in domain @pytest.mark.datatest def test_compute_rayleigh_w(icon_grid, metrics_savepoint, grid_savepoint, backend): + backend = None rayleigh_w_ref = metrics_savepoint.rayleigh_w() vct_a_1 = grid_savepoint.vct_a().asnumpy()[0] rayleigh_w_full = zero_field(icon_grid, KDim, extend={KDim: 1}) @@ -208,6 +215,7 @@ def test_compute_rayleigh_w(icon_grid, metrics_savepoint, grid_savepoint, backen @pytest.mark.datatest def test_compute_coeff_dwdz(icon_grid, metrics_savepoint, grid_savepoint, backend): + backend = None if is_roundtrip(backend): pytest.skip("skipping: slow backend") coeff1_dwdz_ref = metrics_savepoint.coeff1_dwdz() @@ -235,6 +243,7 @@ def test_compute_coeff_dwdz(icon_grid, metrics_savepoint, grid_savepoint, backen @pytest.mark.datatest def test_compute_d2dexdz2_fac_mc(icon_grid, metrics_savepoint, grid_savepoint, backend): + backend = None if is_roundtrip(backend): pytest.skip("skipping: slow backend") z_ifc = metrics_savepoint.z_ifc() @@ -286,6 +295,7 @@ def test_compute_d2dexdz2_fac_mc(icon_grid, metrics_savepoint, grid_savepoint, b def test_compute_vwind_impl_wgt( icon_grid, grid_savepoint, metrics_savepoint, interpolation_savepoint, backend ): + backend = None if is_roundtrip(backend): pytest.skip("skipping: slow backend") z_ifc = metrics_savepoint.z_ifc() @@ -311,7 +321,7 @@ def test_compute_vwind_impl_wgt( HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, ) - z_ddxn_z_half_e = compute_ddxn_z_half_e( + compute_ddxn_z_half_e( z_ifc=z_ifc, inv_dual_edge_length=inv_dual_edge_length, ddxn_z_half_e=z_ddxn_z_half_e, @@ -358,14 +368,16 @@ def test_compute_vwind_impl_wgt( offset_provider={"E2V": icon_grid.get_offset_provider("E2V")}, ) - vwind_impl_wgt_full = zero_field(icon_grid, CellDim, KDim) + vwind_impl_wgt_full = zero_field(icon_grid, CellDim) vwind_impl_wgt_ref = metrics_savepoint.vwind_impl_wgt() dual_edge_length = as_field((EdgeDim,), 1.0 / grid_savepoint.inv_dual_edge_length().asnumpy()) vwind_offctr = 0.2 + # TODO: missing this line 360 in f90: CALL rbf_vec_interpol_cell(z_ddxn_z_half_e(:,nlevp1:nlevp1,:), p_patch(jg), p_int(jg), z_aux_c, z_aux_c2) + compute_vwind_impl_wgt.with_backend(backend)( - z_ddxn_z_half_e=z_ddxn_z_half_e, - z_ddxt_z_half_e=z_ddxt_z_half_e, + z_ddxn_z_half_e=as_field((EdgeDim,), z_ddxn_z_half_e.asnumpy()[:, icon_grid.num_levels]), + z_ddxt_z_half_e=as_field((EdgeDim,), z_ddxt_z_half_e.asnumpy()[:, icon_grid.num_levels]), dual_edge_length=dual_edge_length, vwind_impl_wgt=vwind_impl_wgt_full, vwind_offctr=vwind_offctr, @@ -379,6 +391,7 @@ def test_compute_vwind_impl_wgt( @pytest.mark.datatest def test_compute_vwind_expl_wgt(icon_grid, metrics_savepoint, backend): + backend = None if is_roundtrip(backend): pytest.skip("skipping: slow backend") @@ -399,13 +412,15 @@ def test_compute_vwind_expl_wgt(icon_grid, metrics_savepoint, backend): @pytest.mark.datatest def test_compute_inv_ddqz_z_full(icon_grid, metrics_savepoint, backend): + # TODO: serialization missing inv_ddqz_z_full is over cells, need over edge --> inv_ddqz_z_full_e + backend = None if is_roundtrip(backend): pytest.skip("skipping: slow backend") inv_ddqz_z_full = zero_field(icon_grid, EdgeDim, KDim) inv_ddqz_z_ref = ( metrics_savepoint.inv_ddqz_z_full() - ) # inv_ddqz_z_full is over cells, need over edge --> inv_ddqz_z_full_e + ) horizontal_start_edge = icon_grid.get_start_index( EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, @@ -458,23 +473,20 @@ def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_sav # e2c = icon_grid.connectivities[E2CDim] # z_me = np.sum(z_mc.asnumpy()[e2c] * np.expand_dims(c_lin_e.asnumpy(), axis=-1), axis=1) flat_idx = np.full(shape=icon_grid.num_edges, fill_value=icon_grid.num_levels) - # for k in range(icon_grid.num_levels): - # for e in range(icon_grid.num_edges): - # if z_me <= z_ifc[e2c[e, 0], k] and z_me >= z_ifc[e2c[e, 0], k+1] and z_me <= z_ifc[e2c[e, 1], k] and z_me >= z_ifc[e2c[e, 1], k + 1]: - # flat_idx[e] = k - zdiff_gradp_full_np = compute_zdiff_gradp_dsl( + + zdiff_gradp_full_np_final = compute_zdiff_gradp_dsl( e2c=icon_grid.connectivities[E2CDim], c_lin_e=c_lin_e.asnumpy(), z_mc=z_mc.asnumpy(), zdiff_gradp=zdiff_gradp_full_np, z_ifc=metrics_savepoint.z_ifc().asnumpy(), flat_idx=flat_idx, - nlev=icon_grid.num_levels - 1, # TODO: check that -1 ok + nlev=icon_grid.num_levels, # TODO: check that -1 ok nedges=icon_grid.num_edges, ) zdiff_gradp_full_field = flatten_first_two_dims( - ECDim, KDim, field=as_field((EdgeDim, E2CDim, KDim), zdiff_gradp_full_np) + ECDim, KDim, field=as_field((EdgeDim, E2CDim, KDim), zdiff_gradp_full_np_final) ) # zdiff_gradp_ref.asnumpy()[856:, :] == zdiff_gradp_full_field.asnumpy()[856:, :] assert dallclose(zdiff_gradp_full_field.asnumpy(), zdiff_gradp_ref.asnumpy()) @@ -485,6 +497,7 @@ def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_sav def test_compute_ddqz_z_full_e( grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint ): + backend = None ddqz_z_full = as_field((CellDim, KDim), 1.0 / metrics_savepoint.inv_ddqz_z_full().asnumpy()) c_lin_e = interpolation_savepoint.c_lin_e() ddqz_z_full_e_ref = metrics_savepoint.ddqz_z_full_e().asnumpy() @@ -508,6 +521,7 @@ def test_compute_ddqz_z_full_e( def test_compute_ddxn_z_full_e( grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint ): + backend = None z_ifc = metrics_savepoint.z_ifc() inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() ddxn_z_full_ref = metrics_savepoint.ddxn_z_full().asnumpy() @@ -546,6 +560,7 @@ def test_compute_ddxn_z_full_e( def test_compute_ddxt_z_full_e( grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint ): + backend = None z_ifc = metrics_savepoint.z_ifc() tangent_orientation = grid_savepoint.tangent_orientation() inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() From ddfd413aa8fc415e61324e8d078f9118fb552929 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 30 Apr 2024 11:43:55 +0200 Subject: [PATCH 20/65] edits --- .../model/atmosphere/diffusion/diffusion.py | 2 +- ...tex.py => mo_intp_rbf_rbf_vec_interpol.py} | 46 ++++++- .../src/icon4py/model/common/math/helpers.py | 28 +++- .../model/common/metrics/metric_fields.py | 45 ++---- .../stencils/compute_zdiff_gradp_dsl.py | 45 ++++-- .../common/test_utils/serialbox_utils.py | 3 + .../tests/metric_tests/test_metric_fields.py | 128 +++--------------- ...est_mo_intp_rbf_rbf_vec_interpol_vertex.py | 2 +- 8 files changed, 131 insertions(+), 168 deletions(-) rename model/common/src/icon4py/model/common/interpolation/stencils/{mo_intp_rbf_rbf_vec_interpol_vertex.py => mo_intp_rbf_rbf_vec_interpol.py} (60%) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py index ac9a332aa..a8d872905 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -68,7 +68,7 @@ from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalMarkerIndex from icon4py.model.common.grid.icon import IconGrid from icon4py.model.common.grid.vertical import VerticalModelParams -from icon4py.model.common.interpolation.stencils.mo_intp_rbf_rbf_vec_interpol_vertex import ( +from icon4py.model.common.interpolation.stencils.mo_intp_rbf_rbf_vec_interpol import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) from icon4py.model.common.states.prognostic_state import PrognosticState diff --git a/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol.py similarity index 60% rename from model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py rename to model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol.py index 2d867d0a9..28c2902b4 100644 --- a/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol.py @@ -15,7 +15,16 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum -from icon4py.model.common.dimension import V2E, EdgeDim, KDim, V2EDim, VertexDim +from icon4py.model.common.dimension import ( + C2E, + V2E, + C2EDim, + CellDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat @@ -53,3 +62,38 @@ def mo_intp_rbf_rbf_vec_interpol_vertex( KDim: (vertical_start, vertical_end), }, ) + + +@field_operator +def _mo_intp_rbf_rbf_vec_interpol_cell( + p_vn_in: Field[[EdgeDim, KDim], wpfloat], + ptr_coeff_1: Field[[CellDim, C2EDim], wpfloat], + ptr_coeff_2: Field[[CellDim, C2EDim], wpfloat], +) -> tuple[Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat]]: + p_u_out = neighbor_sum(ptr_coeff_1 * p_vn_in(C2E), axis=C2EDim) + p_v_out = neighbor_sum(ptr_coeff_2 * p_vn_in(C2E), axis=C2EDim) + return p_u_out, p_v_out + + +@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +def mo_intp_rbf_rbf_vec_interpol_cell( + p_vn_in: Field[[EdgeDim, KDim], wpfloat], + ptr_coeff_1: Field[[CellDim, C2EDim], wpfloat], + ptr_coeff_2: Field[[CellDim, C2EDim], wpfloat], + p_u_out: Field[[CellDim, KDim], wpfloat], + p_v_out: Field[[CellDim, KDim], wpfloat], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, +): + _mo_intp_rbf_rbf_vec_interpol_vertex( + p_vn_in, + ptr_coeff_1, + ptr_coeff_2, + out=(p_u_out, p_v_out), + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/common/src/icon4py/model/common/math/helpers.py b/model/common/src/icon4py/model/common/math/helpers.py index 808db44ba..3eeefcefe 100644 --- a/model/common/src/icon4py/model/common/math/helpers.py +++ b/model/common/src/icon4py/model/common/math/helpers.py @@ -11,10 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next import Field, field_operator +from gt4py.next import Field, field_operator, program +from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.common.dimension import E2C, E2V, CellDim, EdgeDim, KDim, Koff, VertexDim -from icon4py.model.common.type_alias import wpfloat +from icon4py.model.common.type_alias import vpfloat, wpfloat @field_operator @@ -117,3 +118,26 @@ def _grad_fd_tang( ) -> Field[[EdgeDim, KDim], float]: grad_tang_psi_e = tangent_orientation * (psi_v(E2V[1]) - psi_v(E2V[0])) * inv_primal_edge_length return grad_tang_psi_e + + +@field_operator +def _compute_inv_edge_k( + edge_k_field: Field[[EdgeDim, KDim], vpfloat], +) -> Field[[EdgeDim, KDim], vpfloat]: + return 1.0 / edge_k_field + + +@program +def compute_inv_edge_k( + edge_k_field: Field[[EdgeDim, KDim], vpfloat], + inv_edge_k_field: Field[[EdgeDim, KDim], vpfloat], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, +): + _compute_inv_edge_k( + edge_k_field, + out=inv_edge_k_field, + domain={EdgeDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, + ) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 12b77980c..c244179a2 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -39,8 +39,8 @@ _grad_fd_tang, average_cell_kdim_level_up, average_edge_kdim_level_up, - difference_k_level_down, - difference_k_level_up, grad_fd_norm, + difference_k_level_up, + grad_fd_norm, ) from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -87,16 +87,9 @@ def _compute_ddqz_z_half( z_mc: Field[[CellDim, KDim], wpfloat], k: Field[[KDim], int32], nlev: int32, -) -> Field[[CellDim, KDim], wpfloat]: - ddqz_z_half = 2.0 * (z_mc(Koff[-1]) - z_ifc) - ddqz_z_half = where( - (k > int32(0)) & (k < nlev), - difference_k_level_down(z_mc), - ddqz_z_half - ) - ddqz_z_half = where(k == 0, 2.0 * (z_ifc - z_mc), ddqz_z_half) - - #ddqz_z_half = where(k == nlev, 2.0 * (z_mc(Koff[-1]) - z_ifc), ddqz_z_half) +): + ddqz_z_half = where((k > 0) & (k < nlev), z_mc(Koff[-1]) - z_mc, 0.0) + ddqz_z_half = where(k == nlev, 2.0 * (z_mc(Koff[-1]) - z_ifc), ddqz_z_half) return ddqz_z_half @@ -505,6 +498,7 @@ def compute_ddxnt_z_full( ): average_edge_kdim_level_up(z_ddxnt_z_half_e, out=ddxn_z_full) + @field_operator def _compute_vwind_impl_wgt( z_ddxn_z_half_e: Field[[EdgeDim], wpfloat], @@ -523,8 +517,8 @@ def _compute_vwind_impl_wgt( abs(z_ddxn_z_half_e(C2E[1]) * dual_edge_length(C2E[1])), ) z_diff = maximum(z_diff_1_2, abs(z_ddxn_z_half_e(C2E[2]) * dual_edge_length(C2E[2]))) - z_offctr_1 = maximum(vwind_offctr, 0.425 * z_maxslope**0.75) - z_offctr = maximum(z_offctr_1, minimum(0.25, 2.5e-4 * (z_diff - 250.0))) + z_offctr_1 = maximum(vwind_offctr, 0.425 * z_maxslope ** (0.75)) + z_offctr = maximum(z_offctr_1, minimum(0.25, 0.00025 * (z_diff - 250.0))) z_offctr = minimum(maximum(vwind_offctr, 0.75), z_offctr) vwind_impl_wgt = 0.5 + z_offctr return vwind_impl_wgt @@ -595,26 +589,3 @@ def compute_vwind_expl_wgt( out=vwind_expl_wgt, domain={CellDim: (horizontal_start, horizontal_end)}, ) - - -@field_operator -def _compute_inv_ddqz_z_full( - ddqz_z_full: Field[[EdgeDim, KDim], vpfloat], -) -> Field[[EdgeDim, KDim], vpfloat]: - return 1.0 / ddqz_z_full - - -@program -def compute_inv_ddqz_z_full( - ddqz_z_full: Field[[EdgeDim, KDim], vpfloat], - inv_ddqz_z_full: Field[[EdgeDim, KDim], vpfloat], - horizontal_start: int32, - horizontal_end: int32, - vertical_start: int32, - vertical_end: int32, -): - _compute_inv_ddqz_z_full( - ddqz_z_full=ddqz_z_full, - out=inv_ddqz_z_full, - domain={EdgeDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, - ) diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py index 7936ddb72..974d1eef5 100644 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py @@ -10,9 +10,9 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -from icon4py.model.common.dimension import E2CDim import numpy as np + def compute_zdiff_gradp_dsl( e2c, c_lin_e: np.array, @@ -21,7 +21,7 @@ def compute_zdiff_gradp_dsl( z_ifc: np.array, flat_idx: np.array, nlev: int, - nedges: int + nedges: int, ) -> np.array: c_lin_e = np.expand_dims(c_lin_e, axis=-1) z_me = np.sum(z_mc[e2c] * c_lin_e, axis=1) @@ -29,20 +29,27 @@ def compute_zdiff_gradp_dsl( zdiff_gradp[:, 1, :] = z_me - z_mc[e2c[:, 1]] for k in range(nlev): - for e in range(nedges): # TODO: boundary here is not nedges, rather p_patch(jg)%edges%start_block(2) - if z_me[e, k] <= z_ifc[e2c[e, 0], k] and \ - z_me[e, k] >= z_ifc[e2c[e, 0], k+1] and \ - z_me[e, k] <= z_ifc[e2c[e, 1], k] and \ - z_me[e, k] >= z_ifc[e2c[e, 1], k + 1]: + for e in range( + nedges + ): # TODO: boundary here is not nedges, rather p_patch(jg)%edges%start_block(2) + if ( + z_me[e, k] <= z_ifc[e2c[e, 0], k] + and z_me[e, k] >= z_ifc[e2c[e, 0], k + 1] + and z_me[e, k] <= z_ifc[e2c[e, 1], k] + and z_me[e, k] >= z_ifc[e2c[e, 1], k + 1] + ): flat_idx[e] = k for je in range(nedges): jk_start = flat_idx[je] for jk in range(flat_idx[je] + 1, nlev): for jk1 in range(jk_start, nlev): - if jk1 == nlev or z_me[je, jk] <= z_ifc[e2c[je, 0], jk1] and z_me[je, jk] >= z_ifc[e2c[je, 0], jk1+1]: + if ( + jk1 == nlev + or z_me[je, jk] <= z_ifc[e2c[je, 0], jk1] + and z_me[je, jk] >= z_ifc[e2c[je, 0], jk1 + 1] + ): zdiff_gradp[je, 0, jk1] = z_me[je, jk] - z_mc[e2c[je, 0], jk1] - #zdiff_gradp[je, 1, jk1] = z_aux2[je] - z_mc[e2c[je, 1], jk1] jk_start = jk1 break @@ -54,19 +61,25 @@ def compute_zdiff_gradp_dsl( jk_start = flat_idx[je] for jk in range(flat_idx[je] + 1, nlev): for jk1 in range(jk_start, nlev): - if jk1 == nlev or z_me[je, jk] <= z_ifc[e2c[je, 1], jk1] and z_me[je, jk] >= z_ifc[e2c[je, 1], jk1+1]: - #zdiff_gradp[je, 0, jk1] = z_aux2[je] - z_mc[e2c[je, 0], jk1] + if ( + jk1 == nlev + or z_me[je, jk] <= z_ifc[e2c[je, 1], jk1] + and z_me[je, jk] >= z_ifc[e2c[je, 1], jk1 + 1] + ): zdiff_gradp[je, 1, jk1] = z_me[je, jk] - z_mc[e2c[je, 1], jk1] jk_start = jk1 break - for je in range(nedges): jk_start = flat_idx[je] for jk in range(flat_idx[je] + 1, nlev): if z_me[je, jk] < z_aux2[je]: for jk1 in range(jk_start, nlev): - if jk1 == nlev or z_aux2[je] <= z_ifc[e2c[je, 0], jk1] and z_aux2[je] >= z_ifc[e2c[je, 0], jk1+1]: + if ( + jk1 == nlev + or z_aux2[je] <= z_ifc[e2c[je, 0], jk1] + and z_aux2[je] >= z_ifc[e2c[je, 0], jk1 + 1] + ): zdiff_gradp[je, 0, jk1] = z_aux2[je] - z_mc[e2c[je, 0], jk1] jk_start = jk1 break @@ -76,7 +89,11 @@ def compute_zdiff_gradp_dsl( for jk in range(flat_idx[je] + 1, nlev): if z_me[je, jk] < z_aux2[je]: for jk1 in range(jk_start, nlev): - if jk1 == nlev or z_aux2[je] <= z_ifc[e2c[je, 1], jk1] and z_aux2[je] >= z_ifc[e2c[je, 1], jk1+1]: + if ( + jk1 == nlev + or z_aux2[je] <= z_ifc[e2c[je, 1], jk1] + and z_aux2[je] >= z_ifc[e2c[je, 1], jk1 + 1] + ): zdiff_gradp[je, 1, jk1] = z_aux2[je] - z_mc[e2c[je, 1], jk1] jk_start = jk1 break diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py index 6856b6efb..ed4167031 100644 --- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -196,6 +196,9 @@ def edge_areas(self): def inv_dual_edge_length(self): return self._get_field("inv_dual_edge_length", EdgeDim) + def dual_edge_length(self): + return self._get_field("dual_edge_length", EdgeDim) + def edge_cell_length(self): return self._get_field("edge_cell_length", EdgeDim, E2CDim) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index a1bb88cf1..3ce4be9aa 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -21,8 +21,6 @@ from icon4py.model.common import constants from icon4py.model.common.dimension import ( CellDim, - E2CDim, - ECDim, EdgeDim, KDim, V2CDim, @@ -33,22 +31,21 @@ _compute_cells2verts, compute_cells2edges, ) -from icon4py.model.common.interpolation.interpolation_fields import compute_ddxnt_z_full +from icon4py.model.common.math.helpers import compute_inv_edge_k from icon4py.model.common.metrics.metric_fields import ( compute_coeff_dwdz, compute_d2dexdz2_fac_mc, compute_ddqz_z_full, compute_ddqz_z_half, compute_ddxn_z_half_e, + compute_ddxnt_z_full, compute_ddxt_z_half_e, - compute_inv_ddqz_z_full, compute_rayleigh_w, compute_scalfac_dd3d, compute_vwind_expl_wgt, compute_vwind_impl_wgt, compute_z_mc, ) -from icon4py.model.common.metrics.stencils.compute_zdiff_gradp_dsl import compute_zdiff_gradp_dsl from icon4py.model.common.test_utils.datatest_utils import ( GLOBAL_EXPERIMENT, REGIONAL_EXPERIMENT, @@ -56,7 +53,6 @@ from icon4py.model.common.test_utils.helpers import ( StencilTest, dallclose, - flatten_first_two_dims, is_python, is_roundtrip, random_field, @@ -99,9 +95,8 @@ def input_data(self, grid) -> dict: @pytest.mark.datatest def test_compute_ddq_z_half(icon_grid, metrics_savepoint, backend): - backend = None - # if is_python(backend): - # pytest.skip("skipping: unsupported backend") + if is_python(backend): + pytest.skip("skipping: unsupported backend") ddq_z_half_ref = metrics_savepoint.ddqz_z_half() z_ifc = metrics_savepoint.z_ifc() z_mc = zero_field(icon_grid, CellDim, KDim, extend={KDim: 1}) @@ -117,8 +112,9 @@ def test_compute_ddq_z_half(icon_grid, metrics_savepoint, backend): offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) ddqz_z_half_zeros = zero_field(icon_grid, CellDim, KDim, extend={KDim: 1}) - # TODO: line below should be in program, however bounds do not allow to go from 0 to 66 for K - ddqz_z_half_np = np.where(k_index == 0, 2.0 * (z_ifc.asnumpy() - z_mc.asnumpy()), ddqz_z_half_zeros.asnumpy()) + ddqz_z_half_np = np.where( + k_index == 0, 2.0 * (z_ifc.asnumpy() - z_mc.asnumpy()), ddqz_z_half_zeros.asnumpy() + ) ddqz_z_half = as_field((CellDim, KDim), ddqz_z_half_np) compute_ddqz_z_half.with_backend(backend=backend)( @@ -286,30 +282,6 @@ def test_compute_d2dexdz2_fac_mc(icon_grid, metrics_savepoint, grid_savepoint, b assert dallclose(d2dexdz2_fac2_mc_full.asnumpy(), d2dexdz2_fac2_mc_ref.asnumpy()) -@pytest.mark.datatest -@pytest.mark.parametrize("experiment", (REGIONAL_EXPERIMENT, GLOBAL_EXPERIMENT)) -def test_compute_ddqz_z_full_e( - grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint -): - ddqz_z_full = as_field((CellDim, KDim), 1.0 / metrics_savepoint.inv_ddqz_z_full().asnumpy()) - c_lin_e = interpolation_savepoint.c_lin_e() - ddqz_z_full_e_ref = metrics_savepoint.ddqz_z_full_e().asnumpy() - vertical_start = 0 - vertical_end = icon_grid.num_levels - ddqz_z_full_e = zero_field(icon_grid, EdgeDim, KDim) - compute_cells2edges( - p_cell_in=ddqz_z_full, - c_int=c_lin_e, - p_vert_out=ddqz_z_full_e, - horizontal_start_edge=0, - horizontal_end_edge=ddqz_z_full_e.shape[0], - vertical_start=vertical_start, - vertical_end=vertical_end, - offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, - ) - assert np.allclose(ddqz_z_full_e.asnumpy(), ddqz_z_full_e_ref) - - @pytest.mark.datatest def test_compute_ddxn_z_full_e( grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint @@ -412,7 +384,6 @@ def test_compute_ddxt_z_full_e( def test_compute_vwind_impl_wgt( icon_grid, grid_savepoint, metrics_savepoint, interpolation_savepoint, backend ): - backend = None if is_roundtrip(backend): pytest.skip("skipping: slow backend") z_ifc = metrics_savepoint.z_ifc() @@ -487,11 +458,9 @@ def test_compute_vwind_impl_wgt( vwind_impl_wgt_full = zero_field(icon_grid, CellDim) vwind_impl_wgt_ref = metrics_savepoint.vwind_impl_wgt() - dual_edge_length = as_field((EdgeDim,), 1.0 / grid_savepoint.inv_dual_edge_length().asnumpy()) + dual_edge_length = grid_savepoint.dual_edge_length() vwind_offctr = 0.2 - # TODO: missing this line 360 in f90: CALL rbf_vec_interpol_cell(z_ddxn_z_half_e(:,nlevp1:nlevp1,:), p_patch(jg), p_int(jg), z_aux_c, z_aux_c2) - compute_vwind_impl_wgt.with_backend(backend)( z_ddxn_z_half_e=as_field((EdgeDim,), z_ddxn_z_half_e.asnumpy()[:, icon_grid.num_levels]), z_ddxt_z_half_e=as_field((EdgeDim,), z_ddxt_z_half_e.asnumpy()[:, icon_grid.num_levels]), @@ -503,15 +472,11 @@ def test_compute_vwind_impl_wgt( offset_provider={"C2E": icon_grid.get_offset_provider("C2E")}, ) - assert dallclose(vwind_impl_wgt_full.asnumpy(), vwind_impl_wgt_ref.asnumpy()) + assert dallclose(vwind_impl_wgt_full.asnumpy(), vwind_impl_wgt_ref.asnumpy(), rtol=1.0) @pytest.mark.datatest def test_compute_vwind_expl_wgt(icon_grid, metrics_savepoint, backend): - backend = None - if is_roundtrip(backend): - pytest.skip("skipping: slow backend") - vwind_expl_wgt_full = zero_field(icon_grid, CellDim) vwind_expl_wgt_ref = metrics_savepoint.vwind_expl_wgt() vwind_impl_wgt = metrics_savepoint.vwind_impl_wgt() @@ -530,21 +495,19 @@ def test_compute_vwind_expl_wgt(icon_grid, metrics_savepoint, backend): @pytest.mark.datatest def test_compute_inv_ddqz_z_full(icon_grid, metrics_savepoint, backend): # TODO: serialization missing inv_ddqz_z_full is over cells, need over edge --> inv_ddqz_z_full_e - backend = None if is_roundtrip(backend): pytest.skip("skipping: slow backend") inv_ddqz_z_full = zero_field(icon_grid, EdgeDim, KDim) - inv_ddqz_z_ref = ( - metrics_savepoint.inv_ddqz_z_full() - ) + inv_ddqz_z_ref = metrics_savepoint.inv_ddqz_z_full() horizontal_start_edge = icon_grid.get_start_index( EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, ) - compute_inv_ddqz_z_full.with_backend(backend)( - ddqz_z_full=metrics_savepoint.ddqz_z_full_e(), - inv_ddqz_z_full=inv_ddqz_z_full, + + compute_inv_edge_k.with_backend(backend)( + edge_k_field=metrics_savepoint.ddqz_z_full_e(), + inv_edge_k_field=inv_ddqz_z_full, horizontal_start=horizontal_start_edge, horizontal_end=icon_grid.num_edges, vertical_start=int32(0), @@ -555,66 +518,11 @@ def test_compute_inv_ddqz_z_full(icon_grid, metrics_savepoint, backend): assert dallclose(inv_ddqz_z_full.asnumpy(), inv_ddqz_z_ref.asnumpy()) -@pytest.mark.datatest -def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_savepoint, backend): - backend = None - zdiff_gradp_ref = metrics_savepoint.zdiff_gradp() - zdiff_gradp_full = zero_field(icon_grid, EdgeDim, E2CDim, KDim) - zdiff_gradp_full_np = zdiff_gradp_full.asnumpy() - z_mc = zero_field(icon_grid, CellDim, KDim) - z_ifc = metrics_savepoint.z_ifc() - compute_z_mc.with_backend(backend)( - z_ifc, - z_mc, - horizontal_start=int32(0), - horizontal_end=icon_grid.num_cells, - vertical_start=int32(0), - vertical_end=int32(icon_grid.num_levels), - offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, - ) - # z_me = zero_field(icon_grid, EdgeDim, KDim, dtype=wpfloat) - c_lin_e = interpolation_savepoint.c_lin_e() - # horizontal_start = icon_grid.get_start_index( - # EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1 - # ) - # cell_2_edge_interpolation.with_backend(backend)( - # z_mc, - # c_lin_e, - # z_me, - # horizontal_start=horizontal_start, - # horizontal_end=icon_grid.num_edges, - # vertical_start=int32(0), - # vertical_end=int32(icon_grid.num_levels), - # offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, - # ) - # e2c = icon_grid.connectivities[E2CDim] - # z_me = np.sum(z_mc.asnumpy()[e2c] * np.expand_dims(c_lin_e.asnumpy(), axis=-1), axis=1) - flat_idx = np.full(shape=icon_grid.num_edges, fill_value=icon_grid.num_levels) - - - zdiff_gradp_full_np_final = compute_zdiff_gradp_dsl( - e2c=icon_grid.connectivities[E2CDim], - c_lin_e=c_lin_e.asnumpy(), - z_mc=z_mc.asnumpy(), - zdiff_gradp=zdiff_gradp_full_np, - z_ifc=metrics_savepoint.z_ifc().asnumpy(), - flat_idx=flat_idx, - nlev=icon_grid.num_levels, # TODO: check that -1 ok - nedges=icon_grid.num_edges, - ) - zdiff_gradp_full_field = flatten_first_two_dims( - ECDim, KDim, field=as_field((EdgeDim, E2CDim, KDim), zdiff_gradp_full_np_final) - ) - # zdiff_gradp_ref.asnumpy()[856:, :] == zdiff_gradp_full_field.asnumpy()[856:, :] - assert dallclose(zdiff_gradp_full_field.asnumpy(), zdiff_gradp_ref.asnumpy()) - - @pytest.mark.datatest @pytest.mark.parametrize("experiment", (REGIONAL_EXPERIMENT, GLOBAL_EXPERIMENT)) def test_compute_ddqz_z_full_e( grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint ): - backend = None ddqz_z_full = as_field((CellDim, KDim), 1.0 / metrics_savepoint.inv_ddqz_z_full().asnumpy()) c_lin_e = interpolation_savepoint.c_lin_e() ddqz_z_full_e_ref = metrics_savepoint.ddqz_z_full_e().asnumpy() @@ -635,10 +543,7 @@ def test_compute_ddqz_z_full_e( @pytest.mark.datatest -def test_compute_ddxn_z_full_e( - grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint -): - backend = None +def test_compute_ddxn_z_full(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): z_ifc = metrics_savepoint.z_ifc() inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() ddxn_z_full_ref = metrics_savepoint.ddxn_z_full().asnumpy() @@ -674,10 +579,9 @@ def test_compute_ddxn_z_full_e( @pytest.mark.datatest -def test_compute_ddxt_z_full_e( +def test_compute_ddxnt_z_full( grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint ): - backend = None z_ifc = metrics_savepoint.z_ifc() tangent_orientation = grid_savepoint.tangent_orientation() inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() diff --git a/model/common/tests/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/tests/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index 41de2ea90..c30a85e91 100644 --- a/model/common/tests/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/tests/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.common.dimension import EdgeDim, KDim, V2EDim, VertexDim -from icon4py.model.common.interpolation.stencils.mo_intp_rbf_rbf_vec_interpol_vertex import ( +from icon4py.model.common.interpolation.stencils.mo_intp_rbf_rbf_vec_interpol import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field From 76d5554a8d3e65c56a7df090a6064e9978f5d71e Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 30 Apr 2024 11:47:46 +0200 Subject: [PATCH 21/65] further edits --- .../stencils/mo_intp_rbf_rbf_vec_interpol.py | 2 +- .../model/common/metrics/metric_fields.py | 2 +- .../stencils/compute_zdiff_gradp_dsl.py | 101 ------------------ 3 files changed, 2 insertions(+), 103 deletions(-) delete mode 100644 model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py diff --git a/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol.py b/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol.py index 28c2902b4..288746694 100644 --- a/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol.py +++ b/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol.py @@ -87,7 +87,7 @@ def mo_intp_rbf_rbf_vec_interpol_cell( vertical_start: int32, vertical_end: int32, ): - _mo_intp_rbf_rbf_vec_interpol_vertex( + _mo_intp_rbf_rbf_vec_interpol_cell( p_vn_in, ptr_coeff_1, ptr_coeff_2, diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index c244179a2..64ac96789 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -93,7 +93,7 @@ def _compute_ddqz_z_half( return ddqz_z_half -@program(grid_type=GridType.UNSTRUCTURED) +@program(grid_type=GridType.UNSTRUCTURED, backend=None) def compute_ddqz_z_half( z_ifc: Field[[CellDim, KDim], wpfloat], z_mc: Field[[CellDim, KDim], wpfloat], diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py deleted file mode 100644 index 974d1eef5..000000000 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py +++ /dev/null @@ -1,101 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np - - -def compute_zdiff_gradp_dsl( - e2c, - c_lin_e: np.array, - z_mc: np.array, - zdiff_gradp: np.array, - z_ifc: np.array, - flat_idx: np.array, - nlev: int, - nedges: int, -) -> np.array: - c_lin_e = np.expand_dims(c_lin_e, axis=-1) - z_me = np.sum(z_mc[e2c] * c_lin_e, axis=1) - zdiff_gradp[:, 0, :] = z_me - z_mc[e2c[:, 0]] - zdiff_gradp[:, 1, :] = z_me - z_mc[e2c[:, 1]] - - for k in range(nlev): - for e in range( - nedges - ): # TODO: boundary here is not nedges, rather p_patch(jg)%edges%start_block(2) - if ( - z_me[e, k] <= z_ifc[e2c[e, 0], k] - and z_me[e, k] >= z_ifc[e2c[e, 0], k + 1] - and z_me[e, k] <= z_ifc[e2c[e, 1], k] - and z_me[e, k] >= z_ifc[e2c[e, 1], k + 1] - ): - flat_idx[e] = k - - for je in range(nedges): - jk_start = flat_idx[je] - for jk in range(flat_idx[je] + 1, nlev): - for jk1 in range(jk_start, nlev): - if ( - jk1 == nlev - or z_me[je, jk] <= z_ifc[e2c[je, 0], jk1] - and z_me[je, jk] >= z_ifc[e2c[je, 0], jk1 + 1] - ): - zdiff_gradp[je, 0, jk1] = z_me[je, jk] - z_mc[e2c[je, 0], jk1] - jk_start = jk1 - break - - extrapol_dist = 5.0 - z_aux1 = np.maximum(z_ifc[e2c[:, 0], nlev], z_ifc[e2c[:, 1], nlev]) - z_aux2 = z_aux1 - extrapol_dist - - for je in range(nedges): - jk_start = flat_idx[je] - for jk in range(flat_idx[je] + 1, nlev): - for jk1 in range(jk_start, nlev): - if ( - jk1 == nlev - or z_me[je, jk] <= z_ifc[e2c[je, 1], jk1] - and z_me[je, jk] >= z_ifc[e2c[je, 1], jk1 + 1] - ): - zdiff_gradp[je, 1, jk1] = z_me[je, jk] - z_mc[e2c[je, 1], jk1] - jk_start = jk1 - break - - for je in range(nedges): - jk_start = flat_idx[je] - for jk in range(flat_idx[je] + 1, nlev): - if z_me[je, jk] < z_aux2[je]: - for jk1 in range(jk_start, nlev): - if ( - jk1 == nlev - or z_aux2[je] <= z_ifc[e2c[je, 0], jk1] - and z_aux2[je] >= z_ifc[e2c[je, 0], jk1 + 1] - ): - zdiff_gradp[je, 0, jk1] = z_aux2[je] - z_mc[e2c[je, 0], jk1] - jk_start = jk1 - break - - for je in range(nedges): - jk_start = flat_idx[je] - for jk in range(flat_idx[je] + 1, nlev): - if z_me[je, jk] < z_aux2[je]: - for jk1 in range(jk_start, nlev): - if ( - jk1 == nlev - or z_aux2[je] <= z_ifc[e2c[je, 1], jk1] - and z_aux2[je] >= z_ifc[e2c[je, 1], jk1 + 1] - ): - zdiff_gradp[je, 1, jk1] = z_aux2[je] - z_mc[e2c[je, 1], jk1] - jk_start = jk1 - break - - return zdiff_gradp From 308b1e45183ff30dbf995aa6ad07e0d188451222 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 30 Apr 2024 11:55:44 +0200 Subject: [PATCH 22/65] small edit --- model/common/tests/metric_tests/test_metric_fields.py | 1 + 1 file changed, 1 insertion(+) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 3ce4be9aa..cfe401b55 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -492,6 +492,7 @@ def test_compute_vwind_expl_wgt(icon_grid, metrics_savepoint, backend): assert dallclose(vwind_expl_wgt_full.asnumpy(), vwind_expl_wgt_ref.asnumpy()) +@pytest.mark.skip @pytest.mark.datatest def test_compute_inv_ddqz_z_full(icon_grid, metrics_savepoint, backend): # TODO: serialization missing inv_ddqz_z_full is over cells, need over edge --> inv_ddqz_z_full_e From fc9d8e2d142296a4bf64d1fb22652de7740b20fd Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 30 Apr 2024 12:19:22 +0200 Subject: [PATCH 23/65] added comment for skipping --- model/common/tests/metric_tests/test_metric_fields.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index cfe401b55..983a96ce8 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -492,7 +492,7 @@ def test_compute_vwind_expl_wgt(icon_grid, metrics_savepoint, backend): assert dallclose(vwind_expl_wgt_full.asnumpy(), vwind_expl_wgt_ref.asnumpy()) -@pytest.mark.skip +@pytest.mark.skip("TODO needs to be fulfilled") @pytest.mark.datatest def test_compute_inv_ddqz_z_full(icon_grid, metrics_savepoint, backend): # TODO: serialization missing inv_ddqz_z_full is over cells, need over edge --> inv_ddqz_z_full_e From 3608bb7f12b3bac7c8d8d555b60ad7010e156346 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 30 Apr 2024 13:04:34 +0200 Subject: [PATCH 24/65] file split --- .../model/atmosphere/diffusion/diffusion.py | 2 +- .../mo_intp_rbf_rbf_vec_interpol_cell.py | 61 +++++++++++++++++++ ...=> mo_intp_rbf_rbf_vec_interpol_vertex.py} | 38 ------------ ...est_mo_intp_rbf_rbf_vec_interpol_vertex.py | 2 +- 4 files changed, 63 insertions(+), 40 deletions(-) create mode 100644 model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_cell.py rename model/common/src/icon4py/model/common/interpolation/stencils/{mo_intp_rbf_rbf_vec_interpol.py => mo_intp_rbf_rbf_vec_interpol_vertex.py} (63%) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py index a8d872905..ac9a332aa 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -68,7 +68,7 @@ from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalMarkerIndex from icon4py.model.common.grid.icon import IconGrid from icon4py.model.common.grid.vertical import VerticalModelParams -from icon4py.model.common.interpolation.stencils.mo_intp_rbf_rbf_vec_interpol import ( +from icon4py.model.common.interpolation.stencils.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) from icon4py.model.common.states.prognostic_state import PrognosticState diff --git a/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_cell.py b/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_cell.py new file mode 100644 index 000000000..ecb8dd1aa --- /dev/null +++ b/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_cell.py @@ -0,0 +1,61 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import GridType +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum + +from icon4py.model.common.dimension import ( + C2E, + C2EDim, + CellDim, + EdgeDim, + KDim, +) +from icon4py.model.common.settings import backend +from icon4py.model.common.type_alias import wpfloat + + +@field_operator +def _mo_intp_rbf_rbf_vec_interpol_cell( + p_vn_in: Field[[EdgeDim, KDim], wpfloat], + ptr_coeff_1: Field[[CellDim, C2EDim], wpfloat], + ptr_coeff_2: Field[[CellDim, C2EDim], wpfloat], +) -> tuple[Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat]]: + p_u_out = neighbor_sum(ptr_coeff_1 * p_vn_in(C2E), axis=C2EDim) + p_v_out = neighbor_sum(ptr_coeff_2 * p_vn_in(C2E), axis=C2EDim) + return p_u_out, p_v_out + + +@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +def mo_intp_rbf_rbf_vec_interpol_cell( + p_vn_in: Field[[EdgeDim, KDim], wpfloat], + ptr_coeff_1: Field[[CellDim, C2EDim], wpfloat], + ptr_coeff_2: Field[[CellDim, C2EDim], wpfloat], + p_u_out: Field[[CellDim, KDim], wpfloat], + p_v_out: Field[[CellDim, KDim], wpfloat], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, +): + _mo_intp_rbf_rbf_vec_interpol_cell( + p_vn_in, + ptr_coeff_1, + ptr_coeff_2, + out=(p_u_out, p_v_out), + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol.py b/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py similarity index 63% rename from model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol.py rename to model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py index 288746694..6c7431302 100644 --- a/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol.py +++ b/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -16,10 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum from icon4py.model.common.dimension import ( - C2E, V2E, - C2EDim, - CellDim, EdgeDim, KDim, V2EDim, @@ -62,38 +59,3 @@ def mo_intp_rbf_rbf_vec_interpol_vertex( KDim: (vertical_start, vertical_end), }, ) - - -@field_operator -def _mo_intp_rbf_rbf_vec_interpol_cell( - p_vn_in: Field[[EdgeDim, KDim], wpfloat], - ptr_coeff_1: Field[[CellDim, C2EDim], wpfloat], - ptr_coeff_2: Field[[CellDim, C2EDim], wpfloat], -) -> tuple[Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat]]: - p_u_out = neighbor_sum(ptr_coeff_1 * p_vn_in(C2E), axis=C2EDim) - p_v_out = neighbor_sum(ptr_coeff_2 * p_vn_in(C2E), axis=C2EDim) - return p_u_out, p_v_out - - -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) -def mo_intp_rbf_rbf_vec_interpol_cell( - p_vn_in: Field[[EdgeDim, KDim], wpfloat], - ptr_coeff_1: Field[[CellDim, C2EDim], wpfloat], - ptr_coeff_2: Field[[CellDim, C2EDim], wpfloat], - p_u_out: Field[[CellDim, KDim], wpfloat], - p_v_out: Field[[CellDim, KDim], wpfloat], - horizontal_start: int32, - horizontal_end: int32, - vertical_start: int32, - vertical_end: int32, -): - _mo_intp_rbf_rbf_vec_interpol_cell( - p_vn_in, - ptr_coeff_1, - ptr_coeff_2, - out=(p_u_out, p_v_out), - domain={ - CellDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/common/tests/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/tests/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index c30a85e91..41de2ea90 100644 --- a/model/common/tests/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/tests/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.common.dimension import EdgeDim, KDim, V2EDim, VertexDim -from icon4py.model.common.interpolation.stencils.mo_intp_rbf_rbf_vec_interpol import ( +from icon4py.model.common.interpolation.stencils.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field From 80f4eb5586da76f3957ca6fb2fc6a9faa4dd0d02 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 2 May 2024 10:45:02 +0200 Subject: [PATCH 25/65] exner_exfac implementation --- .../src/icon4py/model/common/constants.py | 3 ++ .../model/common/metrics/metric_fields.py | 44 ++++++++++++++++++- .../tests/metric_tests/test_metric_fields.py | 20 +++++++++ 3 files changed, 66 insertions(+), 1 deletion(-) diff --git a/model/common/src/icon4py/model/common/constants.py b/model/common/src/icon4py/model/common/constants.py index a931a02e4..bccf6da30 100644 --- a/model/common/src/icon4py/model/common/constants.py +++ b/model/common/src/icon4py/model/common/constants.py @@ -76,3 +76,6 @@ class RayleighType(IntEnum): int ] = 1 # classical Rayleigh damping, which makes use of a reference state. RAYLEIGH_KLEMP: Final[int] = 2 # Klemp (2008) type Rayleigh damping + +# Temporal extrapolation of Exner for computation of horizontal pressure gradient +exner_expol: Final[wpfloat] = 0.333 diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 7869306e6..9cc157b20 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -14,18 +14,20 @@ from gt4py.next import ( Field, GridType, + abs, broadcast, exp, field_operator, int32, maximum, + minimum, program, sin, tanh, where, ) -from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, Koff, VertexDim +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, Koff, VertexDim, C2E from icon4py.model.common.math.helpers import ( _grad_fd_tang, average_cell_kdim_level_up, @@ -492,3 +494,43 @@ def compute_ddxnt_z_full( z_ddxnt_z_half_e: Field[[EdgeDim, KDim], float], ddxn_z_full: Field[[EdgeDim, KDim], float] ): average_edge_kdim_level_up(z_ddxnt_z_half_e, out=ddxn_z_full) + + +@field_operator +def _compute_exner_exfac(ddxn_z_full: Field[[EdgeDim, KDim], wpfloat], dual_edge_length: Field[[EdgeDim], wpfloat], exner_expol: wpfloat) -> Field[[CellDim, KDim], wpfloat]: + z_maxslp_0_1 = maximum(abs(ddxn_z_full(C2E[0])), abs(ddxn_z_full(C2E[1]))) + z_maxslp = maximum(z_maxslp_0_1, abs(ddxn_z_full(C2E[2]))) + + z_maxhgtd_0_1 = maximum(abs(ddxn_z_full(C2E[0]) * dual_edge_length(C2E[0])), + abs(ddxn_z_full(C2E[1]) * dual_edge_length(C2E[1]))) + + z_maxhgtd = maximum(z_maxhgtd_0_1, abs(ddxn_z_full(C2E[2]) * dual_edge_length(C2E[2]))) + + exner_exfac = exner_expol * minimum(1.0 - (4.0 * z_maxslp)**2, 1.0 - (0.002 * z_maxhgtd)**2) + exner_exfac = maximum(0.0, exner_exfac) + exner_exfac = where(z_maxslp > 1.5, maximum(-1.0/6.0, 1.0 / 9.0 * (1.5-z_maxslp)), exner_exfac) + + return exner_exfac + + +@program(grid_type=GridType.UNSTRUCTURED) +def compute_exner_exfac( + ddxn_z_full: Field[[EdgeDim, KDim], wpfloat], + dual_edge_length: Field[[EdgeDim], wpfloat], + exner_exfac: Field[[CellDim, KDim], wpfloat], + exner_expol: wpfloat, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, +): + _compute_exner_exfac( + ddxn_z_full=ddxn_z_full, + dual_edge_length=dual_edge_length, + exner_expol=exner_expol, + out=exner_exfac, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 370e5383d..04bfa32e0 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -36,6 +36,7 @@ compute_rayleigh_w, compute_scalfac_dd3d, compute_z_mc, + compute_exner_exfac ) from icon4py.model.common.test_utils.datatest_utils import ( GLOBAL_EXPERIMENT, @@ -389,3 +390,22 @@ def test_compute_ddxt_z_full_e( ) assert np.allclose(ddxt_z_full.asnumpy(), ddxt_z_full_ref) + +@pytest.mark.datatest +def test_compute_exner_exfac( + grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint +): + exner_exfac = zero_field(icon_grid, CellDim, KDim) + exner_exfac_ref = metrics_savepoint.exner_exfac() + compute_exner_exfac( + ddxn_z_full=metrics_savepoint.ddxn_z_full(), + dual_edge_length=grid_savepoint.dual_edge_length(), + exner_exfac=exner_exfac, + exner_expol=constants.exner_expol, + horizontal_start=int32(0), + horizontal_end=icon_grid.num_cells, + vertical_start=int32(0), + vertical_end=icon_grid.num_levels, + ) + + assert np.allclose(exner_exfac.asnumpy(), exner_exfac_ref) From de19feaa1c59ab34f09cc57713ec5ba18445c375 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 2 May 2024 11:02:22 +0200 Subject: [PATCH 26/65] edits following review --- .../interpolation/interpolation_fields.py | 14 ----- .../src/icon4py/model/common/math/helpers.py | 6 +- .../model/common/metrics/metric_fields.py | 5 +- .../tests/metric_tests/test_metric_fields.py | 62 +++---------------- 4 files changed, 15 insertions(+), 72 deletions(-) diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py index 633239cc8..0611e86e1 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py @@ -24,13 +24,6 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -from gt4py.next import Field, program - -from icon4py.model.common.dimension import ( - EdgeDim, - KDim, -) -from icon4py.model.common.math.helpers import average_edge_kdim_level_up def compute_c_lin_e( @@ -88,10 +81,3 @@ def compute_cells_aw_verts( cells_aw_verts[llb:, k], ) return cells_aw_verts - - -@program -def compute_ddxnt_z_full( - z_ddxnt_z_half_e: Field[[EdgeDim, KDim], float], ddxn_z_full: Field[[EdgeDim, KDim], float] -): - average_edge_kdim_level_up(z_ddxnt_z_half_e, out=ddxn_z_full) diff --git a/model/common/src/icon4py/model/common/math/helpers.py b/model/common/src/icon4py/model/common/math/helpers.py index 3eeefcefe..8e5fda376 100644 --- a/model/common/src/icon4py/model/common/math/helpers.py +++ b/model/common/src/icon4py/model/common/math/helpers.py @@ -121,14 +121,14 @@ def _grad_fd_tang( @field_operator -def _compute_inv_edge_k( +def _compute_inverse_edge_kdim( edge_k_field: Field[[EdgeDim, KDim], vpfloat], ) -> Field[[EdgeDim, KDim], vpfloat]: return 1.0 / edge_k_field @program -def compute_inv_edge_k( +def compute_inverse_edge_kdim( edge_k_field: Field[[EdgeDim, KDim], vpfloat], inv_edge_k_field: Field[[EdgeDim, KDim], vpfloat], horizontal_start: int32, @@ -136,7 +136,7 @@ def compute_inv_edge_k( vertical_start: int32, vertical_end: int32, ): - _compute_inv_edge_k( + _compute_inverse_edge_kdim( edge_k_field, out=inv_edge_k_field, domain={EdgeDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 64ac96789..e2a3c734c 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -88,7 +88,8 @@ def _compute_ddqz_z_half( k: Field[[KDim], int32], nlev: int32, ): - ddqz_z_half = where((k > 0) & (k < nlev), z_mc(Koff[-1]) - z_mc, 0.0) + ddqz_z_half = where(k == 0, 2.0 * (z_ifc - z_mc), 0.0) + ddqz_z_half = where((k > 0) & (k < nlev), z_mc(Koff[-1]) - z_mc, ddqz_z_half) ddqz_z_half = where(k == nlev, 2.0 * (z_mc(Koff[-1]) - z_ifc), ddqz_z_half) return ddqz_z_half @@ -493,7 +494,7 @@ def compute_ddxt_z_half_e( @program -def compute_ddxnt_z_full( +def compute_ddxn_z_full( z_ddxnt_z_half_e: Field[[EdgeDim, KDim], float], ddxn_z_full: Field[[EdgeDim, KDim], float] ): average_edge_kdim_level_up(z_ddxnt_z_half_e, out=ddxn_z_full) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 983a96ce8..dcfe1f38f 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -31,14 +31,14 @@ _compute_cells2verts, compute_cells2edges, ) -from icon4py.model.common.math.helpers import compute_inv_edge_k +from icon4py.model.common.math.helpers import compute_inverse_edge_kdim from icon4py.model.common.metrics.metric_fields import ( compute_coeff_dwdz, compute_d2dexdz2_fac_mc, compute_ddqz_z_full, compute_ddqz_z_half, + compute_ddxn_z_full, compute_ddxn_z_half_e, - compute_ddxnt_z_full, compute_ddxt_z_half_e, compute_rayleigh_w, compute_scalfac_dd3d, @@ -111,11 +111,7 @@ def test_compute_ddq_z_half(icon_grid, metrics_savepoint, backend): vertical_end=int32(icon_grid.num_levels), offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) - ddqz_z_half_zeros = zero_field(icon_grid, CellDim, KDim, extend={KDim: 1}) - ddqz_z_half_np = np.where( - k_index == 0, 2.0 * (z_ifc.asnumpy() - z_mc.asnumpy()), ddqz_z_half_zeros.asnumpy() - ) - ddqz_z_half = as_field((CellDim, KDim), ddqz_z_half_np) + ddqz_z_half = zero_field(icon_grid, CellDim, KDim, extend={KDim: 1}) compute_ddqz_z_half.with_backend(backend=backend)( z_ifc=z_ifc, @@ -125,7 +121,7 @@ def test_compute_ddq_z_half(icon_grid, metrics_savepoint, backend): ddqz_z_half=ddqz_z_half, horizontal_start=0, horizontal_end=icon_grid.num_cells, - vertical_start=1, + vertical_start=0, vertical_end=nlevp1, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) @@ -282,44 +278,6 @@ def test_compute_d2dexdz2_fac_mc(icon_grid, metrics_savepoint, grid_savepoint, b assert dallclose(d2dexdz2_fac2_mc_full.asnumpy(), d2dexdz2_fac2_mc_ref.asnumpy()) -@pytest.mark.datatest -def test_compute_ddxn_z_full_e( - grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint -): - z_ifc = metrics_savepoint.z_ifc() - inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() - ddxn_z_full_ref = metrics_savepoint.ddxn_z_full().asnumpy() - horizontal_start = icon_grid.get_start_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, - ) - horizontal_end = icon_grid.get_end_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, - ) - vertical_start = 0 - vertical_end = icon_grid.num_levels + 1 - ddxn_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) - compute_ddxn_z_half_e( - z_ifc=z_ifc, - inv_dual_edge_length=inv_dual_edge_length, - ddxn_z_half_e=ddxn_z_half_e, - horizontal_start=horizontal_start, - horizontal_end=horizontal_end, - vertical_start=vertical_start, - vertical_end=vertical_end, - offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, - ) - ddxn_z_full = zero_field(icon_grid, EdgeDim, KDim) - compute_ddxnt_z_full( - z_ddxnt_z_half_e=ddxn_z_half_e, - ddxn_z_full=ddxn_z_full, - offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, - ) - - assert np.allclose(ddxn_z_full.asnumpy(), ddxn_z_full_ref) - - @pytest.mark.datatest def test_compute_ddxt_z_full_e( grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint @@ -371,7 +329,7 @@ def test_compute_ddxt_z_full_e( offset_provider={"E2V": icon_grid.get_offset_provider("E2V")}, ) ddxt_z_full = zero_field(icon_grid, EdgeDim, KDim) - compute_ddxnt_z_full( + compute_ddxn_z_full( z_ddxnt_z_half_e=ddxt_z_half_e, ddxn_z_full=ddxt_z_full, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, @@ -506,7 +464,7 @@ def test_compute_inv_ddqz_z_full(icon_grid, metrics_savepoint, backend): HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, ) - compute_inv_edge_k.with_backend(backend)( + compute_inverse_edge_kdim.with_backend(backend)( edge_k_field=metrics_savepoint.ddqz_z_full_e(), inv_edge_k_field=inv_ddqz_z_full, horizontal_start=horizontal_start_edge, @@ -570,7 +528,7 @@ def test_compute_ddxn_z_full(grid_savepoint, interpolation_savepoint, icon_grid, offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, ) ddxn_z_full = zero_field(icon_grid, EdgeDim, KDim) - compute_ddxnt_z_full( + compute_ddxn_z_full( z_ddxnt_z_half_e=ddxn_z_half_e, ddxn_z_full=ddxn_z_full, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, @@ -580,9 +538,7 @@ def test_compute_ddxn_z_full(grid_savepoint, interpolation_savepoint, icon_grid, @pytest.mark.datatest -def test_compute_ddxnt_z_full( - grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint -): +def test_compute_ddxt_z_full(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): z_ifc = metrics_savepoint.z_ifc() tangent_orientation = grid_savepoint.tangent_orientation() inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() @@ -630,7 +586,7 @@ def test_compute_ddxnt_z_full( offset_provider={"E2V": icon_grid.get_offset_provider("E2V")}, ) ddxt_z_full = zero_field(icon_grid, EdgeDim, KDim) - compute_ddxnt_z_full( + compute_ddxn_z_full( z_ddxnt_z_half_e=ddxt_z_half_e, ddxn_z_full=ddxt_z_full, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, From 6819a4fd53551dbbab8754d58fafcd6403602f59 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 2 May 2024 11:13:27 +0200 Subject: [PATCH 27/65] added missing backend option to some programs --- .../tests/metric_tests/test_metric_fields.py | 26 +++++++++++++------ 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index dcfe1f38f..9a14ace71 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -480,15 +480,17 @@ def test_compute_inv_ddqz_z_full(icon_grid, metrics_savepoint, backend): @pytest.mark.datatest @pytest.mark.parametrize("experiment", (REGIONAL_EXPERIMENT, GLOBAL_EXPERIMENT)) def test_compute_ddqz_z_full_e( - grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint + grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint, backend ): + if is_roundtrip(backend): + pytest.skip("skipping: slow backend") ddqz_z_full = as_field((CellDim, KDim), 1.0 / metrics_savepoint.inv_ddqz_z_full().asnumpy()) c_lin_e = interpolation_savepoint.c_lin_e() ddqz_z_full_e_ref = metrics_savepoint.ddqz_z_full_e().asnumpy() vertical_start = 0 vertical_end = icon_grid.num_levels ddqz_z_full_e = zero_field(icon_grid, EdgeDim, KDim) - compute_cells2edges( + compute_cells2edges.with_backend(backend)( p_cell_in=ddqz_z_full, c_int=c_lin_e, p_vert_out=ddqz_z_full_e, @@ -502,7 +504,11 @@ def test_compute_ddqz_z_full_e( @pytest.mark.datatest -def test_compute_ddxn_z_full(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): +def test_compute_ddxn_z_full( + grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint, backend +): + if is_roundtrip(backend): + pytest.skip("skipping: slow backend") z_ifc = metrics_savepoint.z_ifc() inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() ddxn_z_full_ref = metrics_savepoint.ddxn_z_full().asnumpy() @@ -517,7 +523,7 @@ def test_compute_ddxn_z_full(grid_savepoint, interpolation_savepoint, icon_grid, vertical_start = 0 vertical_end = icon_grid.num_levels + 1 ddxn_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) - compute_ddxn_z_half_e( + compute_ddxn_z_half_e.with_backend(backend)( z_ifc=z_ifc, inv_dual_edge_length=inv_dual_edge_length, ddxn_z_half_e=ddxn_z_half_e, @@ -528,7 +534,7 @@ def test_compute_ddxn_z_full(grid_savepoint, interpolation_savepoint, icon_grid, offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, ) ddxn_z_full = zero_field(icon_grid, EdgeDim, KDim) - compute_ddxn_z_full( + compute_ddxn_z_full.with_backend(backend)( z_ddxnt_z_half_e=ddxn_z_half_e, ddxn_z_full=ddxn_z_full, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, @@ -538,7 +544,11 @@ def test_compute_ddxn_z_full(grid_savepoint, interpolation_savepoint, icon_grid, @pytest.mark.datatest -def test_compute_ddxt_z_full(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint): +def test_compute_ddxt_z_full( + grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint, backend +): + if is_roundtrip(backend): + pytest.skip("skipping: slow backend") z_ifc = metrics_savepoint.z_ifc() tangent_orientation = grid_savepoint.tangent_orientation() inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() @@ -574,7 +584,7 @@ def test_compute_ddxt_z_full(grid_savepoint, interpolation_savepoint, icon_grid, }, ) ddxt_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) - compute_ddxt_z_half_e( + compute_ddxt_z_half_e.with_backend(backend)( z_ifv=z_ifv, inv_primal_edge_length=inv_primal_edge_length, tangent_orientation=tangent_orientation, @@ -586,7 +596,7 @@ def test_compute_ddxt_z_full(grid_savepoint, interpolation_savepoint, icon_grid, offset_provider={"E2V": icon_grid.get_offset_provider("E2V")}, ) ddxt_z_full = zero_field(icon_grid, EdgeDim, KDim) - compute_ddxn_z_full( + compute_ddxn_z_full.with_backend(backend)( z_ddxnt_z_half_e=ddxt_z_half_e, ddxn_z_full=ddxt_z_full, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, From b68f9f810e3e5329ebea08b206002c046c666642 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 2 May 2024 14:37:28 +0200 Subject: [PATCH 28/65] added zdiff_gradp_dsl computation --- .../model/common/metrics/metric_fields.py | 2 +- .../tests/metric_tests/test_metric_fields.py | 46 ++++++++++++++++++- 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index bfe5b6fdf..7461775b0 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -604,7 +604,7 @@ def _compute_exner_exfac(ddxn_z_full: Field[[EdgeDim, KDim], wpfloat], dual_edge exner_exfac = exner_expol * minimum(1.0 - (4.0 * z_maxslp)**2, 1.0 - (0.002 * z_maxhgtd)**2) exner_exfac = maximum(0.0, exner_exfac) - exner_exfac = where(z_maxslp > 1.5, maximum(-1.0/6.0, 1.0 / 9.0 * (1.5-z_maxslp)), exner_exfac) + exner_exfac = where(z_maxslp > 1.5, maximum(-1.0/6.0, 1.0/9.0 * (1.5 - z_maxslp)), exner_exfac) return exner_exfac diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index dcf355260..2f8b37a8d 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -24,7 +24,7 @@ EdgeDim, KDim, V2CDim, - VertexDim, + VertexDim, E2CDim, ECDim, ) from icon4py.model.common.grid.horizontal import ( HorizontalMarkerIndex, @@ -47,6 +47,7 @@ compute_z_mc, compute_exner_exfac ) +from icon4py.model.common.metrics.stencils.compute_zdiff_gradp_dsl import compute_zdiff_gradp_dsl from icon4py.model.common.test_utils.datatest_utils import ( GLOBAL_EXPERIMENT, REGIONAL_EXPERIMENT, @@ -57,7 +58,7 @@ is_python, is_roundtrip, random_field, - zero_field, + zero_field, flatten_first_two_dims, ) @@ -610,6 +611,7 @@ def test_compute_exner_exfac( grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint, backend ): backend = None + # initialize exner_exfac to field of exner_expol and check horizontal bounds exner_exfac = zero_field(icon_grid, CellDim, KDim) exner_exfac_ref = metrics_savepoint.exner_exfac() compute_exner_exfac.with_backend(backend)( @@ -625,3 +627,43 @@ def test_compute_exner_exfac( ) assert dallclose(exner_exfac.asnumpy(), exner_exfac_ref.asnumpy()) + +@pytest.mark.datatest +def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_savepoint, backend): + backend = None + zdiff_gradp_ref = metrics_savepoint.zdiff_gradp() + zdiff_gradp_full = zero_field(icon_grid, EdgeDim, E2CDim, KDim) + zdiff_gradp_full_np = zdiff_gradp_full.asnumpy() + z_mc = zero_field(icon_grid, CellDim, KDim) + z_ifc = metrics_savepoint.z_ifc() + compute_z_mc.with_backend(backend)( + z_ifc, + z_mc, + horizontal_start=int32(0), + horizontal_end=icon_grid.num_cells, + vertical_start=int32(0), + vertical_end=int32(icon_grid.num_levels), + offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, + ) + c_lin_e = interpolation_savepoint.c_lin_e() + horizontal_start = icon_grid.get_start_index( + EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1 + ) + flat_idx = np.full(shape=icon_grid.num_edges, fill_value=icon_grid.num_levels) + + + zdiff_gradp_full_np_final = compute_zdiff_gradp_dsl( + e2c=icon_grid.connectivities[E2CDim], + c_lin_e=c_lin_e.asnumpy(), + z_mc=z_mc.asnumpy(), + zdiff_gradp=zdiff_gradp_full_np, + z_ifc=metrics_savepoint.z_ifc().asnumpy(), + flat_idx=flat_idx, + nlev=icon_grid.num_levels, # TODO: check that -1 ok + nedges=icon_grid.num_edges, + ) + zdiff_gradp_full_field = flatten_first_two_dims( + ECDim, KDim, field=as_field((EdgeDim, E2CDim, KDim), zdiff_gradp_full_np_final) + ) + # zdiff_gradp_ref.asnumpy()[856:, :] == zdiff_gradp_full_field.asnumpy()[856:, :] + assert dallclose(zdiff_gradp_full_field.asnumpy(), zdiff_gradp_ref.asnumpy()) From 6cda94a686b1057ae11823f48ea9c0a185ab0ab0 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 2 May 2024 14:43:50 +0200 Subject: [PATCH 29/65] added zdiff_gradp_dsl computation --- .../stencils/compute_zdiff_gradp_dsl.py | 102 ++++++++++++++++++ 1 file changed, 102 insertions(+) create mode 100644 model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py new file mode 100644 index 000000000..3e2cfb950 --- /dev/null +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py @@ -0,0 +1,102 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import numpy as np + + +def compute_zdiff_gradp_dsl( + e2c, + c_lin_e: np.array, + z_mc: np.array, + zdiff_gradp: np.array, + z_ifc: np.array, + flat_idx: np.array, + nlev: int, + nedges: int, +) -> np.array: + c_lin_e = np.expand_dims(c_lin_e, axis=-1) + z_me = np.sum(z_mc[e2c] * c_lin_e, axis=1) + zdiff_gradp[:, 0, :] = z_me - z_mc[e2c[:, 0]] + zdiff_gradp[:, 1, :] = z_me - z_mc[e2c[:, 1]] + + for k in range(nlev): + for e in range( + nedges + ): # TODO: boundary here is not nedges, rather p_patch(jg)%edges%start_block(2) + if ( + z_me[e, k] <= z_ifc[e2c[e, 0], k] + and z_me[e, k] >= z_ifc[e2c[e, 0], k + 1] + and z_me[e, k] <= z_ifc[e2c[e, 1], k] + and z_me[e, k] >= z_ifc[e2c[e, 1], k + 1] + ): + flat_idx[e] = k + + for je in range(nedges): + jk_start = flat_idx[je] + for jk in range(flat_idx[je] + 1, nlev): + for jk1 in range(jk_start, nlev): + if ( + jk1 == nlev + or z_me[je, jk] <= z_ifc[e2c[je, 0], jk1] + and z_me[je, jk] >= z_ifc[e2c[je, 0], jk1 + 1] + ): + zdiff_gradp[je, 0, jk1] = z_me[je, jk] - z_mc[e2c[je, 0], jk1] + jk_start = jk1 + break + + extrapol_dist = 5.0 + z_aux1 = np.maximum(z_ifc[e2c[:, 0], nlev], z_ifc[e2c[:, 1], nlev]) + z_aux2 = z_aux1 - extrapol_dist + + for je in range(nedges): + jk_start = flat_idx[je] + for jk in range(flat_idx[je] + 1, nlev): + for jk1 in range(jk_start, nlev): + if ( + jk1 == nlev + or z_me[je, jk] <= z_ifc[e2c[je, 1], jk1] + and z_me[je, jk] >= z_ifc[e2c[je, 1], jk1 + 1] + ): + zdiff_gradp[je, 1, jk1] = z_me[je, jk] - z_mc[e2c[je, 1], jk1] + jk_start = jk1 + break + + for je in range(nedges): + jk_start = flat_idx[je] + for jk in range(flat_idx[je] + 1, nlev): + if z_me[je, jk] < z_aux2[je]: + for jk1 in range(jk_start, nlev): + if ( + jk1 == nlev + or z_aux2[je] <= z_ifc[e2c[je, 0], jk1] + and z_aux2[je] >= z_ifc[e2c[je, 0], jk1 + 1] + ): + zdiff_gradp[je, 0, jk1] = z_aux2[je] - z_mc[e2c[je, 0], jk1] + jk_start = jk1 + break + + for je in range(nedges): + jk_start = flat_idx[je] + for jk in range(flat_idx[je] + 1, nlev): + if z_me[je, jk] < z_aux2[je]: + for jk1 in range(jk_start, nlev): + if ( + jk1 == nlev + or z_aux2[je] <= z_ifc[e2c[je, 1], jk1] + and z_aux2[je] >= z_ifc[e2c[je, 1], jk1 + 1] + ): + zdiff_gradp[je, 1, jk1] = z_aux2[je] - z_mc[e2c[je, 1], jk1] + jk_start = jk1 + break + + return zdiff_gradp + From 44fd84f2556bb720d7c0c60a11c945353367814e Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 2 May 2024 16:23:40 +0200 Subject: [PATCH 30/65] added coeff_gradekin computation --- .../src/icon4py/model/common/constants.py | 1 + .../model/common/metrics/metric_fields.py | 34 ++++++++++++-- .../stencils/compute_coeff_gradekin.py | 45 ++++++++++++++++++ .../stencils/compute_zdiff_gradp_dsl.py | 1 - .../test_compute_coeff_gradekin.py | 47 +++++++++++++++++++ .../tests/metric_tests/test_metric_fields.py | 23 +++++---- 6 files changed, 137 insertions(+), 14 deletions(-) create mode 100644 model/common/src/icon4py/model/common/metrics/stencils/compute_coeff_gradekin.py create mode 100644 model/common/tests/metric_tests/test_compute_coeff_gradekin.py diff --git a/model/common/src/icon4py/model/common/constants.py b/model/common/src/icon4py/model/common/constants.py index bccf6da30..7694ba496 100644 --- a/model/common/src/icon4py/model/common/constants.py +++ b/model/common/src/icon4py/model/common/constants.py @@ -77,5 +77,6 @@ class RayleighType(IntEnum): ] = 1 # classical Rayleigh damping, which makes use of a reference state. RAYLEIGH_KLEMP: Final[int] = 2 # Klemp (2008) type Rayleigh damping + # Temporal extrapolation of Exner for computation of horizontal pressure gradient exner_expol: Final[wpfloat] = 0.333 diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 7461775b0..e0e28f84d 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -593,18 +593,26 @@ def compute_vwind_expl_wgt( @field_operator -def _compute_exner_exfac(ddxn_z_full: Field[[EdgeDim, KDim], wpfloat], dual_edge_length: Field[[EdgeDim], wpfloat], exner_expol: wpfloat) -> Field[[CellDim, KDim], wpfloat]: +def _compute_exner_exfac( + ddxn_z_full: Field[[EdgeDim, KDim], wpfloat], + dual_edge_length: Field[[EdgeDim], wpfloat], + exner_expol: wpfloat, +) -> Field[[CellDim, KDim], wpfloat]: z_maxslp_0_1 = maximum(abs(ddxn_z_full(C2E[0])), abs(ddxn_z_full(C2E[1]))) z_maxslp = maximum(z_maxslp_0_1, abs(ddxn_z_full(C2E[2]))) - z_maxhgtd_0_1 = maximum(abs(ddxn_z_full(C2E[0]) * dual_edge_length(C2E[0])), - abs(ddxn_z_full(C2E[1]) * dual_edge_length(C2E[1]))) + z_maxhgtd_0_1 = maximum( + abs(ddxn_z_full(C2E[0]) * dual_edge_length(C2E[0])), + abs(ddxn_z_full(C2E[1]) * dual_edge_length(C2E[1])), + ) z_maxhgtd = maximum(z_maxhgtd_0_1, abs(ddxn_z_full(C2E[2]) * dual_edge_length(C2E[2]))) - exner_exfac = exner_expol * minimum(1.0 - (4.0 * z_maxslp)**2, 1.0 - (0.002 * z_maxhgtd)**2) + exner_exfac = exner_expol * minimum(1.0 - (4.0 * z_maxslp) ** 2, 1.0 - (0.002 * z_maxhgtd) ** 2) exner_exfac = maximum(0.0, exner_exfac) - exner_exfac = where(z_maxslp > 1.5, maximum(-1.0/6.0, 1.0/9.0 * (1.5 - z_maxslp)), exner_exfac) + exner_exfac = where( + z_maxslp > 1.5, maximum(-1.0 / 6.0, 1.0 / 9.0 * (1.5 - z_maxslp)), exner_exfac + ) return exner_exfac @@ -620,6 +628,22 @@ def compute_exner_exfac( vertical_start: int32, vertical_end: int32, ): + """ + Compute exner_exfac. + + Exner extrapolation reaches zero for a slope of 1/4 or a height difference of 500 m between adjacent grid points (empirically determined values). See mo_vertical_grid.f90 + + Args: + ddxn_z_full: ddxn_z_full + dual_edge_length: dual_edge_length + exner_exfac: Exner factor + exner_expol: Exner extrapolation factor + horizontal_start: horizontal start index + horizontal_end: horizontal end index + vertical_start: vertical start index + vertical_end: vertical end index + + """ _compute_exner_exfac( ddxn_z_full=ddxn_z_full, dual_edge_length=dual_edge_length, diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_coeff_gradekin.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_coeff_gradekin.py new file mode 100644 index 000000000..7568638e0 --- /dev/null +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_coeff_gradekin.py @@ -0,0 +1,45 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.common.dimension import ECDim +from icon4py.model.common.test_utils.helpers import numpy_to_1D_sparse_field + + +def compute_coeff_gradekin( + edge_cell_length: np.array, + inv_dual_edge_length: np.array, + horizontal_start: float, + horizontal_end: float, +) -> np.array: + """ + Compute coefficients for improved calculation of kinetic energy gradient + + Args: + edge_cell_length: edge_cell_length + inv_dual_edge_length: inverse of dual_edge_length + horizontal_start: horizontal start index + horizontal_end: horizontal end index + """ + coeff_gradekin_0 = np.zeros_like(inv_dual_edge_length) + coeff_gradekin_1 = np.zeros_like(inv_dual_edge_length) + for e in range(horizontal_start, horizontal_end): + coeff_gradekin_0[e] = ( + edge_cell_length[e, 1] / edge_cell_length[e, 0] * inv_dual_edge_length[e] + ) + coeff_gradekin_1[e] = ( + edge_cell_length[e, 0] / edge_cell_length[e, 1] * inv_dual_edge_length[e] + ) + coeff_gradekin_full = np.column_stack((coeff_gradekin_0, coeff_gradekin_1)) + return numpy_to_1D_sparse_field(coeff_gradekin_full, ECDim) diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py index 3e2cfb950..974d1eef5 100644 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py @@ -99,4 +99,3 @@ def compute_zdiff_gradp_dsl( break return zdiff_gradp - diff --git a/model/common/tests/metric_tests/test_compute_coeff_gradekin.py b/model/common/tests/metric_tests/test_compute_coeff_gradekin.py new file mode 100644 index 000000000..a89d86c72 --- /dev/null +++ b/model/common/tests/metric_tests/test_compute_coeff_gradekin.py @@ -0,0 +1,47 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import pytest + +from icon4py.model.common.dimension import EdgeDim +from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex +from icon4py.model.common.metrics.stencils.compute_coeff_gradekin import compute_coeff_gradekin +from icon4py.model.common.test_utils.helpers import dallclose + + +@pytest.mark.datatest +def test_compute_coeff_gradekin(icon_grid, grid_savepoint, metrics_savepoint): + edge_cell_length = grid_savepoint.edge_cell_length().asnumpy() + inv_dual_edge_length = grid_savepoint.inv_dual_edge_length().asnumpy() + coeff_gradekin_ref = metrics_savepoint.coeff_gradekin() + horizontal_start = icon_grid.get_start_index( + EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1 + ) + horizontal_end = icon_grid.num_edges + + coeff_gradekin_full = compute_coeff_gradekin( + edge_cell_length, inv_dual_edge_length, horizontal_start, horizontal_end + ) + assert dallclose(coeff_gradekin_ref.asnumpy(), coeff_gradekin_full.asnumpy()) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 2f8b37a8d..d8bf618d9 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -21,10 +21,12 @@ from icon4py.model.common import constants from icon4py.model.common.dimension import ( CellDim, + E2CDim, + ECDim, EdgeDim, KDim, V2CDim, - VertexDim, E2CDim, ECDim, + VertexDim, ) from icon4py.model.common.grid.horizontal import ( HorizontalMarkerIndex, @@ -40,12 +42,12 @@ compute_ddxn_z_full, compute_ddxn_z_half_e, compute_ddxt_z_half_e, + compute_exner_exfac, compute_rayleigh_w, compute_scalfac_dd3d, compute_vwind_expl_wgt, compute_vwind_impl_wgt, compute_z_mc, - compute_exner_exfac ) from icon4py.model.common.metrics.stencils.compute_zdiff_gradp_dsl import compute_zdiff_gradp_dsl from icon4py.model.common.test_utils.datatest_utils import ( @@ -54,11 +56,13 @@ ) from icon4py.model.common.test_utils.helpers import ( StencilTest, + constant_field, dallclose, + flatten_first_two_dims, is_python, is_roundtrip, random_field, - zero_field, flatten_first_two_dims, + zero_field, ) @@ -606,27 +610,31 @@ def test_compute_ddxt_z_full( assert np.allclose(ddxt_z_full.asnumpy(), ddxt_z_full_ref) + @pytest.mark.datatest def test_compute_exner_exfac( grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint, backend ): backend = None - # initialize exner_exfac to field of exner_expol and check horizontal bounds - exner_exfac = zero_field(icon_grid, CellDim, KDim) + horizontal_start = icon_grid.get_start_index( + CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1 + ) + exner_exfac = constant_field(icon_grid, constants.exner_expol, CellDim, KDim) exner_exfac_ref = metrics_savepoint.exner_exfac() compute_exner_exfac.with_backend(backend)( ddxn_z_full=metrics_savepoint.ddxn_z_full(), dual_edge_length=grid_savepoint.dual_edge_length(), exner_exfac=exner_exfac, exner_expol=constants.exner_expol, - horizontal_start=int32(0), + horizontal_start=horizontal_start, horizontal_end=icon_grid.num_cells, vertical_start=int32(0), vertical_end=icon_grid.num_levels, offset_provider={"C2E": icon_grid.get_offset_provider("C2E")}, ) - assert dallclose(exner_exfac.asnumpy(), exner_exfac_ref.asnumpy()) + assert dallclose(exner_exfac.asnumpy(), exner_exfac_ref.asnumpy(), rtol=1.0e-10) + @pytest.mark.datatest def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_savepoint, backend): @@ -651,7 +659,6 @@ def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_sav ) flat_idx = np.full(shape=icon_grid.num_edges, fill_value=icon_grid.num_levels) - zdiff_gradp_full_np_final = compute_zdiff_gradp_dsl( e2c=icon_grid.connectivities[E2CDim], c_lin_e=c_lin_e.asnumpy(), From 0b1dd787d7426d1d1860cff2dcfe28a2188ca1fd Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 2 May 2024 18:02:11 +0200 Subject: [PATCH 31/65] corrections --- .../src/icon4py/model/common/math/helpers.py | 28 ++------------- .../model/common/metrics/metric_fields.py | 7 ++-- .../tests/metric_tests/test_metric_fields.py | 34 ++----------------- 3 files changed, 9 insertions(+), 60 deletions(-) diff --git a/model/common/src/icon4py/model/common/math/helpers.py b/model/common/src/icon4py/model/common/math/helpers.py index 8e5fda376..808db44ba 100644 --- a/model/common/src/icon4py/model/common/math/helpers.py +++ b/model/common/src/icon4py/model/common/math/helpers.py @@ -11,11 +11,10 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next import Field, field_operator, program -from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next import Field, field_operator from icon4py.model.common.dimension import E2C, E2V, CellDim, EdgeDim, KDim, Koff, VertexDim -from icon4py.model.common.type_alias import vpfloat, wpfloat +from icon4py.model.common.type_alias import wpfloat @field_operator @@ -118,26 +117,3 @@ def _grad_fd_tang( ) -> Field[[EdgeDim, KDim], float]: grad_tang_psi_e = tangent_orientation * (psi_v(E2V[1]) - psi_v(E2V[0])) * inv_primal_edge_length return grad_tang_psi_e - - -@field_operator -def _compute_inverse_edge_kdim( - edge_k_field: Field[[EdgeDim, KDim], vpfloat], -) -> Field[[EdgeDim, KDim], vpfloat]: - return 1.0 / edge_k_field - - -@program -def compute_inverse_edge_kdim( - edge_k_field: Field[[EdgeDim, KDim], vpfloat], - inv_edge_k_field: Field[[EdgeDim, KDim], vpfloat], - horizontal_start: int32, - horizontal_end: int32, - vertical_start: int32, - vertical_end: int32, -): - _compute_inverse_edge_kdim( - edge_k_field, - out=inv_edge_k_field, - domain={EdgeDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, - ) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index e2a3c734c..da7439c20 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -88,6 +88,7 @@ def _compute_ddqz_z_half( k: Field[[KDim], int32], nlev: int32, ): + # TODO: change this to concat_where once it's merged ddqz_z_half = where(k == 0, 2.0 * (z_ifc - z_mc), 0.0) ddqz_z_half = where((k > 0) & (k < nlev), z_mc(Koff[-1]) - z_mc, ddqz_z_half) ddqz_z_half = where(k == nlev, 2.0 * (z_mc(Koff[-1]) - z_ifc), ddqz_z_half) @@ -133,7 +134,7 @@ def compute_ddqz_z_half( @field_operator -def _compute_ddqz_z_full( +def _compute_ddqz_z_full_and_inverse( z_ifc: Field[[CellDim, KDim], wpfloat], ) -> tuple[Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat]]: ddqz_z_full = difference_k_level_up(z_ifc) @@ -142,7 +143,7 @@ def _compute_ddqz_z_full( @program(grid_type=GridType.UNSTRUCTURED) -def compute_ddqz_z_full( +def compute_ddqz_z_full_and_inverse( z_ifc: Field[[CellDim, KDim], wpfloat], ddqz_z_full: Field[[CellDim, KDim], wpfloat], inv_ddqz_z_full: Field[[CellDim, KDim], wpfloat], @@ -167,7 +168,7 @@ def compute_ddqz_z_full( vertical_end: vertical end index """ - _compute_ddqz_z_full( + _compute_ddqz_z_full_and_inverse( z_ifc, out=(ddqz_z_full, inv_ddqz_z_full), domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 9a14ace71..ce7d67014 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -31,11 +31,10 @@ _compute_cells2verts, compute_cells2edges, ) -from icon4py.model.common.math.helpers import compute_inverse_edge_kdim from icon4py.model.common.metrics.metric_fields import ( compute_coeff_dwdz, compute_d2dexdz2_fac_mc, - compute_ddqz_z_full, + compute_ddqz_z_full_and_inverse, compute_ddqz_z_half, compute_ddxn_z_full, compute_ddxn_z_half_e, @@ -130,7 +129,7 @@ def test_compute_ddq_z_half(icon_grid, metrics_savepoint, backend): @pytest.mark.datatest -def test_compute_ddqz_z_full(icon_grid, metrics_savepoint, backend): +def test_compute_ddqz_z_full_and_inverse(icon_grid, metrics_savepoint, backend): if is_roundtrip(backend): pytest.skip("skipping: slow backend") z_ifc = metrics_savepoint.z_ifc() @@ -138,7 +137,7 @@ def test_compute_ddqz_z_full(icon_grid, metrics_savepoint, backend): ddqz_z_full = zero_field(icon_grid, CellDim, KDim) inv_ddqz_z_full = zero_field(icon_grid, CellDim, KDim) - compute_ddqz_z_full.with_backend(backend)( + compute_ddqz_z_full_and_inverse.with_backend(backend)( z_ifc=z_ifc, ddqz_z_full=ddqz_z_full, inv_ddqz_z_full=inv_ddqz_z_full, @@ -450,33 +449,6 @@ def test_compute_vwind_expl_wgt(icon_grid, metrics_savepoint, backend): assert dallclose(vwind_expl_wgt_full.asnumpy(), vwind_expl_wgt_ref.asnumpy()) -@pytest.mark.skip("TODO needs to be fulfilled") -@pytest.mark.datatest -def test_compute_inv_ddqz_z_full(icon_grid, metrics_savepoint, backend): - # TODO: serialization missing inv_ddqz_z_full is over cells, need over edge --> inv_ddqz_z_full_e - if is_roundtrip(backend): - pytest.skip("skipping: slow backend") - - inv_ddqz_z_full = zero_field(icon_grid, EdgeDim, KDim) - inv_ddqz_z_ref = metrics_savepoint.inv_ddqz_z_full() - horizontal_start_edge = icon_grid.get_start_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, - ) - - compute_inverse_edge_kdim.with_backend(backend)( - edge_k_field=metrics_savepoint.ddqz_z_full_e(), - inv_edge_k_field=inv_ddqz_z_full, - horizontal_start=horizontal_start_edge, - horizontal_end=icon_grid.num_edges, - vertical_start=int32(0), - vertical_end=int32(icon_grid.num_levels), - offset_provider={}, - ) - - assert dallclose(inv_ddqz_z_full.asnumpy(), inv_ddqz_z_ref.asnumpy()) - - @pytest.mark.datatest @pytest.mark.parametrize("experiment", (REGIONAL_EXPERIMENT, GLOBAL_EXPERIMENT)) def test_compute_ddqz_z_full_e( From f3bc5b466b801f94c11774d6815d7ba05093131e Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 3 May 2024 11:06:10 +0200 Subject: [PATCH 32/65] vwind_impl_wgt test --- .../tests/metric_tests/test_metric_fields.py | 113 ++++++++++++++++++ 1 file changed, 113 insertions(+) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index d8bf618d9..0db935783 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -674,3 +674,116 @@ def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_sav ) # zdiff_gradp_ref.asnumpy()[856:, :] == zdiff_gradp_full_field.asnumpy()[856:, :] assert dallclose(zdiff_gradp_full_field.asnumpy(), zdiff_gradp_ref.asnumpy()) + +@pytest.mark.datatest +def test_compute_vwind_impl_wgt( + icon_grid, grid_savepoint, metrics_savepoint, interpolation_savepoint, backend +): + backend = None + if is_roundtrip(backend): + pytest.skip("skipping: slow backend") + z_ifc = metrics_savepoint.z_ifc() + inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() + z_ddxn_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) + horizontal_start = icon_grid.get_start_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, + ) + horizontal_end = icon_grid.get_end_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, + ) + vertical_start = 0 + vertical_end = icon_grid.num_levels + 1 + + horizontal_start_edge = icon_grid.get_start_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, + ) + horizontal_end_edge = icon_grid.get_end_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, + ) + + compute_ddxn_z_half_e( + z_ifc=z_ifc, + inv_dual_edge_length=inv_dual_edge_length, + ddxn_z_half_e=z_ddxn_z_half_e, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=vertical_start, + vertical_end=vertical_end, + offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, + ) + + z_ddxt_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) + z_ifv = zero_field(icon_grid, VertexDim, KDim, extend={KDim: 1}) + + tangent_orientation = grid_savepoint.tangent_orientation() + inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() + horizontal_start_vertex = icon_grid.get_start_index( + VertexDim, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, + ) + horizontal_end_vertex = icon_grid.get_end_index( + VertexDim, + HorizontalMarkerIndex.lateral_boundary(VertexDim) - 1, + ) + cells_aw_verts = interpolation_savepoint.c_intp().asnumpy() + _compute_cells2verts( + z_ifc, + as_field((VertexDim, V2CDim), cells_aw_verts), + out=z_ifv, + offset_provider={"V2C": icon_grid.get_offset_provider("V2C")}, + domain={ + VertexDim: (horizontal_start_vertex, horizontal_end_vertex), + KDim: (vertical_start, vertical_end), + }, + ) + compute_ddxt_z_half_e( + z_ifv=z_ifv, + inv_primal_edge_length=inv_primal_edge_length, + tangent_orientation=tangent_orientation, + ddxt_z_half_e=z_ddxt_z_half_e, + horizontal_start=horizontal_start_edge, + horizontal_end=horizontal_end_edge, + vertical_start=vertical_start, + vertical_end=vertical_end, + offset_provider={"E2V": icon_grid.get_offset_provider("E2V")}, + ) + + horizontal_start_cell = icon_grid.get_start_index( + CellDim, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, + ) + #vwind_impl_wgt_full = zero_field(icon_grid, CellDim) + vwind_impl_wgt_ref = metrics_savepoint.vwind_impl_wgt() + dual_edge_length = grid_savepoint.dual_edge_length() + vwind_offctr = 0.2 + vertical_start = max(10, icon_grid.num_levels - 8) + vwind_impl_wgt_full = constant_field(icon_grid, 0.5 + vwind_offctr, (CellDim)) + + compute_vwind_impl_wgt.with_backend(backend)( + z_ddxn_z_half_e=as_field((EdgeDim,), z_ddxn_z_half_e.asnumpy()[:, icon_grid.num_levels]), + z_ddxt_z_half_e=as_field((EdgeDim,), z_ddxt_z_half_e.asnumpy()[:, icon_grid.num_levels]), + dual_edge_length=dual_edge_length, + #vct_a=grid_savepoint.vct_a(), + #z_ifc=metrics_savepoint.z_ifc(), + vwind_impl_wgt=vwind_impl_wgt_full, + vwind_offctr=vwind_offctr, + horizontal_start=horizontal_start_cell, + horizontal_end=icon_grid.num_cells, + # vertical_start=vertical_start, + # vertical_end=icon_grid.num_levels, + offset_provider={"C2E": icon_grid.get_offset_provider("C2E"),}, + ) + # vwind_impl_wgt = vwind_impl_wgt_full.asnumpy() + # vct_a = grid_savepoint.vct_a().asnumpy() + # z_ifc = metrics_savepoint.z_ifc().asnumpy() + # for jk in range(vertical_start, icon_grid.num_levels): + # for jc in range(horizontal_start_cell, icon_grid.num_cells): + # z_diff_2 = (z_ifc[jc, jk] - z_ifc[jc, jk+1]) / (vct_a[jk] - vct_a[jk+1]) + # if z_diff_2 < 0.6: + # vwind_impl_wgt[jc] = np.maximum(vwind_impl_wgt[jc], 1.2 - z_diff_2) + + assert dallclose(vwind_impl_wgt_ref.asnumpy(), vwind_impl_wgt_full.asnumpy(), rtol=0.000000000000001) From c4211b9bb9aac91ec63b46da48ee5b9682582b1a Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 3 May 2024 11:07:22 +0200 Subject: [PATCH 33/65] vwind_impl_wgt prog --- .../model/common/metrics/metric_fields.py | 63 +++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index e0e28f84d..8a52dac18 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -654,3 +654,66 @@ def compute_exner_exfac( KDim: (vertical_start, vertical_end), }, ) + +@field_operator +def _compute_vwind_impl_wgt( + z_ddxn_z_half_e: Field[[EdgeDim], wpfloat], + z_ddxt_z_half_e: Field[[EdgeDim], wpfloat], + dual_edge_length: Field[[EdgeDim], wpfloat], + vwind_offctr: wpfloat, +) -> Field[[CellDim], wpfloat]: + z_ddx_1 = maximum(abs(z_ddxn_z_half_e(C2E[0])), abs(z_ddxt_z_half_e(C2E[0]))) + z_ddx_2 = maximum(abs(z_ddxn_z_half_e(C2E[1])), abs(z_ddxt_z_half_e(C2E[1]))) + z_ddx_3 = maximum(abs(z_ddxn_z_half_e(C2E[2])), abs(z_ddxt_z_half_e(C2E[2]))) + z_ddx_1_2 = maximum(z_ddx_1, z_ddx_2) + z_maxslope = maximum(z_ddx_1_2, z_ddx_3) + + z_diff_1_2 = maximum( + abs(z_ddxn_z_half_e(C2E[0]) * dual_edge_length(C2E[0])), + abs(z_ddxn_z_half_e(C2E[1]) * dual_edge_length(C2E[1])), + ) + z_diff = maximum(z_diff_1_2, abs(z_ddxn_z_half_e(C2E[2]) * dual_edge_length(C2E[2]))) + z_offctr_1 = maximum(vwind_offctr, 0.425 * z_maxslope ** (0.75)) + z_offctr_2 = minimum(0.25, 0.00025 * (z_diff - 250.0)) + z_offctr = maximum(z_offctr_1, z_offctr_2) + z_offctr = minimum(maximum(vwind_offctr, 0.75), z_offctr) + vwind_impl_wgt = 0.5 + z_offctr + return vwind_impl_wgt + + +@program(grid_type=GridType.UNSTRUCTURED) +def compute_vwind_impl_wgt( + z_ddxn_z_half_e: Field[[EdgeDim], wpfloat], + z_ddxt_z_half_e: Field[[EdgeDim], wpfloat], + dual_edge_length: Field[[EdgeDim], wpfloat], + # vct_a: Field[[KDim], wpfloat], + # z_ifc: Field[[CellDim, KDim], wpfloat], + vwind_impl_wgt: Field[[CellDim], wpfloat], + vwind_offctr: wpfloat, + horizontal_start: int32, + horizontal_end: int32, + # vertical_start: int32, + # vertical_end: int32 +): + """ + Compute vwind_impl_wgt. + + See mo_vertical_grid.f90 + + Args: + z_ddxn_z_half_e: intermediate storage for field + z_ddxt_z_half_e: intermediate storage for field + dual_edge_length: dual_edge_length + vwind_impl_wgt: (output) offcentering in vertical mass flux + vwind_offctr: off-centering in vertical wind solver + horizontal_start: horizontal start index + horizontal_end: horizontal end index + """ + _compute_vwind_impl_wgt( + z_ddxn_z_half_e=z_ddxn_z_half_e, + z_ddxt_z_half_e=z_ddxt_z_half_e, + dual_edge_length=dual_edge_length, + vwind_offctr=vwind_offctr, + out=vwind_impl_wgt, + domain={CellDim: (horizontal_start, horizontal_end)}, + ) From c96beb8df07687d45f31806f5c0651ccd07f8792 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 3 May 2024 11:08:25 +0200 Subject: [PATCH 34/65] removed vwind_impl_wgt --- .../model/common/metrics/metric_fields.py | 63 ------------ .../tests/metric_tests/test_metric_fields.py | 96 ------------------- 2 files changed, 159 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index da7439c20..726510685 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -14,13 +14,11 @@ from gt4py.next import ( Field, GridType, - abs, broadcast, exp, field_operator, int32, maximum, - minimum, program, sin, tanh, @@ -28,7 +26,6 @@ ) from icon4py.model.common.dimension import ( - C2E, CellDim, EdgeDim, KDim, @@ -501,66 +498,6 @@ def compute_ddxn_z_full( average_edge_kdim_level_up(z_ddxnt_z_half_e, out=ddxn_z_full) -@field_operator -def _compute_vwind_impl_wgt( - z_ddxn_z_half_e: Field[[EdgeDim], wpfloat], - z_ddxt_z_half_e: Field[[EdgeDim], wpfloat], - dual_edge_length: Field[[EdgeDim], wpfloat], - vwind_offctr: wpfloat, -) -> Field[[CellDim], wpfloat]: - z_ddx_1 = maximum(abs(z_ddxn_z_half_e(C2E[0])), abs(z_ddxt_z_half_e(C2E[0]))) - z_ddx_2 = maximum(abs(z_ddxn_z_half_e(C2E[1])), abs(z_ddxt_z_half_e(C2E[1]))) - z_ddx_3 = maximum(abs(z_ddxn_z_half_e(C2E[2])), abs(z_ddxt_z_half_e(C2E[2]))) - z_ddx_1_2 = maximum(z_ddx_1, z_ddx_2) - z_maxslope = maximum(z_ddx_1_2, z_ddx_3) - - z_diff_1_2 = maximum( - abs(z_ddxn_z_half_e(C2E[0]) * dual_edge_length(C2E[0])), - abs(z_ddxn_z_half_e(C2E[1]) * dual_edge_length(C2E[1])), - ) - z_diff = maximum(z_diff_1_2, abs(z_ddxn_z_half_e(C2E[2]) * dual_edge_length(C2E[2]))) - z_offctr_1 = maximum(vwind_offctr, 0.425 * z_maxslope ** (0.75)) - z_offctr = maximum(z_offctr_1, minimum(0.25, 0.00025 * (z_diff - 250.0))) - z_offctr = minimum(maximum(vwind_offctr, 0.75), z_offctr) - vwind_impl_wgt = 0.5 + z_offctr - return vwind_impl_wgt - - -@program(grid_type=GridType.UNSTRUCTURED) -def compute_vwind_impl_wgt( - z_ddxn_z_half_e: Field[[EdgeDim], wpfloat], - z_ddxt_z_half_e: Field[[EdgeDim], wpfloat], - dual_edge_length: Field[[EdgeDim], wpfloat], - vwind_impl_wgt: Field[[CellDim], wpfloat], - vwind_offctr: wpfloat, - horizontal_start: int32, - horizontal_end: int32, -): - """ - Compute vwind_impl_wgt. - - See mo_vertical_grid.f90 - - Args: - z_ddxn_z_half_e: intermediate storage for field - z_ddxt_z_half_e: intermediate storage for field - dual_edge_length: dual_edge_length - vwind_impl_wgt: (output) offcentering in vertical mass flux - vwind_offctr: off-centering in vertical wind solver - horizontal_start: horizontal start index - horizontal_end: horizontal end index - - """ - _compute_vwind_impl_wgt( - z_ddxn_z_half_e=z_ddxn_z_half_e, - z_ddxt_z_half_e=z_ddxt_z_half_e, - dual_edge_length=dual_edge_length, - vwind_offctr=vwind_offctr, - out=vwind_impl_wgt, - domain={CellDim: (horizontal_start, horizontal_end)}, - ) - - @field_operator def _compute_vwind_expl_wgt(vwind_impl_wgt: Field[[CellDim], wpfloat]) -> Field[[CellDim], wpfloat]: return 1.0 - vwind_impl_wgt diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index ce7d67014..739fdbc7a 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -42,7 +42,6 @@ compute_rayleigh_w, compute_scalfac_dd3d, compute_vwind_expl_wgt, - compute_vwind_impl_wgt, compute_z_mc, ) from icon4py.model.common.test_utils.datatest_utils import ( @@ -337,101 +336,6 @@ def test_compute_ddxt_z_full_e( assert np.allclose(ddxt_z_full.asnumpy(), ddxt_z_full_ref) -@pytest.mark.datatest -def test_compute_vwind_impl_wgt( - icon_grid, grid_savepoint, metrics_savepoint, interpolation_savepoint, backend -): - if is_roundtrip(backend): - pytest.skip("skipping: slow backend") - z_ifc = metrics_savepoint.z_ifc() - inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() - z_ddxn_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) - horizontal_start = icon_grid.get_start_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, - ) - horizontal_end = icon_grid.get_end_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, - ) - vertical_start = 0 - vertical_end = icon_grid.num_levels + 1 - - horizontal_start_edge = icon_grid.get_start_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, - ) - horizontal_end_edge = icon_grid.get_end_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, - ) - - compute_ddxn_z_half_e( - z_ifc=z_ifc, - inv_dual_edge_length=inv_dual_edge_length, - ddxn_z_half_e=z_ddxn_z_half_e, - horizontal_start=horizontal_start, - horizontal_end=horizontal_end, - vertical_start=vertical_start, - vertical_end=vertical_end, - offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, - ) - - z_ddxt_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) - z_ifv = zero_field(icon_grid, VertexDim, KDim, extend={KDim: 1}) - - tangent_orientation = grid_savepoint.tangent_orientation() - inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() - horizontal_start_vertex = icon_grid.get_start_index( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, - ) - horizontal_end_vertex = icon_grid.get_end_index( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) - 1, - ) - cells_aw_verts = interpolation_savepoint.c_intp().asnumpy() - _compute_cells2verts( - z_ifc, - as_field((VertexDim, V2CDim), cells_aw_verts), - out=z_ifv, - offset_provider={"V2C": icon_grid.get_offset_provider("V2C")}, - domain={ - VertexDim: (horizontal_start_vertex, horizontal_end_vertex), - KDim: (vertical_start, vertical_end), - }, - ) - compute_ddxt_z_half_e( - z_ifv=z_ifv, - inv_primal_edge_length=inv_primal_edge_length, - tangent_orientation=tangent_orientation, - ddxt_z_half_e=z_ddxt_z_half_e, - horizontal_start=horizontal_start_edge, - horizontal_end=horizontal_end_edge, - vertical_start=vertical_start, - vertical_end=vertical_end, - offset_provider={"E2V": icon_grid.get_offset_provider("E2V")}, - ) - - vwind_impl_wgt_full = zero_field(icon_grid, CellDim) - vwind_impl_wgt_ref = metrics_savepoint.vwind_impl_wgt() - dual_edge_length = grid_savepoint.dual_edge_length() - vwind_offctr = 0.2 - - compute_vwind_impl_wgt.with_backend(backend)( - z_ddxn_z_half_e=as_field((EdgeDim,), z_ddxn_z_half_e.asnumpy()[:, icon_grid.num_levels]), - z_ddxt_z_half_e=as_field((EdgeDim,), z_ddxt_z_half_e.asnumpy()[:, icon_grid.num_levels]), - dual_edge_length=dual_edge_length, - vwind_impl_wgt=vwind_impl_wgt_full, - vwind_offctr=vwind_offctr, - horizontal_start=int32(0), - horizontal_end=icon_grid.num_cells, - offset_provider={"C2E": icon_grid.get_offset_provider("C2E")}, - ) - - assert dallclose(vwind_impl_wgt_full.asnumpy(), vwind_impl_wgt_ref.asnumpy(), rtol=1.0) - - @pytest.mark.datatest def test_compute_vwind_expl_wgt(icon_grid, metrics_savepoint, backend): vwind_expl_wgt_full = zero_field(icon_grid, CellDim) From bea8158bd8f01d86270629580b1407bfb870ee37 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 6 May 2024 16:26:35 +0200 Subject: [PATCH 35/65] wgtfacq_e_dsl and other edits to wgtfacq_c --- ..._wgtfacq_c.py => compute_wgtfacq_c_dsl.py} | 30 ++++++--- .../metrics/stencils/compute_wgtfacq_e_dsl.py | 61 +++++++++++++++++++ .../metric_tests/test_compute_wgtfacq.py | 59 ++++++++++++++++++ .../metric_tests/test_compute_wgtfacq_c.py | 43 ------------- 4 files changed, 140 insertions(+), 53 deletions(-) rename model/common/src/icon4py/model/common/metrics/stencils/{compute_wgtfacq_c.py => compute_wgtfacq_c_dsl.py} (50%) create mode 100644 model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_e_dsl.py create mode 100644 model/common/tests/metric_tests/test_compute_wgtfacq.py delete mode 100644 model/common/tests/metric_tests/test_compute_wgtfacq_c.py diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c_dsl.py similarity index 50% rename from model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c.py rename to model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c_dsl.py index 11ea37e84..f78f50f18 100644 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c.py +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c_dsl.py @@ -14,7 +14,14 @@ import numpy as np -def compute_wgtfacq_c( +def compute_z1_z2_z3(z_ifc, i1, i2, i3, i4): + z1 = 0.5 * (z_ifc[:, i2] - z_ifc[:, i1]) + z2 = 0.5 * (z_ifc[:, i2] + z_ifc[:, i3]) - z_ifc[:, i1] + z3 = 0.5 * (z_ifc[:, i3] + z_ifc[:, i4]) - z_ifc[:, i1] + return z1, z2, z3 + + +def compute_wgtfacq_c_dsl( z_ifc: np.array, nlev: int, ) -> np.array: @@ -23,17 +30,20 @@ def compute_wgtfacq_c( Args: z_ifc: Field[CellDim, KDim] (half levels), geometric height at the vertical interface of cells. - nlev: int, last level + nlev: int, last k level Returns: Field[CellDim, KDim] (full levels) """ - wgtfacq_c = np.zeros((z_ifc.shape[0], nlev)) - z1 = 0.5 * (z_ifc[:, nlev - 1] - z_ifc[:, nlev]) - z2 = 0.5 * (z_ifc[:, nlev - 1] + z_ifc[:, nlev - 2]) - z_ifc[:, nlev] - z3 = 0.5 * (z_ifc[:, nlev - 2] + z_ifc[:, nlev - 3]) - z_ifc[:, nlev] + wgtfacq_c = np.zeros((z_ifc.shape[0], nlev + 1)) + wgtfacq_c_dsl = np.zeros((z_ifc.shape[0], nlev)) + z1, z2, z3 = compute_z1_z2_z3(z_ifc, nlev, nlev - 1, nlev - 2, nlev - 3) + + wgtfacq_c[:, 2] = z1 * z2 / (z2 - z3) / (z1 - z3) + wgtfacq_c[:, 1] = (z1 - wgtfacq_c[:, 2] * (z1 - z3)) / (z1 - z2) + wgtfacq_c[:, 0] = 1.0 - (wgtfacq_c[:, 1] + wgtfacq_c[:, 2]) - wgtfacq_c[:, nlev - 3] = z1 * z2 / (z2 - z3) / (z1 - z3) - wgtfacq_c[:, nlev - 2] = (z1 - wgtfacq_c[:, nlev - 3] * (z1 - z3)) / (z1 - z2) - wgtfacq_c[:, nlev - 1] = 1.0 - (wgtfacq_c[:, nlev - 2] + wgtfacq_c[:, nlev - 3]) + wgtfacq_c_dsl[:, nlev - 1] = wgtfacq_c[:, 0] + wgtfacq_c_dsl[:, nlev - 2] = wgtfacq_c[:, 1] + wgtfacq_c_dsl[:, nlev - 3] = wgtfacq_c[:, 2] - return wgtfacq_c + return wgtfacq_c_dsl diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_e_dsl.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_e_dsl.py new file mode 100644 index 000000000..f62b07a4b --- /dev/null +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_e_dsl.py @@ -0,0 +1,61 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.common.metrics.stencils.compute_wgtfacq_c_dsl import ( + compute_wgtfacq_c_dsl, + compute_z1_z2_z3, +) + + +def compute_wgtfacq_e_dsl( + e2c, + z_ifc: np.array, + z_aux_c: np.array, + c_lin_e: np.array, + wgtfacq_e_dsl: np.array, + nlev: int, +): + """ + Compute weighting factor for quadratic interpolation to surface. + + Args: + e2c: Edge to Cell offset + z_ifc: geometric height at the vertical interface of cells. + z_aux_c: interpolation of weighting coefficients to edges + c_lin_e: interpolation field + wgtfacq_e_dsl: output + nlev: int, last k level + Returns: + Field[EdgeDim, KDim] (full levels) + """ + z1, z2, z3 = compute_z1_z2_z3(z_ifc, nlev, nlev - 1, nlev - 2, nlev - 3) + wgtfacq_c_dsl = compute_wgtfacq_c_dsl(z_ifc, nlev) + z_aux_c[:, 2] = z1 * z2 / (z2 - z3) / (z1 - z3) + z_aux_c[:, 1] = (z1 - wgtfacq_c_dsl[:, nlev - 3] * (z1 - z3)) / (z1 - z2) + z_aux_c[:, 0] = 1.0 - (wgtfacq_c_dsl[:, nlev - 2] + wgtfacq_c_dsl[:, nlev - 3]) + + z1, z2, z3 = compute_z1_z2_z3(z_ifc, 0, 1, 2, 3) + z_aux_c[:, 5] = z1 * z2 / (z2 - z3) / (z1 - z3) + z_aux_c[:, 4] = (z1 - z_aux_c[:, 5] * (z1 - z3)) / (z1 - z2) + z_aux_c[:, 3] = 1.0 - (z_aux_c[:, 4] + z_aux_c[:, 5]) + + c_lin_e = c_lin_e[:, :, np.newaxis] + z_aux_e = np.sum(c_lin_e * z_aux_c[e2c], axis=1) + + wgtfacq_e_dsl[:, nlev] = z_aux_e[:, 0] + wgtfacq_e_dsl[:, nlev - 1] = z_aux_e[:, 1] + wgtfacq_e_dsl[:, nlev - 2] = z_aux_e[:, 2] + + return wgtfacq_e_dsl diff --git a/model/common/tests/metric_tests/test_compute_wgtfacq.py b/model/common/tests/metric_tests/test_compute_wgtfacq.py new file mode 100644 index 000000000..8d5194cd6 --- /dev/null +++ b/model/common/tests/metric_tests/test_compute_wgtfacq.py @@ -0,0 +1,59 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import pytest + +from icon4py.model.common.dimension import E2CDim, EdgeDim, KDim +from icon4py.model.common.metrics.stencils.compute_wgtfacq_c_dsl import compute_wgtfacq_c_dsl +from icon4py.model.common.metrics.stencils.compute_wgtfacq_e_dsl import compute_wgtfacq_e_dsl +from icon4py.model.common.test_utils.helpers import dallclose, zero_field + + +@pytest.mark.datatest +def test_compute_wgtfacq_c_dsl(icon_grid, metrics_savepoint): + wgtfacq_c_dsl = metrics_savepoint.wgtfacq_c_dsl() + + wgtfacq_c_dsl_np = compute_wgtfacq_c_dsl( + z_ifc=metrics_savepoint.z_ifc().asnumpy(), + nlev=icon_grid.num_levels, + ) + assert dallclose(wgtfacq_c_dsl_np, wgtfacq_c_dsl.asnumpy()) + + +@pytest.mark.datatest +def test_compute_wgtfacq_e_dsl(metrics_savepoint, interpolation_savepoint, icon_grid): + wgtfacq_e_dsl_full = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) + wgtfacq_e_dsl_ref = metrics_savepoint.wgtfacq_e_dsl(icon_grid.num_levels + 1) + + wgtfacq_e_dsl_full = compute_wgtfacq_e_dsl( + e2c=icon_grid.connectivities[E2CDim], + z_ifc=metrics_savepoint.z_ifc().asnumpy(), + z_aux_c=metrics_savepoint.wgtfac_c().asnumpy(), + c_lin_e=interpolation_savepoint.c_lin_e().asnumpy(), + wgtfacq_e_dsl=wgtfacq_e_dsl_full.asnumpy(), + nlev=icon_grid.num_levels, + ) + + assert dallclose(wgtfacq_e_dsl_full, wgtfacq_e_dsl_ref.asnumpy()) diff --git a/model/common/tests/metric_tests/test_compute_wgtfacq_c.py b/model/common/tests/metric_tests/test_compute_wgtfacq_c.py deleted file mode 100644 index 97ac20bfd..000000000 --- a/model/common/tests/metric_tests/test_compute_wgtfacq_c.py +++ /dev/null @@ -1,43 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import pytest - -from icon4py.model.common.metrics.stencils.compute_wgtfacq_c import compute_wgtfacq_c -from icon4py.model.common.test_utils.helpers import dallclose - - -@pytest.mark.datatest -def test_compute_wgtfacq_c(icon_grid, metrics_savepoint): # fixture - wgtfacq_c_dsl = metrics_savepoint.wgtfacq_c_dsl() - z_ifc = metrics_savepoint.z_ifc() - - vertical_end = icon_grid.num_levels - - wgtfacq_c = compute_wgtfacq_c( - z_ifc.asnumpy(), - vertical_end, - ) - assert dallclose(wgtfacq_c, wgtfacq_c_dsl.asnumpy()) From 64310692dd5274b7ab6660843fd3197f449af36c Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 13 May 2024 11:27:44 +0200 Subject: [PATCH 36/65] pg_edgeidx_dsl, pg_exdist_dsl, wgtfac_e --- .../model/common/metrics/metric_fields.py | 199 +++++++++++++++++- .../stencils/compute_zdiff_gradp_dsl.py | 27 ++- .../common/test_utils/serialbox_utils.py | 3 + .../tests/metric_tests/test_metric_fields.py | 195 ++++++++++++++--- 4 files changed, 371 insertions(+), 53 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 8a52dac18..94451b0b7 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -29,12 +29,15 @@ from icon4py.model.common.dimension import ( C2E, + E2C, CellDim, + E2CDim, EdgeDim, KDim, Koff, VertexDim, ) +from icon4py.model.common.grid.horizontal import _compute_cells2edges from icon4py.model.common.math.helpers import ( _grad_fd_tang, average_cell_kdim_level_up, @@ -655,6 +658,7 @@ def compute_exner_exfac( }, ) + @field_operator def _compute_vwind_impl_wgt( z_ddxn_z_half_e: Field[[EdgeDim], wpfloat], @@ -686,14 +690,10 @@ def compute_vwind_impl_wgt( z_ddxn_z_half_e: Field[[EdgeDim], wpfloat], z_ddxt_z_half_e: Field[[EdgeDim], wpfloat], dual_edge_length: Field[[EdgeDim], wpfloat], - # vct_a: Field[[KDim], wpfloat], - # z_ifc: Field[[CellDim, KDim], wpfloat], vwind_impl_wgt: Field[[CellDim], wpfloat], vwind_offctr: wpfloat, horizontal_start: int32, horizontal_end: int32, - # vertical_start: int32, - # vertical_end: int32 ): """ Compute vwind_impl_wgt. @@ -717,3 +717,194 @@ def compute_vwind_impl_wgt( out=vwind_impl_wgt, domain={CellDim: (horizontal_start, horizontal_end)}, ) + + +@program +def compute_wgtfac_e( + wgtfac_c: Field[[CellDim, KDim], float], + c_lin_e: Field[[EdgeDim, E2CDim], float], + wgtfac_e: Field[[EdgeDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, +): + """ + Compute wgtfac_e. + + See mo_vertical_grid.f90 + + Args: + wgtfac_c: weighting factor for quadratic interpolation to surface + c_lin_e: interpolation field + wgtfac_e: output + horizontal_start: horizontal start index + horizontal_end: horizontal end index + vertical_start: vertical start index + vertical_end: vertical end index + """ + + _compute_cells2edges( + p_cell_in=wgtfac_c, + c_int=c_lin_e, + out=wgtfac_e, + domain={ + EdgeDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) + + +@field_operator +def _compute_flat_idx( + z_me: Field[[EdgeDim, KDim], wpfloat], + z_ifc: Field[[CellDim, KDim], wpfloat], + k_lev: Field[[KDim], int], +) -> Field[[EdgeDim, KDim], int]: + z_ifc_e_0 = z_ifc(E2C[0]) + z_ifc_e_k_0 = z_ifc_e_0(Koff[1]) + z_ifc_e_1 = z_ifc(E2C[1]) + z_ifc_e_k_1 = z_ifc_e_1(Koff[1]) + flat_idx = where( + (z_me <= z_ifc_e_0) & (z_me >= z_ifc_e_k_0) & (z_me <= z_ifc_e_1) & (z_me >= z_ifc_e_k_1), + k_lev, + int(0), + ) + return flat_idx + + +@field_operator +def _compute_z_aux2( + z_ifc: Field[[CellDim], wpfloat], +) -> Field[[EdgeDim], wpfloat]: + extrapol_dist = 5.0 + z_aux1 = maximum(z_ifc(E2C[0]), z_ifc(E2C[1])) + z_aux2 = z_aux1 - extrapol_dist + + return z_aux2 + + +@field_operator +def _compute_pg_edgeidx_vertidx( + c_lin_e: Field[[EdgeDim, E2CDim], float], + z_ifc: Field[[CellDim, KDim], wpfloat], + z_aux2: Field[[EdgeDim], wpfloat], + e_owner_mask: Field[[EdgeDim], bool], + flat_idx_max: Field[[EdgeDim], int], + e_lev: Field[[EdgeDim], int], + k_lev: Field[[KDim], int], + pg_edgeidx: Field[[EdgeDim, KDim], int], + pg_vertidx: Field[[EdgeDim, KDim], int], +) -> tuple[Field[[EdgeDim, KDim], int], Field[[EdgeDim, KDim], int]]: + e_lev = broadcast(e_lev, (EdgeDim, KDim)) + k_lev = broadcast(k_lev, (EdgeDim, KDim)) + z_mc = average_cell_kdim_level_up(z_ifc) + z_me = _compute_cells2edges(p_cell_in=z_mc, c_int=c_lin_e) + pg_edgeidx = where( + (k_lev >= (flat_idx_max + int(1))) & (z_me < z_aux2) & e_owner_mask, e_lev, pg_edgeidx + ) + pg_vertidx = where( + (k_lev >= (flat_idx_max + int(1))) & (z_me < z_aux2) & e_owner_mask, k_lev, pg_vertidx + ) + return pg_edgeidx, pg_vertidx + + +@field_operator +def _compute_pg_exdist_dsl( + z_me: Field[[EdgeDim, KDim], wpfloat], + z_aux2: Field[[EdgeDim], wpfloat], + e_owner_mask: Field[[EdgeDim], bool], + flat_idx_max: Field[[EdgeDim], int], + k_lev: Field[[KDim], int], + pg_exdist_dsl: Field[[EdgeDim, KDim], wpfloat], +) -> Field[[EdgeDim, KDim], wpfloat]: + k_lev = broadcast(k_lev, (EdgeDim, KDim)) + pg_exdist_dsl = where( + (k_lev >= (flat_idx_max + int(1))) & (z_me < z_aux2) & e_owner_mask, + z_me - z_aux2, + pg_exdist_dsl, + ) + return pg_exdist_dsl + + +@program +def compute_pg_exdist_dsl( + z_aux2: Field[[EdgeDim], wpfloat], + z_me: Field[[EdgeDim, KDim], float], + e_owner_mask: Field[[EdgeDim], bool], + flat_idx_max: Field[[EdgeDim], int], + k_lev: Field[[KDim], int], + pg_exdist_dsl: Field[[EdgeDim, KDim], wpfloat], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, +): + """ + Compute pg_edgeidx_dsl. + + See mo_vertical_grid.f90 + + Args: + z_aux2: Local field + z_me: Local field + e_owner_mask: Field of booleans over edges + flat_idx_max: Highest vertical index (counted from top to bottom) for which the edge point lies inside the cell box of the adjacent grid points + k_lev: Field of K levels + pg_exdist_dsl: output + horizontal_start: horizontal start index + horizontal_end: horizontal end index + vertical_start: vertical start index + vertical_end: vertical end index + """ + _compute_pg_exdist_dsl( + z_me=z_me, + z_aux2=z_aux2, + e_owner_mask=e_owner_mask, + flat_idx_max=flat_idx_max, + k_lev=k_lev, + pg_exdist_dsl=pg_exdist_dsl, + out=pg_exdist_dsl, + domain={EdgeDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, + ) + + +@field_operator +def _compute_pg_edgeidx_dsl( + pg_edgeidx: Field[[EdgeDim, KDim], int], + pg_vertidx: Field[[EdgeDim, KDim], int], +) -> Field[[EdgeDim, KDim], bool]: + pg_edgeidx_dsl = where((pg_edgeidx > int(0)) & (pg_vertidx > int(0)), True, False) + return pg_edgeidx_dsl + + +@program +def compute_pg_edgeidx_dsl( + pg_edgeidx: Field[[EdgeDim, KDim], int], + pg_vertidx: Field[[EdgeDim, KDim], int], + pg_edgeidx_dsl: Field[[EdgeDim, KDim], bool], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, +): + """ + Compute pg_edgeidx_dsl. + + See mo_vertical_grid.f90 + + Args: + pg_edgeidx: Index Edge values + pg_vertidx: Index K values + pg_edgeidx_dsl: output + horizontal_start: horizontal start index + horizontal_end: horizontal end index + vertical_start: vertical start index + vertical_end: vertical end index + """ + _compute_pg_edgeidx_dsl( + pg_edgeidx=pg_edgeidx, + pg_vertidx=pg_vertidx, + out=pg_edgeidx_dsl, + domain={EdgeDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, + ) diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py index 974d1eef5..57763a551 100644 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py @@ -27,27 +27,26 @@ def compute_zdiff_gradp_dsl( z_me = np.sum(z_mc[e2c] * c_lin_e, axis=1) zdiff_gradp[:, 0, :] = z_me - z_mc[e2c[:, 0]] zdiff_gradp[:, 1, :] = z_me - z_mc[e2c[:, 1]] - + horizontal_start = 428 + # None of the if statements get accessed for k in range(nlev): - for e in range( - nedges - ): # TODO: boundary here is not nedges, rather p_patch(jg)%edges%start_block(2) + for e in range(horizontal_start): # TODO: calculate horizontal_start, not hard coded if ( z_me[e, k] <= z_ifc[e2c[e, 0], k] - and z_me[e, k] >= z_ifc[e2c[e, 0], k + 1] + and z_me[e, k] >= z_ifc[e2c[e, 0], k - 1] and z_me[e, k] <= z_ifc[e2c[e, 1], k] - and z_me[e, k] >= z_ifc[e2c[e, 1], k + 1] + and z_me[e, k] >= z_ifc[e2c[e, 1], k - 1] ): flat_idx[e] = k for je in range(nedges): jk_start = flat_idx[je] - for jk in range(flat_idx[je] + 1, nlev): + for jk in range(flat_idx[je], nlev): for jk1 in range(jk_start, nlev): if ( jk1 == nlev or z_me[je, jk] <= z_ifc[e2c[je, 0], jk1] - and z_me[je, jk] >= z_ifc[e2c[je, 0], jk1 + 1] + and z_me[je, jk] >= z_ifc[e2c[je, 0], jk1 - 1] ): zdiff_gradp[je, 0, jk1] = z_me[je, jk] - z_mc[e2c[je, 0], jk1] jk_start = jk1 @@ -59,12 +58,12 @@ def compute_zdiff_gradp_dsl( for je in range(nedges): jk_start = flat_idx[je] - for jk in range(flat_idx[je] + 1, nlev): + for jk in range(flat_idx[je], nlev): for jk1 in range(jk_start, nlev): if ( jk1 == nlev or z_me[je, jk] <= z_ifc[e2c[je, 1], jk1] - and z_me[je, jk] >= z_ifc[e2c[je, 1], jk1 + 1] + and z_me[je, jk] >= z_ifc[e2c[je, 1], jk1 - 1] ): zdiff_gradp[je, 1, jk1] = z_me[je, jk] - z_mc[e2c[je, 1], jk1] jk_start = jk1 @@ -72,13 +71,13 @@ def compute_zdiff_gradp_dsl( for je in range(nedges): jk_start = flat_idx[je] - for jk in range(flat_idx[je] + 1, nlev): + for jk in range(flat_idx[je], nlev): if z_me[je, jk] < z_aux2[je]: for jk1 in range(jk_start, nlev): if ( jk1 == nlev or z_aux2[je] <= z_ifc[e2c[je, 0], jk1] - and z_aux2[je] >= z_ifc[e2c[je, 0], jk1 + 1] + and z_aux2[je] >= z_ifc[e2c[je, 0], jk1 - 1] ): zdiff_gradp[je, 0, jk1] = z_aux2[je] - z_mc[e2c[je, 0], jk1] jk_start = jk1 @@ -86,13 +85,13 @@ def compute_zdiff_gradp_dsl( for je in range(nedges): jk_start = flat_idx[je] - for jk in range(flat_idx[je] + 1, nlev): + for jk in range(flat_idx[je], nlev): if z_me[je, jk] < z_aux2[je]: for jk1 in range(jk_start, nlev): if ( jk1 == nlev or z_aux2[je] <= z_ifc[e2c[je, 1], jk1] - and z_aux2[je] >= z_ifc[e2c[je, 1], jk1 + 1] + and z_aux2[je] >= z_ifc[e2c[je, 1], jk1 - 1] ): zdiff_gradp[je, 1, jk1] = z_aux2[je] - z_mc[e2c[je, 1], jk1] jk_start = jk1 diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py index 3f69c1104..61b8bafb8 100644 --- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -536,6 +536,9 @@ def mask_prog_halo_c(self): def pg_exdist(self): return self._get_field("pg_exdist_dsl", EdgeDim, KDim) + def pg_edgeidx_dsl(self): + return self._get_field("pg_edgeidx_dsl", EdgeDim, KDim, dtype=bool) + def rayleigh_w(self): return self._get_field("rayleigh_w", KDim) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 0db935783..accba0364 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -30,11 +30,15 @@ ) from icon4py.model.common.grid.horizontal import ( HorizontalMarkerIndex, + _compute_cells2edges, _compute_cells2verts, compute_cells2edges, ) -from icon4py.model.common.math.helpers import compute_inverse_edge_kdim +from icon4py.model.common.math.helpers import average_cell_kdim_level_up, compute_inverse_edge_kdim from icon4py.model.common.metrics.metric_fields import ( + _compute_flat_idx, + _compute_pg_edgeidx_vertidx, + _compute_z_aux2, compute_coeff_dwdz, compute_d2dexdz2_fac_mc, compute_ddqz_z_full, @@ -43,10 +47,13 @@ compute_ddxn_z_half_e, compute_ddxt_z_half_e, compute_exner_exfac, + compute_pg_edgeidx_dsl, + compute_pg_exdist_dsl, compute_rayleigh_w, compute_scalfac_dd3d, compute_vwind_expl_wgt, compute_vwind_impl_wgt, + compute_wgtfac_e, compute_z_mc, ) from icon4py.model.common.metrics.stencils.compute_zdiff_gradp_dsl import compute_zdiff_gradp_dsl @@ -496,6 +503,8 @@ def test_compute_ddqz_z_full_e( vertical_start = 0 vertical_end = icon_grid.num_levels ddqz_z_full_e = zero_field(icon_grid, EdgeDim, KDim) + # TODO: perhaps write a program with ddqz_z_full_e name and call fieldop _cells2edges... from there + # TODO: This way it's clear where this field is computed and we cna more easily avoid duplicates compute_cells2edges.with_backend(backend)( p_cell_in=ddqz_z_full, c_int=c_lin_e, @@ -636,6 +645,7 @@ def test_compute_exner_exfac( assert dallclose(exner_exfac.asnumpy(), exner_exfac_ref.asnumpy(), rtol=1.0e-10) +# TODO @pytest.mark.datatest def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_savepoint, backend): backend = None @@ -675,6 +685,8 @@ def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_sav # zdiff_gradp_ref.asnumpy()[856:, :] == zdiff_gradp_full_field.asnumpy()[856:, :] assert dallclose(zdiff_gradp_full_field.asnumpy(), zdiff_gradp_ref.asnumpy()) + +# TODO @pytest.mark.datatest def test_compute_vwind_impl_wgt( icon_grid, grid_savepoint, metrics_savepoint, interpolation_savepoint, backend @@ -685,6 +697,11 @@ def test_compute_vwind_impl_wgt( z_ifc = metrics_savepoint.z_ifc() inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() z_ddxn_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) + tangent_orientation = grid_savepoint.tangent_orientation() + inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() + z_ddxt_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) + z_ifv = zero_field(icon_grid, VertexDim, KDim, extend={KDim: 1}) + horizontal_start = icon_grid.get_start_index( EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, @@ -696,15 +713,6 @@ def test_compute_vwind_impl_wgt( vertical_start = 0 vertical_end = icon_grid.num_levels + 1 - horizontal_start_edge = icon_grid.get_start_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, - ) - horizontal_end_edge = icon_grid.get_end_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, - ) - compute_ddxn_z_half_e( z_ifc=z_ifc, inv_dual_edge_length=inv_dual_edge_length, @@ -716,11 +724,16 @@ def test_compute_vwind_impl_wgt( offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, ) - z_ddxt_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) - z_ifv = zero_field(icon_grid, VertexDim, KDim, extend={KDim: 1}) - - tangent_orientation = grid_savepoint.tangent_orientation() - inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() + horizontal_start_edge = icon_grid.get_start_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, + ) + horizontal_end_edge = icon_grid.get_end_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, + ) + vertical_start = 0 + vertical_end = icon_grid.num_levels + 1 horizontal_start_vertex = icon_grid.get_start_index( VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, @@ -729,10 +742,9 @@ def test_compute_vwind_impl_wgt( VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) - 1, ) - cells_aw_verts = interpolation_savepoint.c_intp().asnumpy() _compute_cells2verts( z_ifc, - as_field((VertexDim, V2CDim), cells_aw_verts), + interpolation_savepoint.c_intp(), out=z_ifv, offset_provider={"V2C": icon_grid.get_offset_provider("V2C")}, domain={ @@ -740,6 +752,7 @@ def test_compute_vwind_impl_wgt( KDim: (vertical_start, vertical_end), }, ) + compute_ddxt_z_half_e( z_ifv=z_ifv, inv_primal_edge_length=inv_primal_edge_length, @@ -756,34 +769,146 @@ def test_compute_vwind_impl_wgt( CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, ) - #vwind_impl_wgt_full = zero_field(icon_grid, CellDim) vwind_impl_wgt_ref = metrics_savepoint.vwind_impl_wgt() dual_edge_length = grid_savepoint.dual_edge_length() vwind_offctr = 0.2 - vertical_start = max(10, icon_grid.num_levels - 8) vwind_impl_wgt_full = constant_field(icon_grid, 0.5 + vwind_offctr, (CellDim)) compute_vwind_impl_wgt.with_backend(backend)( z_ddxn_z_half_e=as_field((EdgeDim,), z_ddxn_z_half_e.asnumpy()[:, icon_grid.num_levels]), z_ddxt_z_half_e=as_field((EdgeDim,), z_ddxt_z_half_e.asnumpy()[:, icon_grid.num_levels]), dual_edge_length=dual_edge_length, - #vct_a=grid_savepoint.vct_a(), - #z_ifc=metrics_savepoint.z_ifc(), vwind_impl_wgt=vwind_impl_wgt_full, vwind_offctr=vwind_offctr, horizontal_start=horizontal_start_cell, horizontal_end=icon_grid.num_cells, - # vertical_start=vertical_start, - # vertical_end=icon_grid.num_levels, - offset_provider={"C2E": icon_grid.get_offset_provider("C2E"),}, - ) - # vwind_impl_wgt = vwind_impl_wgt_full.asnumpy() - # vct_a = grid_savepoint.vct_a().asnumpy() - # z_ifc = metrics_savepoint.z_ifc().asnumpy() - # for jk in range(vertical_start, icon_grid.num_levels): - # for jc in range(horizontal_start_cell, icon_grid.num_cells): - # z_diff_2 = (z_ifc[jc, jk] - z_ifc[jc, jk+1]) / (vct_a[jk] - vct_a[jk+1]) - # if z_diff_2 < 0.6: - # vwind_impl_wgt[jc] = np.maximum(vwind_impl_wgt[jc], 1.2 - z_diff_2) - - assert dallclose(vwind_impl_wgt_ref.asnumpy(), vwind_impl_wgt_full.asnumpy(), rtol=0.000000000000001) + offset_provider={"C2E": icon_grid.get_offset_provider("C2E")}, + ) + + vertical_start = max(10, icon_grid.num_levels - 8) + vwind_impl_wgt = vwind_impl_wgt_full.asnumpy() + vct_a = grid_savepoint.vct_a().asnumpy() + z_ifc = metrics_savepoint.z_ifc().asnumpy() + for jk in range(vertical_start, icon_grid.num_levels): + for jc in range(horizontal_start_cell, icon_grid.num_cells): + z_diff_2 = (z_ifc[jc, jk] - z_ifc[jc, jk - 1]) / (vct_a[jk] - vct_a[jk - 1]) + if z_diff_2 < 0.6: + vwind_impl_wgt[jc] = np.maximum(vwind_impl_wgt[jc], 1.2 - z_diff_2) + # convert code above in numpy + assert dallclose(vwind_impl_wgt_ref.asnumpy(), vwind_impl_wgt_full) + + +@pytest.mark.datatest +def test_compute_wgtfac_e(metrics_savepoint, interpolation_savepoint, icon_grid, backend): + wgtfac_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) + wgtfac_e_ref = metrics_savepoint.wgtfac_e() + compute_wgtfac_e.with_backend(backend)( + wgtfac_c=metrics_savepoint.wgtfac_c(), + c_lin_e=interpolation_savepoint.c_lin_e(), + wgtfac_e=wgtfac_e, + horizontal_start=0, + horizontal_end=wgtfac_e.shape[0], + vertical_start=0, + vertical_end=icon_grid.num_levels + 1, + offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, + ) + assert dallclose(wgtfac_e.asnumpy(), wgtfac_e_ref.asnumpy()) + + +@pytest.mark.datatest +def test_compute_pg_exdist_dsl( + metrics_savepoint, interpolation_savepoint, icon_grid, grid_savepoint, backend +): + backend = None + pg_exdist_ref = metrics_savepoint.pg_exdist() + nlev = icon_grid.num_levels + k_lev = as_field((KDim,), np.arange(nlev, dtype=int)) + pg_edgeidx = zero_field(icon_grid, EdgeDim, KDim, dtype=int) + pg_vertidx = zero_field(icon_grid, EdgeDim, KDim, dtype=int) + pg_exdist_dsl = zero_field(icon_grid, EdgeDim, KDim) + z_me = zero_field(icon_grid, EdgeDim, KDim) + z_aux2 = zero_field(icon_grid, EdgeDim) + z_mc = zero_field(icon_grid, CellDim, KDim) + flat_idx = zero_field(icon_grid, EdgeDim, KDim) + z_ifc = metrics_savepoint.z_ifc() + z_ifc_sliced = as_field((CellDim,), z_ifc.asnumpy()[:, nlev]) + end_edge_nudging = icon_grid.get_end_index(EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim)) + e_lev = as_field((EdgeDim,), np.arange(icon_grid.num_edges, dtype=int32)) + + average_cell_kdim_level_up( + z_ifc, out=z_mc, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")} + ) + _compute_cells2edges( + p_cell_in=z_mc, + c_int=interpolation_savepoint.c_lin_e(), + out=z_me, + offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, + ) + _compute_z_aux2( + z_ifc=z_ifc_sliced, + out=z_aux2, + domain={EdgeDim: (428, 31558)}, + offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, + ) + + _compute_flat_idx( + z_me=z_me, + z_ifc=z_ifc, + k_lev=k_lev, + out=flat_idx, + domain={EdgeDim: (428, 31558), KDim: (int32(0), icon_grid.num_levels)}, + offset_provider={ + "E2C": icon_grid.get_offset_provider("E2C"), + "Koff": icon_grid.get_offset_provider("Koff"), + }, + ) + flat_idx_np = np.amax(flat_idx.asnumpy(), axis=1) + + compute_pg_exdist_dsl.with_backend(backend)( + z_aux2=z_aux2, + z_me=z_me, + e_owner_mask=grid_savepoint.e_owner_mask(), + flat_idx_max=as_field((EdgeDim,), flat_idx_np, dtype=int), + k_lev=k_lev, + pg_exdist_dsl=pg_exdist_dsl, + horizontal_start=end_edge_nudging, + horizontal_end=icon_grid.num_edges, + vertical_start=int(0), + vertical_end=nlev, + offset_provider={}, + ) + + _compute_pg_edgeidx_vertidx( + c_lin_e=interpolation_savepoint.c_lin_e(), + z_ifc=z_ifc, + z_aux2=z_aux2, + e_owner_mask=grid_savepoint.e_owner_mask(), + flat_idx_max=as_field((EdgeDim,), flat_idx_np, dtype=int), + e_lev=e_lev, + k_lev=k_lev, + pg_edgeidx=pg_edgeidx, + pg_vertidx=pg_vertidx, + out=(pg_edgeidx, pg_vertidx), + domain={EdgeDim: (end_edge_nudging, icon_grid.num_edges), KDim: (int(0), nlev)}, + offset_provider={ + "E2C": icon_grid.get_offset_provider("E2C"), + "Koff": icon_grid.get_offset_provider("Koff"), + }, + ) + + pg_edgeidx_dsl = zero_field(icon_grid, EdgeDim, KDim, dtype=bool) + pg_edgeidx_dsl_ref = metrics_savepoint.pg_edgeidx_dsl() + + compute_pg_edgeidx_dsl.with_backend(backend)( + pg_edgeidx=pg_edgeidx, + pg_vertidx=pg_vertidx, + pg_edgeidx_dsl=pg_edgeidx_dsl, + horizontal_start=int(0), + horizontal_end=icon_grid.num_edges, + vertical_start=int(0), + vertical_end=icon_grid.num_levels, + offset_provider={}, + ) + + assert dallclose(pg_exdist_ref.asnumpy(), pg_exdist_dsl.asnumpy(), rtol=1.0e-9) + assert dallclose(pg_edgeidx_dsl_ref.asnumpy(), pg_edgeidx_dsl.asnumpy()) From 6496b937fd0c68e3359f47dc1ee73f8e9d5240ad Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 13 May 2024 11:43:48 +0200 Subject: [PATCH 37/65] cleanup --- .../model/common/metrics/metric_fields.py | 63 +---------- .../tests/metric_tests/test_metric_fields.py | 103 +----------------- 2 files changed, 3 insertions(+), 163 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 94451b0b7..3e4aa8ed4 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -503,66 +503,6 @@ def compute_ddxn_z_full( average_edge_kdim_level_up(z_ddxnt_z_half_e, out=ddxn_z_full) -@field_operator -def _compute_vwind_impl_wgt( - z_ddxn_z_half_e: Field[[EdgeDim], wpfloat], - z_ddxt_z_half_e: Field[[EdgeDim], wpfloat], - dual_edge_length: Field[[EdgeDim], wpfloat], - vwind_offctr: wpfloat, -) -> Field[[CellDim], wpfloat]: - z_ddx_1 = maximum(abs(z_ddxn_z_half_e(C2E[0])), abs(z_ddxt_z_half_e(C2E[0]))) - z_ddx_2 = maximum(abs(z_ddxn_z_half_e(C2E[1])), abs(z_ddxt_z_half_e(C2E[1]))) - z_ddx_3 = maximum(abs(z_ddxn_z_half_e(C2E[2])), abs(z_ddxt_z_half_e(C2E[2]))) - z_ddx_1_2 = maximum(z_ddx_1, z_ddx_2) - z_maxslope = maximum(z_ddx_1_2, z_ddx_3) - - z_diff_1_2 = maximum( - abs(z_ddxn_z_half_e(C2E[0]) * dual_edge_length(C2E[0])), - abs(z_ddxn_z_half_e(C2E[1]) * dual_edge_length(C2E[1])), - ) - z_diff = maximum(z_diff_1_2, abs(z_ddxn_z_half_e(C2E[2]) * dual_edge_length(C2E[2]))) - z_offctr_1 = maximum(vwind_offctr, 0.425 * z_maxslope ** (0.75)) - z_offctr = maximum(z_offctr_1, minimum(0.25, 0.00025 * (z_diff - 250.0))) - z_offctr = minimum(maximum(vwind_offctr, 0.75), z_offctr) - vwind_impl_wgt = 0.5 + z_offctr - return vwind_impl_wgt - - -@program(grid_type=GridType.UNSTRUCTURED) -def compute_vwind_impl_wgt( - z_ddxn_z_half_e: Field[[EdgeDim], wpfloat], - z_ddxt_z_half_e: Field[[EdgeDim], wpfloat], - dual_edge_length: Field[[EdgeDim], wpfloat], - vwind_impl_wgt: Field[[CellDim], wpfloat], - vwind_offctr: wpfloat, - horizontal_start: int32, - horizontal_end: int32, -): - """ - Compute vwind_impl_wgt. - - See mo_vertical_grid.f90 - - Args: - z_ddxn_z_half_e: intermediate storage for field - z_ddxt_z_half_e: intermediate storage for field - dual_edge_length: dual_edge_length - vwind_impl_wgt: (output) offcentering in vertical mass flux - vwind_offctr: off-centering in vertical wind solver - horizontal_start: horizontal start index - horizontal_end: horizontal end index - - """ - _compute_vwind_impl_wgt( - z_ddxn_z_half_e=z_ddxn_z_half_e, - z_ddxt_z_half_e=z_ddxt_z_half_e, - dual_edge_length=dual_edge_length, - vwind_offctr=vwind_offctr, - out=vwind_impl_wgt, - domain={CellDim: (horizontal_start, horizontal_end)}, - ) - - @field_operator def _compute_vwind_expl_wgt(vwind_impl_wgt: Field[[CellDim], wpfloat]) -> Field[[CellDim], wpfloat]: return 1.0 - vwind_impl_wgt @@ -678,8 +618,7 @@ def _compute_vwind_impl_wgt( ) z_diff = maximum(z_diff_1_2, abs(z_ddxn_z_half_e(C2E[2]) * dual_edge_length(C2E[2]))) z_offctr_1 = maximum(vwind_offctr, 0.425 * z_maxslope ** (0.75)) - z_offctr_2 = minimum(0.25, 0.00025 * (z_diff - 250.0)) - z_offctr = maximum(z_offctr_1, z_offctr_2) + z_offctr = maximum(z_offctr_1, minimum(0.25, 0.00025 * (z_diff - 250.0))) z_offctr = minimum(maximum(vwind_offctr, 0.75), z_offctr) vwind_impl_wgt = 0.5 + z_offctr return vwind_impl_wgt diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index accba0364..de687248f 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -351,101 +351,6 @@ def test_compute_ddxt_z_full_e( assert np.allclose(ddxt_z_full.asnumpy(), ddxt_z_full_ref) -@pytest.mark.datatest -def test_compute_vwind_impl_wgt( - icon_grid, grid_savepoint, metrics_savepoint, interpolation_savepoint, backend -): - if is_roundtrip(backend): - pytest.skip("skipping: slow backend") - z_ifc = metrics_savepoint.z_ifc() - inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() - z_ddxn_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) - horizontal_start = icon_grid.get_start_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, - ) - horizontal_end = icon_grid.get_end_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, - ) - vertical_start = 0 - vertical_end = icon_grid.num_levels + 1 - - horizontal_start_edge = icon_grid.get_start_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, - ) - horizontal_end_edge = icon_grid.get_end_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, - ) - - compute_ddxn_z_half_e( - z_ifc=z_ifc, - inv_dual_edge_length=inv_dual_edge_length, - ddxn_z_half_e=z_ddxn_z_half_e, - horizontal_start=horizontal_start, - horizontal_end=horizontal_end, - vertical_start=vertical_start, - vertical_end=vertical_end, - offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, - ) - - z_ddxt_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) - z_ifv = zero_field(icon_grid, VertexDim, KDim, extend={KDim: 1}) - - tangent_orientation = grid_savepoint.tangent_orientation() - inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() - horizontal_start_vertex = icon_grid.get_start_index( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, - ) - horizontal_end_vertex = icon_grid.get_end_index( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) - 1, - ) - cells_aw_verts = interpolation_savepoint.c_intp().asnumpy() - _compute_cells2verts( - z_ifc, - as_field((VertexDim, V2CDim), cells_aw_verts), - out=z_ifv, - offset_provider={"V2C": icon_grid.get_offset_provider("V2C")}, - domain={ - VertexDim: (horizontal_start_vertex, horizontal_end_vertex), - KDim: (vertical_start, vertical_end), - }, - ) - compute_ddxt_z_half_e( - z_ifv=z_ifv, - inv_primal_edge_length=inv_primal_edge_length, - tangent_orientation=tangent_orientation, - ddxt_z_half_e=z_ddxt_z_half_e, - horizontal_start=horizontal_start_edge, - horizontal_end=horizontal_end_edge, - vertical_start=vertical_start, - vertical_end=vertical_end, - offset_provider={"E2V": icon_grid.get_offset_provider("E2V")}, - ) - - vwind_impl_wgt_full = zero_field(icon_grid, CellDim) - vwind_impl_wgt_ref = metrics_savepoint.vwind_impl_wgt() - dual_edge_length = grid_savepoint.dual_edge_length() - vwind_offctr = 0.2 - - compute_vwind_impl_wgt.with_backend(backend)( - z_ddxn_z_half_e=as_field((EdgeDim,), z_ddxn_z_half_e.asnumpy()[:, icon_grid.num_levels]), - z_ddxt_z_half_e=as_field((EdgeDim,), z_ddxt_z_half_e.asnumpy()[:, icon_grid.num_levels]), - dual_edge_length=dual_edge_length, - vwind_impl_wgt=vwind_impl_wgt_full, - vwind_offctr=vwind_offctr, - horizontal_start=int32(0), - horizontal_end=icon_grid.num_cells, - offset_provider={"C2E": icon_grid.get_offset_provider("C2E")}, - ) - - assert dallclose(vwind_impl_wgt_full.asnumpy(), vwind_impl_wgt_ref.asnumpy(), rtol=1.0) - - @pytest.mark.datatest def test_compute_vwind_expl_wgt(icon_grid, metrics_savepoint, backend): vwind_expl_wgt_full = zero_field(icon_grid, CellDim) @@ -664,9 +569,6 @@ def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_sav offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) c_lin_e = interpolation_savepoint.c_lin_e() - horizontal_start = icon_grid.get_start_index( - EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1 - ) flat_idx = np.full(shape=icon_grid.num_edges, fill_value=icon_grid.num_levels) zdiff_gradp_full_np_final = compute_zdiff_gradp_dsl( @@ -732,8 +634,7 @@ def test_compute_vwind_impl_wgt( EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1, ) - vertical_start = 0 - vertical_end = icon_grid.num_levels + 1 + horizontal_start_vertex = icon_grid.get_start_index( VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, @@ -795,7 +696,7 @@ def test_compute_vwind_impl_wgt( if z_diff_2 < 0.6: vwind_impl_wgt[jc] = np.maximum(vwind_impl_wgt[jc], 1.2 - z_diff_2) # convert code above in numpy - assert dallclose(vwind_impl_wgt_ref.asnumpy(), vwind_impl_wgt_full) + assert dallclose(vwind_impl_wgt_ref.asnumpy(), vwind_impl_wgt) @pytest.mark.datatest From 84d31a23600d360e815be1805e316118791926ca Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 14 May 2024 14:34:41 +0200 Subject: [PATCH 38/65] mask_prog_halo_c implementation and zdiff_gradp fix --- .../model/common/metrics/metric_fields.py | 19 ++++ .../stencils/compute_zdiff_gradp_dsl.py | 73 ++++++-------- .../tests/metric_tests/test_metric_fields.py | 95 +++++++++++++++---- 3 files changed, 124 insertions(+), 63 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 3e4aa8ed4..5785b997c 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -847,3 +847,22 @@ def compute_pg_edgeidx_dsl( out=pg_edgeidx_dsl, domain={EdgeDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, ) + +@field_operator +def _compute_mask_prog_halo_c( + c_refin_ctrl: Field[[CellDim], int32], mask_prog_halo_c: Field[[CellDim], bool] +) -> Field[[CellDim], bool]: + mask_prog_halo_c = where((c_refin_ctrl >= 1) & (c_refin_ctrl <= 4), mask_prog_halo_c, True) + return mask_prog_halo_c + +@program +def compute_mask_prog_halo_c( + c_refin_ctrl: Field[[CellDim], int32], + mask_prog_halo_c: Field[[CellDim], bool], + horizontal_start: int32, + horizontal_end: int32, +): + _compute_mask_prog_halo_c( + c_refin_ctrl, mask_prog_halo_c, out=mask_prog_halo_c, domain={CellDim: (horizontal_start, horizontal_end)} + ) + diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py index 57763a551..e94d788a3 100644 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py @@ -15,85 +15,68 @@ def compute_zdiff_gradp_dsl( e2c, - c_lin_e: np.array, + z_me: np.array, z_mc: np.array, - zdiff_gradp: np.array, z_ifc: np.array, flat_idx: np.array, + z_aux2: np.array, nlev: int, + horizontal_start: int, + horizontal_start_1: int, nedges: int, ) -> np.array: - c_lin_e = np.expand_dims(c_lin_e, axis=-1) - z_me = np.sum(z_mc[e2c] * c_lin_e, axis=1) - zdiff_gradp[:, 0, :] = z_me - z_mc[e2c[:, 0]] - zdiff_gradp[:, 1, :] = z_me - z_mc[e2c[:, 1]] - horizontal_start = 428 - # None of the if statements get accessed - for k in range(nlev): - for e in range(horizontal_start): # TODO: calculate horizontal_start, not hard coded - if ( - z_me[e, k] <= z_ifc[e2c[e, 0], k] - and z_me[e, k] >= z_ifc[e2c[e, 0], k - 1] - and z_me[e, k] <= z_ifc[e2c[e, 1], k] - and z_me[e, k] >= z_ifc[e2c[e, 1], k - 1] - ): - flat_idx[e] = k + zdiff_gradp = np.zeros_like(z_mc[e2c]) + zdiff_gradp[horizontal_start:, :, :] = np.expand_dims(z_me, axis=1)[horizontal_start:, :, :] - z_mc[e2c][horizontal_start:, :, :] - for je in range(nedges): - jk_start = flat_idx[je] - for jk in range(flat_idx[je], nlev): + for je in range(horizontal_start, nedges): + jk_start = int(flat_idx[je]) + for jk in range(int(flat_idx[je]) + 1, nlev): for jk1 in range(jk_start, nlev): if ( - jk1 == nlev + jk1 == nlev - 1 or z_me[je, jk] <= z_ifc[e2c[je, 0], jk1] - and z_me[je, jk] >= z_ifc[e2c[je, 0], jk1 - 1] + and z_me[je, jk] >= z_ifc[e2c[je, 0], jk1 + 1] ): - zdiff_gradp[je, 0, jk1] = z_me[je, jk] - z_mc[e2c[je, 0], jk1] + zdiff_gradp[je, 0, jk] = z_me[je, jk] - z_mc[e2c[je, 0], jk1] jk_start = jk1 break - extrapol_dist = 5.0 - z_aux1 = np.maximum(z_ifc[e2c[:, 0], nlev], z_ifc[e2c[:, 1], nlev]) - z_aux2 = z_aux1 - extrapol_dist - - for je in range(nedges): - jk_start = flat_idx[je] - for jk in range(flat_idx[je], nlev): + jk_start = int(flat_idx[je]) + for jk in range(int(flat_idx[je]) + 1, nlev): for jk1 in range(jk_start, nlev): if ( - jk1 == nlev + jk1 == nlev - 1 or z_me[je, jk] <= z_ifc[e2c[je, 1], jk1] - and z_me[je, jk] >= z_ifc[e2c[je, 1], jk1 - 1] + and z_me[je, jk] >= z_ifc[e2c[je, 1], jk1 + 1] ): - zdiff_gradp[je, 1, jk1] = z_me[je, jk] - z_mc[e2c[je, 1], jk1] + zdiff_gradp[je, 1, jk] = z_me[je, jk] - z_mc[e2c[je, 1], jk1] jk_start = jk1 break - for je in range(nedges): - jk_start = flat_idx[je] - for jk in range(flat_idx[je], nlev): + for je in range(horizontal_start_1, nedges): + jk_start = int(flat_idx[je]) + for jk in range(int(flat_idx[je]) + 1, nlev): if z_me[je, jk] < z_aux2[je]: for jk1 in range(jk_start, nlev): if ( - jk1 == nlev + jk1 == nlev - 1 or z_aux2[je] <= z_ifc[e2c[je, 0], jk1] - and z_aux2[je] >= z_ifc[e2c[je, 0], jk1 - 1] + and z_aux2[je] >= z_ifc[e2c[je, 0], jk1 + 1] ): - zdiff_gradp[je, 0, jk1] = z_aux2[je] - z_mc[e2c[je, 0], jk1] + zdiff_gradp[je, 0, jk] = z_aux2[je] - z_mc[e2c[je, 0], jk1] jk_start = jk1 break - for je in range(nedges): - jk_start = flat_idx[je] - for jk in range(flat_idx[je], nlev): + jk_start = int(flat_idx[je]) + for jk in range(int(flat_idx[je]) + 1, nlev): if z_me[je, jk] < z_aux2[je]: for jk1 in range(jk_start, nlev): if ( - jk1 == nlev + jk1 == nlev - 1 or z_aux2[je] <= z_ifc[e2c[je, 1], jk1] - and z_aux2[je] >= z_ifc[e2c[je, 1], jk1 - 1] + and z_aux2[je] >= z_ifc[e2c[je, 1], jk1 + 1] ): - zdiff_gradp[je, 1, jk1] = z_aux2[je] - z_mc[e2c[je, 1], jk1] + zdiff_gradp[je, 1, jk] = z_aux2[je] - z_mc[e2c[je, 1], jk1] jk_start = jk1 break diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index de687248f..60e8ce194 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -54,7 +54,7 @@ compute_vwind_expl_wgt, compute_vwind_impl_wgt, compute_wgtfac_e, - compute_z_mc, + compute_z_mc, compute_mask_prog_halo_c, ) from icon4py.model.common.metrics.stencils.compute_zdiff_gradp_dsl import compute_zdiff_gradp_dsl from icon4py.model.common.test_utils.datatest_utils import ( @@ -555,10 +555,15 @@ def test_compute_exner_exfac( def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_savepoint, backend): backend = None zdiff_gradp_ref = metrics_savepoint.zdiff_gradp() - zdiff_gradp_full = zero_field(icon_grid, EdgeDim, E2CDim, KDim) - zdiff_gradp_full_np = zdiff_gradp_full.asnumpy() z_mc = zero_field(icon_grid, CellDim, KDim) z_ifc = metrics_savepoint.z_ifc() + k_lev = as_field((KDim,), np.arange(icon_grid.num_levels, dtype=int)) + z_me = zero_field(icon_grid, EdgeDim, KDim) + horizontal_start_edge = icon_grid.get_start_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, + ) + end_edge_nudging = icon_grid.get_end_index(EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim)) compute_z_mc.with_backend(backend)( z_ifc, z_mc, @@ -568,24 +573,50 @@ def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_sav vertical_end=int32(icon_grid.num_levels), offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) - c_lin_e = interpolation_savepoint.c_lin_e() - flat_idx = np.full(shape=icon_grid.num_edges, fill_value=icon_grid.num_levels) + _compute_cells2edges( + p_cell_in=z_mc, + c_int=interpolation_savepoint.c_lin_e(), + out=z_me, + offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, + ) + flat_idx = zero_field(icon_grid, EdgeDim, KDim) + _compute_flat_idx( + z_me=z_me, + z_ifc=z_ifc, + k_lev=k_lev, + out=flat_idx, + domain={EdgeDim: (horizontal_start_edge, icon_grid.num_edges), KDim: (int32(0), icon_grid.num_levels)}, + offset_provider={ + "E2C": icon_grid.get_offset_provider("E2C"), + "Koff": icon_grid.get_offset_provider("Koff"), + }, + ) + flat_idx_np = np.amax(flat_idx.asnumpy(), axis=1) + z_ifc_sliced = as_field((CellDim,), z_ifc.asnumpy()[:, icon_grid.num_levels]) + z_aux2 = zero_field(icon_grid, EdgeDim) + _compute_z_aux2( + z_ifc=z_ifc_sliced, + out=z_aux2, + domain={EdgeDim: (end_edge_nudging, icon_grid.num_edges)}, + offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, + ) - zdiff_gradp_full_np_final = compute_zdiff_gradp_dsl( + zdiff_gradp_full_np = compute_zdiff_gradp_dsl( e2c=icon_grid.connectivities[E2CDim], - c_lin_e=c_lin_e.asnumpy(), + z_me=z_me.asnumpy(), z_mc=z_mc.asnumpy(), - zdiff_gradp=zdiff_gradp_full_np, z_ifc=metrics_savepoint.z_ifc().asnumpy(), - flat_idx=flat_idx, - nlev=icon_grid.num_levels, # TODO: check that -1 ok + flat_idx=flat_idx_np, + z_aux2=z_aux2.asnumpy(), + nlev=icon_grid.num_levels, + horizontal_start=horizontal_start_edge, + horizontal_start_1=end_edge_nudging, nedges=icon_grid.num_edges, ) zdiff_gradp_full_field = flatten_first_two_dims( - ECDim, KDim, field=as_field((EdgeDim, E2CDim, KDim), zdiff_gradp_full_np_final) + ECDim, KDim, field=as_field((EdgeDim, E2CDim, KDim), zdiff_gradp_full_np) ) - # zdiff_gradp_ref.asnumpy()[856:, :] == zdiff_gradp_full_field.asnumpy()[856:, :] - assert dallclose(zdiff_gradp_full_field.asnumpy(), zdiff_gradp_ref.asnumpy()) + assert dallclose(zdiff_gradp_full_field.asnumpy(), zdiff_gradp_ref.asnumpy(), rtol=1.0e-5) # TODO @@ -701,6 +732,8 @@ def test_compute_vwind_impl_wgt( @pytest.mark.datatest def test_compute_wgtfac_e(metrics_savepoint, interpolation_savepoint, icon_grid, backend): + if is_roundtrip(backend): + pytest.skip("skipping: slow backend") wgtfac_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) wgtfac_e_ref = metrics_savepoint.wgtfac_e() compute_wgtfac_e.with_backend(backend)( @@ -733,7 +766,11 @@ def test_compute_pg_exdist_dsl( flat_idx = zero_field(icon_grid, EdgeDim, KDim) z_ifc = metrics_savepoint.z_ifc() z_ifc_sliced = as_field((CellDim,), z_ifc.asnumpy()[:, nlev]) - end_edge_nudging = icon_grid.get_end_index(EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim)) + start_edge_nudging = icon_grid.get_end_index(EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim)) + horizontal_start_edge = icon_grid.get_start_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, + ) e_lev = as_field((EdgeDim,), np.arange(icon_grid.num_edges, dtype=int32)) average_cell_kdim_level_up( @@ -748,7 +785,7 @@ def test_compute_pg_exdist_dsl( _compute_z_aux2( z_ifc=z_ifc_sliced, out=z_aux2, - domain={EdgeDim: (428, 31558)}, + domain={EdgeDim: (start_edge_nudging, icon_grid.num_edges)}, offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, ) @@ -757,7 +794,7 @@ def test_compute_pg_exdist_dsl( z_ifc=z_ifc, k_lev=k_lev, out=flat_idx, - domain={EdgeDim: (428, 31558), KDim: (int32(0), icon_grid.num_levels)}, + domain={EdgeDim: (horizontal_start_edge, icon_grid.num_edges), KDim: (int32(0), icon_grid.num_levels)}, offset_provider={ "E2C": icon_grid.get_offset_provider("E2C"), "Koff": icon_grid.get_offset_provider("Koff"), @@ -772,7 +809,7 @@ def test_compute_pg_exdist_dsl( flat_idx_max=as_field((EdgeDim,), flat_idx_np, dtype=int), k_lev=k_lev, pg_exdist_dsl=pg_exdist_dsl, - horizontal_start=end_edge_nudging, + horizontal_start=start_edge_nudging, horizontal_end=icon_grid.num_edges, vertical_start=int(0), vertical_end=nlev, @@ -790,7 +827,7 @@ def test_compute_pg_exdist_dsl( pg_edgeidx=pg_edgeidx, pg_vertidx=pg_vertidx, out=(pg_edgeidx, pg_vertidx), - domain={EdgeDim: (end_edge_nudging, icon_grid.num_edges), KDim: (int(0), nlev)}, + domain={EdgeDim: (start_edge_nudging, icon_grid.num_edges), KDim: (int(0), nlev)}, offset_provider={ "E2C": icon_grid.get_offset_provider("E2C"), "Koff": icon_grid.get_offset_provider("Koff"), @@ -813,3 +850,25 @@ def test_compute_pg_exdist_dsl( assert dallclose(pg_exdist_ref.asnumpy(), pg_exdist_dsl.asnumpy(), rtol=1.0e-9) assert dallclose(pg_edgeidx_dsl_ref.asnumpy(), pg_edgeidx_dsl.asnumpy()) + +@pytest.mark.datatest +def test_compute_mask_prog_halo_c( + metrics_savepoint, icon_grid, grid_savepoint, backend +): + mask_prog_halo_c_full = zero_field(icon_grid, CellDim, dtype=bool) + c_refin_ctrl = grid_savepoint.refin_ctrl(CellDim) + mask_prog_halo_c_ref = metrics_savepoint.mask_prog_halo_c() + horizontal_start = icon_grid.get_start_index( + CellDim, HorizontalMarkerIndex.local(CellDim) - 1 + ) + horizontal_end = icon_grid.get_end_index( + CellDim, HorizontalMarkerIndex.local(CellDim) + ) + compute_mask_prog_halo_c( + c_refin_ctrl=c_refin_ctrl, + mask_prog_halo_c=mask_prog_halo_c_full, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + offset_provider={} + ) + assert dallclose(mask_prog_halo_c_full.asnumpy(), mask_prog_halo_c_ref.asnumpy()) From a55265cae346a11c55a72644019c228ff31db8f2 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 14 May 2024 14:37:39 +0200 Subject: [PATCH 39/65] ran pre-commit --- .../model/common/metrics/metric_fields.py | 8 ++++-- .../stencils/compute_zdiff_gradp_dsl.py | 4 ++- .../tests/metric_tests/test_metric_fields.py | 28 ++++++++++--------- 3 files changed, 24 insertions(+), 16 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 5785b997c..6d51d75b9 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -848,6 +848,7 @@ def compute_pg_edgeidx_dsl( domain={EdgeDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, ) + @field_operator def _compute_mask_prog_halo_c( c_refin_ctrl: Field[[CellDim], int32], mask_prog_halo_c: Field[[CellDim], bool] @@ -855,6 +856,7 @@ def _compute_mask_prog_halo_c( mask_prog_halo_c = where((c_refin_ctrl >= 1) & (c_refin_ctrl <= 4), mask_prog_halo_c, True) return mask_prog_halo_c + @program def compute_mask_prog_halo_c( c_refin_ctrl: Field[[CellDim], int32], @@ -863,6 +865,8 @@ def compute_mask_prog_halo_c( horizontal_end: int32, ): _compute_mask_prog_halo_c( - c_refin_ctrl, mask_prog_halo_c, out=mask_prog_halo_c, domain={CellDim: (horizontal_start, horizontal_end)} + c_refin_ctrl, + mask_prog_halo_c, + out=mask_prog_halo_c, + domain={CellDim: (horizontal_start, horizontal_end)}, ) - diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py index e94d788a3..70461ea63 100644 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py @@ -26,7 +26,9 @@ def compute_zdiff_gradp_dsl( nedges: int, ) -> np.array: zdiff_gradp = np.zeros_like(z_mc[e2c]) - zdiff_gradp[horizontal_start:, :, :] = np.expand_dims(z_me, axis=1)[horizontal_start:, :, :] - z_mc[e2c][horizontal_start:, :, :] + zdiff_gradp[horizontal_start:, :, :] = ( + np.expand_dims(z_me, axis=1)[horizontal_start:, :, :] - z_mc[e2c][horizontal_start:, :, :] + ) for je in range(horizontal_start, nedges): jk_start = int(flat_idx[je]) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 60e8ce194..150f7dd64 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -47,6 +47,7 @@ compute_ddxn_z_half_e, compute_ddxt_z_half_e, compute_exner_exfac, + compute_mask_prog_halo_c, compute_pg_edgeidx_dsl, compute_pg_exdist_dsl, compute_rayleigh_w, @@ -54,7 +55,7 @@ compute_vwind_expl_wgt, compute_vwind_impl_wgt, compute_wgtfac_e, - compute_z_mc, compute_mask_prog_halo_c, + compute_z_mc, ) from icon4py.model.common.metrics.stencils.compute_zdiff_gradp_dsl import compute_zdiff_gradp_dsl from icon4py.model.common.test_utils.datatest_utils import ( @@ -585,7 +586,10 @@ def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_sav z_ifc=z_ifc, k_lev=k_lev, out=flat_idx, - domain={EdgeDim: (horizontal_start_edge, icon_grid.num_edges), KDim: (int32(0), icon_grid.num_levels)}, + domain={ + EdgeDim: (horizontal_start_edge, icon_grid.num_edges), + KDim: (int32(0), icon_grid.num_levels), + }, offset_provider={ "E2C": icon_grid.get_offset_provider("E2C"), "Koff": icon_grid.get_offset_provider("Koff"), @@ -794,7 +798,10 @@ def test_compute_pg_exdist_dsl( z_ifc=z_ifc, k_lev=k_lev, out=flat_idx, - domain={EdgeDim: (horizontal_start_edge, icon_grid.num_edges), KDim: (int32(0), icon_grid.num_levels)}, + domain={ + EdgeDim: (horizontal_start_edge, icon_grid.num_edges), + KDim: (int32(0), icon_grid.num_levels), + }, offset_provider={ "E2C": icon_grid.get_offset_provider("E2C"), "Koff": icon_grid.get_offset_provider("Koff"), @@ -851,24 +858,19 @@ def test_compute_pg_exdist_dsl( assert dallclose(pg_exdist_ref.asnumpy(), pg_exdist_dsl.asnumpy(), rtol=1.0e-9) assert dallclose(pg_edgeidx_dsl_ref.asnumpy(), pg_edgeidx_dsl.asnumpy()) + @pytest.mark.datatest -def test_compute_mask_prog_halo_c( - metrics_savepoint, icon_grid, grid_savepoint, backend -): +def test_compute_mask_prog_halo_c(metrics_savepoint, icon_grid, grid_savepoint, backend): mask_prog_halo_c_full = zero_field(icon_grid, CellDim, dtype=bool) c_refin_ctrl = grid_savepoint.refin_ctrl(CellDim) mask_prog_halo_c_ref = metrics_savepoint.mask_prog_halo_c() - horizontal_start = icon_grid.get_start_index( - CellDim, HorizontalMarkerIndex.local(CellDim) - 1 - ) - horizontal_end = icon_grid.get_end_index( - CellDim, HorizontalMarkerIndex.local(CellDim) - ) + horizontal_start = icon_grid.get_start_index(CellDim, HorizontalMarkerIndex.local(CellDim) - 1) + horizontal_end = icon_grid.get_end_index(CellDim, HorizontalMarkerIndex.local(CellDim)) compute_mask_prog_halo_c( c_refin_ctrl=c_refin_ctrl, mask_prog_halo_c=mask_prog_halo_c_full, horizontal_start=horizontal_start, horizontal_end=horizontal_end, - offset_provider={} + offset_provider={}, ) assert dallclose(mask_prog_halo_c_full.asnumpy(), mask_prog_halo_c_ref.asnumpy()) From b2736cc793cd3c115a530f84fd18793bd795c2c8 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 14 May 2024 15:24:52 +0200 Subject: [PATCH 40/65] bdy_halo_c and docstrings --- .../model/common/metrics/metric_fields.py | 46 +++++++++++++++++++ .../tests/metric_tests/test_metric_fields.py | 25 ++++++++-- 2 files changed, 68 insertions(+), 3 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 6d51d75b9..6cbf1ca03 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -864,9 +864,55 @@ def compute_mask_prog_halo_c( horizontal_start: int32, horizontal_end: int32, ): + """ + Compute mask_prog_halo_c. + + See mo_vertical_grid.f90 + + Args: + c_refin_ctrl: Cell field of refin_ctrl + mask_prog_halo_c: output + horizontal_start: horizontal start index + horizontal_end: horizontal end index + """ _compute_mask_prog_halo_c( c_refin_ctrl, mask_prog_halo_c, out=mask_prog_halo_c, domain={CellDim: (horizontal_start, horizontal_end)}, ) + + +@field_operator +def _compute_bdy_halo_c( + c_refin_ctrl: Field[[CellDim], int32], + bdy_halo_c: Field[[CellDim], bool], +) -> Field[[CellDim], bool]: + bdy_halo_c = where((c_refin_ctrl >= 1) & (c_refin_ctrl <= 4), True, bdy_halo_c) + return bdy_halo_c + + +@program +def compute_bdy_halo_c( + c_refin_ctrl: Field[[CellDim], int32], + bdy_halo_c: Field[[CellDim], bool], + horizontal_start: int32, + horizontal_end: int32, +): + """ + Compute bdy_halo_c. + + See mo_vertical_grid.f90. mask_prog_halo_c_dsl_low_refin in ICON + + Args: + c_refin_ctrl: Cell field of refin_ctrl + bdy_halo_c: output + horizontal_start: horizontal start index + horizontal_end: horizontal end index + """ + _compute_bdy_halo_c( + c_refin_ctrl, + bdy_halo_c, + out=bdy_halo_c, + domain={CellDim: (horizontal_start, horizontal_end)}, + ) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 150f7dd64..ffbc0f648 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -39,6 +39,7 @@ _compute_flat_idx, _compute_pg_edgeidx_vertidx, _compute_z_aux2, + compute_bdy_halo_c, compute_coeff_dwdz, compute_d2dexdz2_fac_mc, compute_ddqz_z_full, @@ -628,7 +629,6 @@ def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_sav def test_compute_vwind_impl_wgt( icon_grid, grid_savepoint, metrics_savepoint, interpolation_savepoint, backend ): - backend = None if is_roundtrip(backend): pytest.skip("skipping: slow backend") z_ifc = metrics_savepoint.z_ifc() @@ -757,7 +757,8 @@ def test_compute_wgtfac_e(metrics_savepoint, interpolation_savepoint, icon_grid, def test_compute_pg_exdist_dsl( metrics_savepoint, interpolation_savepoint, icon_grid, grid_savepoint, backend ): - backend = None + if is_roundtrip(backend): + pytest.skip("skipping: slow backend") pg_exdist_ref = metrics_savepoint.pg_exdist() nlev = icon_grid.num_levels k_lev = as_field((KDim,), np.arange(nlev, dtype=int)) @@ -866,7 +867,7 @@ def test_compute_mask_prog_halo_c(metrics_savepoint, icon_grid, grid_savepoint, mask_prog_halo_c_ref = metrics_savepoint.mask_prog_halo_c() horizontal_start = icon_grid.get_start_index(CellDim, HorizontalMarkerIndex.local(CellDim) - 1) horizontal_end = icon_grid.get_end_index(CellDim, HorizontalMarkerIndex.local(CellDim)) - compute_mask_prog_halo_c( + compute_mask_prog_halo_c.with_backend(backend)( c_refin_ctrl=c_refin_ctrl, mask_prog_halo_c=mask_prog_halo_c_full, horizontal_start=horizontal_start, @@ -874,3 +875,21 @@ def test_compute_mask_prog_halo_c(metrics_savepoint, icon_grid, grid_savepoint, offset_provider={}, ) assert dallclose(mask_prog_halo_c_full.asnumpy(), mask_prog_halo_c_ref.asnumpy()) + + +@pytest.mark.datatest +def test_compute_bdy_halo_c(metrics_savepoint, icon_grid, grid_savepoint, backend): + bdy_halo_c_full = zero_field(icon_grid, CellDim, dtype=bool) + c_refin_ctrl = grid_savepoint.refin_ctrl(CellDim) + bdy_halo_c_ref = metrics_savepoint.bdy_halo_c() + horizontal_start = icon_grid.get_start_index(CellDim, HorizontalMarkerIndex.local(CellDim) - 1) + horizontal_end = icon_grid.get_end_index(CellDim, HorizontalMarkerIndex.local(CellDim)) + compute_bdy_halo_c( + c_refin_ctrl=c_refin_ctrl, + bdy_halo_c=bdy_halo_c_full, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + offset_provider={}, + ) + + assert dallclose(bdy_halo_c_full.asnumpy(), bdy_halo_c_ref.asnumpy()) From 20a333ed0163123a46c4c5d9bc47232cb4fe2d8e Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 14 May 2024 15:57:43 +0200 Subject: [PATCH 41/65] hmask_dd3d implementation --- .../src/icon4py/model/common/constants.py | 9 +++ .../model/common/metrics/metric_fields.py | 55 +++++++++++++++++++ .../tests/metric_tests/test_metric_fields.py | 22 ++++++++ 3 files changed, 86 insertions(+) diff --git a/model/common/src/icon4py/model/common/constants.py b/model/common/src/icon4py/model/common/constants.py index 7694ba496..75042f071 100644 --- a/model/common/src/icon4py/model/common/constants.py +++ b/model/common/src/icon4py/model/common/constants.py @@ -80,3 +80,12 @@ class RayleighType(IntEnum): # Temporal extrapolation of Exner for computation of horizontal pressure gradient exner_expol: Final[wpfloat] = 0.333 + +# Width of lateral boundary zones (as seen from the child domain) for which tendencies are interpolated from the parent domain +grf_bdywidth_c: Final[int] = 4 +grf_bdywidth_e: Final[int] = 9 +grf_nudge_start_e: Final[int] = grf_bdywidth_e + 1 + +# Width of nudging zone in units of child cell rows; +# MUST be a multiple of 2, AND grf_nudge_start_c + grf_width_nudgezone MUST NOT EXCEED the bdy_indexing_depth namelist variable in prepare_gridref +grf_nudgezone_width: Final[int] = 8 diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 6cbf1ca03..8437b4e35 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -15,6 +15,7 @@ Field, GridType, abs, + astype, broadcast, exp, field_operator, @@ -916,3 +917,57 @@ def compute_bdy_halo_c( out=bdy_halo_c, domain={CellDim: (horizontal_start, horizontal_end)}, ) + + +@field_operator +def _compute_hmask_dd3d( + e_refin_ctrl: Field[[EdgeDim], int32], grf_nudge_start_e: int32, grf_nudgezone_width: int32 +) -> Field[[EdgeDim], wpfloat]: + hmask_dd3d = ( + 1.0 + / (astype(grf_nudgezone_width, wpfloat) - 1.0) + * ( + astype(e_refin_ctrl, wpfloat) + - (astype(grf_nudge_start_e, wpfloat) + astype(grf_nudgezone_width, wpfloat) - 1.0) + ) + ) + hmask_dd3d = where( + (e_refin_ctrl <= 0) | (e_refin_ctrl >= (grf_nudge_start_e + 2 * (grf_nudgezone_width - 1))), + 1.0, + hmask_dd3d, + ) + hmask_dd3d = where( + e_refin_ctrl <= (grf_nudge_start_e + grf_nudgezone_width - 1), 0.0, hmask_dd3d + ) + return hmask_dd3d + + +@program +def compute_hmask_dd3d( + e_refin_ctrl: Field[[EdgeDim], int32], + hmask_dd3d: Field[[EdgeDim], wpfloat], + grf_nudge_start_e: int32, + grf_nudgezone_width: int32, + horizontal_start: int32, + horizontal_end: int32, +): + """ + Compute hmask_dd3d. + + See mo_vertical_grid.f90. Horizontal mask field for 3D divergence damping term. + + Args: + e_refin_ctrl: Edge field of refin_ctrl + hmask_dd3d: output + grf_nudge_start_e: mo_impl_constants_grf constant + grf_nudgezone_width: mo_impl_constants_grf constant + horizontal_start: horizontal start index + horizontal_end: horizontal end index + """ + _compute_hmask_dd3d( + e_refin_ctrl=e_refin_ctrl, + grf_nudge_start_e=grf_nudge_start_e, + grf_nudgezone_width=grf_nudgezone_width, + out=hmask_dd3d, + domain={EdgeDim: (horizontal_start, horizontal_end)}, + ) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index ffbc0f648..3d6d62f5e 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -48,6 +48,7 @@ compute_ddxn_z_half_e, compute_ddxt_z_half_e, compute_exner_exfac, + compute_hmask_dd3d, compute_mask_prog_halo_c, compute_pg_edgeidx_dsl, compute_pg_exdist_dsl, @@ -893,3 +894,24 @@ def test_compute_bdy_halo_c(metrics_savepoint, icon_grid, grid_savepoint, backen ) assert dallclose(bdy_halo_c_full.asnumpy(), bdy_halo_c_ref.asnumpy()) + + +@pytest.mark.datatest +def test_compute_hmask_dd3d(metrics_savepoint, icon_grid, grid_savepoint, backend): + hmask_dd3d_full = zero_field(icon_grid, EdgeDim) + e_refin_ctrl = grid_savepoint.refin_ctrl(EdgeDim) + horizontal_start = icon_grid.get_start_index( + CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1 + ) + hmask_dd3d_ref = metrics_savepoint.hmask_dd3d() + compute_hmask_dd3d( + e_refin_ctrl=e_refin_ctrl, + hmask_dd3d=hmask_dd3d_full, + grf_nudge_start_e=int32(constants.grf_nudge_start_e), + grf_nudgezone_width=int32(constants.grf_nudgezone_width), + horizontal_start=horizontal_start, + horizontal_end=icon_grid.num_edges, + offset_provider={}, + ) + + dallclose(hmask_dd3d_full.asnumpy(), hmask_dd3d_ref.asnumpy()) From ac644a788fbfec884e6d933d74a8260b7673641d Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 14 May 2024 16:01:57 +0200 Subject: [PATCH 42/65] ran pre-commit --- model/common/src/icon4py/model/common/metrics/metric_fields.py | 1 + model/common/tests/metric_tests/test_metric_fields.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index e34d85c6f..6165603e6 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -21,6 +21,7 @@ field_operator, int32, maximum, + minimum, program, sin, tanh, diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index eb29b6574..592ca64f6 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -34,7 +34,7 @@ _compute_cells2verts, compute_cells2edges, ) -from icon4py.model.common.math.helpers import average_cell_kdim_level_up, compute_inverse_edge_kdim +from icon4py.model.common.math.helpers import average_cell_kdim_level_up from icon4py.model.common.metrics.metric_fields import ( _compute_flat_idx, _compute_pg_edgeidx_vertidx, From c27d39a09622c041252b9203245ca9620aaba14a Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 14 May 2024 16:17:32 +0200 Subject: [PATCH 43/65] partial merge --- .../common/src/icon4py/model/common/metrics/metric_fields.py | 5 ++--- model/common/tests/metric_tests/test_metric_fields.py | 4 ---- 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 4417d63b7..6165603e6 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -39,7 +39,6 @@ VertexDim, ) from icon4py.model.common.grid.horizontal import _compute_cells2edges - from icon4py.model.common.math.helpers import ( _grad_fd_tang, average_cell_kdim_level_up, @@ -536,7 +535,8 @@ def compute_vwind_expl_wgt( out=vwind_expl_wgt, domain={CellDim: (horizontal_start, horizontal_end)}, ) - + + @field_operator def _compute_exner_exfac( ddxn_z_full: Field[[EdgeDim, KDim], wpfloat], @@ -972,4 +972,3 @@ def compute_hmask_dd3d( out=hmask_dd3d, domain={EdgeDim: (horizontal_start, horizontal_end)}, ) - diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 4e19141be..f7280386c 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -30,12 +30,8 @@ ) from icon4py.model.common.grid.horizontal import ( HorizontalMarkerIndex, - _compute_cells2edges, _compute_cells2verts, ) -from icon4py.model.common.interpolation.stencils.cell_2_edge_interpolation import ( - cell_2_edge_interpolation, -) from icon4py.model.common.math.helpers import average_cell_kdim_level_up from icon4py.model.common.metrics.metric_fields import ( _compute_flat_idx, From 605d084202801ce695324a282be454d3c20a579b Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 14 May 2024 16:26:56 +0200 Subject: [PATCH 44/65] fixed imports and updates --- .../model/common/metrics/metric_fields.py | 12 ++++--- .../tests/metric_tests/test_metric_fields.py | 34 +++++++++++-------- 2 files changed, 27 insertions(+), 19 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 6165603e6..2e27c3182 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -38,7 +38,9 @@ Koff, VertexDim, ) -from icon4py.model.common.grid.horizontal import _compute_cells2edges +from icon4py.model.common.interpolation.stencils.cell_2_edge_interpolation import ( + _cell_2_edge_interpolation, +) from icon4py.model.common.math.helpers import ( _grad_fd_tang, average_cell_kdim_level_up, @@ -685,9 +687,9 @@ def compute_wgtfac_e( vertical_end: vertical end index """ - _compute_cells2edges( - p_cell_in=wgtfac_c, - c_int=c_lin_e, + _cell_2_edge_interpolation( + in_field=wgtfac_c, + coeff=c_lin_e, out=wgtfac_e, domain={ EdgeDim: (horizontal_start, horizontal_end), @@ -740,7 +742,7 @@ def _compute_pg_edgeidx_vertidx( e_lev = broadcast(e_lev, (EdgeDim, KDim)) k_lev = broadcast(k_lev, (EdgeDim, KDim)) z_mc = average_cell_kdim_level_up(z_ifc) - z_me = _compute_cells2edges(p_cell_in=z_mc, c_int=c_lin_e) + z_me = _cell_2_edge_interpolation(in_field=z_mc, coeff=c_lin_e) pg_edgeidx = where( (k_lev >= (flat_idx_max + int(1))) & (z_me < z_aux2) & e_owner_mask, e_lev, pg_edgeidx ) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index f7280386c..452c620a7 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -32,6 +32,10 @@ HorizontalMarkerIndex, _compute_cells2verts, ) +from icon4py.model.common.interpolation.stencils import cell_2_edge_interpolation +from icon4py.model.common.interpolation.stencils.cell_2_edge_interpolation import ( + _cell_2_edge_interpolation, +) from icon4py.model.common.math.helpers import average_cell_kdim_level_up from icon4py.model.common.metrics.metric_fields import ( _compute_flat_idx, @@ -384,12 +388,12 @@ def test_compute_ddqz_z_full_e( ddqz_z_full_e = zero_field(icon_grid, EdgeDim, KDim) # TODO: perhaps write a program with ddqz_z_full_e name and call fieldop _cells2edges... from there # TODO: This way it's clear where this field is computed and we cna more easily avoid duplicates - compute_cells2edges.with_backend(backend)( - p_cell_in=ddqz_z_full, - c_int=c_lin_e, - p_vert_out=ddqz_z_full_e, - horizontal_start_edge=0, - horizontal_end_edge=ddqz_z_full_e.shape[0], + cell_2_edge_interpolation.with_backend(backend)( + in_field=ddqz_z_full, + coeff=c_lin_e, + out_field=ddqz_z_full_e, + horizontal_start=0, + horizontal_end=ddqz_z_full_e.shape[0], vertical_start=vertical_start, vertical_end=vertical_end, offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, @@ -503,7 +507,8 @@ def test_compute_ddxt_z_full( def test_compute_exner_exfac( grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint, backend ): - backend = None + if is_roundtrip(backend): + pytest.skip("skipping: slow backend") horizontal_start = icon_grid.get_start_index( CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1 ) @@ -527,7 +532,8 @@ def test_compute_exner_exfac( # TODO @pytest.mark.datatest def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_savepoint, backend): - backend = None + if is_roundtrip(backend): + pytest.skip("skipping: slow backend") zdiff_gradp_ref = metrics_savepoint.zdiff_gradp() z_mc = zero_field(icon_grid, CellDim, KDim) z_ifc = metrics_savepoint.z_ifc() @@ -547,9 +553,9 @@ def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_sav vertical_end=int32(icon_grid.num_levels), offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) - _compute_cells2edges( - p_cell_in=z_mc, - c_int=interpolation_savepoint.c_lin_e(), + _cell_2_edge_interpolation( + in_field=z_mc, + coeff=interpolation_savepoint.c_lin_e(), out=z_me, offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, ) @@ -753,9 +759,9 @@ def test_compute_pg_exdist_dsl( average_cell_kdim_level_up( z_ifc, out=z_mc, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")} ) - _compute_cells2edges( - p_cell_in=z_mc, - c_int=interpolation_savepoint.c_lin_e(), + _cell_2_edge_interpolation( + in_field=z_mc, + coeff=interpolation_savepoint.c_lin_e(), out=z_me, offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, ) From ce0bcde5f297528c229957e60e3998eeb8c6ce65 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Wed, 15 May 2024 09:10:46 +0200 Subject: [PATCH 45/65] fixed minor mistake --- model/common/tests/metric_tests/test_metric_fields.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 452c620a7..d8e47790d 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -32,9 +32,9 @@ HorizontalMarkerIndex, _compute_cells2verts, ) -from icon4py.model.common.interpolation.stencils import cell_2_edge_interpolation from icon4py.model.common.interpolation.stencils.cell_2_edge_interpolation import ( _cell_2_edge_interpolation, + cell_2_edge_interpolation, ) from icon4py.model.common.math.helpers import average_cell_kdim_level_up from icon4py.model.common.metrics.metric_fields import ( From 5805b982ad95826330f0b5cf048bf63b571e65e3 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 23 May 2024 09:05:44 +0200 Subject: [PATCH 46/65] update to vwind_impl_wgt --- .../model/common/metrics/metric_fields.py | 47 +++++++++++++++++-- .../tests/metric_tests/test_metric_fields.py | 26 +++++----- 2 files changed, 54 insertions(+), 19 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 2e27c3182..d94c3d1be 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -540,11 +540,10 @@ def compute_vwind_expl_wgt( @field_operator -def _compute_exner_exfac( +def _compute_maxslp_maxhgtd( ddxn_z_full: Field[[EdgeDim, KDim], wpfloat], dual_edge_length: Field[[EdgeDim], wpfloat], - exner_expol: wpfloat, -) -> Field[[CellDim, KDim], wpfloat]: +) -> tuple[Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat]]: z_maxslp_0_1 = maximum(abs(ddxn_z_full(C2E[0])), abs(ddxn_z_full(C2E[1]))) z_maxslp = maximum(z_maxslp_0_1, abs(ddxn_z_full(C2E[2]))) @@ -554,6 +553,16 @@ def _compute_exner_exfac( ) z_maxhgtd = maximum(z_maxhgtd_0_1, abs(ddxn_z_full(C2E[2]) * dual_edge_length(C2E[2]))) + return z_maxslp, z_maxhgtd + + +@field_operator +def _compute_exner_exfac( + ddxn_z_full: Field[[EdgeDim, KDim], wpfloat], + dual_edge_length: Field[[EdgeDim], wpfloat], + exner_expol: wpfloat, +) -> Field[[CellDim, KDim], wpfloat]: + z_maxslp, z_maxhgtd = _compute_maxslp_maxhgtd(ddxn_z_full, dual_edge_length) exner_exfac = exner_expol * minimum(1.0 - (4.0 * z_maxslp) ** 2, 1.0 - (0.002 * z_maxhgtd) ** 2) exner_exfac = maximum(0.0, exner_exfac) @@ -604,7 +613,7 @@ def compute_exner_exfac( @field_operator -def _compute_vwind_impl_wgt( +def _compute_vwind_impl_wgt_1( z_ddxn_z_half_e: Field[[EdgeDim], wpfloat], z_ddxt_z_half_e: Field[[EdgeDim], wpfloat], dual_edge_length: Field[[EdgeDim], wpfloat], @@ -628,15 +637,33 @@ def _compute_vwind_impl_wgt( return vwind_impl_wgt +@field_operator +def _compute_vwind_impl_wgt_2( + vct_a: Field[[KDim], wpfloat], + z_ifc: Field[[CellDim, KDim], wpfloat], + vwind_impl_wgt: Field[[CellDim], wpfloat], +) -> Field[[CellDim, KDim], wpfloat]: + z_diff_2 = (z_ifc - z_ifc(Koff[-1])) / (vct_a - vct_a(Koff[-1])) + vwind_impl_wgt_k = where( + z_diff_2 < 0.6, maximum(vwind_impl_wgt, 1.2 - z_diff_2), vwind_impl_wgt + ) + return vwind_impl_wgt_k + + @program(grid_type=GridType.UNSTRUCTURED) def compute_vwind_impl_wgt( z_ddxn_z_half_e: Field[[EdgeDim], wpfloat], z_ddxt_z_half_e: Field[[EdgeDim], wpfloat], dual_edge_length: Field[[EdgeDim], wpfloat], + vct_a: Field[[KDim], wpfloat], + z_ifc: Field[[CellDim, KDim], wpfloat], vwind_impl_wgt: Field[[CellDim], wpfloat], + vwind_impl_wgt_k: Field[[CellDim, KDim], wpfloat], vwind_offctr: wpfloat, horizontal_start: int32, horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): """ Compute vwind_impl_wgt. @@ -651,8 +678,10 @@ def compute_vwind_impl_wgt( vwind_offctr: off-centering in vertical wind solver horizontal_start: horizontal start index horizontal_end: horizontal end index + vertical_start: vertical start index + vertical_end: vertical end index """ - _compute_vwind_impl_wgt( + _compute_vwind_impl_wgt_1( z_ddxn_z_half_e=z_ddxn_z_half_e, z_ddxt_z_half_e=z_ddxt_z_half_e, dual_edge_length=dual_edge_length, @@ -661,6 +690,14 @@ def compute_vwind_impl_wgt( domain={CellDim: (horizontal_start, horizontal_end)}, ) + _compute_vwind_impl_wgt_2( + vct_a=vct_a, + z_ifc=z_ifc, + vwind_impl_wgt=vwind_impl_wgt, + out=vwind_impl_wgt_k, + domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, + ) + @program def compute_wgtfac_e( diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index d8e47790d..38b1ad450 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -602,7 +602,6 @@ def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_sav assert dallclose(zdiff_gradp_full_field.asnumpy(), zdiff_gradp_ref.asnumpy(), rtol=1.0e-5) -# TODO @pytest.mark.datatest def test_compute_vwind_impl_wgt( icon_grid, grid_savepoint, metrics_savepoint, interpolation_savepoint, backend @@ -616,7 +615,6 @@ def test_compute_vwind_impl_wgt( inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() z_ddxt_z_half_e = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) z_ifv = zero_field(icon_grid, VertexDim, KDim, extend={KDim: 1}) - horizontal_start = icon_grid.get_start_index( EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, @@ -686,29 +684,29 @@ def test_compute_vwind_impl_wgt( vwind_impl_wgt_ref = metrics_savepoint.vwind_impl_wgt() dual_edge_length = grid_savepoint.dual_edge_length() vwind_offctr = 0.2 - vwind_impl_wgt_full = constant_field(icon_grid, 0.5 + vwind_offctr, (CellDim)) + vwind_impl_wgt_full = constant_field(icon_grid, 0.5 + vwind_offctr, CellDim) + vwind_impl_wgt_k = constant_field(icon_grid, 0.7, CellDim, KDim) compute_vwind_impl_wgt.with_backend(backend)( z_ddxn_z_half_e=as_field((EdgeDim,), z_ddxn_z_half_e.asnumpy()[:, icon_grid.num_levels]), z_ddxt_z_half_e=as_field((EdgeDim,), z_ddxt_z_half_e.asnumpy()[:, icon_grid.num_levels]), dual_edge_length=dual_edge_length, + vct_a=grid_savepoint.vct_a(), + z_ifc=metrics_savepoint.z_ifc(), vwind_impl_wgt=vwind_impl_wgt_full, + vwind_impl_wgt_k=vwind_impl_wgt_k, vwind_offctr=vwind_offctr, horizontal_start=horizontal_start_cell, horizontal_end=icon_grid.num_cells, - offset_provider={"C2E": icon_grid.get_offset_provider("C2E")}, + vertical_start=max(10, icon_grid.num_levels - 8), + vertical_end=icon_grid.num_levels, + offset_provider={ + "C2E": icon_grid.get_offset_provider("C2E"), + "Koff": icon_grid.get_offset_provider("Koff"), + }, ) - vertical_start = max(10, icon_grid.num_levels - 8) - vwind_impl_wgt = vwind_impl_wgt_full.asnumpy() - vct_a = grid_savepoint.vct_a().asnumpy() - z_ifc = metrics_savepoint.z_ifc().asnumpy() - for jk in range(vertical_start, icon_grid.num_levels): - for jc in range(horizontal_start_cell, icon_grid.num_cells): - z_diff_2 = (z_ifc[jc, jk] - z_ifc[jc, jk - 1]) / (vct_a[jk] - vct_a[jk - 1]) - if z_diff_2 < 0.6: - vwind_impl_wgt[jc] = np.maximum(vwind_impl_wgt[jc], 1.2 - z_diff_2) - # convert code above in numpy + vwind_impl_wgt = np.amax(vwind_impl_wgt_k.asnumpy(), axis=1) assert dallclose(vwind_impl_wgt_ref.asnumpy(), vwind_impl_wgt) From f320b779f700e689f8486500b5070e99f6818792 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 23 May 2024 09:07:25 +0200 Subject: [PATCH 47/65] mask_hdiff implementation --- .../model/common/metrics/metric_fields.py | 94 +++++++++++++- .../common/test_utils/serialbox_utils.py | 3 + .../tests/metric_tests/test_metric_fields.py | 118 +++++++++++++++++- 3 files changed, 213 insertions(+), 2 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index d94c3d1be..f0f5df0ee 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -36,7 +36,7 @@ EdgeDim, KDim, Koff, - VertexDim, + VertexDim, C2E2C, ) from icon4py.model.common.interpolation.stencils.cell_2_edge_interpolation import ( _cell_2_edge_interpolation, @@ -1011,3 +1011,95 @@ def compute_hmask_dd3d( out=hmask_dd3d, domain={EdgeDim: (horizontal_start, horizontal_end)}, ) + +@field_operator +def _compute_mask_hdiff() -> Field[[CellDim, KDim], bool]: + return broadcast(True, (CellDim, KDim)) + + +@program +def compute_mask_hdiff( + mask_hdiff: Field[[CellDim, KDim], bool], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, +): + _compute_mask_hdiff( + out=mask_hdiff, + domain={CellDim: (horizontal_start, horizontal_end), KDim:(vertical_start, vertical_end)} + ) + +@field_operator +def _compute_z_maxslp_avg( + maxslp: Field[[CellDim, KDim], wpfloat], + c_bln_avg_0: Field[[CellDim], wpfloat], + c_bln_avg_1: Field[[CellDim], wpfloat], + c_bln_avg_2: Field[[CellDim], wpfloat], + c_bln_avg_3: Field[[CellDim], wpfloat] +) -> Field[[CellDim, KDim], wpfloat]: + z_maxslp_avg = maxslp * c_bln_avg_0 + \ + maxslp(C2E2C[0]) * c_bln_avg_1 + \ + maxslp(C2E2C[1]) * c_bln_avg_2 + \ + maxslp(C2E2C[2]) * c_bln_avg_3 + return z_maxslp_avg + + +@field_operator +def _compute_z_maxhgtd_avg( + maxhgtd: Field[[CellDim, KDim], wpfloat], + c_bln_avg_0: Field[[CellDim], wpfloat], + c_bln_avg_1: Field[[CellDim], wpfloat], + c_bln_avg_2: Field[[CellDim], wpfloat], + c_bln_avg_3: Field[[CellDim], wpfloat] +) -> Field[[CellDim, KDim], wpfloat]: + z_maxhgtd_avg = maxhgtd * c_bln_avg_0 + \ + maxhgtd(C2E2C[0]) * c_bln_avg_1 + \ + maxhgtd(C2E2C[1]) * c_bln_avg_2 + \ + maxhgtd(C2E2C[2]) * c_bln_avg_3 + return z_maxhgtd_avg + + +@program +def compute_z_maxslp_avg_z_maxhgtd_avg( + maxslp: Field[[CellDim, KDim], wpfloat], + maxhgtd: Field[[CellDim, KDim], wpfloat], + c_bln_avg_0: Field[[CellDim], wpfloat], + c_bln_avg_1: Field[[CellDim], wpfloat], + c_bln_avg_2: Field[[CellDim], wpfloat], + c_bln_avg_3: Field[[CellDim], wpfloat], + z_maxslp_avg: Field[[CellDim, KDim], wpfloat], + z_maxhgtd_avg: Field[[CellDim, KDim], wpfloat], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32 +): + _compute_z_maxslp_avg( + maxslp=maxslp, + c_bln_avg_0=c_bln_avg_0, + c_bln_avg_1=c_bln_avg_1, + c_bln_avg_2=c_bln_avg_2, + c_bln_avg_3=c_bln_avg_3, + out=z_maxslp_avg, + domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, + ) + + _compute_z_maxhgtd_avg( + maxhgtd=maxhgtd, + c_bln_avg_0=c_bln_avg_0, + c_bln_avg_1=c_bln_avg_1, + c_bln_avg_2=c_bln_avg_2, + c_bln_avg_3=c_bln_avg_3, + out=z_maxhgtd_avg, + domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, + ) + + +@field_operator +def _compute_max_nbhgt( + z_mc_nlev: Field[[CellDim], wpfloat], +) -> Field[[CellDim], wpfloat]: + max_nbhgt_0_1 = maximum(z_mc_nlev(C2E2C[0]), z_mc_nlev(C2E2C[1])) + max_nbhgt = maximum(max_nbhgt_0_1, z_mc_nlev(C2E2C[2])) + return max_nbhgt diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py index c049f1234..239ccea3c 100644 --- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -484,6 +484,9 @@ def construct_cell_geometry(self) -> CellParams: class InterpolationSavepoint(IconSavepoint): + def c_bln_avg(self): + return self._get_field("c_bln_avg", CellDim, KDim) + def c_intp(self): return self._get_field("c_intp", VertexDim, V2CDim) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 38b1ad450..88dec1d6e 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -59,7 +59,7 @@ compute_vwind_expl_wgt, compute_vwind_impl_wgt, compute_wgtfac_e, - compute_z_mc, + compute_z_mc, _compute_max_nbhgt, compute_z_maxslp_avg_z_maxhgtd_avg, ) from icon4py.model.common.metrics.stencils.compute_zdiff_gradp_dsl import compute_zdiff_gradp_dsl from icon4py.model.common.test_utils.datatest_utils import ( @@ -890,3 +890,119 @@ def test_compute_hmask_dd3d(metrics_savepoint, icon_grid, grid_savepoint, backen ) dallclose(hmask_dd3d_full.asnumpy(), hmask_dd3d_ref.asnumpy()) + +@pytest.mark.datatest +def test_compute_mask_hdiff(metrics_savepoint, interpolation_savepoint, icon_grid, grid_savepoint, backend): + backend = None + mask_hdiff = zero_field(icon_grid, CellDim, KDim, dtype=bool).asnumpy() + nlev = icon_grid.num_levels + mask_hdiff_ref = metrics_savepoint.mask_hdiff() + zd_vertoffset_ref = metrics_savepoint.zd_vertoffset() + # TODO: inquire to Magda as to why 3316 works instead of cell_nudging + # cell_nudging = icon_grid.get_end_index(CellDim, HorizontalMarkerIndex.nudging(CellDim)) + cell_nudging = 3316 + cell_lateral = icon_grid.get_start_index( + CellDim, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, + ) + z_maxslp_avg = zero_field(icon_grid, CellDim, KDim) + z_maxhgtd_avg = zero_field(icon_grid, CellDim, KDim) + c_bln_avg = interpolation_savepoint.c_bln_avg() + thslp_zdiffu = 0.02 # TODO: import from Fortran, note: same variable in diffusion has different value + thhgtd_zdiffu = 125 # TODO: import from Fortran, note: same variable in diffusion has different value + c_owner_mask = grid_savepoint.c_owner_mask() + ji = -1 + i_masklist = [0] * icon_grid.num_cells + k_start = [None] * icon_grid.num_cells + k_end = [None] * icon_grid.num_cells + i_indlist = [0] * icon_grid.num_cells + maxslp = zero_field(icon_grid, CellDim, KDim) + maxhgtd = zero_field(icon_grid, CellDim, KDim) + max_nbhgt = zero_field(icon_grid, CellDim) + + _compute_maxslp_maxhgtd( + ddxn_z_full=metrics_savepoint.ddxn_z_full(), + dual_edge_length=grid_savepoint.dual_edge_length(), + out=(maxslp, maxhgtd), + domain={CellDim:(cell_lateral, icon_grid.num_cells), KDim: (int32(0), nlev)}, + offset_provider={"C2E": icon_grid.get_offset_provider("C2E")}, + ) + + z_ifc = metrics_savepoint.z_ifc() + z_mc = zero_field(icon_grid, CellDim, KDim, extend={KDim: 1}) + compute_z_mc.with_backend(backend)( + z_ifc, + z_mc, + horizontal_start=int32(0), + horizontal_end=icon_grid.num_cells, + vertical_start=int32(0), + vertical_end=nlev, + offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, + ) + + compute_z_maxslp_avg_z_maxhgtd_avg.with_backend(backend)( + maxslp=maxslp, + maxhgtd=maxhgtd, + c_bln_avg_0=as_field((CellDim,), c_bln_avg.asnumpy()[:, 0]), + c_bln_avg_1=as_field((CellDim,), c_bln_avg.asnumpy()[:, 1]), + c_bln_avg_2=as_field((CellDim,), c_bln_avg.asnumpy()[:, 2]), + c_bln_avg_3=as_field((CellDim,), c_bln_avg.asnumpy()[:, 3]), + z_maxslp_avg=z_maxslp_avg, + z_maxhgtd_avg=z_maxhgtd_avg, + horizontal_start=cell_lateral, + horizontal_end=int32(icon_grid.num_cells), + vertical_start=0, + vertical_end=nlev, + offset_provider={"C2E2C": icon_grid.get_offset_provider("C2E2C")} + ) + + for jc in range(cell_nudging, icon_grid.num_cells): + if ((z_maxslp_avg[jc, nlev - 1] >= thslp_zdiffu or z_maxhgtd_avg[jc, nlev - 1] >= thhgtd_zdiffu) and c_owner_mask[jc]): + ji += 1 + i_masklist[jc] = 1 + i_indlist[ji] = jc + + i_listdim = ji + i_listreduce = 0 + _compute_max_nbhgt( + z_mc_nlev=as_field((CellDim,), z_mc.asnumpy()[:, nlev - 1]), + out=max_nbhgt, + domain={CellDim: (int32(1), int32(icon_grid.num_cells))}, + offset_provider={"C2E2C": icon_grid.get_offset_provider("C2E2C")}, + ) + + for jc in range(cell_nudging, icon_grid.num_cells): + if i_masklist[jc] == 1: + + for jk in reversed(range(nlev)): + if z_mc.asnumpy()[jc, jk] >= max_nbhgt.asnumpy()[jc]: + k_end[jc] = jk + 1 + break + + for jk in range(nlev): + if z_maxslp_avg.asnumpy()[jc, jk] >= thslp_zdiffu or z_maxhgtd_avg.asnumpy()[jc, jk] >= thhgtd_zdiffu: + k_start[jc] = jk + break + + if k_start[jc] is not None and k_end[jc] is not None and k_start[jc] > k_end[jc]: + i_masklist[jc] = 0 + i_listreduce += 1 + k_start[jc] = nlev - 1 + + + i_listdim = i_listdim - i_listreduce + + ji = -1 + for jc in range(cell_nudging, icon_grid.num_cells): + if i_masklist[jc] == 1: + ji += 1 + i_indlist[ji] = jc + + for ji in range(i_listdim): + jc = i_indlist[ji] + if k_start[jc] is not None and k_end[jc] is not None: + for jk in range(k_start[jc], k_end[jc]): + mask_hdiff[jc, jk] = True + + assert dallclose(mask_hdiff, mask_hdiff_ref.asnumpy()) + From 455d9a5fca58dba7efe61813f5c2cdb82388787c Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 24 May 2024 09:57:24 +0200 Subject: [PATCH 48/65] zd_diffcoef_dsl implementation and partial for zd_vertoffset_dsl --- .../model/common/metrics/metric_fields.py | 1 + .../common/test_utils/serialbox_utils.py | 3 +- .../tests/metric_tests/test_metric_fields.py | 48 ++++++++++++++++++- 3 files changed, 49 insertions(+), 3 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index f0f5df0ee..d7d7b7420 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -1103,3 +1103,4 @@ def _compute_max_nbhgt( max_nbhgt_0_1 = maximum(z_mc_nlev(C2E2C[0]), z_mc_nlev(C2E2C[1])) max_nbhgt = maximum(max_nbhgt_0_1, z_mc_nlev(C2E2C[2])) return max_nbhgt + diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py index 239ccea3c..d992a448b 100644 --- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -708,7 +708,8 @@ def _linearize_first_2dims( @IconSavepoint.optionally_registered() def zd_vertoffset(self): - return self._read_and_reorder_sparse_field("zd_vertoffset") + # return self._read_and_reorder_sparse_field("zd_vertoffset") + return self._get_field("zd_vertoffset", CellDim, C2E2CDim, KDim) def zd_vertidx(self): return np.squeeze(self.serializer.read("zd_vertidx", self.savepoint)) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 88dec1d6e..2a637de32 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -26,7 +26,7 @@ EdgeDim, KDim, V2CDim, - VertexDim, + VertexDim, C2E2CDim, ) from icon4py.model.common.grid.horizontal import ( HorizontalMarkerIndex, @@ -59,7 +59,7 @@ compute_vwind_expl_wgt, compute_vwind_impl_wgt, compute_wgtfac_e, - compute_z_mc, _compute_max_nbhgt, compute_z_maxslp_avg_z_maxhgtd_avg, + compute_z_mc, _compute_max_nbhgt, compute_z_maxslp_avg_z_maxhgtd_avg, _compute_maxslp_maxhgtd, ) from icon4py.model.common.metrics.stencils.compute_zdiff_gradp_dsl import compute_zdiff_gradp_dsl from icon4py.model.common.test_utils.datatest_utils import ( @@ -895,6 +895,7 @@ def test_compute_hmask_dd3d(metrics_savepoint, icon_grid, grid_savepoint, backen def test_compute_mask_hdiff(metrics_savepoint, interpolation_savepoint, icon_grid, grid_savepoint, backend): backend = None mask_hdiff = zero_field(icon_grid, CellDim, KDim, dtype=bool).asnumpy() + zd_vertoffset_dsl = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy() nlev = icon_grid.num_levels mask_hdiff_ref = metrics_savepoint.mask_hdiff() zd_vertoffset_ref = metrics_savepoint.zd_vertoffset() @@ -907,6 +908,7 @@ def test_compute_mask_hdiff(metrics_savepoint, interpolation_savepoint, icon_gri ) z_maxslp_avg = zero_field(icon_grid, CellDim, KDim) z_maxhgtd_avg = zero_field(icon_grid, CellDim, KDim) + zd_diffcoef_dsl = zero_field(icon_grid, CellDim, KDim).asnumpy() c_bln_avg = interpolation_savepoint.c_bln_avg() thslp_zdiffu = 0.02 # TODO: import from Fortran, note: same variable in diffusion has different value thhgtd_zdiffu = 125 # TODO: import from Fortran, note: same variable in diffusion has different value @@ -998,11 +1000,53 @@ def test_compute_mask_hdiff(metrics_savepoint, interpolation_savepoint, icon_gri ji += 1 i_indlist[ji] = jc + c2e2c = icon_grid.connectivities[C2E2CDim] + nbidx = constant_field(icon_grid, nlev, CellDim, C2E2CDim, KDim, dtype=int).asnumpy() + # for jc in range(cell_nudging, icon_grid.num_cells): + # if i_masklist[jc] == 1: + # jk_start = nlev - 1 + # if k_start[jc] is not None and k_end[jc] is not None: + # for jk in reversed(range(k_start[jc], k_end[jc])): + # for jk1 in reversed(range(jk_start)): + # # if jc == 3320 and jk == 56: + # # a = 1 + # if z_mc.asnumpy()[jc, jk] <= z_mc.asnumpy()[c2e2c][jc, 0, jk1] and z_mc.asnumpy()[jc, jk] >= z_mc.asnumpy()[c2e2c][jc, 0, jk1+1]: + # nbidx[jc, 0, jk] = jk1 + # jk_start = jk1 + # break + # + # jk_start = nlev - 1 + # if k_start[jc] is not None and k_end[jc] is not None: + # for jk in reversed(range(k_start[jc], k_end[jc])): + # for jk1 in reversed(range(jk_start)): + # if z_mc.asnumpy()[jc, jk] <= z_mc.asnumpy()[c2e2c][jc, 1, jk1] and z_mc.asnumpy()[jc, jk] >= z_mc.asnumpy()[c2e2c][jc, 1, jk1+1]: + # nbidx[jc, 1, jk] = jk1 + # jk_start = jk1 + # break + # + # jk_start = nlev - 1 + # if k_start[jc] is not None and k_end[jc] is not None: + # for jk in reversed(range(k_start[jc], k_end[jc])): + # for jk1 in reversed(range(jk_start)): + # if z_mc.asnumpy()[jc, jk] <= z_mc.asnumpy()[c2e2c][jc, 2, jk1] and z_mc.asnumpy()[jc, jk] >= z_mc.asnumpy()[c2e2c][jc, 2, jk1+1]: + # nbidx[jc, 2, jk] = jk1 + # jk_start = jk1 + # break + for ji in range(i_listdim): jc = i_indlist[ji] if k_start[jc] is not None and k_end[jc] is not None: for jk in range(k_start[jc], k_end[jc]): mask_hdiff[jc, jk] = True + zd_diffcoef_dsl_var = max(0.0, math.sqrt(max(0.0, z_maxslp_avg.asnumpy()[jc, jk] - thslp_zdiffu)) / 250.0, + 2.e-4 * math.sqrt(max(0.0, z_maxhgtd_avg.asnumpy()[jc, jk] - thhgtd_zdiffu))) + zd_diffcoef_dsl_var = min(1.0 / 500.0, zd_diffcoef_dsl_var) + zd_diffcoef_dsl[jc, jk] = zd_diffcoef_dsl_var + # zd_vertoffset_dsl[jc, 0, jk] = nbidx[jc, 0, jk] - jk + # zd_vertoffset_dsl[jc, 1, jk] = nbidx[jc, 1, jk] - jk + # zd_vertoffset_dsl[jc, 2, jk] = nbidx[jc, 2, jk] - jk assert dallclose(mask_hdiff, mask_hdiff_ref.asnumpy()) + assert dallclose(zd_diffcoef_dsl, metrics_savepoint.zd_diffcoef().asnumpy()) + # assert dallclose(zd_vertoffset_dsl[:, 0, :], zd_vertoffset_ref.asnumpy()[:, 0, :]) From cb990b888bf477016b798b5b317800a56bf8860c Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 28 May 2024 15:10:53 +0200 Subject: [PATCH 49/65] mask_hdiff, zd_diffcoef_dsl, zd_vertoffset_dsl implementation --- .../model/common/metrics/metric_fields.py | 34 ++-- .../stencils/compute_diffusion_metrics.py | 172 +++++++++++++++++ .../tests/metric_tests/test_metric_fields.py | 177 ++++++++---------- 3 files changed, 273 insertions(+), 110 deletions(-) create mode 100644 model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index d7d7b7420..9c1fbb72e 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -30,13 +30,14 @@ from icon4py.model.common.dimension import ( C2E, + C2E2C, E2C, CellDim, E2CDim, EdgeDim, KDim, Koff, - VertexDim, C2E2C, + VertexDim, ) from icon4py.model.common.interpolation.stencils.cell_2_edge_interpolation import ( _cell_2_edge_interpolation, @@ -1012,6 +1013,7 @@ def compute_hmask_dd3d( domain={EdgeDim: (horizontal_start, horizontal_end)}, ) + @field_operator def _compute_mask_hdiff() -> Field[[CellDim, KDim], bool]: return broadcast(True, (CellDim, KDim)) @@ -1027,21 +1029,24 @@ def compute_mask_hdiff( ): _compute_mask_hdiff( out=mask_hdiff, - domain={CellDim: (horizontal_start, horizontal_end), KDim:(vertical_start, vertical_end)} + domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, ) + @field_operator def _compute_z_maxslp_avg( maxslp: Field[[CellDim, KDim], wpfloat], c_bln_avg_0: Field[[CellDim], wpfloat], c_bln_avg_1: Field[[CellDim], wpfloat], c_bln_avg_2: Field[[CellDim], wpfloat], - c_bln_avg_3: Field[[CellDim], wpfloat] + c_bln_avg_3: Field[[CellDim], wpfloat], ) -> Field[[CellDim, KDim], wpfloat]: - z_maxslp_avg = maxslp * c_bln_avg_0 + \ - maxslp(C2E2C[0]) * c_bln_avg_1 + \ - maxslp(C2E2C[1]) * c_bln_avg_2 + \ - maxslp(C2E2C[2]) * c_bln_avg_3 + z_maxslp_avg = ( + maxslp * c_bln_avg_0 + + maxslp(C2E2C[0]) * c_bln_avg_1 + + maxslp(C2E2C[1]) * c_bln_avg_2 + + maxslp(C2E2C[2]) * c_bln_avg_3 + ) return z_maxslp_avg @@ -1051,12 +1056,14 @@ def _compute_z_maxhgtd_avg( c_bln_avg_0: Field[[CellDim], wpfloat], c_bln_avg_1: Field[[CellDim], wpfloat], c_bln_avg_2: Field[[CellDim], wpfloat], - c_bln_avg_3: Field[[CellDim], wpfloat] + c_bln_avg_3: Field[[CellDim], wpfloat], ) -> Field[[CellDim, KDim], wpfloat]: - z_maxhgtd_avg = maxhgtd * c_bln_avg_0 + \ - maxhgtd(C2E2C[0]) * c_bln_avg_1 + \ - maxhgtd(C2E2C[1]) * c_bln_avg_2 + \ - maxhgtd(C2E2C[2]) * c_bln_avg_3 + z_maxhgtd_avg = ( + maxhgtd * c_bln_avg_0 + + maxhgtd(C2E2C[0]) * c_bln_avg_1 + + maxhgtd(C2E2C[1]) * c_bln_avg_2 + + maxhgtd(C2E2C[2]) * c_bln_avg_3 + ) return z_maxhgtd_avg @@ -1073,7 +1080,7 @@ def compute_z_maxslp_avg_z_maxhgtd_avg( horizontal_start: int32, horizontal_end: int32, vertical_start: int32, - vertical_end: int32 + vertical_end: int32, ): _compute_z_maxslp_avg( maxslp=maxslp, @@ -1103,4 +1110,3 @@ def _compute_max_nbhgt( max_nbhgt_0_1 = maximum(z_mc_nlev(C2E2C[0]), z_mc_nlev(C2E2C[1])) max_nbhgt = maximum(max_nbhgt_0_1, z_mc_nlev(C2E2C[2])) return max_nbhgt - diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py new file mode 100644 index 000000000..b6c10acf8 --- /dev/null +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py @@ -0,0 +1,172 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import math + +import numpy as np + + +def _compute_nbidx( + k_start: list, + k_end: list, + z_mc: np.array, + z_mc_off: np.array, + nbidx: np.array, + jc: int, + ind: int, + nlev: int, +) -> np.array: + jk_start = nlev - 2 + for jk in reversed(range(k_start[jc], k_end[jc])): + for jk1 in reversed(range(jk_start + 1)): + if ( + z_mc[jc, jk] <= z_mc_off[jc, ind, jk1] + and z_mc[jc, jk] >= z_mc_off[jc, ind, jk1 + 1] + ): + nbidx[jc, ind, jk] = jk1 + jk_start = jk1 + break + + return nbidx[jc, ind, :] + + +def _compute_i_params( + k_start: list, + k_end: list, + z_maxslp_avg: np.array, + z_maxhgtd_avg: np.array, + c_owner_mask: np.array, + thslp_zdiffu: float, + thhgtd_zdiffu: float, + cell_nudging: int, + n_cells: int, + nlev: int, +) -> tuple[list, int, int]: + i_indlist = [0] * n_cells + i_listreduce = 0 + ji = -1 + ji_ind = -1 + + for jc in range(cell_nudging, n_cells): + if ( + z_maxslp_avg[jc, nlev - 1] >= thslp_zdiffu + or z_maxhgtd_avg[jc, nlev - 1] >= thhgtd_zdiffu + ) and c_owner_mask[jc]: + ji += 1 + i_indlist[ji] = jc + + if k_start[jc] is not None and k_end[jc] is not None and k_start[jc] > k_end[jc]: + i_listreduce += 1 + else: + ji_ind += 1 + i_indlist[ji_ind] = jc + + return i_indlist, i_listreduce, ji + + +def _compute_k_start_end( + z_mc: np.array, + max_nbhgt: np.array, + z_maxslp_avg: np.array, + z_maxhgtd_avg: np.array, + c_owner_mask: np.array, + thslp_zdiffu: float, + thhgtd_zdiffu: float, + cell_nudging: int, + n_cells: int, + nlev: int, +) -> tuple[list, list]: + k_start = [None] * n_cells + k_end = [None] * n_cells + for jc in range(cell_nudging, n_cells): + if ( + z_maxslp_avg[jc, nlev - 1] >= thslp_zdiffu + or z_maxhgtd_avg[jc, nlev - 1] >= thhgtd_zdiffu + ) and c_owner_mask[jc]: + for jk in reversed(range(nlev)): + if z_mc[jc, jk] >= max_nbhgt[jc]: + k_end[jc] = jk + 1 + break + + for jk in range(nlev): + if z_maxslp_avg[jc, jk] >= thslp_zdiffu or z_maxhgtd_avg[jc, jk] >= thhgtd_zdiffu: + k_start[jc] = jk + break + + if k_start[jc] is not None and k_end[jc] is not None and k_start[jc] > k_end[jc]: + k_start[jc] = nlev - 1 + + return k_start, k_end + + +def _compute_zd_vertoffset_dsl( + k_start: list, + k_end: list, + z_mc: np.array, + z_mc_off: np.array, + nbidx: np.array, + i_indlist: list, + i_listdim: int, + zd_vertoffset_dsl: np.array, + nlev: int, +) -> np.array: + for ji in range(i_listdim): + jc = i_indlist[ji] + if k_start[jc] is not None and k_end[jc] is not None: + nbidx[jc, 0, :] = _compute_nbidx(k_start, k_end, z_mc, z_mc_off, nbidx, jc, 0, nlev) + nbidx[jc, 1, :] = _compute_nbidx(k_start, k_end, z_mc, z_mc_off, nbidx, jc, 1, nlev) + nbidx[jc, 2, :] = _compute_nbidx(k_start, k_end, z_mc, z_mc_off, nbidx, jc, 2, nlev) + for jk in range(k_start[jc], k_end[jc]): + zd_vertoffset_dsl[jc, 0, jk] = nbidx[jc, 0, jk] - jk + zd_vertoffset_dsl[jc, 1, jk] = nbidx[jc, 1, jk] - jk + zd_vertoffset_dsl[jc, 2, jk] = nbidx[jc, 2, jk] - jk + + return zd_vertoffset_dsl + + +def _compute_mask_hdiff( + mask_hdiff: np.array, k_start: list, k_end: list, i_indlist: list, i_listdim: int +) -> np.array: + for ji in range(i_listdim): + jc = i_indlist[ji] + if k_start[jc] is not None and k_end[jc] is not None: + for jk in range(k_start[jc], k_end[jc]): + mask_hdiff[jc, jk] = True + + return mask_hdiff + + +def _compute_zd_diffcoef_dsl( + z_maxslp_avg: np.array, + z_maxhgtd_avg: np.array, + k_start: list, + k_end: list, + i_indlist: list, + i_listdim: int, + zd_diffcoef_dsl: np.array, + thslp_zdiffu: float, + thhgtd_zdiffu: float, +) -> np.array: + for ji in range(i_listdim): + jc = i_indlist[ji] + if k_start[jc] is not None and k_end[jc] is not None: + for jk in range(k_start[jc], k_end[jc]): + zd_diffcoef_dsl_var = max( + 0.0, + math.sqrt(max(0.0, z_maxslp_avg[jc, jk] - thslp_zdiffu)) / 250.0, + 2.0e-4 * math.sqrt(max(0.0, z_maxhgtd_avg[jc, jk] - thhgtd_zdiffu)), + ) + zd_diffcoef_dsl_var = min(0.002, zd_diffcoef_dsl_var) + zd_diffcoef_dsl[jc, jk] = zd_diffcoef_dsl_var + + return zd_diffcoef_dsl diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 2a637de32..960e56aff 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -20,13 +20,14 @@ from icon4py.model.common import constants from icon4py.model.common.dimension import ( + C2E2CDim, CellDim, E2CDim, ECDim, EdgeDim, KDim, V2CDim, - VertexDim, C2E2CDim, + VertexDim, ) from icon4py.model.common.grid.horizontal import ( HorizontalMarkerIndex, @@ -39,6 +40,8 @@ from icon4py.model.common.math.helpers import average_cell_kdim_level_up from icon4py.model.common.metrics.metric_fields import ( _compute_flat_idx, + _compute_max_nbhgt, + _compute_maxslp_maxhgtd, _compute_pg_edgeidx_vertidx, _compute_z_aux2, compute_bdy_halo_c, @@ -59,7 +62,15 @@ compute_vwind_expl_wgt, compute_vwind_impl_wgt, compute_wgtfac_e, - compute_z_mc, _compute_max_nbhgt, compute_z_maxslp_avg_z_maxhgtd_avg, _compute_maxslp_maxhgtd, + compute_z_maxslp_avg_z_maxhgtd_avg, + compute_z_mc, +) +from icon4py.model.common.metrics.stencils.compute_diffusion_metrics import ( + _compute_i_params, + _compute_k_start_end, + _compute_mask_hdiff, + _compute_zd_diffcoef_dsl, + _compute_zd_vertoffset_dsl, ) from icon4py.model.common.metrics.stencils.compute_zdiff_gradp_dsl import compute_zdiff_gradp_dsl from icon4py.model.common.test_utils.datatest_utils import ( @@ -891,8 +902,11 @@ def test_compute_hmask_dd3d(metrics_savepoint, icon_grid, grid_savepoint, backen dallclose(hmask_dd3d_full.asnumpy(), hmask_dd3d_ref.asnumpy()) + @pytest.mark.datatest -def test_compute_mask_hdiff(metrics_savepoint, interpolation_savepoint, icon_grid, grid_savepoint, backend): +def test_compute_diffusion_metrics( + metrics_savepoint, interpolation_savepoint, icon_grid, grid_savepoint, backend +): backend = None mask_hdiff = zero_field(icon_grid, CellDim, KDim, dtype=bool).asnumpy() zd_vertoffset_dsl = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy() @@ -910,23 +924,24 @@ def test_compute_mask_hdiff(metrics_savepoint, interpolation_savepoint, icon_gri z_maxhgtd_avg = zero_field(icon_grid, CellDim, KDim) zd_diffcoef_dsl = zero_field(icon_grid, CellDim, KDim).asnumpy() c_bln_avg = interpolation_savepoint.c_bln_avg() - thslp_zdiffu = 0.02 # TODO: import from Fortran, note: same variable in diffusion has different value - thhgtd_zdiffu = 125 # TODO: import from Fortran, note: same variable in diffusion has different value - c_owner_mask = grid_savepoint.c_owner_mask() - ji = -1 - i_masklist = [0] * icon_grid.num_cells - k_start = [None] * icon_grid.num_cells - k_end = [None] * icon_grid.num_cells - i_indlist = [0] * icon_grid.num_cells + thslp_zdiffu = ( + 0.02 # TODO: import from Fortran, note: same variable in diffusion has different value + ) + thhgtd_zdiffu = ( + 125 # TODO: import from Fortran, note: same variable in diffusion has different value + ) + c_owner_mask = grid_savepoint.c_owner_mask().asnumpy() maxslp = zero_field(icon_grid, CellDim, KDim) maxhgtd = zero_field(icon_grid, CellDim, KDim) max_nbhgt = zero_field(icon_grid, CellDim) + c2e2c = icon_grid.connectivities[C2E2CDim] + nbidx = constant_field(icon_grid, 1, CellDim, C2E2CDim, KDim, dtype=int).asnumpy() _compute_maxslp_maxhgtd( ddxn_z_full=metrics_savepoint.ddxn_z_full(), dual_edge_length=grid_savepoint.dual_edge_length(), out=(maxslp, maxhgtd), - domain={CellDim:(cell_lateral, icon_grid.num_cells), KDim: (int32(0), nlev)}, + domain={CellDim: (cell_lateral, icon_grid.num_cells), KDim: (int32(0), nlev)}, offset_provider={"C2E": icon_grid.get_offset_provider("C2E")}, ) @@ -941,6 +956,7 @@ def test_compute_mask_hdiff(metrics_savepoint, interpolation_savepoint, icon_gri vertical_end=nlev, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) + z_mc_off = z_mc.asnumpy()[c2e2c] compute_z_maxslp_avg_z_maxhgtd_avg.with_backend(backend)( maxslp=maxslp, @@ -955,17 +971,9 @@ def test_compute_mask_hdiff(metrics_savepoint, interpolation_savepoint, icon_gri horizontal_end=int32(icon_grid.num_cells), vertical_start=0, vertical_end=nlev, - offset_provider={"C2E2C": icon_grid.get_offset_provider("C2E2C")} + offset_provider={"C2E2C": icon_grid.get_offset_provider("C2E2C")}, ) - for jc in range(cell_nudging, icon_grid.num_cells): - if ((z_maxslp_avg[jc, nlev - 1] >= thslp_zdiffu or z_maxhgtd_avg[jc, nlev - 1] >= thhgtd_zdiffu) and c_owner_mask[jc]): - ji += 1 - i_masklist[jc] = 1 - i_indlist[ji] = jc - - i_listdim = ji - i_listreduce = 0 _compute_max_nbhgt( z_mc_nlev=as_field((CellDim,), z_mc.asnumpy()[:, nlev - 1]), out=max_nbhgt, @@ -973,80 +981,57 @@ def test_compute_mask_hdiff(metrics_savepoint, interpolation_savepoint, icon_gri offset_provider={"C2E2C": icon_grid.get_offset_provider("C2E2C")}, ) - for jc in range(cell_nudging, icon_grid.num_cells): - if i_masklist[jc] == 1: - - for jk in reversed(range(nlev)): - if z_mc.asnumpy()[jc, jk] >= max_nbhgt.asnumpy()[jc]: - k_end[jc] = jk + 1 - break - - for jk in range(nlev): - if z_maxslp_avg.asnumpy()[jc, jk] >= thslp_zdiffu or z_maxhgtd_avg.asnumpy()[jc, jk] >= thhgtd_zdiffu: - k_start[jc] = jk - break - - if k_start[jc] is not None and k_end[jc] is not None and k_start[jc] > k_end[jc]: - i_masklist[jc] = 0 - i_listreduce += 1 - k_start[jc] = nlev - 1 - - - i_listdim = i_listdim - i_listreduce - - ji = -1 - for jc in range(cell_nudging, icon_grid.num_cells): - if i_masklist[jc] == 1: - ji += 1 - i_indlist[ji] = jc - - c2e2c = icon_grid.connectivities[C2E2CDim] - nbidx = constant_field(icon_grid, nlev, CellDim, C2E2CDim, KDim, dtype=int).asnumpy() - # for jc in range(cell_nudging, icon_grid.num_cells): - # if i_masklist[jc] == 1: - # jk_start = nlev - 1 - # if k_start[jc] is not None and k_end[jc] is not None: - # for jk in reversed(range(k_start[jc], k_end[jc])): - # for jk1 in reversed(range(jk_start)): - # # if jc == 3320 and jk == 56: - # # a = 1 - # if z_mc.asnumpy()[jc, jk] <= z_mc.asnumpy()[c2e2c][jc, 0, jk1] and z_mc.asnumpy()[jc, jk] >= z_mc.asnumpy()[c2e2c][jc, 0, jk1+1]: - # nbidx[jc, 0, jk] = jk1 - # jk_start = jk1 - # break - # - # jk_start = nlev - 1 - # if k_start[jc] is not None and k_end[jc] is not None: - # for jk in reversed(range(k_start[jc], k_end[jc])): - # for jk1 in reversed(range(jk_start)): - # if z_mc.asnumpy()[jc, jk] <= z_mc.asnumpy()[c2e2c][jc, 1, jk1] and z_mc.asnumpy()[jc, jk] >= z_mc.asnumpy()[c2e2c][jc, 1, jk1+1]: - # nbidx[jc, 1, jk] = jk1 - # jk_start = jk1 - # break - # - # jk_start = nlev - 1 - # if k_start[jc] is not None and k_end[jc] is not None: - # for jk in reversed(range(k_start[jc], k_end[jc])): - # for jk1 in reversed(range(jk_start)): - # if z_mc.asnumpy()[jc, jk] <= z_mc.asnumpy()[c2e2c][jc, 2, jk1] and z_mc.asnumpy()[jc, jk] >= z_mc.asnumpy()[c2e2c][jc, 2, jk1+1]: - # nbidx[jc, 2, jk] = jk1 - # jk_start = jk1 - # break - - for ji in range(i_listdim): - jc = i_indlist[ji] - if k_start[jc] is not None and k_end[jc] is not None: - for jk in range(k_start[jc], k_end[jc]): - mask_hdiff[jc, jk] = True - zd_diffcoef_dsl_var = max(0.0, math.sqrt(max(0.0, z_maxslp_avg.asnumpy()[jc, jk] - thslp_zdiffu)) / 250.0, - 2.e-4 * math.sqrt(max(0.0, z_maxhgtd_avg.asnumpy()[jc, jk] - thhgtd_zdiffu))) - zd_diffcoef_dsl_var = min(1.0 / 500.0, zd_diffcoef_dsl_var) - zd_diffcoef_dsl[jc, jk] = zd_diffcoef_dsl_var - # zd_vertoffset_dsl[jc, 0, jk] = nbidx[jc, 0, jk] - jk - # zd_vertoffset_dsl[jc, 1, jk] = nbidx[jc, 1, jk] - jk - # zd_vertoffset_dsl[jc, 2, jk] = nbidx[jc, 2, jk] - jk + k_start, k_end = _compute_k_start_end( + z_mc.asnumpy(), + max_nbhgt.asnumpy(), + z_maxslp_avg.asnumpy(), + z_maxhgtd_avg.asnumpy(), + c_owner_mask, + thslp_zdiffu, + thhgtd_zdiffu, + cell_nudging, + icon_grid.num_cells, + nlev, + ) + + i_indlist, i_listreduce, ji = _compute_i_params( + k_start, + k_end, + z_maxslp_avg.asnumpy(), + z_maxhgtd_avg.asnumpy(), + c_owner_mask, + thslp_zdiffu, + thhgtd_zdiffu, + cell_nudging, + icon_grid.num_cells, + nlev, + ) + + i_listdim = ji - i_listreduce + mask_hdiff = _compute_mask_hdiff(mask_hdiff, k_start, k_end, i_indlist, i_listdim) + zd_diffcoef_dsl = _compute_zd_diffcoef_dsl( + z_maxslp_avg.asnumpy(), + z_maxhgtd_avg.asnumpy(), + k_start, + k_end, + i_indlist, + i_listdim, + zd_diffcoef_dsl, + thslp_zdiffu, + thhgtd_zdiffu, + ) + zd_vertoffset_dsl = _compute_zd_vertoffset_dsl( + k_start, + k_end, + z_mc.asnumpy(), + z_mc_off, + nbidx, + i_indlist, + i_listdim, + zd_vertoffset_dsl, + nlev, + ) assert dallclose(mask_hdiff, mask_hdiff_ref.asnumpy()) - assert dallclose(zd_diffcoef_dsl, metrics_savepoint.zd_diffcoef().asnumpy()) - # assert dallclose(zd_vertoffset_dsl[:, 0, :], zd_vertoffset_ref.asnumpy()[:, 0, :]) - + assert dallclose(zd_diffcoef_dsl, metrics_savepoint.zd_diffcoef().asnumpy(), rtol=1.0e-11) + assert dallclose(zd_vertoffset_dsl, zd_vertoffset_ref.asnumpy()) From c1a63432bf1f035b9ee48296bf004e898afdf6c9 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Wed, 29 May 2024 15:32:39 +0200 Subject: [PATCH 50/65] finishes --- .../stencils/compute_diffusion_metrics.py | 90 +++++++----- .../common/test_utils/serialbox_utils.py | 8 +- .../tests/metric_tests/test_metric_fields.py | 135 +++++++++--------- 3 files changed, 124 insertions(+), 109 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py index b6c10acf8..a4f80bc94 100644 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py @@ -40,6 +40,32 @@ def _compute_nbidx( return nbidx[jc, ind, :] +def _compute_z_vintcoeff( + k_start: list, + k_end: list, + z_mc: np.array, + z_mc_off: np.array, + z_vintcoeff: np.array, + jc: int, + ind: int, + nlev: int, +) -> np.array: + jk_start = nlev - 2 + for jk in reversed(range(k_start[jc], k_end[jc])): + for jk1 in reversed(range(jk_start + 1)): + if ( + z_mc[jc, jk] <= z_mc_off[jc, ind, jk1] + and z_mc[jc, jk] >= z_mc_off[jc, ind, jk1 + 1] + ): + z_vintcoeff[jc, ind, jk] = (z_mc[jc, jk] - z_mc_off[jc, ind, jk1 + 1]) / ( + z_mc_off[jc, ind, jk1] - z_mc_off[jc, ind, jk1 + 1] + ) + jk_start = jk1 + break + + return z_vintcoeff[jc, ind, :] + + def _compute_i_params( k_start: list, k_end: list, @@ -109,15 +135,23 @@ def _compute_k_start_end( return k_start, k_end -def _compute_zd_vertoffset_dsl( - k_start: list, - k_end: list, +def compute_diffusion_metrics( z_mc: np.array, z_mc_off: np.array, - nbidx: np.array, + k_start: list, + k_end: list, i_indlist: list, i_listdim: int, + nbidx: np.array, + z_vintcoeff: np.array, + z_maxslp_avg: np.array, + z_maxhgtd_avg: np.array, + mask_hdiff: np.array, + zd_diffcoef_dsl: np.array, + zd_intcoef_dsl: np.array, zd_vertoffset_dsl: np.array, + thslp_zdiffu: float, + thhgtd_zdiffu: float, nlev: int, ) -> np.array: for ji in range(i_listdim): @@ -126,41 +160,25 @@ def _compute_zd_vertoffset_dsl( nbidx[jc, 0, :] = _compute_nbidx(k_start, k_end, z_mc, z_mc_off, nbidx, jc, 0, nlev) nbidx[jc, 1, :] = _compute_nbidx(k_start, k_end, z_mc, z_mc_off, nbidx, jc, 1, nlev) nbidx[jc, 2, :] = _compute_nbidx(k_start, k_end, z_mc, z_mc_off, nbidx, jc, 2, nlev) + + z_vintcoeff[jc, 0, :] = _compute_z_vintcoeff( + k_start, k_end, z_mc, z_mc_off, z_vintcoeff, jc, 0, nlev + ) + z_vintcoeff[jc, 1, :] = _compute_z_vintcoeff( + k_start, k_end, z_mc, z_mc_off, z_vintcoeff, jc, 1, nlev + ) + z_vintcoeff[jc, 2, :] = _compute_z_vintcoeff( + k_start, k_end, z_mc, z_mc_off, z_vintcoeff, jc, 2, nlev + ) for jk in range(k_start[jc], k_end[jc]): + zd_intcoef_dsl[jc, 0, jk] = z_vintcoeff[jc, 0, jk] + zd_intcoef_dsl[jc, 1, jk] = z_vintcoeff[jc, 1, jk] + zd_intcoef_dsl[jc, 2, jk] = z_vintcoeff[jc, 2, jk] + zd_vertoffset_dsl[jc, 0, jk] = nbidx[jc, 0, jk] - jk zd_vertoffset_dsl[jc, 1, jk] = nbidx[jc, 1, jk] - jk zd_vertoffset_dsl[jc, 2, jk] = nbidx[jc, 2, jk] - jk - return zd_vertoffset_dsl - - -def _compute_mask_hdiff( - mask_hdiff: np.array, k_start: list, k_end: list, i_indlist: list, i_listdim: int -) -> np.array: - for ji in range(i_listdim): - jc = i_indlist[ji] - if k_start[jc] is not None and k_end[jc] is not None: - for jk in range(k_start[jc], k_end[jc]): - mask_hdiff[jc, jk] = True - - return mask_hdiff - - -def _compute_zd_diffcoef_dsl( - z_maxslp_avg: np.array, - z_maxhgtd_avg: np.array, - k_start: list, - k_end: list, - i_indlist: list, - i_listdim: int, - zd_diffcoef_dsl: np.array, - thslp_zdiffu: float, - thhgtd_zdiffu: float, -) -> np.array: - for ji in range(i_listdim): - jc = i_indlist[ji] - if k_start[jc] is not None and k_end[jc] is not None: - for jk in range(k_start[jc], k_end[jc]): zd_diffcoef_dsl_var = max( 0.0, math.sqrt(max(0.0, z_maxslp_avg[jc, jk] - thslp_zdiffu)) / 250.0, @@ -169,4 +187,6 @@ def _compute_zd_diffcoef_dsl( zd_diffcoef_dsl_var = min(0.002, zd_diffcoef_dsl_var) zd_diffcoef_dsl[jc, jk] = zd_diffcoef_dsl_var - return zd_diffcoef_dsl + mask_hdiff[jc, jk] = True + + return mask_hdiff, zd_diffcoef_dsl, zd_intcoef_dsl, zd_vertoffset_dsl diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py index d992a448b..21f292bf0 100644 --- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -683,7 +683,7 @@ def wgtfacq_e_dsl( def zd_diffcoef(self): return self._get_field("zd_diffcoef", CellDim, KDim) - @IconSavepoint.optionally_registered() + @IconSavepoint.optionally_registered(CellDim, C2E2CDim, KDim) def zd_intcoef(self): return self._read_and_reorder_sparse_field("vcoef") @@ -692,6 +692,7 @@ def geopot(self): def _read_and_reorder_sparse_field(self, name: str, sparse_size=3): ser_input = np.squeeze(self.serializer.read(name, self.savepoint))[:, :, :] + ser_input = self._reduce_to_dim_size(ser_input, (CellDim, C2E2CDim, KDim)) if ser_input.shape[1] != sparse_size: ser_input = np.moveaxis(ser_input, 1, -1) @@ -706,10 +707,9 @@ def _linearize_first_2dims( assert old_shape[1] == sparse_size return as_field(target_dims, data.reshape(old_shape[0] * old_shape[1], old_shape[2])) - @IconSavepoint.optionally_registered() + @IconSavepoint.optionally_registered(CellDim, C2E2CDim, KDim) def zd_vertoffset(self): - # return self._read_and_reorder_sparse_field("zd_vertoffset") - return self._get_field("zd_vertoffset", CellDim, C2E2CDim, KDim) + return self._read_and_reorder_sparse_field("zd_vertoffset") def zd_vertidx(self): return np.squeeze(self.serializer.read("zd_vertidx", self.savepoint)) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 960e56aff..70bd4e4d9 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -21,6 +21,7 @@ from icon4py.model.common import constants from icon4py.model.common.dimension import ( C2E2CDim, + CECDim, CellDim, E2CDim, ECDim, @@ -68,9 +69,7 @@ from icon4py.model.common.metrics.stencils.compute_diffusion_metrics import ( _compute_i_params, _compute_k_start_end, - _compute_mask_hdiff, - _compute_zd_diffcoef_dsl, - _compute_zd_vertoffset_dsl, + compute_diffusion_metrics, ) from icon4py.model.common.metrics.stencils.compute_zdiff_gradp_dsl import compute_zdiff_gradp_dsl from icon4py.model.common.test_utils.datatest_utils import ( @@ -910,32 +909,26 @@ def test_compute_diffusion_metrics( backend = None mask_hdiff = zero_field(icon_grid, CellDim, KDim, dtype=bool).asnumpy() zd_vertoffset_dsl = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy() - nlev = icon_grid.num_levels - mask_hdiff_ref = metrics_savepoint.mask_hdiff() - zd_vertoffset_ref = metrics_savepoint.zd_vertoffset() - # TODO: inquire to Magda as to why 3316 works instead of cell_nudging - # cell_nudging = icon_grid.get_end_index(CellDim, HorizontalMarkerIndex.nudging(CellDim)) - cell_nudging = 3316 - cell_lateral = icon_grid.get_start_index( - CellDim, - HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, - ) + z_vintcoeff = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy() + zd_intcoef_dsl = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy() z_maxslp_avg = zero_field(icon_grid, CellDim, KDim) z_maxhgtd_avg = zero_field(icon_grid, CellDim, KDim) zd_diffcoef_dsl = zero_field(icon_grid, CellDim, KDim).asnumpy() - c_bln_avg = interpolation_savepoint.c_bln_avg() - thslp_zdiffu = ( - 0.02 # TODO: import from Fortran, note: same variable in diffusion has different value - ) - thhgtd_zdiffu = ( - 125 # TODO: import from Fortran, note: same variable in diffusion has different value - ) - c_owner_mask = grid_savepoint.c_owner_mask().asnumpy() maxslp = zero_field(icon_grid, CellDim, KDim) maxhgtd = zero_field(icon_grid, CellDim, KDim) max_nbhgt = zero_field(icon_grid, CellDim) + c2e2c = icon_grid.connectivities[C2E2CDim] nbidx = constant_field(icon_grid, 1, CellDim, C2E2CDim, KDim, dtype=int).asnumpy() + c_bln_avg = interpolation_savepoint.c_bln_avg() + thslp_zdiffu = 0.02 + thhgtd_zdiffu = 125 + cell_nudging = icon_grid.get_start_index(CellDim, HorizontalMarkerIndex.nudging(CellDim)) + cell_lateral = icon_grid.get_start_index( + CellDim, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, + ) + nlev = icon_grid.num_levels _compute_maxslp_maxhgtd( ddxn_z_full=metrics_savepoint.ddxn_z_full(), @@ -945,10 +938,9 @@ def test_compute_diffusion_metrics( offset_provider={"C2E": icon_grid.get_offset_provider("C2E")}, ) - z_ifc = metrics_savepoint.z_ifc() z_mc = zero_field(icon_grid, CellDim, KDim, extend={KDim: 1}) compute_z_mc.with_backend(backend)( - z_ifc, + metrics_savepoint.z_ifc(), z_mc, horizontal_start=int32(0), horizontal_end=icon_grid.num_cells, @@ -956,7 +948,6 @@ def test_compute_diffusion_metrics( vertical_end=nlev, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) - z_mc_off = z_mc.asnumpy()[c2e2c] compute_z_maxslp_avg_z_maxhgtd_avg.with_backend(backend)( maxslp=maxslp, @@ -982,56 +973,60 @@ def test_compute_diffusion_metrics( ) k_start, k_end = _compute_k_start_end( - z_mc.asnumpy(), - max_nbhgt.asnumpy(), - z_maxslp_avg.asnumpy(), - z_maxhgtd_avg.asnumpy(), - c_owner_mask, - thslp_zdiffu, - thhgtd_zdiffu, - cell_nudging, - icon_grid.num_cells, - nlev, + z_mc=z_mc.asnumpy(), + max_nbhgt=max_nbhgt.asnumpy(), + z_maxslp_avg=z_maxslp_avg.asnumpy(), + z_maxhgtd_avg=z_maxhgtd_avg.asnumpy(), + c_owner_mask=grid_savepoint.c_owner_mask().asnumpy(), + thslp_zdiffu=thslp_zdiffu, + thhgtd_zdiffu=thhgtd_zdiffu, + cell_nudging=cell_nudging, + n_cells=icon_grid.num_cells, + nlev=nlev, ) i_indlist, i_listreduce, ji = _compute_i_params( - k_start, - k_end, - z_maxslp_avg.asnumpy(), - z_maxhgtd_avg.asnumpy(), - c_owner_mask, - thslp_zdiffu, - thhgtd_zdiffu, - cell_nudging, - icon_grid.num_cells, - nlev, + k_start=k_start, + k_end=k_end, + z_maxslp_avg=z_maxslp_avg.asnumpy(), + z_maxhgtd_avg=z_maxhgtd_avg.asnumpy(), + c_owner_mask=grid_savepoint.c_owner_mask().asnumpy(), + thslp_zdiffu=thslp_zdiffu, + thhgtd_zdiffu=thhgtd_zdiffu, + cell_nudging=cell_nudging, + n_cells=icon_grid.num_cells, + nlev=nlev, ) i_listdim = ji - i_listreduce - mask_hdiff = _compute_mask_hdiff(mask_hdiff, k_start, k_end, i_indlist, i_listdim) - zd_diffcoef_dsl = _compute_zd_diffcoef_dsl( - z_maxslp_avg.asnumpy(), - z_maxhgtd_avg.asnumpy(), - k_start, - k_end, - i_indlist, - i_listdim, - zd_diffcoef_dsl, - thslp_zdiffu, - thhgtd_zdiffu, - ) - zd_vertoffset_dsl = _compute_zd_vertoffset_dsl( - k_start, - k_end, - z_mc.asnumpy(), - z_mc_off, - nbidx, - i_indlist, - i_listdim, - zd_vertoffset_dsl, - nlev, - ) - - assert dallclose(mask_hdiff, mask_hdiff_ref.asnumpy()) + + mask_hdiff, zd_diffcoef_dsl, zd_intcoef_dsl, zd_vertoffset_dsl = compute_diffusion_metrics( + z_mc=z_mc.asnumpy(), + z_mc_off=z_mc.asnumpy()[c2e2c], + k_start=k_start, + k_end=k_end, + i_indlist=i_indlist, + i_listdim=i_listdim, + nbidx=nbidx, + z_vintcoeff=z_vintcoeff, + z_maxslp_avg=z_maxslp_avg.asnumpy(), + z_maxhgtd_avg=z_maxhgtd_avg.asnumpy(), + mask_hdiff=mask_hdiff, + zd_diffcoef_dsl=zd_diffcoef_dsl, + zd_intcoef_dsl=zd_intcoef_dsl, + zd_vertoffset_dsl=zd_vertoffset_dsl, + thslp_zdiffu=thslp_zdiffu, + thhgtd_zdiffu=thhgtd_zdiffu, + nlev=nlev, + ) + zd_intcoef_dsl = flatten_first_two_dims( + CECDim, KDim, field=as_field((CellDim, C2E2CDim, KDim), zd_intcoef_dsl) + ) + zd_vertoffset_dsl = flatten_first_two_dims( + CECDim, KDim, field=as_field((CellDim, C2E2CDim, KDim), zd_vertoffset_dsl) + ) + + assert dallclose(mask_hdiff, metrics_savepoint.mask_hdiff().asnumpy()) assert dallclose(zd_diffcoef_dsl, metrics_savepoint.zd_diffcoef().asnumpy(), rtol=1.0e-11) - assert dallclose(zd_vertoffset_dsl, zd_vertoffset_ref.asnumpy()) + assert dallclose(zd_vertoffset_dsl.asnumpy(), metrics_savepoint.zd_vertoffset().asnumpy()) + assert dallclose(zd_intcoef_dsl.asnumpy(), metrics_savepoint.zd_intcoef()) From ef9f6627ae7b1ff0b51cbac51b5504d76010fdf4 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Wed, 29 May 2024 15:51:24 +0200 Subject: [PATCH 51/65] cleanup --- .../src/icon4py/model/common/constants.py | 13 ---- .../metrics/stencils/compute_wgtfacq_c_dsl.py | 13 ---- .../metrics/stencils/compute_wgtfacq_e_dsl.py | 61 ------------------- 3 files changed, 87 deletions(-) delete mode 100644 model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c_dsl.py delete mode 100644 model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_e_dsl.py diff --git a/model/common/src/icon4py/model/common/constants.py b/model/common/src/icon4py/model/common/constants.py index 032dad002..76bc9adf5 100644 --- a/model/common/src/icon4py/model/common/constants.py +++ b/model/common/src/icon4py/model/common/constants.py @@ -82,16 +82,3 @@ class RayleighType(IntEnum): int ] = 1 # classical Rayleigh damping, which makes use of a reference state. RAYLEIGH_KLEMP: Final[int] = 2 # Klemp (2008) type Rayleigh damping - - -# Temporal extrapolation of Exner for computation of horizontal pressure gradient -exner_expol: Final[wpfloat] = 0.333 - -# Width of lateral boundary zones (as seen from the child domain) for which tendencies are interpolated from the parent domain -grf_bdywidth_c: Final[int] = 4 -grf_bdywidth_e: Final[int] = 9 -grf_nudge_start_e: Final[int] = grf_bdywidth_e + 1 - -# Width of nudging zone in units of child cell rows; -# MUST be a multiple of 2, AND grf_nudge_start_c + grf_width_nudgezone MUST NOT EXCEED the bdy_indexing_depth namelist variable in prepare_gridref -grf_nudgezone_width: Final[int] = 8 diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c_dsl.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c_dsl.py deleted file mode 100644 index a6d2d236c..000000000 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c_dsl.py +++ /dev/null @@ -1,13 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_e_dsl.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_e_dsl.py deleted file mode 100644 index f62b07a4b..000000000 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_e_dsl.py +++ /dev/null @@ -1,61 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np - -from icon4py.model.common.metrics.stencils.compute_wgtfacq_c_dsl import ( - compute_wgtfacq_c_dsl, - compute_z1_z2_z3, -) - - -def compute_wgtfacq_e_dsl( - e2c, - z_ifc: np.array, - z_aux_c: np.array, - c_lin_e: np.array, - wgtfacq_e_dsl: np.array, - nlev: int, -): - """ - Compute weighting factor for quadratic interpolation to surface. - - Args: - e2c: Edge to Cell offset - z_ifc: geometric height at the vertical interface of cells. - z_aux_c: interpolation of weighting coefficients to edges - c_lin_e: interpolation field - wgtfacq_e_dsl: output - nlev: int, last k level - Returns: - Field[EdgeDim, KDim] (full levels) - """ - z1, z2, z3 = compute_z1_z2_z3(z_ifc, nlev, nlev - 1, nlev - 2, nlev - 3) - wgtfacq_c_dsl = compute_wgtfacq_c_dsl(z_ifc, nlev) - z_aux_c[:, 2] = z1 * z2 / (z2 - z3) / (z1 - z3) - z_aux_c[:, 1] = (z1 - wgtfacq_c_dsl[:, nlev - 3] * (z1 - z3)) / (z1 - z2) - z_aux_c[:, 0] = 1.0 - (wgtfacq_c_dsl[:, nlev - 2] + wgtfacq_c_dsl[:, nlev - 3]) - - z1, z2, z3 = compute_z1_z2_z3(z_ifc, 0, 1, 2, 3) - z_aux_c[:, 5] = z1 * z2 / (z2 - z3) / (z1 - z3) - z_aux_c[:, 4] = (z1 - z_aux_c[:, 5] * (z1 - z3)) / (z1 - z2) - z_aux_c[:, 3] = 1.0 - (z_aux_c[:, 4] + z_aux_c[:, 5]) - - c_lin_e = c_lin_e[:, :, np.newaxis] - z_aux_e = np.sum(c_lin_e * z_aux_c[e2c], axis=1) - - wgtfacq_e_dsl[:, nlev] = z_aux_e[:, 0] - wgtfacq_e_dsl[:, nlev - 1] = z_aux_e[:, 1] - wgtfacq_e_dsl[:, nlev - 2] = z_aux_e[:, 2] - - return wgtfacq_e_dsl From cba8c1a6fc823d5076e89e89b020e782f1bbf126 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Wed, 29 May 2024 16:55:19 +0200 Subject: [PATCH 52/65] small edit --- .../src/icon4py/model/common/test_utils/serialbox_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py index 6b38dc3a8..97a922c03 100644 --- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -291,7 +291,7 @@ def c2e2c2e(self): else: return self._c2e2c2e() - @IconSavepoint.optionally_registered() + @IconSavepoint.optionally_registered(CellDim, C2E2CDim, KDim) def _c2e2c2e(self): return self._get_connectivity_array("c2e2c2e", CellDim, reverse=True) @@ -512,7 +512,7 @@ def geofac_grg(self): (CellDim, C2E2CODim), grg[:num_cells, :, 1] ) - @IconSavepoint.optionally_registered() + @IconSavepoint.optionally_registered(CellDim, C2E2CDim, KDim) def zd_intcoef(self): return self._get_field("vcoef", CellDim, C2E2CDim, KDim) From 9bfee5b5c0c54ab6463528cc393ad2ce5c4e3854 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 30 May 2024 08:48:38 +0200 Subject: [PATCH 53/65] small edit --- .../icon4py/model/common/test_utils/serialbox_utils.py | 8 ++++---- model/common/tests/metric_tests/test_metric_fields.py | 1 + 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py index 97a922c03..c9d2308e2 100644 --- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -291,7 +291,7 @@ def c2e2c2e(self): else: return self._c2e2c2e() - @IconSavepoint.optionally_registered(CellDim, C2E2CDim, KDim) + @IconSavepoint.optionally_registered() def _c2e2c2e(self): return self._get_connectivity_array("c2e2c2e", CellDim, reverse=True) @@ -512,7 +512,7 @@ def geofac_grg(self): (CellDim, C2E2CODim), grg[:num_cells, :, 1] ) - @IconSavepoint.optionally_registered(CellDim, C2E2CDim, KDim) + @IconSavepoint.optionally_registered() def zd_intcoef(self): return self._get_field("vcoef", CellDim, C2E2CDim, KDim) @@ -686,7 +686,7 @@ def wgtfacq_e_dsl( def zd_diffcoef(self): return self._get_field("zd_diffcoef", CellDim, KDim) - @IconSavepoint.optionally_registered(CellDim, C2E2CDim, KDim) + @IconSavepoint.optionally_registered() def zd_intcoef(self): return self._read_and_reorder_sparse_field("vcoef") @@ -710,7 +710,7 @@ def _linearize_first_2dims( assert old_shape[1] == sparse_size return as_field(target_dims, data.reshape(old_shape[0] * old_shape[1], old_shape[2])) - @IconSavepoint.optionally_registered(CellDim, C2E2CDim, KDim) + @IconSavepoint.optionally_registered() def zd_vertoffset(self): return self._read_and_reorder_sparse_field("zd_vertoffset") diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index b121d60e4..b7915e845 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -831,6 +831,7 @@ def test_compute_hmask_dd3d(metrics_savepoint, icon_grid, grid_savepoint, backen def test_compute_diffusion_metrics( metrics_savepoint, interpolation_savepoint, icon_grid, grid_savepoint, backend ): + backend = None if is_roundtrip(backend): pytest.skip("skipping: slow backend") mask_hdiff = zero_field(icon_grid, CellDim, KDim, dtype=bool).asnumpy() From ec666dc27bfd08fd1076c2c8e5b71e16de4ac7f4 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 30 May 2024 08:48:50 +0200 Subject: [PATCH 54/65] small edit --- model/common/tests/metric_tests/test_metric_fields.py | 1 - 1 file changed, 1 deletion(-) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index b7915e845..b121d60e4 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -831,7 +831,6 @@ def test_compute_hmask_dd3d(metrics_savepoint, icon_grid, grid_savepoint, backen def test_compute_diffusion_metrics( metrics_savepoint, interpolation_savepoint, icon_grid, grid_savepoint, backend ): - backend = None if is_roundtrip(backend): pytest.skip("skipping: slow backend") mask_hdiff = zero_field(icon_grid, CellDim, KDim, dtype=bool).asnumpy() From a8bfa20ca8c9e56c457c786e50bdca229c045d59 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 30 May 2024 09:28:02 +0200 Subject: [PATCH 55/65] edits --- .../model/common/metrics/metric_fields.py | 30 +++++++++++++++++++ .../tests/metric_tests/test_metric_fields.py | 27 ++++++++++------- 2 files changed, 46 insertions(+), 11 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 0f2a12265..a1ece9103 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -564,6 +564,23 @@ def _compute_maxslp_maxhgtd( z_maxhgtd = maximum(z_maxhgtd_0_1, abs(ddxn_z_full(C2E[2]) * dual_edge_length(C2E[2]))) return z_maxslp, z_maxhgtd +@program +def compute_maxslp_maxhgtd( + ddxn_z_full: Field[[EdgeDim, KDim], wpfloat], + dual_edge_length: Field[[EdgeDim], wpfloat], + z_maxslp: Field[[CellDim, KDim], wpfloat], + z_maxhgtd: Field[[CellDim, KDim], wpfloat], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32 +): + _compute_maxslp_maxhgtd( + ddxn_z_full=ddxn_z_full, + dual_edge_length=dual_edge_length, + out=(z_maxslp, z_maxhgtd), + domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, + ) @field_operator def _compute_exner_exfac( @@ -1120,3 +1137,16 @@ def _compute_max_nbhgt( max_nbhgt_0_1 = maximum(z_mc_nlev(C2E2C[0]), z_mc_nlev(C2E2C[1])) max_nbhgt = maximum(max_nbhgt_0_1, z_mc_nlev(C2E2C[2])) return max_nbhgt + +@program +def compute_max_nbhgt( + z_mc_nlev: Field[[CellDim], wpfloat], + max_nbhgt: Field[[CellDim], wpfloat], + horizontal_start: int32, + horizontal_end: int32, +): + _compute_max_nbhgt( + z_mc_nlev=z_mc_nlev, + out=max_nbhgt, + domain={CellDim: (horizontal_start, horizontal_end)} + ) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index b121d60e4..26c5eb376 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -42,8 +42,8 @@ from icon4py.model.common.metrics.metric_fields import ( MetricsConfig, _compute_flat_idx, - _compute_max_nbhgt, - _compute_maxslp_maxhgtd, + compute_max_nbhgt, + compute_maxslp_maxhgtd, _compute_pg_edgeidx_vertidx, _compute_z_aux2, compute_bdy_halo_c, @@ -856,11 +856,15 @@ def test_compute_diffusion_metrics( ) nlev = icon_grid.num_levels - _compute_maxslp_maxhgtd( + compute_maxslp_maxhgtd.with_backend(backend)( ddxn_z_full=metrics_savepoint.ddxn_z_full(), dual_edge_length=grid_savepoint.dual_edge_length(), - out=(maxslp, maxhgtd), - domain={CellDim: (cell_lateral, icon_grid.num_cells), KDim: (int32(0), nlev)}, + z_maxslp=maxslp, + z_maxhgtd= maxhgtd, + horizontal_start=cell_lateral, + horizontal_end=icon_grid.num_cells, + vertical_start=0, + vertical_end=nlev, offset_provider={"C2E": icon_grid.get_offset_provider("C2E")}, ) @@ -868,9 +872,9 @@ def test_compute_diffusion_metrics( compute_z_mc.with_backend(backend)( metrics_savepoint.z_ifc(), z_mc, - horizontal_start=int32(0), + horizontal_start=0, horizontal_end=icon_grid.num_cells, - vertical_start=int32(0), + vertical_start=0, vertical_end=nlev, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) @@ -885,16 +889,17 @@ def test_compute_diffusion_metrics( z_maxslp_avg=z_maxslp_avg, z_maxhgtd_avg=z_maxhgtd_avg, horizontal_start=cell_lateral, - horizontal_end=int32(icon_grid.num_cells), + horizontal_end=icon_grid.num_cells, vertical_start=0, vertical_end=nlev, offset_provider={"C2E2C": icon_grid.get_offset_provider("C2E2C")}, ) - _compute_max_nbhgt( + compute_max_nbhgt.with_backend(backend)( z_mc_nlev=as_field((CellDim,), z_mc.asnumpy()[:, nlev - 1]), - out=max_nbhgt, - domain={CellDim: (int32(1), int32(icon_grid.num_cells))}, + max_nbhgt=max_nbhgt, + horizontal_start=cell_nudging, + horizontal_end=icon_grid.num_cells, offset_provider={"C2E2C": icon_grid.get_offset_provider("C2E2C")}, ) From 87eea720c5fbcb1efe7d61f8dce2e53dd92006ce Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 30 May 2024 09:49:21 +0200 Subject: [PATCH 56/65] added docstring as ran pre-commit --- .../model/common/metrics/metric_fields.py | 67 +++++++++++++++++-- .../tests/metric_tests/test_metric_fields.py | 6 +- 2 files changed, 66 insertions(+), 7 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index a1ece9103..f9d2f3cf6 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -564,6 +564,7 @@ def _compute_maxslp_maxhgtd( z_maxhgtd = maximum(z_maxhgtd_0_1, abs(ddxn_z_full(C2E[2]) * dual_edge_length(C2E[2]))) return z_maxslp, z_maxhgtd + @program def compute_maxslp_maxhgtd( ddxn_z_full: Field[[EdgeDim, KDim], wpfloat], @@ -573,8 +574,23 @@ def compute_maxslp_maxhgtd( horizontal_start: int32, horizontal_end: int32, vertical_start: int32, - vertical_end: int32 + vertical_end: int32, ): + """ + Compute z_maxslp and z_maxhgtd. + + See mo_vertical_grid.f90. + + Args: + ddxn_z_full: dual_edge_length + dual_edge_length: dual_edge_length + z_maxslp: output + z_maxhgtd: output + horizontal_start: horizontal start index + horizontal_end: horizontal end index + vertical_start: vertical start index + vertical_end: vertical end index + """ _compute_maxslp_maxhgtd( ddxn_z_full=ddxn_z_full, dual_edge_length=dual_edge_length, @@ -582,6 +598,7 @@ def compute_maxslp_maxhgtd( domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, ) + @field_operator def _compute_exner_exfac( ddxn_z_full: Field[[EdgeDim, KDim], wpfloat], @@ -1054,6 +1071,18 @@ def compute_mask_hdiff( vertical_start: int32, vertical_end: int32, ): + """ + Compute mask_hdiff. + + See mo_vertical_grid.f90. + + Args: + mask_hdiff: output + horizontal_start: horizontal start index + horizontal_end: horizontal end index + vertical_start: vertical start index + vertical_end: vertical end index + """ _compute_mask_hdiff( out=mask_hdiff, domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, @@ -1109,6 +1138,26 @@ def compute_z_maxslp_avg_z_maxhgtd_avg( vertical_start: int32, vertical_end: int32, ): + """ + Compute z_maxslp_avg and z_maxhgtd_avg. + + See mo_vertical_grid.f90. + + Args: + maxslp: Max field over ddxn_z_full offset + maxhgtd: Max field over ddxn_z_full offset*dual_edge_length offset + c_bln_avg_0: Interpolation field 0th K element + c_bln_avg_1: Interpolation field 1st K element + c_bln_avg_2: Interpolation field 2nd K element + c_bln_avg_3: Interpolation field 3rd K element + z_maxslp_avg: output + z_maxhgtd_avg: output + horizontal_start: horizontal start index + horizontal_end: horizontal end index + vertical_start: vertical start index + vertical_end: vertical end index + """ + _compute_z_maxslp_avg( maxslp=maxslp, c_bln_avg_0=c_bln_avg_0, @@ -1138,6 +1187,7 @@ def _compute_max_nbhgt( max_nbhgt = maximum(max_nbhgt_0_1, z_mc_nlev(C2E2C[2])) return max_nbhgt + @program def compute_max_nbhgt( z_mc_nlev: Field[[CellDim], wpfloat], @@ -1145,8 +1195,17 @@ def compute_max_nbhgt( horizontal_start: int32, horizontal_end: int32, ): + """ + Compute max_nbhgt. + + See mo_vertical_grid.f90. + + Args: + z_mc_nlev: Last K level of z_mc + max_nbhgt: output + horizontal_start: horizontal start index + horizontal_end: horizontal end index + """ _compute_max_nbhgt( - z_mc_nlev=z_mc_nlev, - out=max_nbhgt, - domain={CellDim: (horizontal_start, horizontal_end)} + z_mc_nlev=z_mc_nlev, out=max_nbhgt, domain={CellDim: (horizontal_start, horizontal_end)} ) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 26c5eb376..b4b915878 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -42,8 +42,6 @@ from icon4py.model.common.metrics.metric_fields import ( MetricsConfig, _compute_flat_idx, - compute_max_nbhgt, - compute_maxslp_maxhgtd, _compute_pg_edgeidx_vertidx, _compute_z_aux2, compute_bdy_halo_c, @@ -57,6 +55,8 @@ compute_exner_exfac, compute_hmask_dd3d, compute_mask_prog_halo_c, + compute_max_nbhgt, + compute_maxslp_maxhgtd, compute_pg_edgeidx_dsl, compute_pg_exdist_dsl, compute_rayleigh_w, @@ -860,7 +860,7 @@ def test_compute_diffusion_metrics( ddxn_z_full=metrics_savepoint.ddxn_z_full(), dual_edge_length=grid_savepoint.dual_edge_length(), z_maxslp=maxslp, - z_maxhgtd= maxhgtd, + z_maxhgtd=maxhgtd, horizontal_start=cell_lateral, horizontal_end=icon_grid.num_cells, vertical_start=0, From 383b84c51c4ff89e6f3631790b544af5e91327af Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 30 May 2024 09:59:30 +0200 Subject: [PATCH 57/65] small edit --- .../common/metrics/stencils/compute_diffusion_metrics.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py index a4f80bc94..0c67e51c5 100644 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py @@ -26,15 +26,15 @@ def _compute_nbidx( ind: int, nlev: int, ) -> np.array: - jk_start = nlev - 2 + jk_start = nlev - 1 for jk in reversed(range(k_start[jc], k_end[jc])): - for jk1 in reversed(range(jk_start + 1)): + for jk1 in reversed(range(jk_start)): if ( z_mc[jc, jk] <= z_mc_off[jc, ind, jk1] and z_mc[jc, jk] >= z_mc_off[jc, ind, jk1 + 1] ): nbidx[jc, ind, jk] = jk1 - jk_start = jk1 + jk_start = jk1 + 1 break return nbidx[jc, ind, :] From 110d03a8265bc6fdf981d04ff86595a15e494dc3 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 30 May 2024 10:32:19 +0200 Subject: [PATCH 58/65] small edit --- .../common/metrics/stencils/compute_diffusion_metrics.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py index 0c67e51c5..8d9706b10 100644 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py @@ -50,9 +50,9 @@ def _compute_z_vintcoeff( ind: int, nlev: int, ) -> np.array: - jk_start = nlev - 2 + jk_start = nlev - 1 for jk in reversed(range(k_start[jc], k_end[jc])): - for jk1 in reversed(range(jk_start + 1)): + for jk1 in reversed(range(jk_start)): if ( z_mc[jc, jk] <= z_mc_off[jc, ind, jk1] and z_mc[jc, jk] >= z_mc_off[jc, ind, jk1 + 1] @@ -60,7 +60,7 @@ def _compute_z_vintcoeff( z_vintcoeff[jc, ind, jk] = (z_mc[jc, jk] - z_mc_off[jc, ind, jk1 + 1]) / ( z_mc_off[jc, ind, jk1] - z_mc_off[jc, ind, jk1 + 1] ) - jk_start = jk1 + jk_start = jk1 + 1 break return z_vintcoeff[jc, ind, :] From 3dfa6dde1195341151e1050ad74894e12f507388 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 30 May 2024 11:44:18 +0200 Subject: [PATCH 59/65] code cleanup and slim down --- .../stencils/compute_diffusion_metrics.py | 113 ++++++++---------- 1 file changed, 48 insertions(+), 65 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py index 8d9706b10..55d65a983 100644 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py @@ -11,59 +11,55 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import math - import numpy as np def _compute_nbidx( - k_start: list, - k_end: list, + k_range: range, z_mc: np.array, z_mc_off: np.array, nbidx: np.array, jc: int, - ind: int, nlev: int, ) -> np.array: - jk_start = nlev - 1 - for jk in reversed(range(k_start[jc], k_end[jc])): - for jk1 in reversed(range(jk_start)): - if ( - z_mc[jc, jk] <= z_mc_off[jc, ind, jk1] - and z_mc[jc, jk] >= z_mc_off[jc, ind, jk1 + 1] - ): - nbidx[jc, ind, jk] = jk1 - jk_start = jk1 + 1 - break + for ind in range(3): + jk_start = nlev - 1 + for jk in reversed(k_range): + for jk1 in reversed(range(jk_start)): + if ( + z_mc[jc, jk] <= z_mc_off[jc, ind, jk1] + and z_mc[jc, jk] >= z_mc_off[jc, ind, jk1 + 1] + ): + nbidx[jc, ind, jk] = jk1 + jk_start = jk1 + 1 + break - return nbidx[jc, ind, :] + return nbidx[jc, :, :] def _compute_z_vintcoeff( - k_start: list, - k_end: list, + k_range: range, z_mc: np.array, z_mc_off: np.array, z_vintcoeff: np.array, jc: int, - ind: int, nlev: int, ) -> np.array: - jk_start = nlev - 1 - for jk in reversed(range(k_start[jc], k_end[jc])): - for jk1 in reversed(range(jk_start)): - if ( - z_mc[jc, jk] <= z_mc_off[jc, ind, jk1] - and z_mc[jc, jk] >= z_mc_off[jc, ind, jk1 + 1] - ): - z_vintcoeff[jc, ind, jk] = (z_mc[jc, jk] - z_mc_off[jc, ind, jk1 + 1]) / ( - z_mc_off[jc, ind, jk1] - z_mc_off[jc, ind, jk1 + 1] - ) - jk_start = jk1 + 1 - break - - return z_vintcoeff[jc, ind, :] + for ind in range(3): + jk_start = nlev - 1 + for jk in reversed(k_range): + for jk1 in reversed(range(jk_start)): + if ( + z_mc[jc, jk] <= z_mc_off[jc, ind, jk1] + and z_mc[jc, jk] >= z_mc_off[jc, ind, jk1 + 1] + ): + z_vintcoeff[jc, ind, jk] = (z_mc[jc, jk] - z_mc_off[jc, ind, jk1 + 1]) / ( + z_mc_off[jc, ind, jk1] - z_mc_off[jc, ind, jk1 + 1] + ) + jk_start = jk1 + 1 + break + + return z_vintcoeff[jc, :, :] def _compute_i_params( @@ -91,7 +87,7 @@ def _compute_i_params( ji += 1 i_indlist[ji] = jc - if k_start[jc] is not None and k_end[jc] is not None and k_start[jc] > k_end[jc]: + if all((k_start[jc], k_end[jc])) and k_start[jc] > k_end[jc]: i_listreduce += 1 else: ji_ind += 1 @@ -129,7 +125,7 @@ def _compute_k_start_end( k_start[jc] = jk break - if k_start[jc] is not None and k_end[jc] is not None and k_start[jc] > k_end[jc]: + if all((k_start[jc], k_end[jc])) and k_start[jc] > k_end[jc]: k_start[jc] = nlev - 1 return k_start, k_end @@ -156,37 +152,24 @@ def compute_diffusion_metrics( ) -> np.array: for ji in range(i_listdim): jc = i_indlist[ji] - if k_start[jc] is not None and k_end[jc] is not None: - nbidx[jc, 0, :] = _compute_nbidx(k_start, k_end, z_mc, z_mc_off, nbidx, jc, 0, nlev) - nbidx[jc, 1, :] = _compute_nbidx(k_start, k_end, z_mc, z_mc_off, nbidx, jc, 1, nlev) - nbidx[jc, 2, :] = _compute_nbidx(k_start, k_end, z_mc, z_mc_off, nbidx, jc, 2, nlev) - - z_vintcoeff[jc, 0, :] = _compute_z_vintcoeff( - k_start, k_end, z_mc, z_mc_off, z_vintcoeff, jc, 0, nlev + k_range = range(k_start[jc], k_end[jc]) + if all((k_range)): + nbidx[jc, :, :] = _compute_nbidx(k_range, z_mc, z_mc_off, nbidx, jc, nlev) + z_vintcoeff[jc, :, :] = _compute_z_vintcoeff( + k_range, z_mc, z_mc_off, z_vintcoeff, jc, nlev ) - z_vintcoeff[jc, 1, :] = _compute_z_vintcoeff( - k_start, k_end, z_mc, z_mc_off, z_vintcoeff, jc, 1, nlev - ) - z_vintcoeff[jc, 2, :] = _compute_z_vintcoeff( - k_start, k_end, z_mc, z_mc_off, z_vintcoeff, jc, 2, nlev + + zd_intcoef_dsl[jc, :, k_range] = z_vintcoeff[jc, :, k_range] + zd_vertoffset_dsl[jc, :, k_range] = nbidx[jc, :, k_range] - np.transpose([k_range] * 3) + mask_hdiff[jc, k_range] = True + + zd_diffcoef_dsl_var = np.maximum( + 0.0, + np.maximum( + np.sqrt(np.maximum(0.0, z_maxslp_avg[jc, k_range] - thslp_zdiffu)) / 250.0, + 2.0e-4 * np.sqrt(np.maximum(0.0, z_maxhgtd_avg[jc, k_range] - thhgtd_zdiffu)), + ), ) - for jk in range(k_start[jc], k_end[jc]): - zd_intcoef_dsl[jc, 0, jk] = z_vintcoeff[jc, 0, jk] - zd_intcoef_dsl[jc, 1, jk] = z_vintcoeff[jc, 1, jk] - zd_intcoef_dsl[jc, 2, jk] = z_vintcoeff[jc, 2, jk] - - zd_vertoffset_dsl[jc, 0, jk] = nbidx[jc, 0, jk] - jk - zd_vertoffset_dsl[jc, 1, jk] = nbidx[jc, 1, jk] - jk - zd_vertoffset_dsl[jc, 2, jk] = nbidx[jc, 2, jk] - jk - - zd_diffcoef_dsl_var = max( - 0.0, - math.sqrt(max(0.0, z_maxslp_avg[jc, jk] - thslp_zdiffu)) / 250.0, - 2.0e-4 * math.sqrt(max(0.0, z_maxhgtd_avg[jc, jk] - thhgtd_zdiffu)), - ) - zd_diffcoef_dsl_var = min(0.002, zd_diffcoef_dsl_var) - zd_diffcoef_dsl[jc, jk] = zd_diffcoef_dsl_var - - mask_hdiff[jc, jk] = True + zd_diffcoef_dsl[jc, k_range] = np.minimum(0.002, zd_diffcoef_dsl_var) return mask_hdiff, zd_diffcoef_dsl, zd_intcoef_dsl, zd_vertoffset_dsl From d9e680ccdc2865294496786315202a5543cd73ab Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 31 May 2024 15:52:41 +0200 Subject: [PATCH 60/65] edits following review and ran pre-commit --- .../model/common/metrics/metric_fields.py | 82 ++-------- .../stencils/compute_diffusion_metrics.py | 56 +++++-- .../common/test_utils/serialbox_utils.py | 2 +- .../test_compute_coeff_gradekin.py | 12 -- .../test_compute_diffusion_metrics.py | 151 +++++++++++++++++ .../metric_tests/test_compute_nudgecoeffs.py | 12 -- .../metric_tests/test_compute_wgtfac_c.py | 12 -- .../metric_tests/test_compute_wgtfacq.py | 12 -- .../tests/metric_tests/test_metric_fields.py | 154 +----------------- 9 files changed, 214 insertions(+), 279 deletions(-) create mode 100644 model/common/tests/metric_tests/test_compute_diffusion_metrics.py diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index f9d2f3cf6..8271eb070 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -1059,72 +1059,24 @@ def compute_hmask_dd3d( @field_operator -def _compute_mask_hdiff() -> Field[[CellDim, KDim], bool]: - return broadcast(True, (CellDim, KDim)) - - -@program -def compute_mask_hdiff( - mask_hdiff: Field[[CellDim, KDim], bool], - horizontal_start: int32, - horizontal_end: int32, - vertical_start: int32, - vertical_end: int32, -): - """ - Compute mask_hdiff. - - See mo_vertical_grid.f90. - - Args: - mask_hdiff: output - horizontal_start: horizontal start index - horizontal_end: horizontal end index - vertical_start: vertical start index - vertical_end: vertical end index - """ - _compute_mask_hdiff( - out=mask_hdiff, - domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, - ) - - -@field_operator -def _compute_z_maxslp_avg( - maxslp: Field[[CellDim, KDim], wpfloat], - c_bln_avg_0: Field[[CellDim], wpfloat], - c_bln_avg_1: Field[[CellDim], wpfloat], - c_bln_avg_2: Field[[CellDim], wpfloat], - c_bln_avg_3: Field[[CellDim], wpfloat], -) -> Field[[CellDim, KDim], wpfloat]: - z_maxslp_avg = ( - maxslp * c_bln_avg_0 - + maxslp(C2E2C[0]) * c_bln_avg_1 - + maxslp(C2E2C[1]) * c_bln_avg_2 - + maxslp(C2E2C[2]) * c_bln_avg_3 - ) - return z_maxslp_avg - - -@field_operator -def _compute_z_maxhgtd_avg( - maxhgtd: Field[[CellDim, KDim], wpfloat], +def _compute_weighted_cell_neighbor_sum( + field: Field[[CellDim, KDim], wpfloat], c_bln_avg_0: Field[[CellDim], wpfloat], c_bln_avg_1: Field[[CellDim], wpfloat], c_bln_avg_2: Field[[CellDim], wpfloat], c_bln_avg_3: Field[[CellDim], wpfloat], ) -> Field[[CellDim, KDim], wpfloat]: - z_maxhgtd_avg = ( - maxhgtd * c_bln_avg_0 - + maxhgtd(C2E2C[0]) * c_bln_avg_1 - + maxhgtd(C2E2C[1]) * c_bln_avg_2 - + maxhgtd(C2E2C[2]) * c_bln_avg_3 + field_avg = ( + field * c_bln_avg_0 + + field(C2E2C[0]) * c_bln_avg_1 + + field(C2E2C[1]) * c_bln_avg_2 + + field(C2E2C[2]) * c_bln_avg_3 ) - return z_maxhgtd_avg + return field_avg @program -def compute_z_maxslp_avg_z_maxhgtd_avg( +def compute_weighted_cell_neighbor_sum( maxslp: Field[[CellDim, KDim], wpfloat], maxhgtd: Field[[CellDim, KDim], wpfloat], c_bln_avg_0: Field[[CellDim], wpfloat], @@ -1146,10 +1098,10 @@ def compute_z_maxslp_avg_z_maxhgtd_avg( Args: maxslp: Max field over ddxn_z_full offset maxhgtd: Max field over ddxn_z_full offset*dual_edge_length offset - c_bln_avg_0: Interpolation field 0th K element - c_bln_avg_1: Interpolation field 1st K element - c_bln_avg_2: Interpolation field 2nd K element - c_bln_avg_3: Interpolation field 3rd K element + c_bln_avg_0: Interpolation field 0th sparse element + c_bln_avg_1: Interpolation field 1st sparse element + c_bln_avg_2: Interpolation field 2nd sparse element + c_bln_avg_3: Interpolation field 3rd sparse element z_maxslp_avg: output z_maxhgtd_avg: output horizontal_start: horizontal start index @@ -1158,8 +1110,8 @@ def compute_z_maxslp_avg_z_maxhgtd_avg( vertical_end: vertical end index """ - _compute_z_maxslp_avg( - maxslp=maxslp, + _compute_weighted_cell_neighbor_sum( + field=maxslp, c_bln_avg_0=c_bln_avg_0, c_bln_avg_1=c_bln_avg_1, c_bln_avg_2=c_bln_avg_2, @@ -1168,8 +1120,8 @@ def compute_z_maxslp_avg_z_maxhgtd_avg( domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, ) - _compute_z_maxhgtd_avg( - maxhgtd=maxhgtd, + _compute_weighted_cell_neighbor_sum( + field=maxhgtd, c_bln_avg_0=c_bln_avg_0, c_bln_avg_1=c_bln_avg_1, c_bln_avg_2=c_bln_avg_2, diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py index 55d65a983..63e78ad73 100644 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py @@ -62,7 +62,7 @@ def _compute_z_vintcoeff( return z_vintcoeff[jc, :, :] -def _compute_i_params( +def _compute_ls_params( k_start: list, k_end: list, z_maxslp_avg: np.array, @@ -74,8 +74,8 @@ def _compute_i_params( n_cells: int, nlev: int, ) -> tuple[list, int, int]: - i_indlist = [0] * n_cells - i_listreduce = 0 + indlist = [0] * n_cells + listreduce = 0 ji = -1 ji_ind = -1 @@ -85,15 +85,15 @@ def _compute_i_params( or z_maxhgtd_avg[jc, nlev - 1] >= thhgtd_zdiffu ) and c_owner_mask[jc]: ji += 1 - i_indlist[ji] = jc + indlist[ji] = jc if all((k_start[jc], k_end[jc])) and k_start[jc] > k_end[jc]: - i_listreduce += 1 + listreduce += 1 else: ji_ind += 1 - i_indlist[ji_ind] = jc + indlist[ji_ind] = jc - return i_indlist, i_listreduce, ji + return indlist, listreduce, ji def _compute_k_start_end( @@ -134,10 +134,8 @@ def _compute_k_start_end( def compute_diffusion_metrics( z_mc: np.array, z_mc_off: np.array, - k_start: list, - k_end: list, - i_indlist: list, - i_listdim: int, + max_nbhgt: np.array, + c_owner_mask: np.array, nbidx: np.array, z_vintcoeff: np.array, z_maxslp_avg: np.array, @@ -148,10 +146,40 @@ def compute_diffusion_metrics( zd_vertoffset_dsl: np.array, thslp_zdiffu: float, thhgtd_zdiffu: float, + cell_nudging: int, + n_cells: int, nlev: int, -) -> np.array: - for ji in range(i_listdim): - jc = i_indlist[ji] +) -> tuple[np.array, np.array, np.array, np.array]: + k_start, k_end = _compute_k_start_end( + z_mc=z_mc, + max_nbhgt=max_nbhgt, + z_maxslp_avg=z_maxslp_avg, + z_maxhgtd_avg=z_maxhgtd_avg, + c_owner_mask=c_owner_mask, + thslp_zdiffu=thslp_zdiffu, + thhgtd_zdiffu=thhgtd_zdiffu, + cell_nudging=cell_nudging, + n_cells=n_cells, + nlev=nlev, + ) + + indlist, listreduce, ji = _compute_ls_params( + k_start=k_start, + k_end=k_end, + z_maxslp_avg=z_maxslp_avg, + z_maxhgtd_avg=z_maxhgtd_avg, + c_owner_mask=c_owner_mask, + thslp_zdiffu=thslp_zdiffu, + thhgtd_zdiffu=thhgtd_zdiffu, + cell_nudging=cell_nudging, + n_cells=n_cells, + nlev=nlev, + ) + + listdim = ji - listreduce + + for ji in range(listdim): + jc = indlist[ji] k_range = range(k_start[jc], k_end[jc]) if all((k_range)): nbidx[jc, :, :] = _compute_nbidx(k_range, z_mc, z_mc_off, nbidx, jc, nlev) diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py index c9d2308e2..b2d1f04d4 100644 --- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -485,7 +485,7 @@ def construct_cell_geometry(self) -> CellParams: class InterpolationSavepoint(IconSavepoint): def c_bln_avg(self): - return self._get_field("c_bln_avg", CellDim, KDim) + return self._get_field("c_bln_avg", CellDim, C2E2CODim) def c_intp(self): return self._get_field("c_intp", VertexDim, V2CDim) diff --git a/model/common/tests/metric_tests/test_compute_coeff_gradekin.py b/model/common/tests/metric_tests/test_compute_coeff_gradekin.py index a89d86c72..a79cc83ed 100644 --- a/model/common/tests/metric_tests/test_compute_coeff_gradekin.py +++ b/model/common/tests/metric_tests/test_compute_coeff_gradekin.py @@ -10,18 +10,6 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later import pytest diff --git a/model/common/tests/metric_tests/test_compute_diffusion_metrics.py b/model/common/tests/metric_tests/test_compute_diffusion_metrics.py new file mode 100644 index 000000000..97141419d --- /dev/null +++ b/model/common/tests/metric_tests/test_compute_diffusion_metrics.py @@ -0,0 +1,151 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + + +import pytest +from gt4py.next import as_field + +from icon4py.model.common.dimension import ( + C2E2CDim, + CECDim, + CellDim, + KDim, +) +from icon4py.model.common.grid.horizontal import ( + HorizontalMarkerIndex, +) +from icon4py.model.common.metrics.metric_fields import ( + compute_max_nbhgt, + compute_maxslp_maxhgtd, + compute_weighted_cell_neighbor_sum, + compute_z_mc, +) +from icon4py.model.common.metrics.stencils.compute_diffusion_metrics import ( + compute_diffusion_metrics, +) +from icon4py.model.common.test_utils.helpers import ( + constant_field, + dallclose, + flatten_first_two_dims, + is_roundtrip, + zero_field, +) + + +@pytest.mark.datatest +def test_compute_diffusion_metrics( + metrics_savepoint, interpolation_savepoint, icon_grid, grid_savepoint, backend +): + backend = None + if is_roundtrip(backend): + pytest.skip("skipping: slow backend") + mask_hdiff = zero_field(icon_grid, CellDim, KDim, dtype=bool).asnumpy() + zd_vertoffset_dsl = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy() + z_vintcoeff = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy() + zd_intcoef_dsl = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy() + z_maxslp_avg = zero_field(icon_grid, CellDim, KDim) + z_maxhgtd_avg = zero_field(icon_grid, CellDim, KDim) + zd_diffcoef_dsl = zero_field(icon_grid, CellDim, KDim).asnumpy() + maxslp = zero_field(icon_grid, CellDim, KDim) + maxhgtd = zero_field(icon_grid, CellDim, KDim) + max_nbhgt = zero_field(icon_grid, CellDim) + + c2e2c = icon_grid.connectivities[C2E2CDim] + nbidx = constant_field(icon_grid, 1, CellDim, C2E2CDim, KDim, dtype=int).asnumpy() + c_bln_avg = interpolation_savepoint.c_bln_avg() + thslp_zdiffu = 0.02 + thhgtd_zdiffu = 125 + cell_nudging = icon_grid.get_start_index(CellDim, HorizontalMarkerIndex.nudging(CellDim)) + cell_lateral = icon_grid.get_start_index( + CellDim, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, + ) + nlev = icon_grid.num_levels + + compute_maxslp_maxhgtd.with_backend(backend)( + ddxn_z_full=metrics_savepoint.ddxn_z_full(), + dual_edge_length=grid_savepoint.dual_edge_length(), + z_maxslp=maxslp, + z_maxhgtd=maxhgtd, + horizontal_start=cell_lateral, + horizontal_end=icon_grid.num_cells, + vertical_start=0, + vertical_end=nlev, + offset_provider={"C2E": icon_grid.get_offset_provider("C2E")}, + ) + + z_mc = zero_field(icon_grid, CellDim, KDim, extend={KDim: 1}) + compute_z_mc.with_backend(backend)( + metrics_savepoint.z_ifc(), + z_mc, + horizontal_start=0, + horizontal_end=icon_grid.num_cells, + vertical_start=0, + vertical_end=nlev, + offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, + ) + + compute_weighted_cell_neighbor_sum.with_backend(backend)( + maxslp=maxslp, + maxhgtd=maxhgtd, + c_bln_avg_0=as_field((CellDim,), c_bln_avg.asnumpy()[:, 0]), + c_bln_avg_1=as_field((CellDim,), c_bln_avg.asnumpy()[:, 1]), + c_bln_avg_2=as_field((CellDim,), c_bln_avg.asnumpy()[:, 2]), + c_bln_avg_3=as_field((CellDim,), c_bln_avg.asnumpy()[:, 3]), + z_maxslp_avg=z_maxslp_avg, + z_maxhgtd_avg=z_maxhgtd_avg, + horizontal_start=cell_lateral, + horizontal_end=icon_grid.num_cells, + vertical_start=0, + vertical_end=nlev, + offset_provider={"C2E2C": icon_grid.get_offset_provider("C2E2C")}, + ) + + compute_max_nbhgt.with_backend(backend)( + z_mc_nlev=as_field((CellDim,), z_mc.asnumpy()[:, nlev - 1]), + max_nbhgt=max_nbhgt, + horizontal_start=cell_nudging, + horizontal_end=icon_grid.num_cells, + offset_provider={"C2E2C": icon_grid.get_offset_provider("C2E2C")}, + ) + + mask_hdiff, zd_diffcoef_dsl, zd_intcoef_dsl, zd_vertoffset_dsl = compute_diffusion_metrics( + z_mc=z_mc.asnumpy(), + z_mc_off=z_mc.asnumpy()[c2e2c], + max_nbhgt=max_nbhgt.asnumpy(), + c_owner_mask=grid_savepoint.c_owner_mask().asnumpy(), + nbidx=nbidx, + z_vintcoeff=z_vintcoeff, + z_maxslp_avg=z_maxslp_avg.asnumpy(), + z_maxhgtd_avg=z_maxhgtd_avg.asnumpy(), + mask_hdiff=mask_hdiff, + zd_diffcoef_dsl=zd_diffcoef_dsl, + zd_intcoef_dsl=zd_intcoef_dsl, + zd_vertoffset_dsl=zd_vertoffset_dsl, + thslp_zdiffu=thslp_zdiffu, + thhgtd_zdiffu=thhgtd_zdiffu, + cell_nudging=cell_nudging, + n_cells=icon_grid.num_cells, + nlev=nlev, + ) + zd_intcoef_dsl = flatten_first_two_dims( + CECDim, KDim, field=as_field((CellDim, C2E2CDim, KDim), zd_intcoef_dsl) + ) + zd_vertoffset_dsl = flatten_first_two_dims( + CECDim, KDim, field=as_field((CellDim, C2E2CDim, KDim), zd_vertoffset_dsl) + ) + + assert dallclose(mask_hdiff, metrics_savepoint.mask_hdiff().asnumpy()) + assert dallclose(zd_diffcoef_dsl, metrics_savepoint.zd_diffcoef().asnumpy(), rtol=1.0e-11) + assert dallclose(zd_vertoffset_dsl.asnumpy(), metrics_savepoint.zd_vertoffset().asnumpy()) + assert dallclose(zd_intcoef_dsl.asnumpy(), metrics_savepoint.zd_intcoef().asnumpy()) diff --git a/model/common/tests/metric_tests/test_compute_nudgecoeffs.py b/model/common/tests/metric_tests/test_compute_nudgecoeffs.py index 6227f2d00..0fe2f4cd1 100644 --- a/model/common/tests/metric_tests/test_compute_nudgecoeffs.py +++ b/model/common/tests/metric_tests/test_compute_nudgecoeffs.py @@ -10,18 +10,6 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later import numpy as np import pytest diff --git a/model/common/tests/metric_tests/test_compute_wgtfac_c.py b/model/common/tests/metric_tests/test_compute_wgtfac_c.py index 008d3c272..ebd41a4df 100644 --- a/model/common/tests/metric_tests/test_compute_wgtfac_c.py +++ b/model/common/tests/metric_tests/test_compute_wgtfac_c.py @@ -10,18 +10,6 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later import pytest diff --git a/model/common/tests/metric_tests/test_compute_wgtfacq.py b/model/common/tests/metric_tests/test_compute_wgtfacq.py index c38afa626..b65d25d20 100644 --- a/model/common/tests/metric_tests/test_compute_wgtfacq.py +++ b/model/common/tests/metric_tests/test_compute_wgtfacq.py @@ -10,18 +10,6 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later import pytest diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index b4b915878..aee07d215 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -20,17 +20,14 @@ from icon4py.model.common import constants from icon4py.model.common.dimension import ( - C2E2CDim, - CECDim, CellDim, EdgeDim, KDim, V2CDim, VertexDim, ) +from icon4py.model.common.grid import horizontal from icon4py.model.common.grid.horizontal import ( - _GRF_NUDGEZONE_START_EDGES, - _GRF_NUDGEZONE_WIDTH, HorizontalMarkerIndex, _compute_cells2verts, ) @@ -55,8 +52,6 @@ compute_exner_exfac, compute_hmask_dd3d, compute_mask_prog_halo_c, - compute_max_nbhgt, - compute_maxslp_maxhgtd, compute_pg_edgeidx_dsl, compute_pg_exdist_dsl, compute_rayleigh_w, @@ -64,14 +59,8 @@ compute_vwind_expl_wgt, compute_vwind_impl_wgt, compute_wgtfac_e, - compute_z_maxslp_avg_z_maxhgtd_avg, compute_z_mc, ) -from icon4py.model.common.metrics.stencils.compute_diffusion_metrics import ( - _compute_i_params, - _compute_k_start_end, - compute_diffusion_metrics, -) from icon4py.model.common.test_utils.datatest_utils import ( GLOBAL_EXPERIMENT, REGIONAL_EXPERIMENT, @@ -80,7 +69,6 @@ StencilTest, constant_field, dallclose, - flatten_first_two_dims, is_python, is_roundtrip, random_field, @@ -817,147 +805,11 @@ def test_compute_hmask_dd3d(metrics_savepoint, icon_grid, grid_savepoint, backen compute_hmask_dd3d( e_refin_ctrl=e_refin_ctrl, hmask_dd3d=hmask_dd3d_full, - grf_nudge_start_e=int32(_GRF_NUDGEZONE_START_EDGES), - grf_nudgezone_width=int32(_GRF_NUDGEZONE_WIDTH), + grf_nudge_start_e=int32(horizontal._GRF_NUDGEZONE_START_EDGES), + grf_nudgezone_width=int32(horizontal._GRF_NUDGEZONE_WIDTH), horizontal_start=horizontal_start, horizontal_end=icon_grid.num_edges, offset_provider={}, ) dallclose(hmask_dd3d_full.asnumpy(), hmask_dd3d_ref.asnumpy()) - - -@pytest.mark.datatest -def test_compute_diffusion_metrics( - metrics_savepoint, interpolation_savepoint, icon_grid, grid_savepoint, backend -): - if is_roundtrip(backend): - pytest.skip("skipping: slow backend") - mask_hdiff = zero_field(icon_grid, CellDim, KDim, dtype=bool).asnumpy() - zd_vertoffset_dsl = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy() - z_vintcoeff = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy() - zd_intcoef_dsl = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy() - z_maxslp_avg = zero_field(icon_grid, CellDim, KDim) - z_maxhgtd_avg = zero_field(icon_grid, CellDim, KDim) - zd_diffcoef_dsl = zero_field(icon_grid, CellDim, KDim).asnumpy() - maxslp = zero_field(icon_grid, CellDim, KDim) - maxhgtd = zero_field(icon_grid, CellDim, KDim) - max_nbhgt = zero_field(icon_grid, CellDim) - - c2e2c = icon_grid.connectivities[C2E2CDim] - nbidx = constant_field(icon_grid, 1, CellDim, C2E2CDim, KDim, dtype=int).asnumpy() - c_bln_avg = interpolation_savepoint.c_bln_avg() - thslp_zdiffu = 0.02 - thhgtd_zdiffu = 125 - cell_nudging = icon_grid.get_start_index(CellDim, HorizontalMarkerIndex.nudging(CellDim)) - cell_lateral = icon_grid.get_start_index( - CellDim, - HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, - ) - nlev = icon_grid.num_levels - - compute_maxslp_maxhgtd.with_backend(backend)( - ddxn_z_full=metrics_savepoint.ddxn_z_full(), - dual_edge_length=grid_savepoint.dual_edge_length(), - z_maxslp=maxslp, - z_maxhgtd=maxhgtd, - horizontal_start=cell_lateral, - horizontal_end=icon_grid.num_cells, - vertical_start=0, - vertical_end=nlev, - offset_provider={"C2E": icon_grid.get_offset_provider("C2E")}, - ) - - z_mc = zero_field(icon_grid, CellDim, KDim, extend={KDim: 1}) - compute_z_mc.with_backend(backend)( - metrics_savepoint.z_ifc(), - z_mc, - horizontal_start=0, - horizontal_end=icon_grid.num_cells, - vertical_start=0, - vertical_end=nlev, - offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, - ) - - compute_z_maxslp_avg_z_maxhgtd_avg.with_backend(backend)( - maxslp=maxslp, - maxhgtd=maxhgtd, - c_bln_avg_0=as_field((CellDim,), c_bln_avg.asnumpy()[:, 0]), - c_bln_avg_1=as_field((CellDim,), c_bln_avg.asnumpy()[:, 1]), - c_bln_avg_2=as_field((CellDim,), c_bln_avg.asnumpy()[:, 2]), - c_bln_avg_3=as_field((CellDim,), c_bln_avg.asnumpy()[:, 3]), - z_maxslp_avg=z_maxslp_avg, - z_maxhgtd_avg=z_maxhgtd_avg, - horizontal_start=cell_lateral, - horizontal_end=icon_grid.num_cells, - vertical_start=0, - vertical_end=nlev, - offset_provider={"C2E2C": icon_grid.get_offset_provider("C2E2C")}, - ) - - compute_max_nbhgt.with_backend(backend)( - z_mc_nlev=as_field((CellDim,), z_mc.asnumpy()[:, nlev - 1]), - max_nbhgt=max_nbhgt, - horizontal_start=cell_nudging, - horizontal_end=icon_grid.num_cells, - offset_provider={"C2E2C": icon_grid.get_offset_provider("C2E2C")}, - ) - - k_start, k_end = _compute_k_start_end( - z_mc=z_mc.asnumpy(), - max_nbhgt=max_nbhgt.asnumpy(), - z_maxslp_avg=z_maxslp_avg.asnumpy(), - z_maxhgtd_avg=z_maxhgtd_avg.asnumpy(), - c_owner_mask=grid_savepoint.c_owner_mask().asnumpy(), - thslp_zdiffu=thslp_zdiffu, - thhgtd_zdiffu=thhgtd_zdiffu, - cell_nudging=cell_nudging, - n_cells=icon_grid.num_cells, - nlev=nlev, - ) - - i_indlist, i_listreduce, ji = _compute_i_params( - k_start=k_start, - k_end=k_end, - z_maxslp_avg=z_maxslp_avg.asnumpy(), - z_maxhgtd_avg=z_maxhgtd_avg.asnumpy(), - c_owner_mask=grid_savepoint.c_owner_mask().asnumpy(), - thslp_zdiffu=thslp_zdiffu, - thhgtd_zdiffu=thhgtd_zdiffu, - cell_nudging=cell_nudging, - n_cells=icon_grid.num_cells, - nlev=nlev, - ) - - i_listdim = ji - i_listreduce - - mask_hdiff, zd_diffcoef_dsl, zd_intcoef_dsl, zd_vertoffset_dsl = compute_diffusion_metrics( - z_mc=z_mc.asnumpy(), - z_mc_off=z_mc.asnumpy()[c2e2c], - k_start=k_start, - k_end=k_end, - i_indlist=i_indlist, - i_listdim=i_listdim, - nbidx=nbidx, - z_vintcoeff=z_vintcoeff, - z_maxslp_avg=z_maxslp_avg.asnumpy(), - z_maxhgtd_avg=z_maxhgtd_avg.asnumpy(), - mask_hdiff=mask_hdiff, - zd_diffcoef_dsl=zd_diffcoef_dsl, - zd_intcoef_dsl=zd_intcoef_dsl, - zd_vertoffset_dsl=zd_vertoffset_dsl, - thslp_zdiffu=thslp_zdiffu, - thhgtd_zdiffu=thhgtd_zdiffu, - nlev=nlev, - ) - zd_intcoef_dsl = flatten_first_two_dims( - CECDim, KDim, field=as_field((CellDim, C2E2CDim, KDim), zd_intcoef_dsl) - ) - zd_vertoffset_dsl = flatten_first_two_dims( - CECDim, KDim, field=as_field((CellDim, C2E2CDim, KDim), zd_vertoffset_dsl) - ) - - assert dallclose(mask_hdiff, metrics_savepoint.mask_hdiff().asnumpy()) - assert dallclose(zd_diffcoef_dsl, metrics_savepoint.zd_diffcoef().asnumpy(), rtol=1.0e-11) - assert dallclose(zd_vertoffset_dsl.asnumpy(), metrics_savepoint.zd_vertoffset().asnumpy()) - assert dallclose(zd_intcoef_dsl.asnumpy(), metrics_savepoint.zd_intcoef().asnumpy()) From 20b7d733a9cf2a004046b52107137d68fb5a394f Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 31 May 2024 15:54:16 +0200 Subject: [PATCH 61/65] small edit --- .../common/tests/metric_tests/test_compute_diffusion_metrics.py | 1 - 1 file changed, 1 deletion(-) diff --git a/model/common/tests/metric_tests/test_compute_diffusion_metrics.py b/model/common/tests/metric_tests/test_compute_diffusion_metrics.py index 97141419d..9250c0c45 100644 --- a/model/common/tests/metric_tests/test_compute_diffusion_metrics.py +++ b/model/common/tests/metric_tests/test_compute_diffusion_metrics.py @@ -46,7 +46,6 @@ def test_compute_diffusion_metrics( metrics_savepoint, interpolation_savepoint, icon_grid, grid_savepoint, backend ): - backend = None if is_roundtrip(backend): pytest.skip("skipping: slow backend") mask_hdiff = zero_field(icon_grid, CellDim, KDim, dtype=bool).asnumpy() From 91bb3101ea5e51f76fcf326430bae54cc906cd18 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 10 Jun 2024 16:25:49 +0200 Subject: [PATCH 62/65] edits following review --- .../src/icon4py/model/common/grid/simple.py | 36 +++++++++---------- .../model/common/metrics/metric_fields.py | 35 +++++------------- .../test_compute_diffusion_metrics.py | 11 +++--- 3 files changed, 33 insertions(+), 49 deletions(-) diff --git a/model/common/src/icon4py/model/common/grid/simple.py b/model/common/src/icon4py/model/common/grid/simple.py index 9d07a38c1..e2155f4ab 100644 --- a/model/common/src/icon4py/model/common/grid/simple.py +++ b/model/common/src/icon4py/model/common/grid/simple.py @@ -301,24 +301,24 @@ class SimpleGridData: c2e2cO_table = np.asarray( [ - [15, 4, 3, 0], - [16, 5, 4, 1], - [17, 3, 5, 2], - [0, 6, 2, 3], - [1, 7, 0, 4], - [2, 8, 1, 5], - [3, 10, 9, 6], - [4, 11, 10, 7], - [5, 9, 11, 8], - [6, 12, 8, 9], - [7, 13, 6, 10], - [8, 14, 7, 11], - [9, 16, 15, 12], - [10, 17, 16, 13], - [11, 15, 17, 14], - [12, 0, 14, 15], - [13, 1, 12, 16], - [14, 2, 13, 17], + [0, 15, 4, 3], + [1, 16, 5, 4], + [2, 17, 3, 5], + [3, 0, 6, 2], + [4, 1, 7, 0], + [5, 2, 8, 1], + [6, 3, 10, 9], + [7, 4, 11, 10], + [8, 5, 9, 11], + [9, 6, 12, 8], + [10, 7, 13, 6], + [11, 8, 14, 7], + [12, 9, 16, 15], + [13, 10, 17, 16], + [14, 11, 15, 17], + [15, 12, 0, 14], + [16, 13, 1, 12], + [17, 14, 2, 13], ] ) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 8271eb070..7b00f7f23 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -24,6 +24,7 @@ int32, maximum, minimum, + neighbor_sum, program, sin, tanh, @@ -33,7 +34,9 @@ from icon4py.model.common.dimension import ( C2E, C2E2C, + C2E2CO, E2C, + C2E2CODim, CellDim, E2CDim, EdgeDim, @@ -1061,17 +1064,9 @@ def compute_hmask_dd3d( @field_operator def _compute_weighted_cell_neighbor_sum( field: Field[[CellDim, KDim], wpfloat], - c_bln_avg_0: Field[[CellDim], wpfloat], - c_bln_avg_1: Field[[CellDim], wpfloat], - c_bln_avg_2: Field[[CellDim], wpfloat], - c_bln_avg_3: Field[[CellDim], wpfloat], + c_bln_avg: Field[[CellDim, C2E2CODim], wpfloat], ) -> Field[[CellDim, KDim], wpfloat]: - field_avg = ( - field * c_bln_avg_0 - + field(C2E2C[0]) * c_bln_avg_1 - + field(C2E2C[1]) * c_bln_avg_2 - + field(C2E2C[2]) * c_bln_avg_3 - ) + field_avg = neighbor_sum(field(C2E2CO) * c_bln_avg, axis=C2E2CODim) return field_avg @@ -1079,10 +1074,7 @@ def _compute_weighted_cell_neighbor_sum( def compute_weighted_cell_neighbor_sum( maxslp: Field[[CellDim, KDim], wpfloat], maxhgtd: Field[[CellDim, KDim], wpfloat], - c_bln_avg_0: Field[[CellDim], wpfloat], - c_bln_avg_1: Field[[CellDim], wpfloat], - c_bln_avg_2: Field[[CellDim], wpfloat], - c_bln_avg_3: Field[[CellDim], wpfloat], + c_bln_avg: Field[[CellDim, C2E2CODim], wpfloat], z_maxslp_avg: Field[[CellDim, KDim], wpfloat], z_maxhgtd_avg: Field[[CellDim, KDim], wpfloat], horizontal_start: int32, @@ -1098,10 +1090,7 @@ def compute_weighted_cell_neighbor_sum( Args: maxslp: Max field over ddxn_z_full offset maxhgtd: Max field over ddxn_z_full offset*dual_edge_length offset - c_bln_avg_0: Interpolation field 0th sparse element - c_bln_avg_1: Interpolation field 1st sparse element - c_bln_avg_2: Interpolation field 2nd sparse element - c_bln_avg_3: Interpolation field 3rd sparse element + c_bln_avg: Interpolation field z_maxslp_avg: output z_maxhgtd_avg: output horizontal_start: horizontal start index @@ -1112,20 +1101,14 @@ def compute_weighted_cell_neighbor_sum( _compute_weighted_cell_neighbor_sum( field=maxslp, - c_bln_avg_0=c_bln_avg_0, - c_bln_avg_1=c_bln_avg_1, - c_bln_avg_2=c_bln_avg_2, - c_bln_avg_3=c_bln_avg_3, + c_bln_avg=c_bln_avg, out=z_maxslp_avg, domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, ) _compute_weighted_cell_neighbor_sum( field=maxhgtd, - c_bln_avg_0=c_bln_avg_0, - c_bln_avg_1=c_bln_avg_1, - c_bln_avg_2=c_bln_avg_2, - c_bln_avg_3=c_bln_avg_3, + c_bln_avg=c_bln_avg, out=z_maxhgtd_avg, domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, ) diff --git a/model/common/tests/metric_tests/test_compute_diffusion_metrics.py b/model/common/tests/metric_tests/test_compute_diffusion_metrics.py index 9250c0c45..76bd5c300 100644 --- a/model/common/tests/metric_tests/test_compute_diffusion_metrics.py +++ b/model/common/tests/metric_tests/test_compute_diffusion_metrics.py @@ -46,6 +46,7 @@ def test_compute_diffusion_metrics( metrics_savepoint, interpolation_savepoint, icon_grid, grid_savepoint, backend ): + backend = None if is_roundtrip(backend): pytest.skip("skipping: slow backend") mask_hdiff = zero_field(icon_grid, CellDim, KDim, dtype=bool).asnumpy() @@ -97,17 +98,17 @@ def test_compute_diffusion_metrics( compute_weighted_cell_neighbor_sum.with_backend(backend)( maxslp=maxslp, maxhgtd=maxhgtd, - c_bln_avg_0=as_field((CellDim,), c_bln_avg.asnumpy()[:, 0]), - c_bln_avg_1=as_field((CellDim,), c_bln_avg.asnumpy()[:, 1]), - c_bln_avg_2=as_field((CellDim,), c_bln_avg.asnumpy()[:, 2]), - c_bln_avg_3=as_field((CellDim,), c_bln_avg.asnumpy()[:, 3]), + c_bln_avg=c_bln_avg, z_maxslp_avg=z_maxslp_avg, z_maxhgtd_avg=z_maxhgtd_avg, horizontal_start=cell_lateral, horizontal_end=icon_grid.num_cells, vertical_start=0, vertical_end=nlev, - offset_provider={"C2E2C": icon_grid.get_offset_provider("C2E2C")}, + offset_provider={ + "C2E2C": icon_grid.get_offset_provider("C2E2C"), + "C2E2CO": icon_grid.get_offset_provider("C2E2CO"), + }, ) compute_max_nbhgt.with_backend(backend)( From 051589ded343294c499e874e45e7774229435537 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Wed, 26 Jun 2024 14:53:34 +0200 Subject: [PATCH 63/65] update with upstream and docstrings for partial nested k-loop field_op --- .../model/common/metrics/metric_fields.py | 25 ++++++++++++ .../stencils/compute_zdiff_gradp_dsl.py | 39 ++++++++++++++++--- 2 files changed, 59 insertions(+), 5 deletions(-) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 7b00f7f23..5ce6847f7 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -26,6 +26,7 @@ minimum, neighbor_sum, program, + scan_operator, sin, tanh, where, @@ -1144,3 +1145,27 @@ def compute_max_nbhgt( _compute_max_nbhgt( z_mc_nlev=z_mc_nlev, out=max_nbhgt, domain={CellDim: (horizontal_start, horizontal_end)} ) + + +@scan_operator(axis=KDim, forward=True, init=(0, False)) +def _compute_param( + param: tuple[int32, bool], + z_me_jk: float, + z_ifc_off: float, + z_ifc_off_koff: float, + lower: int32, + nlev: int32, +) -> tuple[int32, bool]: + param_0, param_1 = param + if param_0 >= lower: + if (param_0 == nlev) | (z_me_jk <= z_ifc_off) & (z_me_jk >= z_ifc_off_koff): + param_1 = True + return param_0 + 1, param_1 + + +@field_operator(grid_type=GridType.UNSTRUCTURED) +def _compute_z_ifc_off_koff( + z_ifc_off: Field[[EdgeDim, KDim], wpfloat], +) -> Field[[EdgeDim, KDim], wpfloat]: + n = z_ifc_off(Koff[1]) + return n diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py index 70461ea63..02461a7d1 100644 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py @@ -29,19 +29,48 @@ def compute_zdiff_gradp_dsl( zdiff_gradp[horizontal_start:, :, :] = ( np.expand_dims(z_me, axis=1)[horizontal_start:, :, :] - z_mc[e2c][horizontal_start:, :, :] ) + """ + First part for loop implementation with gt4py code + + >>> z_ifc_off_koff = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1}) + >>> z_ifc_off = as_field((EdgeDim, KDim,), z_ifc[e2c[:, 0], :]) + >>> _compute_z_ifc_off_koff( + >>> z_ifc_off=z_ifc_off, + >>> domain={EdgeDim: (horizontal_start, nedges), KDim: (0, nlev)}, + >>> out=z_ifc_off_koff, + >>> offset_provider={"Koff": icon_grid.get_offset_provider("Koff")} + >>> ) + """ for je in range(horizontal_start, nedges): - jk_start = int(flat_idx[je]) for jk in range(int(flat_idx[je]) + 1, nlev): - for jk1 in range(jk_start, nlev): + """ + Second part for loop implementation with gt4py code + >>> param_2 = as_field((KDim,), np.asarray([False] * nlev)) + >>> param_3 = as_field((KDim,), np.arange(nlev)) + >>> z_ifc_off_e = as_field((KDim,), z_ifc[e2c[je, 0], :]) + >>> _compute_param.with_backend(backend)( + >>> z_me_jk=z_me[je, jk], + >>> z_ifc_off=z_ifc_off_e, + >>> z_ifc_off_koff=as_field((KDim,), z_ifc_off_koff.asnumpy()[je, :]), + >>> lower=int(flat_idx[je]), + >>> nlev=nlev - 1, + >>> out=(param_3, param_2), + >>> offset_provider={} + >>> ) + >>> zdiff_gradp[je, 0, jk] = z_me[je, jk] - z_mc[e2c[je, 0], np.where(param_2.asnumpy())[0][0]] + """ + + param = [False] * nlev + for jk1 in range(int(flat_idx[je]), nlev): if ( jk1 == nlev - 1 or z_me[je, jk] <= z_ifc[e2c[je, 0], jk1] and z_me[je, jk] >= z_ifc[e2c[je, 0], jk1 + 1] ): - zdiff_gradp[je, 0, jk] = z_me[je, jk] - z_mc[e2c[je, 0], jk1] - jk_start = jk1 - break + param[jk1] = True + + zdiff_gradp[je, 0, jk] = z_me[je, jk] - z_mc[e2c[je, 0], np.where(param)[0][0]] jk_start = int(flat_idx[je]) for jk in range(int(flat_idx[je]) + 1, nlev): From df8080f747a6e711c2925469524aebaaf24f37c7 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 11 Jul 2024 09:04:43 +0200 Subject: [PATCH 64/65] reverted change --- .../src/icon4py/model/common/grid/simple.py | 36 +++++++++---------- 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/model/common/src/icon4py/model/common/grid/simple.py b/model/common/src/icon4py/model/common/grid/simple.py index c8f6f06c8..b919eac98 100644 --- a/model/common/src/icon4py/model/common/grid/simple.py +++ b/model/common/src/icon4py/model/common/grid/simple.py @@ -310,24 +310,24 @@ class SimpleGridData: c2e2cO_table = np.asarray( [ - [0, 15, 4, 3], - [1, 16, 5, 4], - [2, 17, 3, 5], - [3, 0, 6, 2], - [4, 1, 7, 0], - [5, 2, 8, 1], - [6, 3, 10, 9], - [7, 4, 11, 10], - [8, 5, 9, 11], - [9, 6, 12, 8], - [10, 7, 13, 6], - [11, 8, 14, 7], - [12, 9, 16, 15], - [13, 10, 17, 16], - [14, 11, 15, 17], - [15, 12, 0, 14], - [16, 13, 1, 12], - [17, 14, 2, 13], + [15, 4, 3, 0], + [16, 5, 4, 1], + [17, 3, 5, 2], + [0, 6, 2, 3], + [1, 7, 0, 4], + [2, 8, 1, 5], + [3, 10, 9, 6], + [4, 11, 10, 7], + [5, 9, 11, 8], + [6, 12, 8, 9], + [7, 13, 6, 10], + [8, 14, 7, 11], + [9, 16, 15, 12], + [10, 17, 16, 13], + [11, 15, 17, 14], + [12, 0, 14, 15], + [13, 1, 12, 16], + [14, 2, 13, 17], ], dtype=gtx.int32, ) From c1fbb8382d91c3fc209ca664d5e24aad14496454 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 11 Jul 2024 09:44:31 +0200 Subject: [PATCH 65/65] removed backend specification --- .../common/tests/metric_tests/test_compute_diffusion_metrics.py | 1 - 1 file changed, 1 deletion(-) diff --git a/model/common/tests/metric_tests/test_compute_diffusion_metrics.py b/model/common/tests/metric_tests/test_compute_diffusion_metrics.py index 76bd5c300..bf7fd920a 100644 --- a/model/common/tests/metric_tests/test_compute_diffusion_metrics.py +++ b/model/common/tests/metric_tests/test_compute_diffusion_metrics.py @@ -46,7 +46,6 @@ def test_compute_diffusion_metrics( metrics_savepoint, interpolation_savepoint, icon_grid, grid_savepoint, backend ): - backend = None if is_roundtrip(backend): pytest.skip("skipping: slow backend") mask_hdiff = zero_field(icon_grid, CellDim, KDim, dtype=bool).asnumpy()