From af437713913d46e0b2122517442c4f30b4378667 Mon Sep 17 00:00:00 2001 From: Andreas Date: Fri, 22 Mar 2024 18:09:36 +0100 Subject: [PATCH 01/20] 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 0000000000..0e9e07889c --- /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 0000000000..9281fa444c --- /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/20] 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 0e9e07889c..00c56e69a0 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 ade2946538..9b0ddd9ea5 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 9281fa444c..cf7dfff521 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/20] 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 00c56e69a0..b2eea0460e 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 cf7dfff521..aff90ab004 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/20] 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 b2eea0460e..45f23f54f0 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 9b0ddd9ea5..5100ca68a0 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 aff90ab004..61549c7d68 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/20] 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 61549c7d68..346626cfb7 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/20] 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 45f23f54f0..70399f6cb6 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/20] 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 70399f6cb6..22c3fbf754 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 346626cfb7..1efdfa13f8 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/20] 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 22c3fbf754..0aa4d38e92 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 1efdfa13f8..2b6d99c34c 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/20] 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 0aa4d38e92..ad9ec0bbd1 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 2b6d99c34c..4ea59e2f96 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 3d715780063b7072688c55d52803140c3e963399 Mon Sep 17 00:00:00 2001 From: Andreas Date: Tue, 23 Apr 2024 13:31:04 +0200 Subject: [PATCH 10/20] 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 4ea59e2f96..b67d4f72b8 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 11/20] 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 ad9ec0bbd1..f0176ea3a9 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 12/20] 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 22030c5d78..678670711e 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 f0176ea3a9..c4e20b50fb 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 fec02c9841..2f960e5897 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 b67d4f72b8..5d399c5627 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 13/20] 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 5d399c5627..7d6811aef0 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 14/20] 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 c4e20b50fb..c121c2acc1 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 15/20] 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 678670711e..01d7178d93 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 cd7fa4ab78..2b2618eef0 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 c121c2acc1..0000000000 --- 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 2f960e5897..506fe9ed79 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 449f8c5168..e59465abe8 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 335cd6de9c..cb42c70fb6 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 7d6811aef0..0000000000 --- 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 6ba890fefb..145a08a0aa 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 16/20] 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 506fe9ed79..1487023df3 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 17/20] 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 2b2618eef0..2a767f3895 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 1487023df3..8c0cd3b42f 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 e59465abe8..3600c2bf9e 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 145a08a0aa..30633d8364 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 c4a2c42d22f3462d2dc22ebac412f104dae6c57b Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 26 Apr 2024 20:36:38 +0200 Subject: [PATCH 18/20] Update model/common/src/icon4py/model/common/math/helpers.py Co-authored-by: Magdalena --- model/common/src/icon4py/model/common/math/helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/common/src/icon4py/model/common/math/helpers.py b/model/common/src/icon4py/model/common/math/helpers.py index 8c0cd3b42f..d1083a819d 100644 --- a/model/common/src/icon4py/model/common/math/helpers.py +++ b/model/common/src/icon4py/model/common/math/helpers.py @@ -36,7 +36,7 @@ def cell_kdim( @field_operator -def edge_kdim( +def average_edge_kdim_level_up( half_level_field: Field[[EdgeDim, KDim], wpfloat], ) -> Field[[EdgeDim, KDim], wpfloat]: """ From 0b71c49ec8b82ac97410072140ec017af9c2473d Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 26 Apr 2024 20:36:45 +0200 Subject: [PATCH 19/20] Update model/common/src/icon4py/model/common/math/helpers.py Co-authored-by: Magdalena --- model/common/src/icon4py/model/common/math/helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/common/src/icon4py/model/common/math/helpers.py b/model/common/src/icon4py/model/common/math/helpers.py index d1083a819d..808db44ba7 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 cell_kdim( +def average_cell_kdim_level_up( half_level_field: Field[[CellDim, KDim], wpfloat], ) -> Field[[CellDim, KDim], wpfloat]: """ From 92e366078d3290238f44133b9cdf5880b3552d93 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 26 Apr 2024 20:48:00 +0200 Subject: [PATCH 20/20] edits following review --- .../interpolation/interpolation_fields.py | 14 -------- .../model/common/metrics/metric_fields.py | 36 +++++++++++-------- .../tests/metric_tests/test_metric_fields.py | 10 +++--- 3 files changed, 27 insertions(+), 33 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 2a767f3895..0611e86e11 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 edge_kdim 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] -): - edge_kdim(z_ddxnt_z_half_e, out=ddxn_z_full) 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 3600c2bf9e..7869306e6b 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -28,7 +28,8 @@ from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, Koff, VertexDim from icon4py.model.common.math.helpers import ( _grad_fd_tang, - cell_kdim, + average_cell_kdim_level_up, + average_edge_kdim_level_up, difference_k_level_down, difference_k_level_up, grad_fd_norm, @@ -65,7 +66,7 @@ def compute_z_mc( vertical_end:int32 end index of vertical domain """ - cell_kdim( + average_cell_kdim_level_up( z_ifc, out=z_mc, domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)}, @@ -447,18 +448,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), }, ) @@ -469,10 +470,10 @@ def compute_ddxt_z_half_e( 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, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _grad_fd_tang( z_ifv, @@ -480,7 +481,14 @@ def compute_ddxt_z_half_e( tangent_orientation, out=ddxt_z_half_e, domain={ - EdgeDim: (horizontal_lower, horizontal_upper), - KDim: (vertical_lower, vertical_upper), + EdgeDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), }, ) + + +@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/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 30633d8364..370e5383d0 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -25,13 +25,13 @@ _compute_cells2verts, compute_cells2edges, ) -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_ddxnt_z_full, compute_ddxt_z_half_e, compute_rayleigh_w, compute_scalfac_dd3d, @@ -375,10 +375,10 @@ def test_compute_ddxt_z_full_e( 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, + 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")}, ) ddxt_z_full = zero_field(icon_grid, EdgeDim, KDim)