From 4f635487dcf3254426e0b3b7426fb78cd9c41fbd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Thu, 24 Nov 2022 09:04:58 +0100 Subject: [PATCH 001/263] First version of fused diffusion stencils --- .../fused_mo_nh_diffusion_stencil_02_03.py | 47 ++++++++++ .../fused_mo_nh_diffusion_stencil_04_05_06.py | 82 ++++++++++++++++++ ...sed_mo_nh_diffusion_stencil_07_08_09_10.py | 86 +++++++++++++++++++ .../fused_mo_nh_diffusion_stencil_11_12.py | 40 +++++++++ .../fused_mo_nh_diffusion_stencil_13_14.py | 67 +++++++++++++++ .../mo_nh_diffusion_stencil_10.py | 8 +- 6 files changed, 326 insertions(+), 4 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py new file mode 100644 index 000000000..8c111b5d4 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -0,0 +1,47 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum + +from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim, Koff + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import _mo_nh_diffusion_stencil_02 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import _mo_nh_diffusion_stencil_03 + +@field_operator +def _fused_mo_nh_diffusion_stencil_02_03( + kh_smag_ec: Field[[EdgeDim, KDim], float], + vn: Field[[EdgeDim, KDim], float], + e_bln_c_s: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + diff_multfac_smag: Field[[KDim], float], + wgtfac_c: Field[[CellDim, KDim], float], +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + kh_c, div = _mo_nh_diffusion_stencil_02(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag) + div_ic, hdef_ic = _mo_nh_diffusion_stencil_03(div, kh_c, wgtfac_c) + return div_ic, hdef_ic + + +@program +def fused_mo_nh_diffusion_stencil_02_03( + kh_smag_ec: Field[[EdgeDim, KDim], float], + vn: Field[[EdgeDim, KDim], float], + e_bln_c_s: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + diff_multfac_smag: Field[[KDim], float], + wgtfac_c: Field[[CellDim, KDim], float], + div_ic: Field[[CellDim, KDim], float], + hdef_ic: Field[[CellDim, KDim], float], +): + _fused_mo_nh_diffusion_stencil_02_03(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag, wgtfac_c, out=(div_ic, hdef_ic)) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py new file mode 100644 index 000000000..1438e07c1 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py @@ -0,0 +1,82 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum, maximum, where, int32 + +from icon4py.common.dimension import ( + E2C2V, + E2ECV, + ECVDim, + EdgeDim, + KDim, + VertexDim, +) + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_04 import _mo_nh_diffusion_stencil_04 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_05 import _mo_nh_diffusion_stencil_05 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_06 import _mo_nh_diffusion_stencil_06 + +@field_operator +def _fused_mo_nh_diffusion_stencil_04_05_06( + u_vert: Field[[VertexDim, KDim], float], + v_vert: Field[[VertexDim, KDim], float], + primal_normal_vert_v1: Field[[ECVDim], float], + primal_normal_vert_v2: Field[[ECVDim], float], + z_nabla2_e: Field[[EdgeDim, KDim], float], + inv_vert_vert_length: Field[[EdgeDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + area_edge: Field[[EdgeDim], float], + kh_smag_e: Field[[EdgeDim, KDim], float], + diff_multfac_vn: Field[[KDim], float], + nudgecoeff_e: Field[[EdgeDim], float], + vn: Field[[EdgeDim, KDim], float], + horz_idx: Field[[EdgeDim], int32], + nudgezone_diff: float, + fac_bdydiff_v: float, + start_2nd_nudge_line_idx_e: int32, +) -> Field[[EdgeDim, KDim], float]: + + z_nabla4_e2 = _mo_nh_diffusion_stencil_04(u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2, z_nabla2_e, inv_vert_vert_length, inv_primal_edge_length) + + vn = where( + horz_idx >= start_2nd_nudge_line_idx_e, + _mo_nh_diffusion_stencil_05(area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, nudgecoeff_e, vn, nudgezone_diff), + _mo_nh_diffusion_stencil_06(z_nabla2_e, area_edge, vn, fac_bdydiff_v) + ) + + return vn + + +@program +def fused_mo_nh_diffusion_stencil_04_05_06( + u_vert: Field[[VertexDim, KDim], float], + v_vert: Field[[VertexDim, KDim], float], + primal_normal_vert_v1: Field[[ECVDim], float], + primal_normal_vert_v2: Field[[ECVDim], float], + z_nabla2_e: Field[[EdgeDim, KDim], float], + inv_vert_vert_length: Field[[EdgeDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + area_edge: Field[[EdgeDim], float], + kh_smag_e: Field[[EdgeDim, KDim], float], + diff_multfac_vn: Field[[KDim], float], + nudgecoeff_e: Field[[EdgeDim], float], + vn: Field[[EdgeDim, KDim], float], + horz_idx: Field[[EdgeDim], int32], + nudgezone_diff: float, + fac_bdydiff_v: float, + start_2nd_nudge_line_idx_e: int32, +): + _fused_mo_nh_diffusion_stencil_04_05_06(u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2, z_nabla2_e, inv_vert_vert_length, + inv_primal_edge_length, area_edge, kh_smag_e, diff_multfac_vn, nudgecoeff_e, vn, horz_idx, nudgezone_diff, fac_bdydiff_v, + start_2nd_nudge_line_idx_e, out=vn) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py new file mode 100644 index 000000000..a1e305ae5 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -0,0 +1,86 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum, where, int32, broadcast + +from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_07 import _mo_nh_diffusion_stencil_07 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_08 import _mo_nh_diffusion_stencil_08 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_09 import _mo_nh_diffusion_stencil_09 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_10 import _mo_nh_diffusion_stencil_10 + +@field_operator +def _fused_mo_nh_diffusion_stencil_07_08_09_10( + area: Field[[CellDim], float], + geofac_n2s: Field[[CellDim, C2E2CODim], float], + geofac_grg_x: Field[[CellDim, C2E2CODim], float], + geofac_grg_y: Field[[CellDim, C2E2CODim], float], + w_old: Field[[CellDim, KDim], float], + w: Field[[CellDim, KDim], float], + dwdx: Field[[CellDim, KDim], float], + dwdy: Field[[CellDim, KDim], float], + diff_multfac_w: float, + diff_multfac_n2w: Field[[KDim], float], + vert_idx: Field[[KDim], int32], + horz_idx: Field[[CellDim], int32], + nrdmax: int32, + interior_idx: int32, + halo_idx: int32, +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + + vert_idx = broadcast(vert_idx, (CellDim, KDim)) + + dwdx, dwdy = where( + vert_idx > int32(0), + _mo_nh_diffusion_stencil_08(w_old, geofac_grg_x, geofac_grg_y), + (dwdx, dwdy) + ) + + z_nabla2_c = _mo_nh_diffusion_stencil_07(w_old, geofac_n2s) + + w = where( + (horz_idx >= interior_idx) & (horz_idx < halo_idx), + _mo_nh_diffusion_stencil_09(area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w), + w_old, + ) + + w = where( + (vert_idx > int32(0)) & (vert_idx < nrdmax) & (horz_idx >= interior_idx) & (horz_idx < halo_idx), + _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, area, z_nabla2_c), + w, + ) + + return w, dwdx, dwdy + + +@program +def fused_mo_nh_diffusion_stencil_07_08_09_10( + area: Field[[CellDim], float], + geofac_n2s: Field[[CellDim, C2E2CODim], float], + geofac_grg_x: Field[[CellDim, C2E2CODim], float], + geofac_grg_y: Field[[CellDim, C2E2CODim], float], + w_old: Field[[CellDim, KDim], float], + w: Field[[CellDim, KDim], float], + dwdx: Field[[CellDim, KDim], float], + dwdy: Field[[CellDim, KDim], float], + diff_multfac_w: float, + diff_multfac_n2w: Field[[KDim], float], + vert_idx: Field[[KDim], int32], + horz_idx: Field[[CellDim], int32], + nrdmax: int32, + interior_idx: int32, + halo_idx: int32, +): + _fused_mo_nh_diffusion_stencil_07_08_09_10(area, geofac_n2s, geofac_grg_x, geofac_grg_y, w_old, w, dwdx, dwdy, diff_multfac_w, diff_multfac_n2w, vert_idx, horz_idx, nrdmax, interior_idx, halo_idx, out = (w, dwdx, dwdy)) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py new file mode 100644 index 000000000..02fd00530 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py @@ -0,0 +1,40 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, where, neighbor_sum, max_over, maximum + +from icon4py.common.dimension import E2C, E2CDim, C2E2C, C2E2CDim, CellDim, EdgeDim, KDim + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_11 import _mo_nh_diffusion_stencil_11 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_12 import _mo_nh_diffusion_stencil_12 + +@field_operator +def _fused_mo_nh_diffusion_stencil_11_12( + theta_v: Field[[CellDim, KDim], float], + theta_ref_mc: Field[[CellDim, KDim], float], + thresh_tdiff: float, + kh_smag_e: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + enh_diffu_3d = _mo_nh_diffusion_stencil_11(theta_v, theta_ref_mc, thresh_tdiff) + kh_smag_e = _mo_nh_diffusion_stencil_12(kh_smag_e, enh_diffu_3d) + return kh_smag_e + +@program +def fused_mo_nh_diffusion_stencil_11_12( + theta_v: Field[[CellDim, KDim], float], + theta_ref_mc: Field[[CellDim, KDim], float], + thresh_tdiff: float, + kh_smag_e: Field[[EdgeDim, KDim], float], +): + _fused_mo_nh_diffusion_stencil_11_12(theta_v, theta_ref_mc, thresh_tdiff, kh_smag_e, out=kh_smag_e) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py new file mode 100644 index 000000000..40faa10b5 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -0,0 +1,67 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum + +from icon4py.common.dimension import ( + E2C, + E2CDim, + C2E2C, + C2E2CDim, + C2CE, + C2E, + C2EDim, + CEDim, + CellDim, + EdgeDim, + KDim, +) + +@field_operator +def _mo_nh_diffusion_stencil_13( + kh_smag_e: Field[[EdgeDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + theta_v: Field[[CellDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + z_nabla2_e = kh_smag_e * inv_dual_edge_length * (theta_v(E2C[1]) - theta_v(E2C[0])) + return z_nabla2_e + +@field_operator +def _mo_nh_diffusion_stencil_14( + z_nabla2_e: Field[[EdgeDim, KDim], float], + geofac_div: Field[[CEDim], float], +) -> Field[[CellDim, KDim], float]: + z_temp = neighbor_sum(z_nabla2_e(C2E) * geofac_div(C2CE), axis=C2EDim) + return z_temp + +@field_operator +def _fused_mo_nh_diffusion_stencil_13_14( + kh_smag_e: Field[[EdgeDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + theta_v: Field[[CellDim, KDim], float], + geofac_div: Field[[CEDim], float], +) -> Field[[CellDim, KDim], float]: + z_nabla2_e = _mo_nh_diffusion_stencil_13(kh_smag_e, inv_dual_edge_length, theta_v) + z_temp = _mo_nh_diffusion_stencil_14(z_nabla2_e, geofac_div) + return z_temp + +@program +def fused_mo_nh_diffusion_stencil_13_14( + kh_smag_e: Field[[EdgeDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + theta_v: Field[[CellDim, KDim], float], + geofac_div: Field[[CEDim], float], + z_temp: Field[[CellDim, KDim], float], +): + _fused_mo_nh_diffusion_stencil_13_14(kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py index e439f742e..fed277e85 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py @@ -21,10 +21,10 @@ def _mo_nh_diffusion_stencil_10( w: Field[[CellDim, KDim], float], diff_multfac_n2w: Field[[KDim], float], - cell_area: Field[[CellDim], float], + area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - w = w + diff_multfac_n2w * (cell_area * z_nabla2_c) + w = w + diff_multfac_n2w * (area * z_nabla2_c) return w @@ -32,7 +32,7 @@ def _mo_nh_diffusion_stencil_10( def mo_nh_diffusion_stencil_10( w: Field[[CellDim, KDim], float], diff_multfac_n2w: Field[[KDim], float], - cell_area: Field[[CellDim], float], + area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], ): - _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w) + _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, area, z_nabla2_c, out=w) From eb814e12388e05864c3f982af05888a2b2a4640c Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 15 Nov 2022 10:55:30 +0100 Subject: [PATCH 002/263] add diffusion.py : initialization of diff_multfac_vn --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 49 +++++++++++++++++++ atm_dyn_iconam/tests/test_diffusion.py | 25 ++++++++++ 2 files changed, 74 insertions(+) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py create mode 100644 atm_dyn_iconam/tests/test_diffusion.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py new file mode 100644 index 000000000..e066f7b23 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -0,0 +1,49 @@ +from collections import namedtuple + +from functional.common import Field +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import minimum, broadcast + +from icon4py.common.dimension import KDim + +DiffusionTupleVT = namedtuple('DiffusionParamVT', 'v t') + +@field_operator +def _init_diff_multfac_vn(k4: float, dyn_substeps: float): + con = 1.0/128.0 + dyn = k4 * dyn_substeps / 3.0 + return broadcast(minimum(con, dyn), (KDim,)) + +@program +def init_diff_multfac_vn(k4:float, dyn_substeps: float, diff_multfac_vn:Field[[KDim], float]): + _init_diff_multfac_vn(k4, dyn_substeps, out=diff_multfac_vn) + + + +class DiffusionConfig: + """contains necessary parameter to configure a diffusion run: + - encapsulates namelist parameters + - TODO [ml] derived parameters (k2 to k8) here or in Diffusion class? + """ + ndyn_substeps = 5 + horizontal_diffusion_order = 5 + lhdiff_rcf = True # remove if always true + hdiff_efdt_ratio = 24.0 + + + lateral_boundary_denominator = DiffusionTupleVT( v=200.0, t=135.0) + + +class Diffusion: + """class that configures diffusion and does one diffusion step""" + def __init__(self, config: DiffusionConfig): + pass + #set smag_offset, smag_limit, diff_multfac_vn (depend on configuration: horizontal_diffusion_order, lhdiff_rcf, + # they are smag_offset and + + # MIN(1._wp/128._wp, diffusion_config(jg)%k4*REAL(ndyn_substeps,wp)/3._wp) # different for initial run! + #diff_mult_fac_vn: Field[[KDim] = + #smag_offset = + + + diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py new file mode 100644 index 000000000..76a122660 --- /dev/null +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -0,0 +1,25 @@ +import numpy as np + +from icon4py.atm_dyn_iconam.diffusion import init_diff_multfac_vn +from icon4py.common.dimension import KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import zero_field + + +def test_init_diff_multifac_vn_const(): + mesh = SimpleMesh() + diff_multfac_vn = zero_field(mesh, KDim) + + k4 = 1.0 + substeps = 5.0 + init_diff_multfac_vn(k4, substeps, diff_multfac_vn, offset_provider={}) + assert np.allclose(1.0/128.0 * np.ones(np.asarray(diff_multfac_vn).shape), diff_multfac_vn) + +def test_init_diff_multifac_vn_k4_substeps(): + mesh = SimpleMesh() + diff_multfac_vn = zero_field(mesh, KDim) + + k4 = 0.003 + substeps = 1.0 + init_diff_multfac_vn(k4, substeps, diff_multfac_vn, offset_provider={}) + assert np.allclose(k4*substeps / 3.0 * np.ones(np.asarray(diff_multfac_vn).shape), diff_multfac_vn) From 16cd2def08f9d1a83824f00629f8eb8a454f2c1d Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 15 Nov 2022 12:00:26 +0100 Subject: [PATCH 003/263] add pseudo p_patch: grid.py add initialization of smag_offset --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 104 +++++++++++++----- .../src/icon4py/atm_dyn_iconam/grid.py | 20 ++++ atm_dyn_iconam/tests/test_diffusion.py | 44 +++++++- 3 files changed, 136 insertions(+), 32 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/grid.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index e066f7b23..bce64bf72 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -1,49 +1,101 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + from collections import namedtuple +from typing import Final +import numpy as np from functional.common import Field from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import minimum, broadcast +from functional.ffront.fbuiltins import broadcast, minimum +from functional.iterator.embedded import np_as_located_field +from icon4py.atm_dyn_iconam.grid import GridConfig from icon4py.common.dimension import KDim -DiffusionTupleVT = namedtuple('DiffusionParamVT', 'v t') + +DiffusionTupleVT = namedtuple("DiffusionParamVT", "v t") + +# TODO +# def _setup_initial_diff_multfac_vn + @field_operator -def _init_diff_multfac_vn(k4: float, dyn_substeps: float): - con = 1.0/128.0 +def _setup_runtime_diff_multfac_vn(k4: float, dyn_substeps: float): + con = 1.0 / 128.0 dyn = k4 * dyn_substeps / 3.0 return broadcast(minimum(con, dyn), (KDim,)) -@program -def init_diff_multfac_vn(k4:float, dyn_substeps: float, diff_multfac_vn:Field[[KDim], float]): - _init_diff_multfac_vn(k4, dyn_substeps, out=diff_multfac_vn) +@field_operator +def _setup_smag_limit(diff_multfac_vn: Field[[KDim], float]): + return 0.125 - 4.0 * diff_multfac_vn -class DiffusionConfig: - """contains necessary parameter to configure a diffusion run: - - encapsulates namelist parameters - - TODO [ml] derived parameters (k2 to k8) here or in Diffusion class? - """ - ndyn_substeps = 5 - horizontal_diffusion_order = 5 - lhdiff_rcf = True # remove if always true - hdiff_efdt_ratio = 24.0 +@program +def init_diffusion_local_fields( + k4: float, + dyn_substeps: float, + diff_multfac_vn: Field[[KDim], float], + smag_limit: Field[[KDim], float], +) -> tuple[Field[[KDim], float], Field[[KDim], float]]: + _setup_runtime_diff_multfac_vn(k4, dyn_substeps, out=diff_multfac_vn) + _setup_smag_limit(diff_multfac_vn, out=smag_limit) - lateral_boundary_denominator = DiffusionTupleVT( v=200.0, t=135.0) +class DiffusionConfig: + """contains necessary parameter to configure a diffusion run. + - encapsulates namelist parameters + """ -class Diffusion: - """class that configures diffusion and does one diffusion step""" - def __init__(self, config: DiffusionConfig): - pass - #set smag_offset, smag_limit, diff_multfac_vn (depend on configuration: horizontal_diffusion_order, lhdiff_rcf, - # they are smag_offset and + grid = GridConfig() + ndyn_substeps = 5 # namelist mo_nonhydro_nml + horizontal_diffusion_order = 5 # namelist + lhdiff_rcf = True # namelist, remove if always true + hdiff_efdt_ratio = 24.0 # namelist + lateral_boundary_denominator = DiffusionTupleVT(v=200.0, t=135.0) + + # TODO [ml] keep those derived params in config or move to diffustion.__init__ + K2: Final[float] = 1.0 / (hdiff_efdt_ratio * 8.0) + K4: Final[float] = K2 / 8.0 + K8: Final[float] = K4 / 8.0 - # MIN(1._wp/128._wp, diffusion_config(jg)%k4*REAL(ndyn_substeps,wp)/3._wp) # different for initial run! - #diff_mult_fac_vn: Field[[KDim] = - #smag_offset = + def substep_as_float(self): + return float(self.ndyn_substeps) +class Diffusion: + """class that configures diffusion and does one diffusion step.""" + def __init__(self, config: DiffusionConfig): + """ + TODO [ml]: handle initial run: linit = .TRUE.: smag_offset and diff_multfac_vn are defined + differently. + """ + # different for init call smag_offset = 0 + self.smag_offset = 0.25 * config.K4 * config.substep_as_float() + + # different for initial run!, through diff_multfac_vn + self.diff_multfac_vn = np_as_located_field(KDim)( + np.zeros(config.grid.get_k_size()) + ) + self.smag_limit = np_as_located_field(KDim)(np.zeros(config.grid.get_k_size())) + + init_diffusion_local_fields( + config.K4, + config.substep_as_float(), + self.diff_multfac_vn, + self.smag_limit, + offset_provider={}, + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/grid.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/grid.py new file mode 100644 index 000000000..b6e1636e1 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/grid.py @@ -0,0 +1,20 @@ +# 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 + + +class GridConfig: + def __init__(self): + self._n_lev = 65 + + def get_k_size(self): + return self._n_lev diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 76a122660..1aacab482 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -1,25 +1,57 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + import numpy as np -from icon4py.atm_dyn_iconam.diffusion import init_diff_multfac_vn +from icon4py.atm_dyn_iconam.diffusion import init_diffusion_local_fields from icon4py.common.dimension import KDim from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import zero_field +def _smag_limit_numpy(diff_multfac_vn: np.array): + return 0.125 - 4.0 * diff_multfac_vn + + def test_init_diff_multifac_vn_const(): mesh = SimpleMesh() diff_multfac_vn = zero_field(mesh, KDim) + smag_offset = zero_field(mesh, KDim) + expected_diff_multfac_vn = 1.0 / 128.0 * np.ones(np.asarray(diff_multfac_vn).shape) + expected_smag_limit = _smag_limit_numpy(expected_diff_multfac_vn) k4 = 1.0 substeps = 5.0 - init_diff_multfac_vn(k4, substeps, diff_multfac_vn, offset_provider={}) - assert np.allclose(1.0/128.0 * np.ones(np.asarray(diff_multfac_vn).shape), diff_multfac_vn) + init_diffusion_local_fields( + k4, substeps, diff_multfac_vn, smag_offset, offset_provider={} + ) + assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) + assert np.allclose(expected_smag_limit, smag_offset) + def test_init_diff_multifac_vn_k4_substeps(): mesh = SimpleMesh() diff_multfac_vn = zero_field(mesh, KDim) - + smag_limit = zero_field(mesh, KDim) k4 = 0.003 substeps = 1.0 - init_diff_multfac_vn(k4, substeps, diff_multfac_vn, offset_provider={}) - assert np.allclose(k4*substeps / 3.0 * np.ones(np.asarray(diff_multfac_vn).shape), diff_multfac_vn) + expected_diff_multfac_vn = ( + k4 * substeps / 3.0 * np.ones(np.asarray(diff_multfac_vn).shape) + ) + expected_smag_limit = _smag_limit_numpy(expected_diff_multfac_vn) + + init_diffusion_local_fields( + k4, substeps, diff_multfac_vn, smag_limit, offset_provider={} + ) + assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) + assert np.allclose(expected_smag_limit, smag_limit) From 1002a2d052c234025a77e676da37b4a8aa5a2736 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 15 Nov 2022 17:41:57 +0100 Subject: [PATCH 004/263] - fix program output type - add outline for diffusion step function --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 68 ++++++++++++++++--- atm_dyn_iconam/tests/test_diffusion.py | 11 +++ 2 files changed, 71 insertions(+), 8 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index bce64bf72..26976909a 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -26,19 +26,21 @@ DiffusionTupleVT = namedtuple("DiffusionParamVT", "v t") -# TODO +# TODO [ml] initial RUN linit = TRUE # def _setup_initial_diff_multfac_vn @field_operator -def _setup_runtime_diff_multfac_vn(k4: float, dyn_substeps: float): +def _setup_runtime_diff_multfac_vn( + k4: float, dyn_substeps: float +) -> Field[[KDim], float]: con = 1.0 / 128.0 dyn = k4 * dyn_substeps / 3.0 return broadcast(minimum(con, dyn), (KDim,)) @field_operator -def _setup_smag_limit(diff_multfac_vn: Field[[KDim], float]): +def _setup_smag_limit(diff_multfac_vn: Field[[KDim], float]) -> Field[[KDim], float]: return 0.125 - 4.0 * diff_multfac_vn @@ -48,7 +50,7 @@ def init_diffusion_local_fields( dyn_substeps: float, diff_multfac_vn: Field[[KDim], float], smag_limit: Field[[KDim], float], -) -> tuple[Field[[KDim], float], Field[[KDim], float]]: +): _setup_runtime_diff_multfac_vn(k4, dyn_substeps, out=diff_multfac_vn) _setup_smag_limit(diff_multfac_vn, out=smag_limit) @@ -67,8 +69,10 @@ class DiffusionConfig: lateral_boundary_denominator = DiffusionTupleVT(v=200.0, t=135.0) # TODO [ml] keep those derived params in config or move to diffustion.__init__ + K2: Final[float] = 1.0 / (hdiff_efdt_ratio * 8.0) K4: Final[float] = K2 / 8.0 + K4W: Final[float] = K2 / 4.0 K8: Final[float] = K4 / 8.0 def substep_as_float(self): @@ -76,15 +80,22 @@ def substep_as_float(self): class Diffusion: - """class that configures diffusion and does one diffusion step.""" + """Class that configures diffusion and does one diffusion step.""" def __init__(self, config: DiffusionConfig): """ - TODO [ml]: handle initial run: linit = .TRUE.: smag_offset and diff_multfac_vn are defined - differently. + Initialize Diffusion granule. + + TODO [ml]: initial run: linit = .TRUE.: smag_offset and diff_multfac_vn are defined + differently. """ - # different for init call smag_offset = 0 + if not config: + raise + # different for init call: smag_offset = 0 self.smag_offset = 0.25 * config.K4 * config.substep_as_float() + self.diff_multfac_w = np.minimum( + 1.0 / 48.0, config.K4W * config.substep_as_float() + ) # different for initial run!, through diff_multfac_vn self.diff_multfac_vn = np_as_located_field(KDim)( @@ -99,3 +110,44 @@ def __init__(self, config: DiffusionConfig): self.smag_limit, offset_provider={}, ) + + self.diff_multfac_n2w = np_as_located_field(KDim)( + np.zeros(config.grid.get_k_size()) + ) + # TODO [ml] init diff_multfac_n2w + + def do_step( + self, + diagnostic_state, + prognostic_state, + metric_state, + interpolation_state, + dtime, + ): + """ + Run a diffusion step. + + inputs are + - output fields: + - input fields that have changed from one time step to another: + - simulation parameters: dtime - timestep + """ + # ------- + # OUTLINE + # ------- + # 1. CALL rbf_vec_interpol_vertex + # 2. HALO EXCHANGE -- CALL sync_patch_array_mult + # 3. CALL wrap_run_mo_nh_diffusion_stencil_01, CALL wrap_run_mo_nh_diffusion_stencil_02, + # CALL wrap_run_mo_nh_diffusion_stencil_03 + # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH + # 5. CALL rbf_vec_interpol_vertex_wp + # 6. HALO EXCHANGE -- CALL sync_patch_array_mult + # 7. CALL wrap_run_mo_nh_diffusion_stencil_04, wrap_run_mo_nh_diffusion_stencil_05 + # 7a. IF (l_limited_area .OR. jg > 1) CALL wrap_run_mo_nh_diffusion_stencil_06 + # 7b. call wrap_run_mo_nh_diffusion_stencil_07, wrap_run_mo_nh_diffusion_stencil_08, + # wrap_run_mo_nh_diffusion_stencil_09, wrap_run_mo_nh_diffusion_stencil_10 + # 8. HALO EXCHANGE: CALL sync_patch_array + # 9. call wrap_run_mo_nh_diffusion_stencil_11, wrap_run_mo_nh_diffusion_stencil_12, + # wrap_run_mo_nh_diffusion_stencil_13, wrap_run_mo_nh_diffusion_stencil_14, + # wrap_run_mo_nh_diffusion_stencil_15, wrap_run_mo_nh_diffusion_stencil_16 + # 10. HALO EXCHANGE sync_patch_array diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 1aacab482..8c3efeecf 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -12,6 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.atm_dyn_iconam.diffusion import init_diffusion_local_fields from icon4py.common.dimension import KDim @@ -55,3 +56,13 @@ def test_init_diff_multifac_vn_k4_substeps(): ) assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) assert np.allclose(expected_smag_limit, smag_limit) + + +@pytest.mark.xfail +def test_diffusion_init(): + pytest.fail("not implemented yet") + + +@pytest.mark.xfail +def test_diffusion_run(): + pytest.fail("not implemented yet") From cf6b1b48d8cc58a1b1d7e6223579bf11313ccf6b Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 16 Nov 2022 10:59:28 +0100 Subject: [PATCH 005/263] - extract derived parameters out of config --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 76 ++++++++++++------- atm_dyn_iconam/tests/test_diffusion.py | 27 ++++++- 2 files changed, 73 insertions(+), 30 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index 26976909a..d6e6fadea 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -54,11 +54,19 @@ def init_diffusion_local_fields( _setup_runtime_diff_multfac_vn(k4, dyn_substeps, out=diff_multfac_vn) _setup_smag_limit(diff_multfac_vn, out=smag_limit) +#@field_operator +def _enhanced_smagorinski_factor(): + # aka diff_multfac_smag + pass + +#@field_operator +def _diff_multfac_n2w(shift: int):# -> Field[[KDim], float]: + pass class DiffusionConfig: """contains necessary parameter to configure a diffusion run. - - encapsulates namelist parameters + - encapsulates namelist parameters and derived parameters (for now) """ grid = GridConfig() @@ -67,34 +75,41 @@ class DiffusionConfig: lhdiff_rcf = True # namelist, remove if always true hdiff_efdt_ratio = 24.0 # namelist lateral_boundary_denominator = DiffusionTupleVT(v=200.0, t=135.0) - - # TODO [ml] keep those derived params in config or move to diffustion.__init__ - - K2: Final[float] = 1.0 / (hdiff_efdt_ratio * 8.0) - K4: Final[float] = K2 / 8.0 - K4W: Final[float] = K2 / 4.0 - K8: Final[float] = K4 / 8.0 + hdiff_smag_fac = 0.025 # namelist def substep_as_float(self): return float(self.ndyn_substeps) + +class DiffusionParams: + def __init__(self, config:DiffusionConfig): + #TODO [ml] logging for case KX == 0 + #TODO [ml] calculation for x_dom (jg) > 2..n_dom + + self.K2: Final[float] = 1.0 / (config.hdiff_efdt_ratio * 8.0) \ + if config.hdiff_efdt_ratio > 0.0 else 0.0 + self.K4: Final[float] = self.K2 / 8.0 + self.K8: Final[float] = self.K2 / 64.0 + self.K4W: Final[float] = self.K2 / 4.0 + + class Diffusion: """Class that configures diffusion and does one diffusion step.""" - def __init__(self, config: DiffusionConfig): + def __init__(self, config: DiffusionConfig, params: DiffusionParams): """ Initialize Diffusion granule. TODO [ml]: initial run: linit = .TRUE.: smag_offset and diff_multfac_vn are defined differently. """ - if not config: - raise + self.params: params + # different for init call: smag_offset = 0 - self.smag_offset = 0.25 * config.K4 * config.substep_as_float() + self.smag_offset = 0.25 * params.K4 * config.substep_as_float() self.diff_multfac_w = np.minimum( - 1.0 / 48.0, config.K4W * config.substep_as_float() + 1.0 / 48.0, params.K4W * config.substep_as_float() ) # different for initial run!, through diff_multfac_vn @@ -104,7 +119,7 @@ def __init__(self, config: DiffusionConfig): self.smag_limit = np_as_located_field(KDim)(np.zeros(config.grid.get_k_size())) init_diffusion_local_fields( - config.K4, + params.K4, config.substep_as_float(), self.diff_multfac_vn, self.smag_limit, @@ -116,6 +131,11 @@ def __init__(self, config: DiffusionConfig): ) # TODO [ml] init diff_multfac_n2w + self.enh_smag_fac = np_as_located_field(KDim)( + np.zeros(config.grid.get_k_size()) + ) + # TODO [ml] init enh_smag_fac + def do_step( self, diagnostic_state, @@ -135,19 +155,17 @@ def do_step( # ------- # OUTLINE # ------- - # 1. CALL rbf_vec_interpol_vertex - # 2. HALO EXCHANGE -- CALL sync_patch_array_mult - # 3. CALL wrap_run_mo_nh_diffusion_stencil_01, CALL wrap_run_mo_nh_diffusion_stencil_02, - # CALL wrap_run_mo_nh_diffusion_stencil_03 - # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH - # 5. CALL rbf_vec_interpol_vertex_wp - # 6. HALO EXCHANGE -- CALL sync_patch_array_mult - # 7. CALL wrap_run_mo_nh_diffusion_stencil_04, wrap_run_mo_nh_diffusion_stencil_05 - # 7a. IF (l_limited_area .OR. jg > 1) CALL wrap_run_mo_nh_diffusion_stencil_06 - # 7b. call wrap_run_mo_nh_diffusion_stencil_07, wrap_run_mo_nh_diffusion_stencil_08, - # wrap_run_mo_nh_diffusion_stencil_09, wrap_run_mo_nh_diffusion_stencil_10 - # 8. HALO EXCHANGE: CALL sync_patch_array - # 9. call wrap_run_mo_nh_diffusion_stencil_11, wrap_run_mo_nh_diffusion_stencil_12, - # wrap_run_mo_nh_diffusion_stencil_13, wrap_run_mo_nh_diffusion_stencil_14, - # wrap_run_mo_nh_diffusion_stencil_15, wrap_run_mo_nh_diffusion_stencil_16 + # 1. CALL rbf_vec_interpol_vertex + # 2. HALO EXCHANGE -- CALL sync_patch_array_mult + # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 + # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH + # 5. CALL rbf_vec_interpol_vertex_wp + # 6. HALO EXCHANGE -- CALL sync_patch_array_mult + # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 + # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 + # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, + # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 + # 8. HALO EXCHANGE: CALL sync_patch_array + # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, + # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 # 10. HALO EXCHANGE sync_patch_array diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 8c3efeecf..ceef60b2f 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -14,7 +14,8 @@ import numpy as np import pytest -from icon4py.atm_dyn_iconam.diffusion import init_diffusion_local_fields +from icon4py.atm_dyn_iconam.diffusion import init_diffusion_local_fields, DiffusionConfig, \ + DiffusionParams from icon4py.common.dimension import KDim from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import zero_field @@ -66,3 +67,27 @@ def test_diffusion_init(): @pytest.mark.xfail def test_diffusion_run(): pytest.fail("not implemented yet") + + +def test_diffusion_coefficients_with_hdiff_efdt_ratio(): + config :DiffusionConfig = DiffusionConfig() + config.hdiff_efdt_ratio = 1.0 + + params = DiffusionParams(config) + + assert params.K2 == pytest.approx(0.125, abs=1e-12) + assert params.K4 == pytest.approx(0.125 / 8.0, abs=1e-12) + assert params.K8 == pytest.approx(0.125/64.0, abs=1e-12) + assert params.K4W == pytest.approx(0.125/4.0, abs=1e-12) + + +def test_diffusion_coefficients_without_hdiff_efdt_ratio(): + config: DiffusionConfig = DiffusionConfig() + config.hdiff_efdt_ratio = 0.0 + + params = DiffusionParams(config) + + assert params.K2 == 0.0 + assert params.K4 == 0.0 + assert params.K8 == 0.0 + assert params.K4W == 0.0 From a4a129173bbfdb5b7ff86e78d1470db38bff306d Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 16 Nov 2022 11:48:16 +0100 Subject: [PATCH 006/263] - add preliminarz initialization of diff_multfac_n2w --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 30 ++++++++++++------- .../src/icon4py/atm_dyn_iconam/grid.py | 5 ++++ atm_dyn_iconam/tests/test_diffusion.py | 11 +++++-- 3 files changed, 34 insertions(+), 12 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index d6e6fadea..5839d9e17 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -54,13 +54,17 @@ def init_diffusion_local_fields( _setup_runtime_diff_multfac_vn(k4, dyn_substeps, out=diff_multfac_vn) _setup_smag_limit(diff_multfac_vn, out=smag_limit) -#@field_operator -def _enhanced_smagorinski_factor(): - # aka diff_multfac_smag - pass +@field_operator +def _set_zero_k(): + return broadcast(0.0, (KDim, )) + +@program +def init_nabla2_factor_in_upper_damping_zone(diff_multfac_n2w: Field[[KDim], float]): + # fix missing the IF (nrdmax(jg) > 1) (l.332 following) + _set_zero_k(out=diff_multfac_n2w) #@field_operator -def _diff_multfac_n2w(shift: int):# -> Field[[KDim], float]: +def _diff_multfac_n2w(shift: int, nrdmax: Field[[],float]):# -> Field[[KDim], float]: pass class DiffusionConfig: @@ -68,7 +72,6 @@ class DiffusionConfig: - encapsulates namelist parameters and derived parameters (for now) """ - grid = GridConfig() ndyn_substeps = 5 # namelist mo_nonhydro_nml horizontal_diffusion_order = 5 # namelist @@ -126,15 +129,22 @@ def __init__(self, config: DiffusionConfig, params: DiffusionParams): offset_provider={}, ) - self.diff_multfac_n2w = np_as_located_field(KDim)( + self.enh_smag_fac = np_as_located_field(KDim)( np.zeros(config.grid.get_k_size()) ) - # TODO [ml] init diff_multfac_n2w - self.enh_smag_fac = np_as_located_field(KDim)( + + + + self.diff_multfac_n2w = np_as_located_field(KDim)( np.zeros(config.grid.get_k_size()) ) - # TODO [ml] init enh_smag_fac + ## TODO [ml] missing parts... related to nrdmax + init_enhanced_smagorinski_factor(self.diff_multfac_n2w, offset_provider={}) + # TODO [ml] init diff_multfac_n2w + + + def do_step( self, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/grid.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/grid.py index b6e1636e1..7caa06a29 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/grid.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/grid.py @@ -15,6 +15,11 @@ class GridConfig: def __init__(self): self._n_lev = 65 + # see mo_model_domimp_patches.f90 l. 405 + self.n_shift_total = 0 def get_k_size(self): return self._n_lev + + def get_n_shift(self): + return self.n_shift_total diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index ceef60b2f..4215860a3 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -15,10 +15,10 @@ import pytest from icon4py.atm_dyn_iconam.diffusion import init_diffusion_local_fields, DiffusionConfig, \ - DiffusionParams + DiffusionParams, init_nabla2_factor_in_upper_damping_zone from icon4py.common.dimension import KDim from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import zero_field +from icon4py.testutils.utils import zero_field, random_field def _smag_limit_numpy(diff_multfac_vn: np.array): @@ -59,6 +59,13 @@ def test_init_diff_multifac_vn_k4_substeps(): assert np.allclose(expected_smag_limit, smag_limit) +def test_init_nabla2_factor_in_upper_damping_zone(): + mesh = SimpleMesh() + diff_multfac_n2w = random_field(mesh, KDim) + init_nabla2_factor_in_upper_damping_zone(diff_multfac_n2w, offset_provider={}) + + assert np.allclose(0, diff_multfac_n2w) + @pytest.mark.xfail def test_diffusion_init(): pytest.fail("not implemented yet") From 9284472a3227ae857ebd7b96c52c87554ab8d46b Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 16 Nov 2022 14:12:00 +0100 Subject: [PATCH 007/263] calculate smagorinski factors (I) --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 116 ++++++++++++++---- atm_dyn_iconam/tests/test_diffusion.py | 39 ++++-- 2 files changed, 122 insertions(+), 33 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index 5839d9e17..578263820 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - +import math from collections import namedtuple from typing import Final @@ -26,6 +26,7 @@ DiffusionTupleVT = namedtuple("DiffusionParamVT", "v t") + # TODO [ml] initial RUN linit = TRUE # def _setup_initial_diff_multfac_vn @@ -54,47 +55,115 @@ def init_diffusion_local_fields( _setup_runtime_diff_multfac_vn(k4, dyn_substeps, out=diff_multfac_vn) _setup_smag_limit(diff_multfac_vn, out=smag_limit) + @field_operator def _set_zero_k(): - return broadcast(0.0, (KDim, )) + return broadcast(0.0, (KDim,)) + @program def init_nabla2_factor_in_upper_damping_zone(diff_multfac_n2w: Field[[KDim], float]): # fix missing the IF (nrdmax(jg) > 1) (l.332 following) _set_zero_k(out=diff_multfac_n2w) -#@field_operator -def _diff_multfac_n2w(shift: int, nrdmax: Field[[],float]):# -> Field[[KDim], float]: - pass class DiffusionConfig: """contains necessary parameter to configure a diffusion run. - encapsulates namelist parameters and derived parameters (for now) + + currently we use the MCH r04b09_dsl experiment as constants here. These should + be read from config and the default from mo_diffusion_nml.f90 set as defaults. + + TODO: [ml] read from config + TODO: [ml] handle dependencies on other namelists (below...) """ - grid = GridConfig() - ndyn_substeps = 5 # namelist mo_nonhydro_nml - horizontal_diffusion_order = 5 # namelist + + # from namelist diffusion_nml + diffusion_type = 5 # hdiff_order ! order of nabla operator for diffusion + lhdiff_vn = True # ! diffusion on the horizontal wind field + lhdiff_temp = True # ! diffusion on the temperature field + lhdiff_w = True # ! diffusion on the vertical wind field lhdiff_rcf = True # namelist, remove if always true - hdiff_efdt_ratio = 24.0 # namelist + itype_vn_diffu = 1 # ! reconstruction method used for Smagorinsky diffusion + itype_t_diffu = 2 # ! discretization of temperature diffusion + hdiff_efdt_ratio = 24.0 # ! ratio of e-folding time to time step + hdiff_smag_fac = 0.025 # ! scaling factor for Smagorinsky diffusion + # defaults: + + # TODO [ml]: external stuff, p_patch, other than diffusion namelist + grid = GridConfig() + + # namelist nonhydrostatic_nml + l_zdiffu_t = ( + True # ! l_zdiffu_t: specifies computation of Smagorinsky temperature diffusion + ) + ndyn_substeps = 5 + + # from namelist gridref_nml + # denom_diffu_v = 150 ! denominator for lateral boundary diffusion of velocity lateral_boundary_denominator = DiffusionTupleVT(v=200.0, t=135.0) - hdiff_smag_fac = 0.025 # namelist + + # from namelist grid_nml + l_limited_area = True def substep_as_float(self): return float(self.ndyn_substeps) - class DiffusionParams: - def __init__(self, config:DiffusionConfig): - #TODO [ml] logging for case KX == 0 - #TODO [ml] calculation for x_dom (jg) > 2..n_dom - - self.K2: Final[float] = 1.0 / (config.hdiff_efdt_ratio * 8.0) \ - if config.hdiff_efdt_ratio > 0.0 else 0.0 + def __init__(self, config: DiffusionConfig): + # TODO [ml] logging for case KX == 0 + # TODO [ml] generrally calculation for x_dom (jg) > 2..n_dom, why is jg special + + self.K2: Final[float] = ( + 1.0 / (config.hdiff_efdt_ratio * 8.0) + if config.hdiff_efdt_ratio > 0.0 + else 0.0 + ) self.K4: Final[float] = self.K2 / 8.0 self.K8: Final[float] = self.K2 / 64.0 self.K4W: Final[float] = self.K2 / 4.0 + ( + self.smagorinski_factor, + self.smagorinski_height, + ) = self.determine_enhanced_smagorinski_factor(config) + + def determine_enhanced_smagorinski_factor(self, config: DiffusionConfig): + """Enhanced Smagorinsky diffusion factor. + + Smagorinsky diffusion factor is defined as a profile in height + above sea level with 4 height sections. + + It is calculated/used only in the case of diffusion_type 3 or 5 + """ + match config.diffusion_type: + case 5: + ( + smagorinski_factor, + smagorinski_height, + ) = self._calculate_enhanced_smagorinski_factor(config) + case 4: + # according to mo_nh_diffusion.f90 this isn't used anywhere the factor is only + # used for diffusion_type (3,5) but the defaults are only defined for iequations=3 + smagorinski_factor = ( + config.hdiff_smag_fac if config.hdiff_smag_fac else 0.15, + ) + smagorinski_height = None + case _: + pass + return smagorinski_factor, smagorinski_height + + @staticmethod + def _calculate_enhanced_smagorinski_factor(config: DiffusionConfig): + magic_sqrt = math.sqrt(1600.0 * (1600 + 50000.0)) + magic_fac2_value = 2e-6 * (1600.0 + 25000.0 + magic_sqrt) + magic_z2 = 1600.0 + 50000.0 + magic_sqrt + initial_smagorinski_factor = (config.hdiff_smag_fac, magic_fac2_value, 0.0, 1.0) + hdiff_smagorinski_heights = (32500.0, magic_z2, 50000.0, 90000) + + enhanced_factor = initial_smagorinski_factor + return enhanced_factor, hdiff_smagorinski_heights class Diffusion: @@ -133,18 +202,13 @@ def __init__(self, config: DiffusionConfig, params: DiffusionParams): np.zeros(config.grid.get_k_size()) ) - - - self.diff_multfac_n2w = np_as_located_field(KDim)( np.zeros(config.grid.get_k_size()) ) - ## TODO [ml] missing parts... related to nrdmax - init_enhanced_smagorinski_factor(self.diff_multfac_n2w, offset_provider={}) - # TODO [ml] init diff_multfac_n2w - - - + # TODO [ml] missing parts... related to nrdmax + init_nabla2_factor_in_upper_damping_zone( + self.diff_multfac_n2w, offset_provider={} + ) def do_step( self, diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 4215860a3..c0f2d7246 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -14,11 +14,15 @@ import numpy as np import pytest -from icon4py.atm_dyn_iconam.diffusion import init_diffusion_local_fields, DiffusionConfig, \ - DiffusionParams, init_nabla2_factor_in_upper_damping_zone +from icon4py.atm_dyn_iconam.diffusion import ( + DiffusionConfig, + DiffusionParams, + init_diffusion_local_fields, + init_nabla2_factor_in_upper_damping_zone, +) from icon4py.common.dimension import KDim from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import zero_field, random_field +from icon4py.testutils.utils import random_field, zero_field def _smag_limit_numpy(diff_multfac_vn: np.array): @@ -66,6 +70,7 @@ def test_init_nabla2_factor_in_upper_damping_zone(): assert np.allclose(0, diff_multfac_n2w) + @pytest.mark.xfail def test_diffusion_init(): pytest.fail("not implemented yet") @@ -77,15 +82,15 @@ def test_diffusion_run(): def test_diffusion_coefficients_with_hdiff_efdt_ratio(): - config :DiffusionConfig = DiffusionConfig() + config: DiffusionConfig = DiffusionConfig() config.hdiff_efdt_ratio = 1.0 - params = DiffusionParams(config) + params = DiffusionParams(config) assert params.K2 == pytest.approx(0.125, abs=1e-12) assert params.K4 == pytest.approx(0.125 / 8.0, abs=1e-12) - assert params.K8 == pytest.approx(0.125/64.0, abs=1e-12) - assert params.K4W == pytest.approx(0.125/4.0, abs=1e-12) + assert params.K8 == pytest.approx(0.125 / 64.0, abs=1e-12) + assert params.K4W == pytest.approx(0.125 / 4.0, abs=1e-12) def test_diffusion_coefficients_without_hdiff_efdt_ratio(): @@ -98,3 +103,23 @@ def test_diffusion_coefficients_without_hdiff_efdt_ratio(): assert params.K4 == 0.0 assert params.K8 == 0.0 assert params.K4W == 0.0 + + +def test_smagorinski_factor_for_diffusion_type_4(): + config: DiffusionConfig = DiffusionConfig() + config.hdiff_smag_fac = 0.15 + config.diffusion_type = 4 + + params = DiffusionParams(config) + assert len(params.smagorinski_factor) == 1 + assert params.smagorinski_factor[0] == pytest.approx(0.15, abs=1e-16) + assert params.smagorinski_height is None + + +def test_smagorinsik_factor_diffusion_type_5(): + config: DiffusionConfig = DiffusionConfig() + config.hdiff_smag_fac = 0.15 + config.diffusion_type = 5 + + params = DiffusionParams(config) + assert len(params.smagorinski_factor) == 4 From f53ad2000203f655fe8e44ba4b100e221c58733e Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 16 Nov 2022 16:36:31 +0100 Subject: [PATCH 008/263] calculate smagorinski factors (II) --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 50 ++++++++++++++++--- atm_dyn_iconam/tests/test_diffusion.py | 30 +++++++++-- 2 files changed, 69 insertions(+), 11 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index 578263820..c8d45b5b4 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -15,14 +15,12 @@ from typing import Final import numpy as np -from functional.common import Field from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import broadcast, minimum +from functional.ffront.fbuiltins import broadcast, minimum, maximum, Field from functional.iterator.embedded import np_as_located_field from icon4py.atm_dyn_iconam.grid import GridConfig -from icon4py.common.dimension import KDim - +from icon4py.common.dimension import KDim, Koff DiffusionTupleVT = namedtuple("DiffusionParamVT", "v t") @@ -55,7 +53,26 @@ def init_diffusion_local_fields( _setup_runtime_diff_multfac_vn(k4, dyn_substeps, out=diff_multfac_vn) _setup_smag_limit(diff_multfac_vn, out=smag_limit) +@field_operator +def _en_smag_fac(hdiff_smag_fac:float, hdiff_smag_fac2:float, hdiff_smag_fac3:float, hdiff_smag_fac4:float, hdiff_smag_z:float, hdiff_smag_z2:float, hdiff_smag_z3:float, hdiff_smag_z4:float, vect_a:Field[[KDim],float])->Field[[KDim], float]: + dz21 = hdiff_smag_z2 - hdiff_smag_z + alin = (hdiff_smag_fac2 - hdiff_smag_fac) / dz21 + df32 = hdiff_smag_fac3 - hdiff_smag_fac2 + df42 = hdiff_smag_fac4 - hdiff_smag_fac2 + dz32 = hdiff_smag_z3 - hdiff_smag_z2 + dz42 = hdiff_smag_z4 - hdiff_smag_z2 + bsq = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32)) + asq = df32 / dz32 - bsq * dz32 + zf = 0.5 * (vect_a + vect_a(Koff[1])) + + dzlin = minimum(dz21, maximum(0., zf - hdiff_smag_z )) + dzqdr = minimum(dz42, maximum(0., zf - hdiff_smag_z2)) + enh_smag_fac = hdiff_smag_fac + (dzlin * alin) + dzqdr * (asq + dzqdr * bsq) + return enh_smag_fac +@program +def enhanced_smagorinski_factor(hdiff_smag_fac:float, hdiff_smag_fac2:float, hdiff_smag_fac3:float, hdiff_smag_fac4:float, hdiff_smag_z:float, hdiff_smag_z2:float, hdiff_smag_z3:float, hdiff_smag_z4:float, vect_a:Field[[KDim],float], enh_smag_fac:Field[[KDim], float]): + _en_smag_fac(hdiff_smag_fac,hdiff_smag_fac2, hdiff_smag_fac3, hdiff_smag_fac4, hdiff_smag_z, hdiff_smag_z2, hdiff_smag_z3, hdiff_smag_z4, vect_a, out=enh_smag_fac) @field_operator def _set_zero_k(): return broadcast(0.0, (KDim,)) @@ -154,16 +171,33 @@ def determine_enhanced_smagorinski_factor(self, config: DiffusionConfig): pass return smagorinski_factor, smagorinski_height + @staticmethod def _calculate_enhanced_smagorinski_factor(config: DiffusionConfig): + # initial values from mo_diffusion_nml.f90 magic_sqrt = math.sqrt(1600.0 * (1600 + 50000.0)) magic_fac2_value = 2e-6 * (1600.0 + 25000.0 + magic_sqrt) magic_z2 = 1600.0 + 50000.0 + magic_sqrt - initial_smagorinski_factor = (config.hdiff_smag_fac, magic_fac2_value, 0.0, 1.0) - hdiff_smagorinski_heights = (32500.0, magic_z2, 50000.0, 90000) + factor = (config.hdiff_smag_fac, magic_fac2_value, 0.0, 1.0) + heights = (32500.0, magic_z2, 50000.0, 90000.0) + # enhance factors according to mo_nh_diffusion.f90/diffusion + + df32 = factor[2] - factor[1] + df42 = factor[3] - factor[1] + + dz32 = heights[2]- heights[1] + dz42 = heights[3] - heights[1] + + alin = (factor[1] - factor[0]) / heights[1] - heights[0] + bsquare = (df42 * dz32 - df32 * dz42) / (dz32*dz42*(dz42-dz32)) + asquare = df32/dz32-bsquare*dz32 + + + + - enhanced_factor = initial_smagorinski_factor - return enhanced_factor, hdiff_smagorinski_heights + enhanced_factor = factor + return enhanced_factor, heights class Diffusion: diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index c0f2d7246..fb922418d 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -13,12 +13,13 @@ import numpy as np import pytest +from numpy import int32 from icon4py.atm_dyn_iconam.diffusion import ( DiffusionConfig, DiffusionParams, init_diffusion_local_fields, - init_nabla2_factor_in_upper_damping_zone, + init_nabla2_factor_in_upper_damping_zone, _en_smag_fac, enhanced_smagorinski_factor, ) from icon4py.common.dimension import KDim from icon4py.testutils.simple_mesh import SimpleMesh @@ -28,7 +29,6 @@ def _smag_limit_numpy(diff_multfac_vn: np.array): return 0.125 - 4.0 * diff_multfac_vn - def test_init_diff_multifac_vn_const(): mesh = SimpleMesh() diff_multfac_vn = zero_field(mesh, KDim) @@ -71,6 +71,18 @@ def test_init_nabla2_factor_in_upper_damping_zone(): assert np.allclose(0, diff_multfac_n2w) +def test_enhanced_smagorinski_factor(): + mesh = SimpleMesh() + a_vec = random_field(mesh, KDim) + result = zero_field(mesh, KDim) + z = (0.1,0.2,0.3, 0.4) + fac = (1./3, 0.5, 1.33,0.077) + enhanced_smagorinski_factor(*fac, *z, a_vec, result, offset_provider={"Koff":KDim}) + #enhanced_smagorinski_factor(nshift, fac[0], fac[1], fac[2], fac[3], z[0], z[1], z[2], z[3], + # a_vec, result, offset_provider={"Koff": KDim}) + + + @pytest.mark.xfail def test_diffusion_init(): pytest.fail("not implemented yet") @@ -116,10 +128,22 @@ def test_smagorinski_factor_for_diffusion_type_4(): assert params.smagorinski_height is None -def test_smagorinsik_factor_diffusion_type_5(): +def test_smagorinski_heights_diffusion_type_5_are_consistent(): + config: DiffusionConfig = DiffusionConfig() + config.hdiff_smag_fac = 0.15 + config.diffusion_type = 5 + + params = DiffusionParams(config) + assert len(params.smagorinski_height) == 4 + assert min(params.smagorinski_height) == params.smagorinski_height[0] + assert max(params.smagorinski_height) == params.smagorinski_height[-1] + assert np.all(params.smagorinski_height[1:] > params.smagorinski_height[:-1]) + +def test_smagorinski_factor_diffusion_type_5(): config: DiffusionConfig = DiffusionConfig() config.hdiff_smag_fac = 0.15 config.diffusion_type = 5 params = DiffusionParams(config) assert len(params.smagorinski_factor) == 4 + From 1e7a85a5c71795b7a1a97519aea62129547bd68b Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 17 Nov 2022 18:01:13 +0100 Subject: [PATCH 009/263] calculate smagorinski factors (III) --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 98 ++++++++++++------- atm_dyn_iconam/tests/test_diffusion.py | 42 +++++--- 2 files changed, 96 insertions(+), 44 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index c8d45b5b4..bd7412109 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -16,12 +16,13 @@ import numpy as np from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import broadcast, minimum, maximum, Field +from functional.ffront.fbuiltins import Field, broadcast, maximum, minimum from functional.iterator.embedded import np_as_located_field from icon4py.atm_dyn_iconam.grid import GridConfig from icon4py.common.dimension import KDim, Koff + DiffusionTupleVT = namedtuple("DiffusionParamVT", "v t") @@ -53,26 +54,63 @@ def init_diffusion_local_fields( _setup_runtime_diff_multfac_vn(k4, dyn_substeps, out=diff_multfac_vn) _setup_smag_limit(diff_multfac_vn, out=smag_limit) + @field_operator -def _en_smag_fac(hdiff_smag_fac:float, hdiff_smag_fac2:float, hdiff_smag_fac3:float, hdiff_smag_fac4:float, hdiff_smag_z:float, hdiff_smag_z2:float, hdiff_smag_z3:float, hdiff_smag_z4:float, vect_a:Field[[KDim],float])->Field[[KDim], float]: +def _en_smag_fac_for_zero_nshift( + hdiff_smag_fac: float, + hdiff_smag_fac2: float, + hdiff_smag_fac3: float, + hdiff_smag_fac4: float, + hdiff_smag_z: float, + hdiff_smag_z2: float, + hdiff_smag_z3: float, + hdiff_smag_z4: float, + vect_a: Field[[KDim], float], +) -> Field[[KDim], float]: dz21 = hdiff_smag_z2 - hdiff_smag_z alin = (hdiff_smag_fac2 - hdiff_smag_fac) / dz21 df32 = hdiff_smag_fac3 - hdiff_smag_fac2 df42 = hdiff_smag_fac4 - hdiff_smag_fac2 dz32 = hdiff_smag_z3 - hdiff_smag_z2 dz42 = hdiff_smag_z4 - hdiff_smag_z2 - bsq = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32)) - asq = df32 / dz32 - bsq * dz32 + + bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32)) + aqdr = df32 / dz32 - bqdr * dz32 zf = 0.5 * (vect_a + vect_a(Koff[1])) - dzlin = minimum(dz21, maximum(0., zf - hdiff_smag_z )) - dzqdr = minimum(dz42, maximum(0., zf - hdiff_smag_z2)) - enh_smag_fac = hdiff_smag_fac + (dzlin * alin) + dzqdr * (asq + dzqdr * bsq) + dzlin = minimum(dz21, maximum(0.0, zf - hdiff_smag_z)) + dzqdr = minimum(dz42, maximum(0.0, zf - hdiff_smag_z2)) + enh_smag_fac = hdiff_smag_fac + (dzlin * alin) + dzqdr * (aqdr + dzqdr * bqdr) return enh_smag_fac + @program -def enhanced_smagorinski_factor(hdiff_smag_fac:float, hdiff_smag_fac2:float, hdiff_smag_fac3:float, hdiff_smag_fac4:float, hdiff_smag_z:float, hdiff_smag_z2:float, hdiff_smag_z3:float, hdiff_smag_z4:float, vect_a:Field[[KDim],float], enh_smag_fac:Field[[KDim], float]): - _en_smag_fac(hdiff_smag_fac,hdiff_smag_fac2, hdiff_smag_fac3, hdiff_smag_fac4, hdiff_smag_z, hdiff_smag_z2, hdiff_smag_z3, hdiff_smag_z4, vect_a, out=enh_smag_fac) +def enhanced_smagorinski_factor( + hdiff_smag_fac: float, + hdiff_smag_fac2: float, + hdiff_smag_fac3: float, + hdiff_smag_fac4: float, + hdiff_smag_z: float, + hdiff_smag_z2: float, + hdiff_smag_z3: float, + hdiff_smag_z4: float, + vect_a: Field[[KDim], float], + enh_smag_fac: Field[[KDim], float], +): + _en_smag_fac_for_zero_nshift( + hdiff_smag_fac, + hdiff_smag_fac2, + hdiff_smag_fac3, + hdiff_smag_fac4, + hdiff_smag_z, + hdiff_smag_z2, + hdiff_smag_z3, + hdiff_smag_z4, + vect_a, + out=enh_smag_fac, + ) + + @field_operator def _set_zero_k(): return broadcast(0.0, (KDim,)) @@ -144,9 +182,9 @@ def __init__(self, config: DiffusionConfig): ( self.smagorinski_factor, self.smagorinski_height, - ) = self.determine_enhanced_smagorinski_factor(config) + ) = self.determine_smagorinski_factor(config) - def determine_enhanced_smagorinski_factor(self, config: DiffusionConfig): + def determine_smagorinski_factor(self, config: DiffusionConfig): """Enhanced Smagorinsky diffusion factor. Smagorinsky diffusion factor is defined as a profile in height @@ -159,7 +197,7 @@ def determine_enhanced_smagorinski_factor(self, config: DiffusionConfig): ( smagorinski_factor, smagorinski_height, - ) = self._calculate_enhanced_smagorinski_factor(config) + ) = self.diffusion_type_5_smagorinski_factor(config) case 4: # according to mo_nh_diffusion.f90 this isn't used anywhere the factor is only # used for diffusion_type (3,5) but the defaults are only defined for iequations=3 @@ -171,39 +209,26 @@ def determine_enhanced_smagorinski_factor(self, config: DiffusionConfig): pass return smagorinski_factor, smagorinski_height - @staticmethod - def _calculate_enhanced_smagorinski_factor(config: DiffusionConfig): + def diffusion_type_5_smagorinski_factor(config: DiffusionConfig): # initial values from mo_diffusion_nml.f90 magic_sqrt = math.sqrt(1600.0 * (1600 + 50000.0)) magic_fac2_value = 2e-6 * (1600.0 + 25000.0 + magic_sqrt) magic_z2 = 1600.0 + 50000.0 + magic_sqrt factor = (config.hdiff_smag_fac, magic_fac2_value, 0.0, 1.0) heights = (32500.0, magic_z2, 50000.0, 90000.0) - # enhance factors according to mo_nh_diffusion.f90/diffusion - - df32 = factor[2] - factor[1] - df42 = factor[3] - factor[1] - - dz32 = heights[2]- heights[1] - dz42 = heights[3] - heights[1] - - alin = (factor[1] - factor[0]) / heights[1] - heights[0] - bsquare = (df42 * dz32 - df32 * dz42) / (dz32*dz42*(dz42-dz32)) - asquare = df32/dz32-bsquare*dz32 - - - - - - enhanced_factor = factor - return enhanced_factor, heights + return factor, heights class Diffusion: """Class that configures diffusion and does one diffusion step.""" - def __init__(self, config: DiffusionConfig, params: DiffusionParams): + def __init__( + self, + config: DiffusionConfig, + params: DiffusionParams, + a_vect: Field[[KDim], float], + ): """ Initialize Diffusion granule. @@ -235,6 +260,13 @@ def __init__(self, config: DiffusionConfig, params: DiffusionParams): self.enh_smag_fac = np_as_located_field(KDim)( np.zeros(config.grid.get_k_size()) ) + enhanced_smagorinski_factor( + *params.smagorinski_factor, + *params.smagorinski_height, + a_vect, + self.enh_smag_fac, + offset_provider={"Koff", KDim}, + ) self.diff_multfac_n2w = np_as_located_field(KDim)( np.zeros(config.grid.get_k_size()) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index fb922418d..c08806913 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -13,13 +13,13 @@ import numpy as np import pytest -from numpy import int32 from icon4py.atm_dyn_iconam.diffusion import ( DiffusionConfig, DiffusionParams, + enhanced_smagorinski_factor, init_diffusion_local_fields, - init_nabla2_factor_in_upper_damping_zone, _en_smag_fac, enhanced_smagorinski_factor, + init_nabla2_factor_in_upper_damping_zone, ) from icon4py.common.dimension import KDim from icon4py.testutils.simple_mesh import SimpleMesh @@ -29,6 +29,7 @@ def _smag_limit_numpy(diff_multfac_vn: np.array): return 0.125 - 4.0 * diff_multfac_vn + def test_init_diff_multifac_vn_const(): mesh = SimpleMesh() diff_multfac_vn = zero_field(mesh, KDim) @@ -72,15 +73,29 @@ def test_init_nabla2_factor_in_upper_damping_zone(): def test_enhanced_smagorinski_factor(): + def enhanced_smagorinski_factor_np(factor_in, heigths_in, a_vec): + alin = (factor_in[1] - factor_in[0]) / (heigths_in[1] - heigths_in[0]) + df32 = factor_in[2] - factor_in[1] + df42 = factor_in[3] - factor_in[1] + dz32 = heigths_in[2] - heigths_in[1] + dz42 = heigths_in[3] - heigths_in[1] + bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32)) + aqdr = df32 / dz32 - bqdr * dz32 + zf = 0.5 * (a_vec[:-1] + a_vec[1:]) + max0 = np.maximum(0.0, zf - heigths_in[0]) + dzlin = np.minimum(heigths_in[1] - heigths_in[0], max0) + max1 = np.maximum(0.0, zf - heigths_in[1]) + dzqdr = np.minimum(heigths_in[3] - heigths_in[1], max1) + return factor_in[0] + dzlin * alin + dzqdr * (aqdr + dzqdr * bqdr) + mesh = SimpleMesh() - a_vec = random_field(mesh, KDim) + a_vec = random_field(mesh, KDim, low=1.0, high=10.0) result = zero_field(mesh, KDim) - z = (0.1,0.2,0.3, 0.4) - fac = (1./3, 0.5, 1.33,0.077) - enhanced_smagorinski_factor(*fac, *z, a_vec, result, offset_provider={"Koff":KDim}) - #enhanced_smagorinski_factor(nshift, fac[0], fac[1], fac[2], fac[3], z[0], z[1], z[2], z[3], - # a_vec, result, offset_provider={"Koff": KDim}) - + fac = (0.67, 0.5, 1.3, 0.8) + z = (0.1, 0.2, 0.3, 0.4) + enhanced_smagorinski_factor(*fac, *z, a_vec, result, offset_provider={"Koff": KDim}) + enhanced_smag_fac_np = enhanced_smagorinski_factor_np(fac, z, np.asarray(a_vec)) + assert np.allclose(enhanced_smag_fac_np, np.asarray(result[:-1])) @pytest.mark.xfail @@ -137,7 +152,11 @@ def test_smagorinski_heights_diffusion_type_5_are_consistent(): assert len(params.smagorinski_height) == 4 assert min(params.smagorinski_height) == params.smagorinski_height[0] assert max(params.smagorinski_height) == params.smagorinski_height[-1] - assert np.all(params.smagorinski_height[1:] > params.smagorinski_height[:-1]) + assert params.smagorinski_height[0] < params.smagorinski_height[1] + assert params.smagorinski_height[1] < params.smagorinski_height[3] + assert params.smagorinski_height[2] != params.smagorinski_height[1] + assert params.smagorinski_height[2] != params.smagorinski_height[3] + def test_smagorinski_factor_diffusion_type_5(): config: DiffusionConfig = DiffusionConfig() @@ -145,5 +164,6 @@ def test_smagorinski_factor_diffusion_type_5(): config.diffusion_type = 5 params = DiffusionParams(config) + assert len(params.smagorinski_factor) == len(params.smagorinski_height) assert len(params.smagorinski_factor) == 4 - + assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor))) From f8ade7ddc2902ed0788c3ced9522c8b154a5d35a Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 18 Nov 2022 09:07:07 +0100 Subject: [PATCH 010/263] - fuse init of diffusion local fields - add calculation of nrdmax --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 48 +++++++++++++++---- .../src/icon4py/atm_dyn_iconam/vertical.py | 33 +++++++++++++ atm_dyn_iconam/tests/test_vertical.py | 34 +++++++++++++ 3 files changed, 107 insertions(+), 8 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/vertical.py create mode 100644 atm_dyn_iconam/tests/test_vertical.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index bd7412109..b01e49eed 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -44,6 +44,15 @@ def _setup_smag_limit(diff_multfac_vn: Field[[KDim], float]) -> Field[[KDim], fl return 0.125 - 4.0 * diff_multfac_vn +@field_operator +def _init_diffusion_local_fields( + k4: float, dyn_substeps: float +) -> tuple[Field[[KDim], float], Field[[KDim], float]]: + diff_multfac_vn = _setup_runtime_diff_multfac_vn(k4, dyn_substeps) + smag_limit = _setup_smag_limit(diff_multfac_vn) + return diff_multfac_vn, smag_limit + + @program def init_diffusion_local_fields( k4: float, @@ -51,8 +60,7 @@ def init_diffusion_local_fields( diff_multfac_vn: Field[[KDim], float], smag_limit: Field[[KDim], float], ): - _setup_runtime_diff_multfac_vn(k4, dyn_substeps, out=diff_multfac_vn) - _setup_smag_limit(diff_multfac_vn, out=smag_limit) + _init_diffusion_local_fields(k4, dyn_substeps, out=(diff_multfac_vn, smag_limit)) @field_operator @@ -116,6 +124,16 @@ def _set_zero_k(): return broadcast(0.0, (KDim,)) +@field_operator +def _init_diff_multifac_n2w( + offset: int, vct_a: Field[[KDim], float] +) -> Field[[KDim], float]: + # offset = nshift + nrdmax, but we assume nshift=0 + init_diff_multifac_n2w = _set_zero_k() + # 1.0 / 12.0 * (vct_a - vct_a(Koff[offset + 1])) / (vct_a(2) - vct_a(Koff[offset]))**4 + return init_diff_multifac_n2w + + @program def init_nabla2_factor_in_upper_damping_zone(diff_multfac_n2w: Field[[KDim], float]): # fix missing the IF (nrdmax(jg) > 1) (l.332 following) @@ -130,10 +148,15 @@ class DiffusionConfig: currently we use the MCH r04b09_dsl experiment as constants here. These should be read from config and the default from mo_diffusion_nml.f90 set as defaults. + + TODO: [ml] read from config TODO: [ml] handle dependencies on other namelists (below...) """ + # TODO [ml]: external stuff grid related, p_patch, + grid = GridConfig() + # from namelist diffusion_nml diffusion_type = 5 # hdiff_order ! order of nabla operator for diffusion lhdiff_vn = True # ! diffusion on the horizontal wind field @@ -146,20 +169,20 @@ class DiffusionConfig: hdiff_smag_fac = 0.025 # ! scaling factor for Smagorinsky diffusion # defaults: - # TODO [ml]: external stuff, p_patch, other than diffusion namelist - grid = GridConfig() - - # namelist nonhydrostatic_nml + # + # from parent namelist nonhydrostatic_nml l_zdiffu_t = ( True # ! l_zdiffu_t: specifies computation of Smagorinsky temperature diffusion ) + damp_height = 12500 ndyn_substeps = 5 - # from namelist gridref_nml + # from other namelists + # namelist gridref_nml # denom_diffu_v = 150 ! denominator for lateral boundary diffusion of velocity lateral_boundary_denominator = DiffusionTupleVT(v=200.0, t=135.0) - # from namelist grid_nml + # namelist grid_nml -> TODO [ml] could go to grid config? l_limited_area = True def substep_as_float(self): @@ -167,6 +190,8 @@ def substep_as_float(self): class DiffusionParams: + """Calculates derived quantities depending on the diffusion config.""" + def __init__(self, config: DiffusionConfig): # TODO [ml] logging for case KX == 0 # TODO [ml] generrally calculation for x_dom (jg) > 2..n_dom, why is jg special @@ -295,6 +320,13 @@ def do_step( # ------- # OUTLINE # ------- + # Oa logging + # 0b call timer start + # 0c. apply dtime to enh_smag_factor + # TODO does not work because the self.enh_smag_fact is a field: do this where the factor is used, ie inside stencil + + timestep_scaled_smagorinski_factor = self.enh_smag_fac * dtime + # 1. CALL rbf_vec_interpol_vertex # 2. HALO EXCHANGE -- CALL sync_patch_array_mult # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/vertical.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/vertical.py new file mode 100644 index 000000000..767bd1c33 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/vertical.py @@ -0,0 +1,33 @@ +# 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 + + +class VerticalGridConfig: + num_k_levels = 65 + + +class VerticalModelConfig: + rayleigh_damping_height = 12500 # height of rayleigh damping in [m] + vct_a = [] # read from vertical_coord_tables + + +class VerticalModelParams: + def __init__(self, vertical_model_config: VerticalModelConfig): + self.index_of_damping_height = np.argmax( + vertical_model_config.vct_a >= vertical_model_config.rayleigh_damping_height + ) + + def get_index_of_damping_layer(self): + return self.index_of_damping_height diff --git a/atm_dyn_iconam/tests/test_vertical.py b/atm_dyn_iconam/tests/test_vertical.py new file mode 100644 index 000000000..a463a3158 --- /dev/null +++ b/atm_dyn_iconam/tests/test_vertical.py @@ -0,0 +1,34 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import math + +import numpy as np +import pytest + +from icon4py.atm_dyn_iconam.vertical import ( + VerticalModelConfig, + VerticalModelParams, +) + + +@pytest.mark.parametrize( + "max_h,damping,delta", + [(60000, 34000, 612), (12000, 10000, 100), (109000, 45000, 123)], +) +def test_nrdmax_calculation(max_h, damping, delta): + vertical_model_config = VerticalModelConfig() + vertical_model_config.rayleigh_damping_height = damping + vertical_model_config.vct_a = np.arange(0, max_h, delta) + vertical_params = VerticalModelParams(vertical_model_config=vertical_model_config) + assert vertical_params.get_index_of_damping_layer() == math.ceil(damping / delta) From 730de3ac0a6e0044b906225bb67760487c6cb5b1 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 18 Nov 2022 18:54:11 +0100 Subject: [PATCH 011/263] - deleted duplicated grid stuff - implement numpy initialization for diff_multfac_n2w --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 36 +++++--------- .../src/icon4py/atm_dyn_iconam/grid.py | 25 ---------- .../src/icon4py/atm_dyn_iconam/vertical.py | 47 ++++++++++++++++--- atm_dyn_iconam/tests/test_diffusion.py | 9 ---- atm_dyn_iconam/tests/test_vertical.py | 12 ++--- 5 files changed, 57 insertions(+), 72 deletions(-) delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/grid.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index b01e49eed..a0de4039c 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -19,7 +19,10 @@ from functional.ffront.fbuiltins import Field, broadcast, maximum, minimum from functional.iterator.embedded import np_as_located_field -from icon4py.atm_dyn_iconam.grid import GridConfig +from icon4py.atm_dyn_iconam.vertical import ( + VerticalGridConfig, + VerticalModelParams, +) from icon4py.common.dimension import KDim, Koff @@ -124,22 +127,6 @@ def _set_zero_k(): return broadcast(0.0, (KDim,)) -@field_operator -def _init_diff_multifac_n2w( - offset: int, vct_a: Field[[KDim], float] -) -> Field[[KDim], float]: - # offset = nshift + nrdmax, but we assume nshift=0 - init_diff_multifac_n2w = _set_zero_k() - # 1.0 / 12.0 * (vct_a - vct_a(Koff[offset + 1])) / (vct_a(2) - vct_a(Koff[offset]))**4 - return init_diff_multifac_n2w - - -@program -def init_nabla2_factor_in_upper_damping_zone(diff_multfac_n2w: Field[[KDim], float]): - # fix missing the IF (nrdmax(jg) > 1) (l.332 following) - _set_zero_k(out=diff_multfac_n2w) - - class DiffusionConfig: """contains necessary parameter to configure a diffusion run. @@ -155,7 +142,10 @@ class DiffusionConfig: """ # TODO [ml]: external stuff grid related, p_patch, - grid = GridConfig() + grid = VerticalGridConfig() + vertical_params = VerticalModelParams( + rayleigh_damping_height=12500, vct_a=np.zeros(grid.get_k_size()) + ) # from namelist diffusion_nml diffusion_type = 5 # hdiff_order ! order of nabla operator for diffusion @@ -293,12 +283,10 @@ def __init__( offset_provider={"Koff", KDim}, ) - self.diff_multfac_n2w = np_as_located_field(KDim)( - np.zeros(config.grid.get_k_size()) - ) - # TODO [ml] missing parts... related to nrdmax - init_nabla2_factor_in_upper_damping_zone( - self.diff_multfac_n2w, offset_provider={} + self.diff_multfac_n2w = ( + config.vertical_params.init_nabla2_factor_in_upper_damping_zone( + k_size=config.grid.get_k_size() + ) ) def do_step( diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/grid.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/grid.py deleted file mode 100644 index 7caa06a29..000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/grid.py +++ /dev/null @@ -1,25 +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 - - -class GridConfig: - def __init__(self): - self._n_lev = 65 - # see mo_model_domimp_patches.f90 l. 405 - self.n_shift_total = 0 - - def get_k_size(self): - return self._n_lev - - def get_n_shift(self): - return self.n_shift_total diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/vertical.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/vertical.py index 767bd1c33..e099d8632 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/vertical.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/vertical.py @@ -12,22 +12,57 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +from functional.common import Field +from functional.iterator.embedded import np_as_located_field + +from icon4py.common.dimension import KDim class VerticalGridConfig: - num_k_levels = 65 + def __init__(self): + self.num_k_levels = 65 + self.n_shift_total = 0 + def get_k_size(self): + return self.num_k_levels -class VerticalModelConfig: - rayleigh_damping_height = 12500 # height of rayleigh damping in [m] - vct_a = [] # read from vertical_coord_tables + def get_n_shift(self): + return self.n_shift_total class VerticalModelParams: - def __init__(self, vertical_model_config: VerticalModelConfig): + def __init__(self, vct_a: np.ndarray, rayleigh_damping_height: float = 12500.0): + """ + Arguments. + + - vct_a: # read from vertical_coord_tables + - rayleigh_damping_height height of rayleigh damping in [m] mo_nonhydro_nml + """ + self.rayleigh_damping_height = rayleigh_damping_height + self.vct_a = vct_a + # TODO klevels in ICON are inversed self.index_of_damping_height = np.argmax( - vertical_model_config.vct_a >= vertical_model_config.rayleigh_damping_height + self.vct_a >= self.rayleigh_damping_height ) def get_index_of_damping_layer(self): return self.index_of_damping_height + + def get_physical_heights(self) -> Field[[KDim], float]: + return np_as_located_field(KDim)(self.vct_a) + + def init_nabla2_factor_in_upper_damping_zone( + self, k_size: int + ) -> Field[[KDim], float]: + # this assumes n_shift == 0 + buffer = np.zeros(k_size) + buffer[2 : self.index_of_damping_height] = ( + 1.0 + / 12.0 + * ( + self.vct_a[2 : self.index_of_damping_height] + - self.vct_a[self.index_of_damping_height + 1] + ) + / (self.vct_a[2] - self.vct_a[self.index_of_damping_height + 1]) + ) + return np_as_located_field(KDim)(buffer) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index c08806913..f418581ab 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -19,7 +19,6 @@ DiffusionParams, enhanced_smagorinski_factor, init_diffusion_local_fields, - init_nabla2_factor_in_upper_damping_zone, ) from icon4py.common.dimension import KDim from icon4py.testutils.simple_mesh import SimpleMesh @@ -64,14 +63,6 @@ def test_init_diff_multifac_vn_k4_substeps(): assert np.allclose(expected_smag_limit, smag_limit) -def test_init_nabla2_factor_in_upper_damping_zone(): - mesh = SimpleMesh() - diff_multfac_n2w = random_field(mesh, KDim) - init_nabla2_factor_in_upper_damping_zone(diff_multfac_n2w, offset_provider={}) - - assert np.allclose(0, diff_multfac_n2w) - - def test_enhanced_smagorinski_factor(): def enhanced_smagorinski_factor_np(factor_in, heigths_in, a_vec): alin = (factor_in[1] - factor_in[0]) / (heigths_in[1] - heigths_in[0]) diff --git a/atm_dyn_iconam/tests/test_vertical.py b/atm_dyn_iconam/tests/test_vertical.py index a463a3158..12ac75aad 100644 --- a/atm_dyn_iconam/tests/test_vertical.py +++ b/atm_dyn_iconam/tests/test_vertical.py @@ -16,19 +16,15 @@ import numpy as np import pytest -from icon4py.atm_dyn_iconam.vertical import ( - VerticalModelConfig, - VerticalModelParams, -) +from icon4py.atm_dyn_iconam.vertical import VerticalModelParams @pytest.mark.parametrize( "max_h,damping,delta", [(60000, 34000, 612), (12000, 10000, 100), (109000, 45000, 123)], ) +# TODO [ml] klevels run from num_lev (ground) to 1 (top most) def test_nrdmax_calculation(max_h, damping, delta): - vertical_model_config = VerticalModelConfig() - vertical_model_config.rayleigh_damping_height = damping - vertical_model_config.vct_a = np.arange(0, max_h, delta) - vertical_params = VerticalModelParams(vertical_model_config=vertical_model_config) + vct_a = np.arange(0, max_h, delta) + vertical_params = VerticalModelParams(rayleigh_damping_height=damping, vct_a=vct_a) assert vertical_params.get_index_of_damping_layer() == math.ceil(damping / delta) From c3f5ea109da3d305c0c4d4f1bfa476622bb0bb92 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 24 Nov 2022 09:24:57 +0100 Subject: [PATCH 012/263] - add data classes for prognostic and diagnostic state - call stencil 01 in diffusion.run - add factor method with default vals (MCH_CH_r04b09) for DiffusionConfig --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 260 +++++++++++++++--- .../src/icon4py/atm_dyn_iconam/horizontal.py | 28 ++ .../src/icon4py/atm_dyn_iconam/icon_mesh.py | 113 ++++++++ .../atm_dyn_iconam/interpolation_state.py | 28 ++ .../src/icon4py/atm_dyn_iconam/prognostic.py | 32 +++ .../src/icon4py/atm_dyn_iconam/vertical.py | 68 ----- atm_dyn_iconam/tests/test_diffusion.py | 20 +- atm_dyn_iconam/tests/test_vertical.py | 2 +- 8 files changed, 431 insertions(+), 120 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_mesh.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/prognostic.py delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/vertical.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index a0de4039c..13c195b04 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -19,14 +19,21 @@ from functional.ffront.fbuiltins import Field, broadcast, maximum, minimum from functional.iterator.embedded import np_as_located_field -from icon4py.atm_dyn_iconam.vertical import ( - VerticalGridConfig, - VerticalModelParams, +from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig +from icon4py.atm_dyn_iconam.icon_mesh import MeshConfig, VerticalModelParams +from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState +from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( + mo_intp_rbf_rbf_vec_interpol_vertex, ) -from icon4py.common.dimension import KDim, Koff +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import ( + _mo_nh_diffusion_stencil_01, +) +from icon4py.atm_dyn_iconam.prognostic import PrognosticState +from icon4py.common.dimension import ECVDim, EdgeDim, KDim, Koff, VertexDim DiffusionTupleVT = namedtuple("DiffusionParamVT", "v t") +CartesianVectorTuple = namedtuple("CartesianVectorTuple", "x y") # TODO [ml] initial RUN linit = TRUE @@ -47,6 +54,85 @@ def _setup_smag_limit(diff_multfac_vn: Field[[KDim], float]) -> Field[[KDim], fl return 0.125 - 4.0 * diff_multfac_vn +# TODO [ml] rename! +@field_operator +def _mo_nh_diffusion_stencil_01_scale_dtime( + enh_smag_fac: Field[[KDim], float], + tangent_orientation: Field[[EdgeDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + inv_vert_vert_length: Field[[EdgeDim], float], + u_vert: Field[[VertexDim, KDim], float], + v_vert: Field[[VertexDim, KDim], float], + primal_normal_vert_x: Field[[ECVDim], float], + primal_normal_vert_y: Field[[ECVDim], float], + dual_normal_vert_x: Field[[ECVDim], float], + dual_normal_vert_y: Field[[ECVDim], float], + vn: Field[[EdgeDim, KDim], float], + smag_limit: Field[[KDim], float], + smag_offset: float, + dtime: float, +) -> tuple[ + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], +]: + diff_multfac_smag = enh_smag_fac * dtime + return _mo_nh_diffusion_stencil_01( + diff_multfac_smag, + tangent_orientation, + inv_primal_edge_length, + inv_vert_vert_length, + u_vert, + v_vert, + primal_normal_vert_x, + primal_normal_vert_y, + dual_normal_vert_x, + dual_normal_vert_y, + vn, + smag_limit, + smag_offset, + ) + + +@program +def mo_nh_diffusion_stencil_01_scaled_dtime( + enh_smag_fac: Field[[KDim], float], + tangent_orientation: Field[[EdgeDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + inv_vert_vert_length: Field[[EdgeDim], float], + u_vert: Field[[VertexDim, KDim], float], + v_vert: Field[[VertexDim, KDim], float], + primal_normal_vert_x: Field[[ECVDim], float], + primal_normal_vert_y: Field[[ECVDim], float], + dual_normal_vert_x: Field[[ECVDim], float], + dual_normal_vert_y: Field[[ECVDim], float], + vn: Field[[EdgeDim, KDim], float], + smag_limit: Field[[KDim], float], + kh_smag_e: Field[[EdgeDim, KDim], float], + kh_smag_ec: Field[[EdgeDim, KDim], float], + z_nabla2_e: Field[[EdgeDim, KDim], float], + smag_offset: float, + dtime: float, +): + _mo_nh_diffusion_stencil_01_scale_dtime( + enh_smag_fac, + tangent_orientation, + inv_primal_edge_length, + inv_vert_vert_length, + u_vert, + v_vert, + primal_normal_vert_x, + primal_normal_vert_y, + dual_normal_vert_x, + dual_normal_vert_y, + vn, + smag_limit, + smag_offset, + dtime, + out=(kh_smag_e, kh_smag_ec, z_nabla2_e), + ) + + @field_operator def _init_diffusion_local_fields( k4: float, dyn_substeps: float @@ -123,10 +209,20 @@ def enhanced_smagorinski_factor( @field_operator -def _set_zero_k(): +def _set_zero_k() -> Field[[KDim], float]: return broadcast(0.0, (KDim,)) +@field_operator +def _set_zero_v_k() -> Field[[VertexDim, KDim], float]: + return broadcast(0.0, (VertexDim, KDim)) + + +@program +def set_zero_v_k(field: Field[[VertexDim, KDim], float]): + _set_zero_v_k(out=field) + + class DiffusionConfig: """contains necessary parameter to configure a diffusion run. @@ -141,39 +237,50 @@ class DiffusionConfig: TODO: [ml] handle dependencies on other namelists (below...) """ - # TODO [ml]: external stuff grid related, p_patch, - grid = VerticalGridConfig() - vertical_params = VerticalModelParams( - rayleigh_damping_height=12500, vct_a=np.zeros(grid.get_k_size()) - ) + @classmethod + def create_with_defaults(cls): + """ + Create DiffusionConfig. - # from namelist diffusion_nml - diffusion_type = 5 # hdiff_order ! order of nabla operator for diffusion - lhdiff_vn = True # ! diffusion on the horizontal wind field - lhdiff_temp = True # ! diffusion on the temperature field - lhdiff_w = True # ! diffusion on the vertical wind field - lhdiff_rcf = True # namelist, remove if always true - itype_vn_diffu = 1 # ! reconstruction method used for Smagorinsky diffusion - itype_t_diffu = 2 # ! discretization of temperature diffusion - hdiff_efdt_ratio = 24.0 # ! ratio of e-folding time to time step - hdiff_smag_fac = 0.025 # ! scaling factor for Smagorinsky diffusion - # defaults: - - # - # from parent namelist nonhydrostatic_nml - l_zdiffu_t = ( - True # ! l_zdiffu_t: specifies computation of Smagorinsky temperature diffusion - ) - damp_height = 12500 - ndyn_substeps = 5 + initialize with values from exp.mch_ch_r04b09_dsl namelist + """ + horizontal = HorizontalMeshConfig( + num_vertices=50000, num_cells=50000, num_edges=50000 + ) + gridConfig = MeshConfig(horizontalMesh=horizontal) + verticalParams = VerticalModelParams( + rayleigh_damping_height=12500, vct_a=np.zeros(gridConfig.get_num_k_levels()) + ) + return cls(grid_config=gridConfig, vertical_params=verticalParams) + + def __init__(self, grid_config: MeshConfig, vertical_params: VerticalModelParams): + # TODO [ml]: external stuff grid related, p_patch, + self.grid = grid_config + self.vertical_params = vertical_params + # from namelist diffusion_nml + self.diffusion_type = 5 # hdiff_order ! order of nabla operator for diffusion + self.lhdiff_vn = True # ! diffusion on the horizontal wind field + self.lhdiff_temp = True # ! diffusion on the temperature field + self.lhdiff_w = True # ! diffusion on the vertical wind field + self.lhdiff_rcf = True # namelist, remove if always true + self.itype_vn_diffu = ( + 1 # ! reconstruction method used for Smagorinsky diffusion + ) + self.itype_t_diffu = 2 # ! discretization of temperature diffusion + self.hdiff_efdt_ratio = 24.0 # ! ratio of e-folding time to time step + self.hdiff_smag_fac = 0.025 # ! scaling factor for Smagorinsky diffusion - # from other namelists - # namelist gridref_nml - # denom_diffu_v = 150 ! denominator for lateral boundary diffusion of velocity - lateral_boundary_denominator = DiffusionTupleVT(v=200.0, t=135.0) + # from other namelists + # from parent namelist nonhydrostatic_nml + self.l_zdiffu_t = True # ! l_zdiffu_t: specifies computation of Smagorinsky temperature diffusion + self.ndyn_substeps = 5 - # namelist grid_nml -> TODO [ml] could go to grid config? - l_limited_area = True + # namelist gridref_nml + # denom_diffu_v = 150 ! denominator for lateral boundary diffusion of velocity + self.lateral_boundary_denominator = DiffusionTupleVT(v=200.0, t=135.0) + + # namelist grid_nml -> TODO [ml] should go to grid config? + self.l_limited_area = True def substep_as_float(self): return float(self.ndyn_substeps) @@ -260,9 +367,11 @@ def __init__( # different for initial run!, through diff_multfac_vn self.diff_multfac_vn = np_as_located_field(KDim)( - np.zeros(config.grid.get_k_size()) + np.zeros(config.grid.get_num_k_levels()) + ) + self.smag_limit = np_as_located_field(KDim)( + np.zeros(config.grid.get_num_k_levels()) ) - self.smag_limit = np_as_located_field(KDim)(np.zeros(config.grid.get_k_size())) init_diffusion_local_fields( params.K4, @@ -273,7 +382,7 @@ def __init__( ) self.enh_smag_fac = np_as_located_field(KDim)( - np.zeros(config.grid.get_k_size()) + np.zeros(config.grid.get_num_k_levels()) ) enhanced_smagorinski_factor( *params.smagorinski_factor, @@ -285,17 +394,38 @@ def __init__( self.diff_multfac_n2w = ( config.vertical_params.init_nabla2_factor_in_upper_damping_zone( - k_size=config.grid.get_k_size() + k_size=config.grid.get_num_k_levels() ) ) + self.u_vert = np_as_located_field(VertexDim, KDim)( + np.zeros(config.grid.get_num_vertices(), config.grid.get_num_k_levels()) + ) + self.v_vert = np_as_located_field(VertexDim, KDim)( + np.zeros(config.grid.get_num_vertices(), config.grid.get_num_k_levels()) + ) + allocate_ek = np_as_located_field(VertexDim, KDim)( + np.zeros(config.grid.get_num_edges(), config.grid.get_num_k_levels()) + ) + self.kh_smag_e = allocate_ek + self.kh_smag_ec = allocate_ek + self.z_nabla2_e = allocate_ek - def do_step( + def run( self, diagnostic_state, - prognostic_state, + prognostic_state: PrognosticState, metric_state, - interpolation_state, + interpolation_state: InterpolationState, dtime, + tangent_orientation: Field[[EdgeDim], float], + inverse_primal_edge_lengths: Field[[EdgeDim], float], + inverse_vertical_vertex_lengths: Field[[EdgeDim], float], + primal_normal_vert: CartesianVectorTuple[ + Field[[ECVDim], float], Field[[ECVDim], float] + ], + dual_normal_vert: CartesianVectorTuple[ + Field[[ECVDim], float], Field[[ECVDim], float] + ], ): """ Run a diffusion step. @@ -310,16 +440,56 @@ def do_step( # ------- # Oa logging # 0b call timer start - # 0c. apply dtime to enh_smag_factor - # TODO does not work because the self.enh_smag_fact is a field: do this where the factor is used, ie inside stencil + # ~~0c. apply dtime to enh_smag_factor~~ done inside stencil - timestep_scaled_smagorinski_factor = self.enh_smag_fac * dtime + # TODO is this needed? + set_zero_v_k(self.u_vert) + set_zero_v_k(self.v_vert) # 1. CALL rbf_vec_interpol_vertex + + mo_intp_rbf_rbf_vec_interpol_vertex( + prognostic_state.normal_wind, + interpolation_state.rbf_coeff_1, + interpolation_state.rbf_coeff_2, + self.u_vert, + self.v_vert, + ) # 2. HALO EXCHANGE -- CALL sync_patch_array_mult # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 + + # tangent_orientation = p_patch % edges % tangent_orientation(:, 1) + # inv_primal_edge_length=p_patch%edges%inv_primal_edge_length(:,1) + # inv_vert_vert_length = p_patch % edges % inv_vert_vert_length(:, 1), + # primal_normal_vert_x = p_patch % edges % primal_normal_vert_x(:,:, 1) + # primal_normal_vert_y = p_patch % edges % primal_normal_vert_y(:,:, 1) + # dual_normal_vert_x=p_patch%edges%dual_normal_vert_x(:,:,1) + # dual_normal_vert_x = p_patch % edges % dual_normal_vert_x(:,:, 1) + + mo_nh_diffusion_stencil_01_scaled_dtime( + self.enh_smag_fac, + tangent_orientation, + inverse_primal_edge_lengths, + inverse_vertical_vertex_lengths, + self.u_vert, + self.v_vert, + primal_normal_vert.x, + primal_normal_vert.y, + dual_normal_vert.x, + dual_normal_vert.y, + prognostic_state.normal_wind, + self.smag_limit, + self.kh_smag_e, + self.kh_smag_ec, + self.z_nabla2_e, + self.smag_offset, + dtime, + ) + # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH + # 5. CALL rbf_vec_interpol_vertex_wp + # 6. HALO EXCHANGE -- CALL sync_patch_array_mult # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py new file mode 100644 index 000000000..8bbbc3c2e --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py @@ -0,0 +1,28 @@ +# 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 + + +class HorizontalMeshConfig: + def __init__(self, num_vertices: int, num_edges: int, num_cells: int): + self._num_vertices = num_vertices + self._num_edges = num_edges + self._num_cells = num_cells + + def get_num_vertices(self): + return self._num_vertices + + def get_num_edges(self): + return self._num_edges + + def get_num_cells(self): + return self._num_cells diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_mesh.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_mesh.py new file mode 100644 index 000000000..435133edc --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_mesh.py @@ -0,0 +1,113 @@ +# 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 functional.common import Field +from functional.iterator.embedded import np_as_located_field + +from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig +from icon4py.common.dimension import EdgeDim, KDim + + +class MeshConfig: + def __init__(self, horizontalMesh: HorizontalMeshConfig): + self.num_k_levels = 65 + self.n_shift_total = 0 + self.horizontal = horizontalMesh + + def get_num_k_levels(self): + return self.num_k_levels + + def get_n_shift(self): + return self.n_shift_total + + def get_num_vertices(self): + return self.horizontal._num_vertices + + def get_num_edges(self): + return self.horizontal._num_edges + + +class IconMesh: + def __init__(self, config: MeshConfig): + self.config = config + # TODO [ml] calculate the content of this stuff? read in? + self._tangent_orientation = np_as_located_field(EdgeDim, KDim)( + np.zeros(self.config.get_num_edges(), self.config.get_num_k_levels()) + ) + self._triangle_edge_inverse_length = np_as_located_field(EdgeDim)( + np.zeros(self.config.get_num_edges()) + ) + # normal to triangle edge projected to the location of the vertices + self._primal_normal_vert = ( + np_as_located_field(EdgeDim)(np.zeros(self.config.get_num_edges())), + np_as_located_field(EdgeDim)(np.zeros(self.config.get_num_edges())), + ) + + # tangent to triangle edge, projected to the location of the vertices + self._dual_normal_vert_x = ( + np_as_located_field(EdgeDim)(np.zeros(self.config.get_num_edges())), + np_as_located_field(EdgeDim)(np.zeros(self.config.get_num_edges())), + ) + + # TODO [ml] geometry + def tangent_orientation(self): + return self._tangent_orientation + + def inv_primal_edge_length(self): + return self._triangle_edge_inverse_length + + def primal_normal_vert(self): + return self._primal_normal_vert + + def dual_normal_vert_x(self): + return self._dual_normal_vert_x + + +class VerticalModelParams: + def __init__(self, vct_a: np.ndarray, rayleigh_damping_height: float = 12500.0): + """ + Contains physical parameters defined on the grid. + + Args: + vct_a: TODO read from vertical_coord_tables! + rayleigh_damping_height height of rayleigh damping in [m] mo_nonhydro_nml + """ + self.rayleigh_damping_height = rayleigh_damping_height + self.vct_a = vct_a + # TODO klevels in ICON are inverted! + self.index_of_damping_height = np.argmax( + self.vct_a >= self.rayleigh_damping_height + ) + + def get_index_of_damping_layer(self): + return self.index_of_damping_height + + def get_physical_heights(self) -> Field[[KDim], float]: + return np_as_located_field(KDim)(self.vct_a) + + def init_nabla2_factor_in_upper_damping_zone( + self, k_size: int + ) -> Field[[KDim], float]: + # this assumes n_shift == 0 + buffer = np.zeros(k_size) + buffer[2 : self.index_of_damping_height] = ( + 1.0 + / 12.0 + * ( + self.vct_a[2 : self.index_of_damping_height] + - self.vct_a[self.index_of_damping_height + 1] + ) + / (self.vct_a[2] - self.vct_a[self.index_of_damping_height + 1]) + ) + return np_as_located_field(KDim)(buffer) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py new file mode 100644 index 000000000..d0425ad92 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py @@ -0,0 +1,28 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from dataclasses import dataclass + +from functional.common import Field + +from icon4py.common.dimension import V2EDim, VertexDim + + +@dataclass +class InterpolationState: + rbf_coeff_1: Field[ + [VertexDim, V2EDim], float + ] # rbf_vec_coeff_v_1(nproma, rbf_vec_dim_v, nblks_v) + rbf_coeff_2: Field[ + [VertexDim, V2EDim], float + ] # rbf_vec_coeff_v_2(nproma, rbf_vec_dim_v, nblks_v) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/prognostic.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/prognostic.py new file mode 100644 index 000000000..c7f9e1b95 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/prognostic.py @@ -0,0 +1,32 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from dataclasses import dataclass + +from functional.common import Field + +from icon4py.common.dimension import CellDim, EdgeDim, KDim + + +@dataclass +class PrognosticState: + """Class that contains the prognostic state. + + corresponds to ICON t_nh_prog + """ + + vertical_wind: Field[[CellDim, KDim], float] # w(nproma, nlevp1, nblks_c) [m/s] + normal_wind: Field[[EdgeDim, KDim], float] # vn(nproma, nlev, nblks_e) [m/s] + density: Field[[CellDim, KDim], float] # rho(nproma, nlev, nblks_c) [kg/m^3] + exner_pressure: Field[[CellDim, KDim], float] # exner(nrpoma, nlev, nblks_c) + theta_v: Field[[CellDim, KDim], float] # (nproma, nlev, nlbks_c) [K] diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/vertical.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/vertical.py deleted file mode 100644 index e099d8632..000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/vertical.py +++ /dev/null @@ -1,68 +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 functional.common import Field -from functional.iterator.embedded import np_as_located_field - -from icon4py.common.dimension import KDim - - -class VerticalGridConfig: - def __init__(self): - self.num_k_levels = 65 - self.n_shift_total = 0 - - def get_k_size(self): - return self.num_k_levels - - def get_n_shift(self): - return self.n_shift_total - - -class VerticalModelParams: - def __init__(self, vct_a: np.ndarray, rayleigh_damping_height: float = 12500.0): - """ - Arguments. - - - vct_a: # read from vertical_coord_tables - - rayleigh_damping_height height of rayleigh damping in [m] mo_nonhydro_nml - """ - self.rayleigh_damping_height = rayleigh_damping_height - self.vct_a = vct_a - # TODO klevels in ICON are inversed - self.index_of_damping_height = np.argmax( - self.vct_a >= self.rayleigh_damping_height - ) - - def get_index_of_damping_layer(self): - return self.index_of_damping_height - - def get_physical_heights(self) -> Field[[KDim], float]: - return np_as_located_field(KDim)(self.vct_a) - - def init_nabla2_factor_in_upper_damping_zone( - self, k_size: int - ) -> Field[[KDim], float]: - # this assumes n_shift == 0 - buffer = np.zeros(k_size) - buffer[2 : self.index_of_damping_height] = ( - 1.0 - / 12.0 - * ( - self.vct_a[2 : self.index_of_damping_height] - - self.vct_a[self.index_of_damping_height + 1] - ) - / (self.vct_a[2] - self.vct_a[self.index_of_damping_height + 1]) - ) - return np_as_located_field(KDim)(buffer) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index f418581ab..b804655bc 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -19,8 +19,9 @@ DiffusionParams, enhanced_smagorinski_factor, init_diffusion_local_fields, + set_zero_v_k, ) -from icon4py.common.dimension import KDim +from icon4py.common.dimension import KDim, VertexDim from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import random_field, zero_field @@ -89,6 +90,13 @@ def enhanced_smagorinski_factor_np(factor_in, heigths_in, a_vec): assert np.allclose(enhanced_smag_fac_np, np.asarray(result[:-1])) +def test_set_zero_vertex_k(): + mesh = SimpleMesh() + f = random_field(mesh, VertexDim, KDim) + set_zero_v_k(f, offset_provider={}) + assert np.allclose(0.0, f) + + @pytest.mark.xfail def test_diffusion_init(): pytest.fail("not implemented yet") @@ -100,7 +108,7 @@ def test_diffusion_run(): def test_diffusion_coefficients_with_hdiff_efdt_ratio(): - config: DiffusionConfig = DiffusionConfig() + config: DiffusionConfig = DiffusionConfig.create_with_defaults() config.hdiff_efdt_ratio = 1.0 params = DiffusionParams(config) @@ -112,7 +120,7 @@ def test_diffusion_coefficients_with_hdiff_efdt_ratio(): def test_diffusion_coefficients_without_hdiff_efdt_ratio(): - config: DiffusionConfig = DiffusionConfig() + config: DiffusionConfig = DiffusionConfig.create_with_defaults() config.hdiff_efdt_ratio = 0.0 params = DiffusionParams(config) @@ -124,7 +132,7 @@ def test_diffusion_coefficients_without_hdiff_efdt_ratio(): def test_smagorinski_factor_for_diffusion_type_4(): - config: DiffusionConfig = DiffusionConfig() + config: DiffusionConfig = DiffusionConfig.create_with_defaults() config.hdiff_smag_fac = 0.15 config.diffusion_type = 4 @@ -135,7 +143,7 @@ def test_smagorinski_factor_for_diffusion_type_4(): def test_smagorinski_heights_diffusion_type_5_are_consistent(): - config: DiffusionConfig = DiffusionConfig() + config: DiffusionConfig = DiffusionConfig.create_with_defaults() config.hdiff_smag_fac = 0.15 config.diffusion_type = 5 @@ -150,7 +158,7 @@ def test_smagorinski_heights_diffusion_type_5_are_consistent(): def test_smagorinski_factor_diffusion_type_5(): - config: DiffusionConfig = DiffusionConfig() + config: DiffusionConfig = DiffusionConfig.create_with_defaults() config.hdiff_smag_fac = 0.15 config.diffusion_type = 5 diff --git a/atm_dyn_iconam/tests/test_vertical.py b/atm_dyn_iconam/tests/test_vertical.py index 12ac75aad..225028d09 100644 --- a/atm_dyn_iconam/tests/test_vertical.py +++ b/atm_dyn_iconam/tests/test_vertical.py @@ -16,7 +16,7 @@ import numpy as np import pytest -from icon4py.atm_dyn_iconam.vertical import VerticalModelParams +from icon4py.atm_dyn_iconam.icon_mesh import VerticalModelParams @pytest.mark.parametrize( From d2c3400d178909904c4059147dad5f8dba1edaf4 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 25 Nov 2022 16:15:53 +0100 Subject: [PATCH 013/263] add pytest for initialization of diffusion --- atm_dyn_iconam/requirements-dev.txt | 3 +- .../src/icon4py/atm_dyn_iconam/diffusion.py | 17 ++-- atm_dyn_iconam/tests/test_diffusion.py | 52 ++++++------ atm_dyn_iconam/tests/test_run_diffusion.py | 85 +++++++++++++++++++ 4 files changed, 126 insertions(+), 31 deletions(-) create mode 100644 atm_dyn_iconam/tests/test_run_diffusion.py diff --git a/atm_dyn_iconam/requirements-dev.txt b/atm_dyn_iconam/requirements-dev.txt index 2cb7bda5e..204703984 100644 --- a/atm_dyn_iconam/requirements-dev.txt +++ b/atm_dyn_iconam/requirements-dev.txt @@ -1,3 +1,4 @@ -r ../base-requirements-dev.txt -e ../common --e . +matplotlib +#git+https://github.com/GridTools/serialbox#egg=serialbox&subdirectory=src/serialbox-python diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index 13c195b04..fdea541f0 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -328,6 +328,9 @@ def determine_smagorinski_factor(self, config: DiffusionConfig): ) smagorinski_height = None case _: + print("not implemented") + smagorinski_factor = None + smagorinski_height = None pass return smagorinski_factor, smagorinski_height @@ -360,8 +363,8 @@ def __init__( self.params: params # different for init call: smag_offset = 0 - self.smag_offset = 0.25 * params.K4 * config.substep_as_float() - self.diff_multfac_w = np.minimum( + self.smag_offset:float = 0.25 * params.K4 * config.substep_as_float() + self.diff_multfac_w :float = min( 1.0 / 48.0, params.K4W * config.substep_as_float() ) @@ -389,7 +392,7 @@ def __init__( *params.smagorinski_height, a_vect, self.enh_smag_fac, - offset_provider={"Koff", KDim}, + offset_provider={"Koff": KDim}, ) self.diff_multfac_n2w = ( @@ -397,14 +400,16 @@ def __init__( k_size=config.grid.get_num_k_levels() ) ) + shape_vk = (config.grid.get_num_vertices(), config.grid.get_num_k_levels()) self.u_vert = np_as_located_field(VertexDim, KDim)( - np.zeros(config.grid.get_num_vertices(), config.grid.get_num_k_levels()) + np.zeros(shape_vk, float) ) self.v_vert = np_as_located_field(VertexDim, KDim)( - np.zeros(config.grid.get_num_vertices(), config.grid.get_num_k_levels()) + np.zeros(shape_vk, float) ) + shape_ek = (config.grid.get_num_edges(), config.grid.get_num_k_levels()) allocate_ek = np_as_located_field(VertexDim, KDim)( - np.zeros(config.grid.get_num_edges(), config.grid.get_num_k_levels()) + np.zeros(shape_ek, float) ) self.kh_smag_e = allocate_ek self.kh_smag_ec = allocate_ek diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index b804655bc..098ecd301 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -26,19 +26,19 @@ from icon4py.testutils.utils import random_field, zero_field -def _smag_limit_numpy(diff_multfac_vn: np.array): - return 0.125 - 4.0 * diff_multfac_vn +def smag_limit_numpy(shape, k4, substeps): + return 0.125 - 4.0 * diff_multfac_vn_numpy(shape, k4, substeps) def test_init_diff_multifac_vn_const(): mesh = SimpleMesh() diff_multfac_vn = zero_field(mesh, KDim) smag_offset = zero_field(mesh, KDim) - expected_diff_multfac_vn = 1.0 / 128.0 * np.ones(np.asarray(diff_multfac_vn).shape) - expected_smag_limit = _smag_limit_numpy(expected_diff_multfac_vn) - k4 = 1.0 substeps = 5.0 + expected_diff_multfac_vn = diff_multfac_vn_numpy(np.asarray(diff_multfac_vn).shape, k4, substeps) + expected_smag_limit = smag_limit_numpy(expected_diff_multfac_vn) + init_diffusion_local_fields( k4, substeps, diff_multfac_vn, smag_offset, offset_provider={} ) @@ -52,10 +52,9 @@ def test_init_diff_multifac_vn_k4_substeps(): smag_limit = zero_field(mesh, KDim) k4 = 0.003 substeps = 1.0 - expected_diff_multfac_vn = ( - k4 * substeps / 3.0 * np.ones(np.asarray(diff_multfac_vn).shape) - ) - expected_smag_limit = _smag_limit_numpy(expected_diff_multfac_vn) + expected_diff_multfac_vn = diff_multfac_vn_numpy(np.asarray(diff_multfac_vn).shape, k4, substeps) + + expected_smag_limit = smag_limit_numpy(expected_diff_multfac_vn) init_diffusion_local_fields( k4, substeps, diff_multfac_vn, smag_limit, offset_provider={} @@ -64,22 +63,27 @@ def test_init_diff_multifac_vn_k4_substeps(): assert np.allclose(expected_smag_limit, smag_limit) -def test_enhanced_smagorinski_factor(): - def enhanced_smagorinski_factor_np(factor_in, heigths_in, a_vec): - alin = (factor_in[1] - factor_in[0]) / (heigths_in[1] - heigths_in[0]) - df32 = factor_in[2] - factor_in[1] - df42 = factor_in[3] - factor_in[1] - dz32 = heigths_in[2] - heigths_in[1] - dz42 = heigths_in[3] - heigths_in[1] - bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32)) - aqdr = df32 / dz32 - bqdr * dz32 - zf = 0.5 * (a_vec[:-1] + a_vec[1:]) - max0 = np.maximum(0.0, zf - heigths_in[0]) - dzlin = np.minimum(heigths_in[1] - heigths_in[0], max0) - max1 = np.maximum(0.0, zf - heigths_in[1]) - dzqdr = np.minimum(heigths_in[3] - heigths_in[1], max1) - return factor_in[0] + dzlin * alin + dzqdr * (aqdr + dzqdr * bqdr) +def diff_multfac_vn_numpy(shape, k4, substeps): + factor = min(1.0/128.0, k4 * substeps / 3.0) + return factor * np.ones(shape) + +def enhanced_smagorinski_factor_np(factor_in, heigths_in, a_vec): + alin = (factor_in[1] - factor_in[0]) / (heigths_in[1] - heigths_in[0]) + df32 = factor_in[2] - factor_in[1] + df42 = factor_in[3] - factor_in[1] + dz32 = heigths_in[2] - heigths_in[1] + dz42 = heigths_in[3] - heigths_in[1] + bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32)) + aqdr = df32 / dz32 - bqdr * dz32 + zf = 0.5 * (a_vec[:-1] + a_vec[1:]) + max0 = np.maximum(0.0, zf - heigths_in[0]) + dzlin = np.minimum(heigths_in[1] - heigths_in[0], max0) + max1 = np.maximum(0.0, zf - heigths_in[1]) + dzqdr = np.minimum(heigths_in[3] - heigths_in[1], max1) + return factor_in[0] + dzlin * alin + dzqdr * (aqdr + dzqdr * bqdr) + +def test_enhanced_smagorinski_factor(): mesh = SimpleMesh() a_vec = random_field(mesh, KDim, low=1.0, high=10.0) result = zero_field(mesh, KDim) diff --git a/atm_dyn_iconam/tests/test_run_diffusion.py b/atm_dyn_iconam/tests/test_run_diffusion.py new file mode 100644 index 000000000..d075f1f2a --- /dev/null +++ b/atm_dyn_iconam/tests/test_run_diffusion.py @@ -0,0 +1,85 @@ +import os +import numpy as np + +try: + import serialbox as ser +except ImportError: + os.system( + "git clone --recursive https://github.com/GridTools/serialbox; CC=`which gcc` CXX=`which g++` pip install serialbox/src/serialbox-python" + ) + import serialbox as ser +from functional.iterator.embedded import np_as_located_field +from atm_dyn_iconam.tests.test_diffusion import diff_multfac_vn_numpy, smag_limit_numpy, \ + enhanced_smagorinski_factor_np +from icon4py.atm_dyn_iconam.diffusion import DiffusionConfig, Diffusion, DiffusionParams +from icon4py.common.dimension import KDim + +data_path = os.path.join(os.path.dirname(__file__), "ser_data") + +def read_from_ser_data(path, metadata=["nproma"], fields=[]): + rank = 0 + serializer = ser.Serializer(ser.OpenModeKind.Read, path, + f"reference_icon_rank{str(rank)}") + save_points = serializer.savepoint_list() + print(save_points) + field_names = serializer.fieldnames() + print(field_names) + savepoint = serializer.savepoint["diffusion-in"].id[0].as_savepoint() + print(type(savepoint)) + print(savepoint) + meta_present={} + meta_absent=[] + for md in metadata: + if savepoint.metainfo.has_key(md): + meta_present[md] = savepoint.metainfo[md] + else: + meta_absent.append(md) + + fields_present = {} + fields_absent = [] + for field_name in fields: + if field_name in field_names: + fields_present[field_name] = serializer.read(field_name, savepoint) + else: + fields_absent.append(field_name) + [print(f"field {f} not present in savepoint") for f in fields_absent] + [print(f"metadata {f} not present in savepoint") for f in meta_absent] + + return fields_present, meta_present + + +def test_diffusion_init(): + meta, fields = read_from_ser_data(data_path, metadata = ["nproma"], fields=["a_vect"]) + + config = DiffusionConfig.create_with_defaults() + additional_parameters = DiffusionParams(config) + vct_a = np_as_located_field(KDim)(np.zeros(config.grid.get_num_k_levels())) + #diffusion = Diffusion(config, additional_parameters, fields["vct_a"]) + diffusion = Diffusion(config, additional_parameters, vct_a) + + ## assert static local fields are initialized and correct: + assert diffusion.smag_offset == 0.25 * additional_parameters.K4 * config.substep_as_float() + assert diffusion.diff_multfac_w == min(1./48., additional_parameters.K4W * config.substep_as_float()) + + assert np.allclose(0.0, np.asarray(diffusion.v_vert)) + assert np.allclose(0.0, np.asarray(diffusion.u_vert)) + assert np.allclose(0.0, np.asarray(diffusion.diff_multfac_n2w)) + assert np.allclose(0.0, np.asarray(diffusion.kh_smag_ec)) + assert np.allclose(0.0, np.asarray(diffusion.kh_smag_e)) + + shape_k = np.asarray(diffusion.diff_multfac_vn.shape) + expected_smag_limit = smag_limit_numpy(shape_k, additional_parameters.K4, config.substep_as_float()) + assert np.allclose(expected_smag_limit, np.asarray(diffusion.smag_limit)) + + expected_diff_multfac_vn = diff_multfac_vn_numpy(shape_k, additional_parameters.K4, config.substep_as_float()) + assert np.allclose(expected_diff_multfac_vn, np.asarray(diffusion.diff_multfac_vn)) + expected_enh_smag_fac = enhanced_smagorinski_factor_np(additional_parameters.smagorinski_factor, additional_parameters.smagorinski_height, vct_a) + assert np.allclose(expected_enh_smag_fac, np.asarray(diffusion.enh_smag_fac)[:-1]) + + + + + + + + From 6bc6cc7b7ca6bcc8206964c718e4e710016aee48 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 25 Nov 2022 18:03:38 +0100 Subject: [PATCH 014/263] clean up tests, fix serialbox installation workaround --- atm_dyn_iconam/requirements-dev.txt | 2 - .../src/icon4py/atm_dyn_iconam/diffusion.py | 16 ++-- atm_dyn_iconam/tests/test_diffusion.py | 65 +++++++++++--- atm_dyn_iconam/tests/test_run_diffusion.py | 85 ------------------- .../src/icon4py/testutils/serialbox_utils.py | 59 +++++++++++++ 5 files changed, 119 insertions(+), 108 deletions(-) delete mode 100644 atm_dyn_iconam/tests/test_run_diffusion.py create mode 100644 testutils/src/icon4py/testutils/serialbox_utils.py diff --git a/atm_dyn_iconam/requirements-dev.txt b/atm_dyn_iconam/requirements-dev.txt index 204703984..8f8d43342 100644 --- a/atm_dyn_iconam/requirements-dev.txt +++ b/atm_dyn_iconam/requirements-dev.txt @@ -1,4 +1,2 @@ -r ../base-requirements-dev.txt -e ../common -matplotlib -#git+https://github.com/GridTools/serialbox#egg=serialbox&subdirectory=src/serialbox-python diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index fdea541f0..0062c6a61 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -363,8 +363,8 @@ def __init__( self.params: params # different for init call: smag_offset = 0 - self.smag_offset:float = 0.25 * params.K4 * config.substep_as_float() - self.diff_multfac_w :float = min( + self.smag_offset: float = 0.25 * params.K4 * config.substep_as_float() + self.diff_multfac_w: float = min( 1.0 / 48.0, params.K4W * config.substep_as_float() ) @@ -401,16 +401,10 @@ def __init__( ) ) shape_vk = (config.grid.get_num_vertices(), config.grid.get_num_k_levels()) - self.u_vert = np_as_located_field(VertexDim, KDim)( - np.zeros(shape_vk, float) - ) - self.v_vert = np_as_located_field(VertexDim, KDim)( - np.zeros(shape_vk, float) - ) + self.u_vert = np_as_located_field(VertexDim, KDim)(np.zeros(shape_vk, float)) + self.v_vert = np_as_located_field(VertexDim, KDim)(np.zeros(shape_vk, float)) shape_ek = (config.grid.get_num_edges(), config.grid.get_num_k_levels()) - allocate_ek = np_as_located_field(VertexDim, KDim)( - np.zeros(shape_ek, float) - ) + allocate_ek = np_as_located_field(VertexDim, KDim)(np.zeros(shape_ek, float)) self.kh_smag_e = allocate_ek self.kh_smag_ec = allocate_ek self.z_nabla2_e = allocate_ek diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 098ecd301..50caa52ea 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -10,11 +10,14 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import os import numpy as np import pytest +from functional.iterator.embedded import np_as_located_field from icon4py.atm_dyn_iconam.diffusion import ( + Diffusion, DiffusionConfig, DiffusionParams, enhanced_smagorinski_factor, @@ -22,6 +25,7 @@ set_zero_v_k, ) from icon4py.common.dimension import KDim, VertexDim +from icon4py.testutils.serialbox_utils import read_from_ser_data from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import random_field, zero_field @@ -36,8 +40,9 @@ def test_init_diff_multifac_vn_const(): smag_offset = zero_field(mesh, KDim) k4 = 1.0 substeps = 5.0 - expected_diff_multfac_vn = diff_multfac_vn_numpy(np.asarray(diff_multfac_vn).shape, k4, substeps) - expected_smag_limit = smag_limit_numpy(expected_diff_multfac_vn) + shape = np.asarray(diff_multfac_vn).shape + expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps) + expected_smag_limit = smag_limit_numpy(shape, k4, substeps) init_diffusion_local_fields( k4, substeps, diff_multfac_vn, smag_offset, offset_provider={} @@ -52,9 +57,10 @@ def test_init_diff_multifac_vn_k4_substeps(): smag_limit = zero_field(mesh, KDim) k4 = 0.003 substeps = 1.0 - expected_diff_multfac_vn = diff_multfac_vn_numpy(np.asarray(diff_multfac_vn).shape, k4, substeps) + shape = np.asarray(diff_multfac_vn).shape + expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps) - expected_smag_limit = smag_limit_numpy(expected_diff_multfac_vn) + expected_smag_limit = smag_limit_numpy(shape, k4, substeps) init_diffusion_local_fields( k4, substeps, diff_multfac_vn, smag_limit, offset_provider={} @@ -64,7 +70,7 @@ def test_init_diff_multifac_vn_k4_substeps(): def diff_multfac_vn_numpy(shape, k4, substeps): - factor = min(1.0/128.0, k4 * substeps / 3.0) + factor = min(1.0 / 128.0, k4 * substeps / 3.0) return factor * np.ones(shape) @@ -83,6 +89,7 @@ def enhanced_smagorinski_factor_np(factor_in, heigths_in, a_vec): dzqdr = np.minimum(heigths_in[3] - heigths_in[1], max1) return factor_in[0] + dzlin * alin + dzqdr * (aqdr + dzqdr * bqdr) + def test_enhanced_smagorinski_factor(): mesh = SimpleMesh() a_vec = random_field(mesh, KDim, low=1.0, high=10.0) @@ -101,11 +108,6 @@ def test_set_zero_vertex_k(): assert np.allclose(0.0, f) -@pytest.mark.xfail -def test_diffusion_init(): - pytest.fail("not implemented yet") - - @pytest.mark.xfail def test_diffusion_run(): pytest.fail("not implemented yet") @@ -170,3 +172,46 @@ def test_smagorinski_factor_diffusion_type_5(): assert len(params.smagorinski_factor) == len(params.smagorinski_height) assert len(params.smagorinski_factor) == 4 assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor))) + + +def test_diffusion_init(): + data_path = os.path.join(os.path.dirname(__file__), "ser_data") + meta, fields = read_from_ser_data(data_path, metadata=["nproma"], fields=["a_vect"]) + + config = DiffusionConfig.create_with_defaults() + additional_parameters = DiffusionParams(config) + vct_a = np_as_located_field(KDim)(np.zeros(config.grid.get_num_k_levels())) + # diffusion = Diffusion(config, additional_parameters, fields["vct_a"]) + diffusion = Diffusion(config, additional_parameters, vct_a) + + # assert static local fields are initialized and correct: + assert ( + diffusion.smag_offset + == 0.25 * additional_parameters.K4 * config.substep_as_float() + ) + assert diffusion.diff_multfac_w == min( + 1.0 / 48.0, additional_parameters.K4W * config.substep_as_float() + ) + + assert np.allclose(0.0, np.asarray(diffusion.v_vert)) + assert np.allclose(0.0, np.asarray(diffusion.u_vert)) + assert np.allclose(0.0, np.asarray(diffusion.diff_multfac_n2w)) + assert np.allclose(0.0, np.asarray(diffusion.kh_smag_ec)) + assert np.allclose(0.0, np.asarray(diffusion.kh_smag_e)) + + shape_k = np.asarray(diffusion.diff_multfac_vn.shape) + expected_smag_limit = smag_limit_numpy( + shape_k, additional_parameters.K4, config.substep_as_float() + ) + assert np.allclose(expected_smag_limit, np.asarray(diffusion.smag_limit)) + + expected_diff_multfac_vn = diff_multfac_vn_numpy( + shape_k, additional_parameters.K4, config.substep_as_float() + ) + assert np.allclose(expected_diff_multfac_vn, np.asarray(diffusion.diff_multfac_vn)) + expected_enh_smag_fac = enhanced_smagorinski_factor_np( + additional_parameters.smagorinski_factor, + additional_parameters.smagorinski_height, + vct_a, + ) + assert np.allclose(expected_enh_smag_fac, np.asarray(diffusion.enh_smag_fac)[:-1]) diff --git a/atm_dyn_iconam/tests/test_run_diffusion.py b/atm_dyn_iconam/tests/test_run_diffusion.py deleted file mode 100644 index d075f1f2a..000000000 --- a/atm_dyn_iconam/tests/test_run_diffusion.py +++ /dev/null @@ -1,85 +0,0 @@ -import os -import numpy as np - -try: - import serialbox as ser -except ImportError: - os.system( - "git clone --recursive https://github.com/GridTools/serialbox; CC=`which gcc` CXX=`which g++` pip install serialbox/src/serialbox-python" - ) - import serialbox as ser -from functional.iterator.embedded import np_as_located_field -from atm_dyn_iconam.tests.test_diffusion import diff_multfac_vn_numpy, smag_limit_numpy, \ - enhanced_smagorinski_factor_np -from icon4py.atm_dyn_iconam.diffusion import DiffusionConfig, Diffusion, DiffusionParams -from icon4py.common.dimension import KDim - -data_path = os.path.join(os.path.dirname(__file__), "ser_data") - -def read_from_ser_data(path, metadata=["nproma"], fields=[]): - rank = 0 - serializer = ser.Serializer(ser.OpenModeKind.Read, path, - f"reference_icon_rank{str(rank)}") - save_points = serializer.savepoint_list() - print(save_points) - field_names = serializer.fieldnames() - print(field_names) - savepoint = serializer.savepoint["diffusion-in"].id[0].as_savepoint() - print(type(savepoint)) - print(savepoint) - meta_present={} - meta_absent=[] - for md in metadata: - if savepoint.metainfo.has_key(md): - meta_present[md] = savepoint.metainfo[md] - else: - meta_absent.append(md) - - fields_present = {} - fields_absent = [] - for field_name in fields: - if field_name in field_names: - fields_present[field_name] = serializer.read(field_name, savepoint) - else: - fields_absent.append(field_name) - [print(f"field {f} not present in savepoint") for f in fields_absent] - [print(f"metadata {f} not present in savepoint") for f in meta_absent] - - return fields_present, meta_present - - -def test_diffusion_init(): - meta, fields = read_from_ser_data(data_path, metadata = ["nproma"], fields=["a_vect"]) - - config = DiffusionConfig.create_with_defaults() - additional_parameters = DiffusionParams(config) - vct_a = np_as_located_field(KDim)(np.zeros(config.grid.get_num_k_levels())) - #diffusion = Diffusion(config, additional_parameters, fields["vct_a"]) - diffusion = Diffusion(config, additional_parameters, vct_a) - - ## assert static local fields are initialized and correct: - assert diffusion.smag_offset == 0.25 * additional_parameters.K4 * config.substep_as_float() - assert diffusion.diff_multfac_w == min(1./48., additional_parameters.K4W * config.substep_as_float()) - - assert np.allclose(0.0, np.asarray(diffusion.v_vert)) - assert np.allclose(0.0, np.asarray(diffusion.u_vert)) - assert np.allclose(0.0, np.asarray(diffusion.diff_multfac_n2w)) - assert np.allclose(0.0, np.asarray(diffusion.kh_smag_ec)) - assert np.allclose(0.0, np.asarray(diffusion.kh_smag_e)) - - shape_k = np.asarray(diffusion.diff_multfac_vn.shape) - expected_smag_limit = smag_limit_numpy(shape_k, additional_parameters.K4, config.substep_as_float()) - assert np.allclose(expected_smag_limit, np.asarray(diffusion.smag_limit)) - - expected_diff_multfac_vn = diff_multfac_vn_numpy(shape_k, additional_parameters.K4, config.substep_as_float()) - assert np.allclose(expected_diff_multfac_vn, np.asarray(diffusion.diff_multfac_vn)) - expected_enh_smag_fac = enhanced_smagorinski_factor_np(additional_parameters.smagorinski_factor, additional_parameters.smagorinski_height, vct_a) - assert np.allclose(expected_enh_smag_fac, np.asarray(diffusion.enh_smag_fac)[:-1]) - - - - - - - - diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py new file mode 100644 index 000000000..c048edbf6 --- /dev/null +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -0,0 +1,59 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import os +from typing import List + + +try: + import serialbox as ser +except ImportError: + external_src = os.path.join(os.path.dirname(__file__), "../../_external_src/") + os.chdir(external_src) + os.system( + "git clone --recursive https://github.com/GridTools/serialbox; CC=`which gcc` CXX=`which g++` pip install serialbox/src/serialbox-python" + ) + import serialbox as ser + + +def read_from_ser_data(path, metadata: List[str], fields: List[str]): + rank = 0 + serializer = ser.Serializer( + ser.OpenModeKind.Read, path, f"reference_icon_rank{str(rank)}" + ) + save_points = serializer.savepoint_list() + print(save_points) + field_names = serializer.fieldnames() + print(field_names) + savepoint = serializer.savepoint["diffusion-in"].id[0].as_savepoint() + print(type(savepoint)) + print(savepoint) + meta_present = {} + meta_absent = [] + for md in metadata: + if md in savepoint.metainfo: + meta_present[md] = savepoint.metainfo[md] + else: + meta_absent.append(md) + + fields_present = {} + fields_absent = [] + for field_name in fields: + if field_name in field_names: + fields_present[field_name] = serializer.read(field_name, savepoint) + else: + fields_absent.append(field_name) + [print(f"field {f} not present in savepoint") for f in fields_absent] + [print(f"metadata {f} not present in savepoint") for f in meta_absent] + + return fields_present, meta_present From f158867a5fb50cdf2cd7e0fe0bca7d1f50ead98d Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 29 Nov 2022 09:05:03 +0100 Subject: [PATCH 015/263] implement initialization test for diffusion reading from serialized data --- atm_dyn_iconam/tests/test_diffusion.py | 23 +++++++++++++------ .../src/icon4py/testutils/serialbox_utils.py | 21 ++++++++++------- 2 files changed, 29 insertions(+), 15 deletions(-) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 50caa52ea..6b8237607 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -25,7 +25,9 @@ set_zero_v_k, ) from icon4py.common.dimension import KDim, VertexDim -from icon4py.testutils.serialbox_utils import read_from_ser_data +from icon4py.testutils.serialbox_utils import ( + read_from_call_diffusion_init_ser_data, +) from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import random_field, zero_field @@ -176,14 +178,21 @@ def test_smagorinski_factor_diffusion_type_5(): def test_diffusion_init(): data_path = os.path.join(os.path.dirname(__file__), "ser_data") - meta, fields = read_from_ser_data(data_path, metadata=["nproma"], fields=["a_vect"]) - + physical_heights_name = "vct_a" + fields, meta = read_from_call_diffusion_init_ser_data( + data_path, + "icon_diffusion", + metadata=["nlev", "linit", "date"], + fields=[physical_heights_name], + ) + print(f"meta {meta}") + assert meta["nlev"] == 65 + assert meta["linit"] is False + assert meta["date"] == "2021-06-20T12:00:10.000" config = DiffusionConfig.create_with_defaults() additional_parameters = DiffusionParams(config) - vct_a = np_as_located_field(KDim)(np.zeros(config.grid.get_num_k_levels())) - # diffusion = Diffusion(config, additional_parameters, fields["vct_a"]) + vct_a = np_as_located_field(KDim)(fields[physical_heights_name]) diffusion = Diffusion(config, additional_parameters, vct_a) - # assert static local fields are initialized and correct: assert ( diffusion.smag_offset @@ -214,4 +223,4 @@ def test_diffusion_init(): additional_parameters.smagorinski_height, vct_a, ) - assert np.allclose(expected_enh_smag_fac, np.asarray(diffusion.enh_smag_fac)[:-1]) + assert np.allclose(expected_enh_smag_fac, np.asarray(diffusion.enh_smag_fac)) diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index c048edbf6..cf9848d90 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -26,22 +26,27 @@ import serialbox as ser -def read_from_ser_data(path, metadata: List[str], fields: List[str]): +def read_from_call_diffusion_init_ser_data( + path, fname_prefix, metadata: List[str], fields: List[str] +): rank = 0 - serializer = ser.Serializer( - ser.OpenModeKind.Read, path, f"reference_icon_rank{str(rank)}" - ) + fname = f"{fname_prefix}_rank{str(rank)}" + serializer = ser.Serializer(ser.OpenModeKind.Read, path, fname) save_points = serializer.savepoint_list() print(save_points) field_names = serializer.fieldnames() print(field_names) - savepoint = serializer.savepoint["diffusion-in"].id[0].as_savepoint() - print(type(savepoint)) - print(savepoint) + savepoint = ( + serializer.savepoint["call-diffusion-init"] + .linit[False] + .date["2021-06-20T12:00:10.000"] + .as_savepoint() + ) + print(savepoint.metainfo) meta_present = {} meta_absent = [] for md in metadata: - if md in savepoint.metainfo: + if md in savepoint.metainfo.to_dict(): meta_present[md] = savepoint.metainfo[md] else: meta_absent.append(md) From f9488ea368ce24d7be90010634924acd49cef2f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Thu, 24 Nov 2022 09:04:58 +0100 Subject: [PATCH 016/263] First version of fused diffusion stencils --- .../fused_mo_nh_diffusion_stencil_02_03.py | 47 ++++++++++ .../fused_mo_nh_diffusion_stencil_04_05_06.py | 82 ++++++++++++++++++ ...sed_mo_nh_diffusion_stencil_07_08_09_10.py | 86 +++++++++++++++++++ .../fused_mo_nh_diffusion_stencil_11_12.py | 40 +++++++++ .../fused_mo_nh_diffusion_stencil_13_14.py | 67 +++++++++++++++ .../mo_nh_diffusion_stencil_10.py | 8 +- 6 files changed, 326 insertions(+), 4 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py new file mode 100644 index 000000000..8c111b5d4 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -0,0 +1,47 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum + +from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim, Koff + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import _mo_nh_diffusion_stencil_02 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import _mo_nh_diffusion_stencil_03 + +@field_operator +def _fused_mo_nh_diffusion_stencil_02_03( + kh_smag_ec: Field[[EdgeDim, KDim], float], + vn: Field[[EdgeDim, KDim], float], + e_bln_c_s: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + diff_multfac_smag: Field[[KDim], float], + wgtfac_c: Field[[CellDim, KDim], float], +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + kh_c, div = _mo_nh_diffusion_stencil_02(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag) + div_ic, hdef_ic = _mo_nh_diffusion_stencil_03(div, kh_c, wgtfac_c) + return div_ic, hdef_ic + + +@program +def fused_mo_nh_diffusion_stencil_02_03( + kh_smag_ec: Field[[EdgeDim, KDim], float], + vn: Field[[EdgeDim, KDim], float], + e_bln_c_s: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + diff_multfac_smag: Field[[KDim], float], + wgtfac_c: Field[[CellDim, KDim], float], + div_ic: Field[[CellDim, KDim], float], + hdef_ic: Field[[CellDim, KDim], float], +): + _fused_mo_nh_diffusion_stencil_02_03(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag, wgtfac_c, out=(div_ic, hdef_ic)) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py new file mode 100644 index 000000000..1438e07c1 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py @@ -0,0 +1,82 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum, maximum, where, int32 + +from icon4py.common.dimension import ( + E2C2V, + E2ECV, + ECVDim, + EdgeDim, + KDim, + VertexDim, +) + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_04 import _mo_nh_diffusion_stencil_04 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_05 import _mo_nh_diffusion_stencil_05 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_06 import _mo_nh_diffusion_stencil_06 + +@field_operator +def _fused_mo_nh_diffusion_stencil_04_05_06( + u_vert: Field[[VertexDim, KDim], float], + v_vert: Field[[VertexDim, KDim], float], + primal_normal_vert_v1: Field[[ECVDim], float], + primal_normal_vert_v2: Field[[ECVDim], float], + z_nabla2_e: Field[[EdgeDim, KDim], float], + inv_vert_vert_length: Field[[EdgeDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + area_edge: Field[[EdgeDim], float], + kh_smag_e: Field[[EdgeDim, KDim], float], + diff_multfac_vn: Field[[KDim], float], + nudgecoeff_e: Field[[EdgeDim], float], + vn: Field[[EdgeDim, KDim], float], + horz_idx: Field[[EdgeDim], int32], + nudgezone_diff: float, + fac_bdydiff_v: float, + start_2nd_nudge_line_idx_e: int32, +) -> Field[[EdgeDim, KDim], float]: + + z_nabla4_e2 = _mo_nh_diffusion_stencil_04(u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2, z_nabla2_e, inv_vert_vert_length, inv_primal_edge_length) + + vn = where( + horz_idx >= start_2nd_nudge_line_idx_e, + _mo_nh_diffusion_stencil_05(area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, nudgecoeff_e, vn, nudgezone_diff), + _mo_nh_diffusion_stencil_06(z_nabla2_e, area_edge, vn, fac_bdydiff_v) + ) + + return vn + + +@program +def fused_mo_nh_diffusion_stencil_04_05_06( + u_vert: Field[[VertexDim, KDim], float], + v_vert: Field[[VertexDim, KDim], float], + primal_normal_vert_v1: Field[[ECVDim], float], + primal_normal_vert_v2: Field[[ECVDim], float], + z_nabla2_e: Field[[EdgeDim, KDim], float], + inv_vert_vert_length: Field[[EdgeDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + area_edge: Field[[EdgeDim], float], + kh_smag_e: Field[[EdgeDim, KDim], float], + diff_multfac_vn: Field[[KDim], float], + nudgecoeff_e: Field[[EdgeDim], float], + vn: Field[[EdgeDim, KDim], float], + horz_idx: Field[[EdgeDim], int32], + nudgezone_diff: float, + fac_bdydiff_v: float, + start_2nd_nudge_line_idx_e: int32, +): + _fused_mo_nh_diffusion_stencil_04_05_06(u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2, z_nabla2_e, inv_vert_vert_length, + inv_primal_edge_length, area_edge, kh_smag_e, diff_multfac_vn, nudgecoeff_e, vn, horz_idx, nudgezone_diff, fac_bdydiff_v, + start_2nd_nudge_line_idx_e, out=vn) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py new file mode 100644 index 000000000..a1e305ae5 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -0,0 +1,86 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum, where, int32, broadcast + +from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_07 import _mo_nh_diffusion_stencil_07 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_08 import _mo_nh_diffusion_stencil_08 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_09 import _mo_nh_diffusion_stencil_09 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_10 import _mo_nh_diffusion_stencil_10 + +@field_operator +def _fused_mo_nh_diffusion_stencil_07_08_09_10( + area: Field[[CellDim], float], + geofac_n2s: Field[[CellDim, C2E2CODim], float], + geofac_grg_x: Field[[CellDim, C2E2CODim], float], + geofac_grg_y: Field[[CellDim, C2E2CODim], float], + w_old: Field[[CellDim, KDim], float], + w: Field[[CellDim, KDim], float], + dwdx: Field[[CellDim, KDim], float], + dwdy: Field[[CellDim, KDim], float], + diff_multfac_w: float, + diff_multfac_n2w: Field[[KDim], float], + vert_idx: Field[[KDim], int32], + horz_idx: Field[[CellDim], int32], + nrdmax: int32, + interior_idx: int32, + halo_idx: int32, +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + + vert_idx = broadcast(vert_idx, (CellDim, KDim)) + + dwdx, dwdy = where( + vert_idx > int32(0), + _mo_nh_diffusion_stencil_08(w_old, geofac_grg_x, geofac_grg_y), + (dwdx, dwdy) + ) + + z_nabla2_c = _mo_nh_diffusion_stencil_07(w_old, geofac_n2s) + + w = where( + (horz_idx >= interior_idx) & (horz_idx < halo_idx), + _mo_nh_diffusion_stencil_09(area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w), + w_old, + ) + + w = where( + (vert_idx > int32(0)) & (vert_idx < nrdmax) & (horz_idx >= interior_idx) & (horz_idx < halo_idx), + _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, area, z_nabla2_c), + w, + ) + + return w, dwdx, dwdy + + +@program +def fused_mo_nh_diffusion_stencil_07_08_09_10( + area: Field[[CellDim], float], + geofac_n2s: Field[[CellDim, C2E2CODim], float], + geofac_grg_x: Field[[CellDim, C2E2CODim], float], + geofac_grg_y: Field[[CellDim, C2E2CODim], float], + w_old: Field[[CellDim, KDim], float], + w: Field[[CellDim, KDim], float], + dwdx: Field[[CellDim, KDim], float], + dwdy: Field[[CellDim, KDim], float], + diff_multfac_w: float, + diff_multfac_n2w: Field[[KDim], float], + vert_idx: Field[[KDim], int32], + horz_idx: Field[[CellDim], int32], + nrdmax: int32, + interior_idx: int32, + halo_idx: int32, +): + _fused_mo_nh_diffusion_stencil_07_08_09_10(area, geofac_n2s, geofac_grg_x, geofac_grg_y, w_old, w, dwdx, dwdy, diff_multfac_w, diff_multfac_n2w, vert_idx, horz_idx, nrdmax, interior_idx, halo_idx, out = (w, dwdx, dwdy)) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py new file mode 100644 index 000000000..02fd00530 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py @@ -0,0 +1,40 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, where, neighbor_sum, max_over, maximum + +from icon4py.common.dimension import E2C, E2CDim, C2E2C, C2E2CDim, CellDim, EdgeDim, KDim + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_11 import _mo_nh_diffusion_stencil_11 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_12 import _mo_nh_diffusion_stencil_12 + +@field_operator +def _fused_mo_nh_diffusion_stencil_11_12( + theta_v: Field[[CellDim, KDim], float], + theta_ref_mc: Field[[CellDim, KDim], float], + thresh_tdiff: float, + kh_smag_e: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + enh_diffu_3d = _mo_nh_diffusion_stencil_11(theta_v, theta_ref_mc, thresh_tdiff) + kh_smag_e = _mo_nh_diffusion_stencil_12(kh_smag_e, enh_diffu_3d) + return kh_smag_e + +@program +def fused_mo_nh_diffusion_stencil_11_12( + theta_v: Field[[CellDim, KDim], float], + theta_ref_mc: Field[[CellDim, KDim], float], + thresh_tdiff: float, + kh_smag_e: Field[[EdgeDim, KDim], float], +): + _fused_mo_nh_diffusion_stencil_11_12(theta_v, theta_ref_mc, thresh_tdiff, kh_smag_e, out=kh_smag_e) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py new file mode 100644 index 000000000..40faa10b5 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -0,0 +1,67 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum + +from icon4py.common.dimension import ( + E2C, + E2CDim, + C2E2C, + C2E2CDim, + C2CE, + C2E, + C2EDim, + CEDim, + CellDim, + EdgeDim, + KDim, +) + +@field_operator +def _mo_nh_diffusion_stencil_13( + kh_smag_e: Field[[EdgeDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + theta_v: Field[[CellDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + z_nabla2_e = kh_smag_e * inv_dual_edge_length * (theta_v(E2C[1]) - theta_v(E2C[0])) + return z_nabla2_e + +@field_operator +def _mo_nh_diffusion_stencil_14( + z_nabla2_e: Field[[EdgeDim, KDim], float], + geofac_div: Field[[CEDim], float], +) -> Field[[CellDim, KDim], float]: + z_temp = neighbor_sum(z_nabla2_e(C2E) * geofac_div(C2CE), axis=C2EDim) + return z_temp + +@field_operator +def _fused_mo_nh_diffusion_stencil_13_14( + kh_smag_e: Field[[EdgeDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + theta_v: Field[[CellDim, KDim], float], + geofac_div: Field[[CEDim], float], +) -> Field[[CellDim, KDim], float]: + z_nabla2_e = _mo_nh_diffusion_stencil_13(kh_smag_e, inv_dual_edge_length, theta_v) + z_temp = _mo_nh_diffusion_stencil_14(z_nabla2_e, geofac_div) + return z_temp + +@program +def fused_mo_nh_diffusion_stencil_13_14( + kh_smag_e: Field[[EdgeDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + theta_v: Field[[CellDim, KDim], float], + geofac_div: Field[[CEDim], float], + z_temp: Field[[CellDim, KDim], float], +): + _fused_mo_nh_diffusion_stencil_13_14(kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py index e439f742e..fed277e85 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py @@ -21,10 +21,10 @@ def _mo_nh_diffusion_stencil_10( w: Field[[CellDim, KDim], float], diff_multfac_n2w: Field[[KDim], float], - cell_area: Field[[CellDim], float], + area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - w = w + diff_multfac_n2w * (cell_area * z_nabla2_c) + w = w + diff_multfac_n2w * (area * z_nabla2_c) return w @@ -32,7 +32,7 @@ def _mo_nh_diffusion_stencil_10( def mo_nh_diffusion_stencil_10( w: Field[[CellDim, KDim], float], diff_multfac_n2w: Field[[KDim], float], - cell_area: Field[[CellDim], float], + area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], ): - _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w) + _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, area, z_nabla2_c, out=w) From c480d77260e9edff0fe1954f510b14271acc2c8e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Tue, 29 Nov 2022 12:30:45 +0100 Subject: [PATCH 017/263] Added prototype of fused_mo_nh_diffusion_stencil_01_02_03_rbf Currently this stencil does not work since it would require elements of a tuple of output variables of a field_operator to have different types. --- ...ed_mo_nh_diffusion_stencil_01_02_03_rbf.py | 92 +++++++++++++++++++ 1 file changed, 92 insertions(+) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py new file mode 100644 index 000000000..ae2b30777 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py @@ -0,0 +1,92 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field + +from icon4py.common.dimension import ECVDim, C2EDim, V2EDim, CellDim, EdgeDim, VertexDim, KDim + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import _mo_nh_diffusion_stencil_01 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import _mo_nh_diffusion_stencil_02 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import _mo_nh_diffusion_stencil_03 +from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import _mo_intp_rbf_rbf_vec_interpol_vertex + +@field_operator +def _fused_mo_nh_diffusion_stencil_01_02_03_rbf( + diff_multfac_smag: Field[[KDim], float], + tangent_orientation: Field[[EdgeDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + inv_vert_vert_length: Field[[EdgeDim], float], + u_vert_old: Field[[VertexDim, KDim], float], + v_vert_old: Field[[VertexDim, KDim], float], + primal_normal_vert_x: Field[[ECVDim], float], + primal_normal_vert_y: Field[[ECVDim], float], + dual_normal_vert_x: Field[[ECVDim], float], + dual_normal_vert_y: Field[[ECVDim], float], + vn: Field[[EdgeDim, KDim], float], + smag_limit: Field[[KDim], float], + smag_offset: float, + e_bln_c_s: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + wgtfac_c: Field[[CellDim, KDim], float], + ptr_coeff_1: Field[[VertexDim, V2EDim], float], + ptr_coeff_2: Field[[VertexDim, V2EDim], float], +) -> tuple[ + Field[[EdgeDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[VertexDim, KDim], float], + Field[[VertexDim, KDim], float], + ]: + + kh_smag_e, kh_smag_ec, z_nabla2_e = _mo_nh_diffusion_stencil_01( diff_multfac_smag, + tangent_orientation, inv_primal_edge_length, inv_vert_vert_length, u_vert_old, v_vert_old, + primal_normal_vert_x, primal_normal_vert_y, dual_normal_vert_x, dual_normal_vert_y, + vn, smag_limit, smag_offset) + + kh_c, div = _mo_nh_diffusion_stencil_02(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag) + + div_ic, hdef_ic = _mo_nh_diffusion_stencil_03(div, kh_c, wgtfac_c) + + u_vert, v_vert = _mo_intp_rbf_rbf_vec_interpol_vertex(z_nabla2_e, ptr_coeff_1, ptr_coeff_2) + + return kh_smag_e, div_ic, hdef_ic, u_vert, v_vert + + +@program +def fused_mo_nh_diffusion_stencil_01_02_03_rbf( + diff_multfac_smag: Field[[KDim], float], + tangent_orientation: Field[[EdgeDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + inv_vert_vert_length: Field[[EdgeDim], float], + u_vert_old: Field[[VertexDim, KDim], float], + v_vert_old: Field[[VertexDim, KDim], float], + primal_normal_vert_x: Field[[ECVDim], float], + primal_normal_vert_y: Field[[ECVDim], float], + dual_normal_vert_x: Field[[ECVDim], float], + dual_normal_vert_y: Field[[ECVDim], float], + vn: Field[[EdgeDim, KDim], float], + smag_limit: Field[[KDim], float], + smag_offset: float, + e_bln_c_s: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + wgtfac_c: Field[[CellDim, KDim], float], + ptr_coeff_1: Field[[VertexDim, V2EDim], float], + ptr_coeff_2: Field[[VertexDim, V2EDim], float], + kh_smag_e: Field[[EdgeDim, KDim], float], + div_ic: Field[[CellDim, KDim], float], + hdef_ic: Field[[CellDim, KDim], float], + u_vert: Field[[VertexDim, KDim], float], + v_vert: Field[[VertexDim, KDim], float], +): + _fused_mo_nh_diffusion_stencil_01_02_03_rbf(diff_multfac_smag, tangent_orientation, inv_primal_edge_length, inv_vert_vert_length, u_vert_old, v_vert_old, primal_normal_vert_x, primal_normal_vert_y, dual_normal_vert_x, dual_normal_vert_y, vn, smag_limit, smag_offset, e_bln_c_s, geofac_div, wgtfac_c, ptr_coeff_1, ptr_coeff_2, out=(kh_smag_e, div_ic, hdef_ic, u_vert, v_vert)) From 0d3ae055d04fcde176b059f525a6d96791387a86 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 1 Dec 2022 16:48:17 +0100 Subject: [PATCH 018/263] add remaining stencils in run function --- .../src/icon4py/atm_dyn_iconam/constants.py | 6 + .../src/icon4py/atm_dyn_iconam/diagnostic.py | 14 ++ .../src/icon4py/atm_dyn_iconam/diffusion.py | 186 +++++++++++++++--- .../src/icon4py/atm_dyn_iconam/horizontal.py | 21 ++ .../src/icon4py/atm_dyn_iconam/icon_mesh.py | 11 ++ .../atm_dyn_iconam/interpolation_state.py | 27 ++- .../icon4py/atm_dyn_iconam/metric_state.py | 24 +++ atm_dyn_iconam/tests/test_diffusion.py | 67 ++++++- .../src/icon4py/testutils/serialbox_utils.py | 83 ++++---- 9 files changed, 363 insertions(+), 76 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/constants.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diagnostic.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/constants.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/constants.py new file mode 100644 index 000000000..ced31729a --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/constants.py @@ -0,0 +1,6 @@ +from typing import Annotated + +GAS_CONSTANT_DRY_AIR: Annotated[float, "gas constant for dry air [J/K/kg], called 'rd' in ICON (mo_physical_constants.f90), https://glossary.ametsoc.org/wiki/Gas_constant"] = 287.04 +CPD: Annotated[float, "specific heat at constant pressure [J/K/kg]"] = 1004.64 +GAS_CONSTANT_WATER_VAPOR: Annotated[float, "gas constant for water vapor [J/K/kg], rv in Icon"] = 461.51 + diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diagnostic.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diagnostic.py new file mode 100644 index 000000000..e0634d202 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diagnostic.py @@ -0,0 +1,14 @@ +from dataclasses import dataclass + +from functional.common import Field + +from icon4py.common.dimension import KDim, CellDim + + +@dataclass +class DiagnosticState: + # fields for 3D elements in turbdiff + hdef_ic: Field[[CellDim, KDim], float] #! divergence at half levels(nproma,nlevp1,nblks_c) [1/s] + div_ic = Field[[CellDim, KDim], float] #! horizontal wind field deformation (nproma,nlevp1,nblks_c) [1/s^2] + dwdx = None #zonal gradient of vertical wind speed (nproma,nlevp1,nblks_c) [1/s] + dwdy = None #meridional gradient of vertical wind speed (nproma,nlevp1,nblks_c) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index 0062c6a61..152682d6c 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import math +import sys from collections import namedtuple from typing import Final @@ -19,18 +20,32 @@ from functional.ffront.fbuiltins import Field, broadcast, maximum, minimum from functional.iterator.embedded import np_as_located_field +from icon4py.atm_dyn_iconam.constants import GAS_CONSTANT_DRY_AIR, CPD +from icon4py.atm_dyn_iconam.diagnostic import DiagnosticState +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import \ + fused_mo_nh_diffusion_stencil_02_03 +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import \ + fused_mo_nh_diffusion_stencil_04_05_06 +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_07_08_09_10 import \ + fused_mo_nh_diffusion_stencil_07_08_09_10 +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_11_12 import \ + fused_mo_nh_diffusion_stencil_11_12 +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_13_14 import \ + fused_mo_nh_diffusion_stencil_13_14 from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig from icon4py.atm_dyn_iconam.icon_mesh import MeshConfig, VerticalModelParams from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState +from icon4py.atm_dyn_iconam.metric_state import MetricState from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import ( - _mo_nh_diffusion_stencil_01, + _mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_01, ) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_16 import mo_nh_diffusion_stencil_16 from icon4py.atm_dyn_iconam.prognostic import PrognosticState -from icon4py.common.dimension import ECVDim, EdgeDim, KDim, Koff, VertexDim - +from icon4py.common.dimension import ECVDim, EdgeDim, KDim, Koff, VertexDim, CellDim, \ + C2E2CDim, C2E2CODim DiffusionTupleVT = namedtuple("DiffusionParamVT", "v t") CartesianVectorTuple = namedtuple("CartesianVectorTuple", "x y") @@ -54,6 +69,16 @@ def _setup_smag_limit(diff_multfac_vn: Field[[KDim], float]) -> Field[[KDim], fl return 0.125 - 4.0 * diff_multfac_vn +@field_operator +def _scale_k(field: Field[[KDim], float], factor:float) -> Field[[KDim], float]: + return field * factor + +@program +def scale_k(field: Field[[KDim], float], factor:float, scaled_field: Field[[KDim],float]): + _scale_k(field, factor, out=scaled_field) + + + # TODO [ml] rename! @field_operator def _mo_nh_diffusion_stencil_01_scale_dtime( @@ -76,7 +101,7 @@ def _mo_nh_diffusion_stencil_01_scale_dtime( Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], ]: - diff_multfac_smag = enh_smag_fac * dtime + diff_multfac_smag = _scale_k(enh_smag_fac, dtime) return _mo_nh_diffusion_stencil_01( diff_multfac_smag, tangent_orientation, @@ -263,9 +288,9 @@ def __init__(self, grid_config: MeshConfig, vertical_params: VerticalModelParams self.lhdiff_temp = True # ! diffusion on the temperature field self.lhdiff_w = True # ! diffusion on the vertical wind field self.lhdiff_rcf = True # namelist, remove if always true - self.itype_vn_diffu = ( - 1 # ! reconstruction method used for Smagorinsky diffusion - ) + self.itype_vn_diffu = 1 # ! reconstruction method used for Smagorinsky diffusion + self.l_smag_d = False #namelist lsmag_d, if `true`, compute 3D Smagorinsky diffusion coefficient. + self.itype_t_diffu = 2 # ! discretization of temperature diffusion self.hdiff_efdt_ratio = 24.0 # ! ratio of e-folding time to time step self.hdiff_smag_fac = 0.025 # ! scaling factor for Smagorinsky diffusion @@ -282,6 +307,9 @@ def __init__(self, grid_config: MeshConfig, vertical_params: VerticalModelParams # namelist grid_nml -> TODO [ml] should go to grid config? self.l_limited_area = True + #name list: interpol_nml + self.nudge_max_coeff = 0.075 + def substep_as_float(self): return float(self.ndyn_substeps) @@ -345,6 +373,18 @@ def diffusion_type_5_smagorinski_factor(config: DiffusionConfig): return factor, heights + +def mo_nh_diffusion_stencil_15_numpy(mask_hdiff: Field[[CellDim, KDim], int], + zd_vertidx: Field[[CellDim, C2E2CDim, KDim], int], + zd_diffcoef: Field[[CellDim, KDim], float], + geofac_n2s: Field[[CellDim, C2E2CODim], float], + vcoef: Field[[C2E2CDim, KDim], float], + theta_v: Field[[CellDim, KDim], float], + z_temp: Field[[CellDim, KDim], float] + ): + pass + + class Diffusion: """Class that configures diffusion and does one diffusion step.""" @@ -360,7 +400,15 @@ def __init__( TODO [ml]: initial run: linit = .TRUE.: smag_offset and diff_multfac_vn are defined differently. """ - self.params: params + + self.params = params + self.config = config + self.rd_o_cvd: float = GAS_CONSTANT_DRY_AIR / (CPD - GAS_CONSTANT_DRY_AIR) + self.nudgezone_diff = 0.04 / (config.nudge_max_coeff + sys.float_info.epsilon) + self.bdy_diff = 0.015/(config.nudge_max_coeff + sys.float_info.epsilon) + self.fac_bdydiff_v = math.sqrt(config.substep_as_float())/config.lateral_boundary_denominator.v if config.lhdiff_rcf else 1. /config.lateral_boundary_denominator.v + self.thresh_tdiff = - 5.0 # threshold temperature deviation from neighboring grid points hat activates extra diffusion against runaway cooling + # different for init call: smag_offset = 0 self.smag_offset: float = 0.25 * params.K4 * config.substep_as_float() @@ -400,24 +448,30 @@ def __init__( k_size=config.grid.get_num_k_levels() ) ) + self.diff_multfac_smag = np_as_located_field(KDim)( + np.zeros(config.grid.get_num_k_levels()) + ) shape_vk = (config.grid.get_num_vertices(), config.grid.get_num_k_levels()) + shape_ck = (config.grid.get_num_cells(), config.grid.get_num_k_levels()) self.u_vert = np_as_located_field(VertexDim, KDim)(np.zeros(shape_vk, float)) self.v_vert = np_as_located_field(VertexDim, KDim)(np.zeros(shape_vk, float)) shape_ek = (config.grid.get_num_edges(), config.grid.get_num_k_levels()) - allocate_ek = np_as_located_field(VertexDim, KDim)(np.zeros(shape_ek, float)) + allocate_ek = np_as_located_field(EdgeDim, KDim)(np.zeros(shape_ek, float)) self.kh_smag_e = allocate_ek self.kh_smag_ec = allocate_ek self.z_nabla2_e = allocate_ek + self.z_temp = np_as_located_field(CellDim, KDim)(np.zeros(shape_ck, float)) def run( self, - diagnostic_state, + diagnostic_state: DiagnosticState, prognostic_state: PrognosticState, - metric_state, + metric_state: MetricState, interpolation_state: InterpolationState, - dtime, + dtime:float, tangent_orientation: Field[[EdgeDim], float], inverse_primal_edge_lengths: Field[[EdgeDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], inverse_vertical_vertex_lengths: Field[[EdgeDim], float], primal_normal_vert: CartesianVectorTuple[ Field[[ECVDim], float], Field[[ECVDim], float] @@ -425,6 +479,8 @@ def run( dual_normal_vert: CartesianVectorTuple[ Field[[ECVDim], float], Field[[ECVDim], float] ], + area_edges: Field[[EdgeDim], float], + cell_areas: Field[[CellDim], float] ): """ Run a diffusion step. @@ -439,8 +495,9 @@ def run( # ------- # Oa logging # 0b call timer start - # ~~0c. apply dtime to enh_smag_factor~~ done inside stencil - + # 0c. dtime dependent stuff: enh_smag_factor, r_dtimensubsteps + r_dtimensubsteps = 1.0 / dtime if self.config.lhdiff_rcf else 1./ (dtime * self.config.substep_as_float) + klevels = self.config.grid.get_num_k_levels() # TODO is this needed? set_zero_v_k(self.u_vert) set_zero_v_k(self.v_vert) @@ -456,17 +513,11 @@ def run( ) # 2. HALO EXCHANGE -- CALL sync_patch_array_mult # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 + # 0c. dtime dependent stuff: enh_smag_factor, ~~r_dtimensubsteps~~ + scale_k(self.enh_smag_fac, dtime, self.diff_multfac_smag) - # tangent_orientation = p_patch % edges % tangent_orientation(:, 1) - # inv_primal_edge_length=p_patch%edges%inv_primal_edge_length(:,1) - # inv_vert_vert_length = p_patch % edges % inv_vert_vert_length(:, 1), - # primal_normal_vert_x = p_patch % edges % primal_normal_vert_x(:,:, 1) - # primal_normal_vert_y = p_patch % edges % primal_normal_vert_y(:,:, 1) - # dual_normal_vert_x=p_patch%edges%dual_normal_vert_x(:,:,1) - # dual_normal_vert_x = p_patch % edges % dual_normal_vert_x(:,:, 1) - - mo_nh_diffusion_stencil_01_scaled_dtime( - self.enh_smag_fac, + mo_nh_diffusion_stencil_01( + self.diff_multfac_smag, tangent_orientation, inverse_primal_edge_lengths, inverse_vertical_vertex_lengths, @@ -482,19 +533,102 @@ def run( self.kh_smag_ec, self.z_nabla2_e, self.smag_offset, - dtime, + dtime, domain={KDim:(0, klevels), EdgeDim:(), VertexDim:()} ) + fused_mo_nh_diffusion_stencil_02_03(self.kh_smag_ec, + prognostic_state.normal_wind, + interpolation_state.e_bln_c_s, #:Field[[CellDim, C2EDim], float], + interpolation_state.geofac_div, # Field[[CellDim, C2EDim], float], + self.diff_multfac_smag, # Field[[KDim], float], + metric_state.wgtfac_c, # Field[[CellDim, KDim], float], + diagnostic_state.div_ic, #Field[[CellDim, KDim], float], + diagnostic_state.hdef_ic, #: Field[[CellDim, KDim], float]) + offset_provider={"C2E": "TODO_C2E"}) + # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH # 5. CALL rbf_vec_interpol_vertex_wp - + mo_intp_rbf_rbf_vec_interpol_vertex(self.z_nabla2_e, + interpolation_state.rbf_coeff_1, + interpolation_state.rbf_coeff_2, + self.u_vert, + self.v_vert) # 6. HALO EXCHANGE -- CALL sync_patch_array_mult + # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 + #TODO where to get index field from + #horz_idx: Field[[EdgeDim], int32], + #start_2nd_nudge_line_idx_e: int32, + + fused_mo_nh_diffusion_stencil_04_05_06(u_vert=self.u_vert, + v_vert=self.v_vert, + primal_normal_vert_v1=primal_normal_vert.x, + primal_normal_vert_v2=primal_normal_vert.y, + z_nabla2_e=self.z_nabla2_e, + inv_vert_vert_length = inverse_vertical_vertex_lengths, + inv_primal_edge_length = inverse_primal_edge_lengths, + area_edges=area_edges, + kh_smag_e=self.kh_smag_e, + diff_multfac_vn=self.diff_multfac_vn, + nudgecoeff_e=interpolation_state.nudgecoeff_e, + vn=prognostic_state.normal_wind, + horz_idx = horizontal_index, + nudgezone_diff=self.nudgezone_diff, + fac_bdydiff_v= self.fac_bdydiff_v, + start_2nd_nudge_line_idx_e = s + ) # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 + + fused_mo_nh_diffusion_stencil_07_08_09_10(area = cell_areas, geofac_n2s=interpolation_state.geofac_n2s, + geofac_grg_x=interpolation_state.geofac_grg_x, + geofac_grg_y=interpolation_state.geofac_grg_y, + w_old=prognostic_state.vertical_wind, + w=prognostic_state.vertical_wind, + dwdx=diagnostic_state.dwdx, + dwdy=diagnostic_state.dwdy, + diff_multfac_w=self.diff_multfac_w, + diff_multfac_n2w=self.diff_multfac_n2w, + vert_idx=, + horz_idx=, + nrdmax=self.config.vertical_params.index_of_damping_height, + interior_idx=, + halo_idx=, + ) # 8. HALO EXCHANGE: CALL sync_patch_array # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 + + # TODO check: kh_smag_e is an out field, should not be calculated in init? + fused_mo_nh_diffusion_stencil_11_12(theta_v=prognostic_state.theta_v, + theta_ref_mc=metric_state.theta_ref_mc, + thresh_tdiff=self.thresh_tdiff, + kh_smag_e=self.kh_smag_e, + ) + + + fused_mo_nh_diffusion_stencil_13_14(kh_smag_e=self.kh_smag_e, + inv_dual_edge_length=inv_dual_edge_length, + theta_v=prognostic_state.theta_v, + geofac_div=interpolation_state.geofac_div, + z_temp=self.z_temp + ) + + mo_nh_diffusion_stencil_15_numpy(mask_hdiff=metric_state.mask_hdiff, + zd_vertidx=metric_state.zd_vertidx_dsl, + vcoef=metric_state.zd_diffcoef, + zd_diffcoef=metric_state.zd_diffcoef, + geofac_n2s=interpolation_state.geofac_n2s, + theta_v=prognostic_state.theta_v, + z_temp=self.z_temp + ) + + mo_nh_diffusion_stencil_16(z_temp = self.z_temp, + area=cell_areas, + theta_v=prognostic_state.theta_v, + exner=prognostic_state.exner_pressure, + rd_o_cvd=self.rd_o_cvd + ) # 10. HALO EXCHANGE sync_patch_array diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py index 8bbbc3c2e..b1d1e5f3e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py @@ -10,7 +10,19 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +from enum import Enum +from functional.common import Dimension, DimensionKind + + +# rather use DSL (dusk) names +class HorizontalIndexMarker(Enum): + INTERIOR = "0" + NUDGING = "grf_bdywidth" + HALO= "min_rl_int" + END="min_rl" + LOCAL_BOUNDARY="1" + MAX_RL="max_rl" class HorizontalMeshConfig: def __init__(self, num_vertices: int, num_edges: int, num_cells: int): @@ -26,3 +38,12 @@ def get_num_edges(self): def get_num_cells(self): return self._num_cells + + def get_index(self, dim:Dimension, marker:HorizontalIndexMarker) -> int: + if dim.kind != DimensionKind.HORIZONTAL: + raise ValueError("only defined for {} dimension kind ", DimensionKind.HORIZONTAL) + return 0 + + + + diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_mesh.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_mesh.py index 435133edc..01a2ffdb0 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_mesh.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_mesh.py @@ -10,6 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +from enum import Enum import numpy as np from functional.common import Field @@ -19,6 +20,10 @@ from icon4py.common.dimension import EdgeDim, KDim + + + + class MeshConfig: def __init__(self, horizontalMesh: HorizontalMeshConfig): self.num_k_levels = 65 @@ -37,6 +42,10 @@ def get_num_vertices(self): def get_num_edges(self): return self.horizontal._num_edges + def get_num_cells(self): + return self.horizontal._num_cells + + class IconMesh: def __init__(self, config: MeshConfig): @@ -111,3 +120,5 @@ def init_nabla2_factor_in_upper_damping_zone( / (self.vct_a[2] - self.vct_a[self.index_of_damping_height + 1]) ) return np_as_located_field(KDim)(buffer) + + diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py index d0425ad92..9e94ea574 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py @@ -15,14 +15,27 @@ from functional.common import Field -from icon4py.common.dimension import V2EDim, VertexDim +from icon4py.common.dimension import V2EDim, VertexDim, CellDim, C2EDim, E2CDim, EdgeDim, C2E2CODim @dataclass class InterpolationState: - rbf_coeff_1: Field[ - [VertexDim, V2EDim], float - ] # rbf_vec_coeff_v_1(nproma, rbf_vec_dim_v, nblks_v) - rbf_coeff_2: Field[ - [VertexDim, V2EDim], float - ] # rbf_vec_coeff_v_2(nproma, rbf_vec_dim_v, nblks_v) + def __init__(self): + self.g = None + + """ + represents + """ + + c_lin_e: Field[[EdgeDim, E2CDim], float] # (nproma, 2, nblks_e) + e_bln_c_s: Field[[CellDim, C2EDim], float] # coefficent for bilinear interpolation from edge to cell () + rbf_coeff_1: Field[[VertexDim, V2EDim], float] # rbf_vec_coeff_v_1(nproma, rbf_vec_dim_v, nblks_v) + rbf_coeff_2: Field[[VertexDim, V2EDim], float] # rbf_vec_coeff_v_2(nproma, rbf_vec_dim_v, nblks_v) + + geofac_div: Field[[CellDim,C2EDim], float] # factor for divergence (nproma,cell_type,nblks_c) + + geofac_n2s: Field[[CellDim, C2E2CODim], float] # factor for nabla2-scalar (nproma,cell_type+1,nblks_c) + geofac_grg_x: Field[[CellDim, C2E2CODim], float] # factor for green gauss gradient (nproma,4,nblks_c,2) + geofac_grg_y: Field[[CellDim, C2E2CODim], float] # TODO combine to tuple + nudgecoeff_e: Field[[EdgeDim],float] #Nudgeing coeffients for edges + nudgecoeff_c: Field[[CellDim], float] # nudgeing coeffiecients for cells diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py new file mode 100644 index 000000000..ecebde53f --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py @@ -0,0 +1,24 @@ +from dataclasses import dataclass + +from functional.common import Field + +from icon4py.common.dimension import KDim, EdgeDim, CellDim, C2E2CDim + + +@dataclass +class MetricState: + theta_ref_mc: Field[[CellDim, KDim], float] + enhfac_diffu: Field[[KDim], float] #Enhancement factor for nabla4 background diffusion TODO check dimension + wgtfac_e: Field[[EdgeDim, KDim], float] #weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_e) + wgtfac_c: Field[[CellDim, KDim], float] #weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_c) + wgtfacq1_e: Field[[EdgeDim, ], float] #weighting factor for quadratic interpolation to model top (nproma,3,nblks_e) + wgtfacq_e: Field[[EdgeDim, ], float] #weighting factor for quadratic interpolation to surface (nproma,3,nblks_e) + ddqz_z_full_e: Field[[EdgeDim, KDim], float] #functional determinant of the metrics [sqrt(gamma)] (nproma,nlev,nblks_e) + mask_hdiff: Field[[CellDim, KDim], int] + zd_vertidx_dsl: Field[[CellDim, C2E2CDim, KDim], int] + zd_diffcoef: Field[[CellDim, KDim], float] + zd_intcoef: Field[[CellDim, KDim], float] + + + + diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 6b8237607..260826d57 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -16,22 +16,34 @@ import pytest from functional.iterator.embedded import np_as_located_field +from icon4py.atm_dyn_iconam.diagnostic import DiagnosticState from icon4py.atm_dyn_iconam.diffusion import ( Diffusion, DiffusionConfig, DiffusionParams, enhanced_smagorinski_factor, init_diffusion_local_fields, - set_zero_v_k, + set_zero_v_k, scale_k, CartesianVectorTuple, ) +from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState +from icon4py.atm_dyn_iconam.metric_state import MetricState +from icon4py.atm_dyn_iconam.prognostic import PrognosticState from icon4py.common.dimension import KDim, VertexDim from icon4py.testutils.serialbox_utils import ( - read_from_call_diffusion_init_ser_data, + read_from_call_diffusion_init_ser_data, SerializedDataProvider, ) from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import random_field, zero_field +def test_scale_k(): + mesh = SimpleMesh() + field = random_field(mesh, KDim) + scaled_field = zero_field(mesh, KDim) + factor = 2.0 + scale_k(field, factor, scaled_field, offset_provider = {}) + assert np.allclose(factor * np.asarray(field), scaled_field) + def smag_limit_numpy(shape, k4, substeps): return 0.125 - 4.0 * diff_multfac_vn_numpy(shape, k4, substeps) @@ -177,14 +189,12 @@ def test_smagorinski_factor_diffusion_type_5(): def test_diffusion_init(): - data_path = os.path.join(os.path.dirname(__file__), "ser_data") + data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") physical_heights_name = "vct_a" - fields, meta = read_from_call_diffusion_init_ser_data( - data_path, - "icon_diffusion", - metadata=["nlev", "linit", "date"], - fields=[physical_heights_name], - ) + serializer = SerializedDataProvider("icon_diffusion", data_path) + serializer.print_info() + meta, fields = serializer.get_fields(metadata=["nlev", "linit", "date"], fields=[physical_heights_name]) + print(f"meta {meta}") assert meta["nlev"] == 65 assert meta["linit"] is False @@ -224,3 +234,42 @@ def test_diffusion_init(): vct_a, ) assert np.allclose(expected_enh_smag_fac, np.asarray(diffusion.enh_smag_fac)) + + +def test_diffusion_run(): + data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") + physical_heights_name = "vct_a" + serializer = SerializedDataProvider("icon_diffusion", data_path) + serializer.print_info() + meta, fields = serializer.get_fields(metadata=["nlev", "linit", "dtime"], + fields=[physical_heights_name, "tangent_orientation"]) + config = DiffusionConfig.create_with_defaults() + additional_parameters = DiffusionParams(config) + vct_a = np_as_located_field(KDim)(fields[physical_heights_name]) + diffusion = Diffusion(config, additional_parameters, vct_a) + + diagnostic_state = DiagnosticState() + prognostic_state = PrognosticState() + interpolation_state = InterpolationState() + metric_state = MetricState() + dtime = meta["dtime"] + orientation = fields["tangent_orientation"] + + # inverse_primal_edge_lengths: Field[[EdgeDim], float], + inverse_vertical_vertex_lengths = fields["inv_vert_vert_length"] + primal_normal_vert: CartesianVectorTuple = (fields["primal_normal_vert_x"], fields["primal_normal_vert_y"]) + dual_normal_vert: CartesianVectorTuple = (fields["dual_normal_vert_x"], fields["dual_normal_vert_y"]) + + + + diffusion.run(diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + metric_state=metric_state, + interpolation_state=interpolation_state, + dtime=dtime, + tangent_orientation=orientation, + inverse_primal_edge_lengths=inverse_primal_edge_lengths, + inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, + primal_normal_vert=primal_normal_vert, + dual_normal_vert=dual_normal_vert + ) diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index cf9848d90..233c81463 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -26,39 +26,54 @@ import serialbox as ser -def read_from_call_diffusion_init_ser_data( - path, fname_prefix, metadata: List[str], fields: List[str] -): - rank = 0 - fname = f"{fname_prefix}_rank{str(rank)}" - serializer = ser.Serializer(ser.OpenModeKind.Read, path, fname) - save_points = serializer.savepoint_list() - print(save_points) - field_names = serializer.fieldnames() - print(field_names) - savepoint = ( - serializer.savepoint["call-diffusion-init"] - .linit[False] - .date["2021-06-20T12:00:10.000"] - .as_savepoint() - ) - print(savepoint.metainfo) - meta_present = {} - meta_absent = [] - for md in metadata: - if md in savepoint.metainfo.to_dict(): - meta_present[md] = savepoint.metainfo[md] - else: - meta_absent.append(md) +class SerializedDataProvider: + + def __init__(self, fname_prefix, path="."): + self.rank = 0 + self.serializer: ser.Serializer = None + self.file_path: str = path + self.fname = f"{fname_prefix}_rank{str(self.rank)}" + self._init_serializer() + + def _init_serializer(self): + if not self.fname: + print(" WARNING: no filename! closing serializer") + self.serializer = ser.Serializer(ser.OpenModeKind.Read, self.file_path, self.fname) - fields_present = {} - fields_absent = [] - for field_name in fields: - if field_name in field_names: - fields_present[field_name] = serializer.read(field_name, savepoint) - else: - fields_absent.append(field_name) - [print(f"field {f} not present in savepoint") for f in fields_absent] - [print(f"metadata {f} not present in savepoint") for f in meta_absent] + def print_info(self): + print(f"SAVEPOINTS: {self.serializer.savepoint_list()}") + print(f"FIELDNAMES: {self.serializer.fieldnames()}") - return fields_present, meta_present + def get_fields(self, metadata: List[str], fields: List[str]): + savepoint = ( + self.serializer.savepoint["call-diffusion-init"] + .linit[False] + .date["2021-06-20T12:00:10.000"] + .as_savepoint() + ) + print(savepoint.metainfo) + meta_present = {} + meta_absent = [] + for md in metadata: + if md in savepoint.metainfo.to_dict(): + meta_present[md] = savepoint.metainfo[md] + else: + meta_absent.append(md) + + fields_present = {} + fields_absent = [] + for field_name in fields: + if field_name in self.serializer.fieldnames(): + fields_present[field_name] = self.serializer.read(field_name, savepoint) + else: + fields_absent.append(field_name) + [print(f"field {f} not present in savepoint") for f in fields_absent] + [print(f"metadata {f} not present in savepoint") for f in meta_absent] + return meta_present, fields_present + +def read_from_call_diffusion_init_ser_data( + path, fname_prefix,metadata: List[str], fields: List[str] +): + serializer = SerializedDataProvider(fname_prefix, path) + serializer.print_info() + return serializer.get_fields(metadata, fields) From bed2a62845cf0d07353ee75f3bd7c18860aa9695 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 2 Dec 2022 18:00:48 +0100 Subject: [PATCH 019/263] add icon grid: start/end indices deserialize connectivity fields --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 49 +++--- .../fused_mo_nh_diffusion_stencil_04_05_06.py | 3 +- .../src/icon4py/atm_dyn_iconam/horizontal.py | 34 ++-- .../src/icon4py/atm_dyn_iconam/icon_grid.py | 145 ++++++++++++++++++ .../src/icon4py/atm_dyn_iconam/icon_mesh.py | 124 --------------- atm_dyn_iconam/tests/test_diffusion.py | 59 +++---- atm_dyn_iconam/tests/test_icon_grid.py | 72 +++++++++ atm_dyn_iconam/tests/test_vertical.py | 2 +- .../src/icon4py/testutils/serialbox_utils.py | 115 +++++++++++++- 9 files changed, 411 insertions(+), 192 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_mesh.py create mode 100644 atm_dyn_iconam/tests/test_icon_grid.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index 152682d6c..ace763af4 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -16,6 +16,7 @@ from typing import Final import numpy as np +from functional.common import Dimension from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, broadcast, maximum, minimum from functional.iterator.embedded import np_as_located_field @@ -33,7 +34,7 @@ from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_13_14 import \ fused_mo_nh_diffusion_stencil_13_14 from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig -from icon4py.atm_dyn_iconam.icon_mesh import MeshConfig, VerticalModelParams +from icon4py.atm_dyn_iconam.icon_grid import MeshConfig, VerticalModelParams, IconGrid from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState from icon4py.atm_dyn_iconam.metric_state import MetricState from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( @@ -272,15 +273,16 @@ def create_with_defaults(cls): horizontal = HorizontalMeshConfig( num_vertices=50000, num_cells=50000, num_edges=50000 ) - gridConfig = MeshConfig(horizontalMesh=horizontal) + + grid = IconGrid().with_config(MeshConfig(horizontalMesh=horizontal)) verticalParams = VerticalModelParams( - rayleigh_damping_height=12500, vct_a=np.zeros(gridConfig.get_num_k_levels()) + rayleigh_damping_height=12500, vct_a=np.zeros(grid.get_num_k_levels()) ) - return cls(grid_config=gridConfig, vertical_params=verticalParams) + return cls(grid=grid, vertical_params=verticalParams) - def __init__(self, grid_config: MeshConfig, vertical_params: VerticalModelParams): + def __init__(self, grid: IconGrid, vertical_params: VerticalModelParams): # TODO [ml]: external stuff grid related, p_patch, - self.grid = grid_config + self.grid = grid self.vertical_params = vertical_params # from namelist diffusion_nml self.diffusion_type = 5 # hdiff_order ! order of nabla operator for diffusion @@ -314,6 +316,8 @@ def substep_as_float(self): return float(self.ndyn_substeps) + + class DiffusionParams: """Calculates derived quantities depending on the diffusion config.""" @@ -403,6 +407,7 @@ def __init__( self.params = params self.config = config + self.grid = config.grid self.rd_o_cvd: float = GAS_CONSTANT_DRY_AIR / (CPD - GAS_CONSTANT_DRY_AIR) self.nudgezone_diff = 0.04 / (config.nudge_max_coeff + sys.float_info.epsilon) self.bdy_diff = 0.015/(config.nudge_max_coeff + sys.float_info.epsilon) @@ -479,7 +484,7 @@ def run( dual_normal_vert: CartesianVectorTuple[ Field[[ECVDim], float], Field[[ECVDim], float] ], - area_edges: Field[[EdgeDim], float], + edge_areas: Field[[EdgeDim], float], cell_areas: Field[[CellDim], float] ): """ @@ -497,7 +502,7 @@ def run( # 0b call timer start # 0c. dtime dependent stuff: enh_smag_factor, r_dtimensubsteps r_dtimensubsteps = 1.0 / dtime if self.config.lhdiff_rcf else 1./ (dtime * self.config.substep_as_float) - klevels = self.config.grid.get_num_k_levels() + klevels = self.config.grid.k_levels() # TODO is this needed? set_zero_v_k(self.u_vert) set_zero_v_k(self.v_vert) @@ -544,7 +549,7 @@ def run( metric_state.wgtfac_c, # Field[[CellDim, KDim], float], diagnostic_state.div_ic, #Field[[CellDim, KDim], float], diagnostic_state.hdef_ic, #: Field[[CellDim, KDim], float]) - offset_provider={"C2E": "TODO_C2E"}) + offset_provider={"C2E": self.grid.get_c2e_connectivity()}) # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH @@ -569,15 +574,16 @@ def run( z_nabla2_e=self.z_nabla2_e, inv_vert_vert_length = inverse_vertical_vertex_lengths, inv_primal_edge_length = inverse_primal_edge_lengths, - area_edges=area_edges, + area_edges=edge_areas, kh_smag_e=self.kh_smag_e, diff_multfac_vn=self.diff_multfac_vn, nudgecoeff_e=interpolation_state.nudgecoeff_e, vn=prognostic_state.normal_wind, - horz_idx = horizontal_index, + horz_idx = None, nudgezone_diff=self.nudgezone_diff, fac_bdydiff_v= self.fac_bdydiff_v, - start_2nd_nudge_line_idx_e = s + start_2nd_nudge_line_idx_e = None, + offset_provider={} ) # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 @@ -591,11 +597,12 @@ def run( dwdy=diagnostic_state.dwdy, diff_multfac_w=self.diff_multfac_w, diff_multfac_n2w=self.diff_multfac_n2w, - vert_idx=, - horz_idx=, + vert_idx= None, + horz_idx=None, nrdmax=self.config.vertical_params.index_of_damping_height, - interior_idx=, - halo_idx=, + interior_idx=None, + halo_idx=None, + offset_provider={"C2E2CO": self.grid.get_c2e2c0_connectivity()} ) # 8. HALO EXCHANGE: CALL sync_patch_array # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, @@ -606,6 +613,7 @@ def run( theta_ref_mc=metric_state.theta_ref_mc, thresh_tdiff=self.thresh_tdiff, kh_smag_e=self.kh_smag_e, + offset_provider={"E2C":self.grid.get_e2c_connectivity(), "C2E2C":self.grid.get_c2e2c_connectivity()} ) @@ -613,7 +621,8 @@ def run( inv_dual_edge_length=inv_dual_edge_length, theta_v=prognostic_state.theta_v, geofac_div=interpolation_state.geofac_div, - z_temp=self.z_temp + z_temp=self.z_temp, + offset_provider={"C2E":self.grid.get_c2e_connectivity(), "E2C":self.grid.get_e2c_connectivity()} ) mo_nh_diffusion_stencil_15_numpy(mask_hdiff=metric_state.mask_hdiff, @@ -622,13 +631,15 @@ def run( zd_diffcoef=metric_state.zd_diffcoef, geofac_n2s=interpolation_state.geofac_n2s, theta_v=prognostic_state.theta_v, - z_temp=self.z_temp + z_temp=self.z_temp, + offset_provider={} ) mo_nh_diffusion_stencil_16(z_temp = self.z_temp, area=cell_areas, theta_v=prognostic_state.theta_v, exner=prognostic_state.exner_pressure, - rd_o_cvd=self.rd_o_cvd + rd_o_cvd=self.rd_o_cvd, + offset_provider={} ) # 10. HALO EXCHANGE sync_patch_array diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py index 1438e07c1..62f4316d5 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py @@ -77,6 +77,7 @@ def fused_mo_nh_diffusion_stencil_04_05_06( fac_bdydiff_v: float, start_2nd_nudge_line_idx_e: int32, ): - _fused_mo_nh_diffusion_stencil_04_05_06(u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2, z_nabla2_e, inv_vert_vert_length, + _fused_mo_nh_diffusion_stencil_04_05_06(u_vert, v_vert, primal_normal_vert_v1, + primal_normal_vert_v2, z_nabla2_e, inv_vert_vert_length, inv_primal_edge_length, area_edge, kh_smag_e, diff_multfac_vn, nudgecoeff_e, vn, horz_idx, nudgezone_diff, fac_bdydiff_v, start_2nd_nudge_line_idx_e, out=vn) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py index b1d1e5f3e..b1b3a4afb 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py @@ -10,19 +10,30 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +from dataclasses import dataclass from enum import Enum +from typing import Final from functional.common import Dimension, DimensionKind - -# rather use DSL (dusk) names -class HorizontalIndexMarker(Enum): - INTERIOR = "0" - NUDGING = "grf_bdywidth" - HALO= "min_rl_int" - END="min_rl" - LOCAL_BOUNDARY="1" - MAX_RL="max_rl" +@dataclass(frozen=True) +class HorizontalMeshParams: + NUM_GHOST_ROWS:Final[int] = 2 + GRF_BOUNDARY_WIDTH_CELL:Final[int] = 12 + MIN_RL_CELL_INT:Final[int] = 4 + MIN_RL_CELL:Final[int] = MIN_RL_CELL_INT - 2 * NUM_GHOST_ROWS + MAX_RL_CELL:Final[int] = 13 + MIN_RL_VERTEX_INT:Final[int] = MIN_RL_CELL_INT + MIN_RL_VERTEX:Final[int] = MIN_RL_VERTEX_INT - (NUM_GHOST_ROWS + 1) + MAX_RL_VERTEX:Final[int] = MAX_RL_CELL + MIN_RL_EDGE_INT:Final[int] = 2 * MIN_RL_CELL_INT + MIN_RL_EDGE:Final[int] = MIN_RL_EDGE_INT - (2 * NUM_GHOST_ROWS + 1) + MAX_RL_EDGE:Final[int] = 2 * MAX_RL_CELL + GRF_BOUNDARY_WIDTH_EDGES: Final[int] = 9 + +class HorizontalMarkerIndex(Enum): + START_PROG_CELL = HorizontalMeshParams.GRF_BOUNDARY_WIDTH_CELL + 1 + END_PROG_CELL = HorizontalMeshParams.MIN_RL_CELL_INT class HorizontalMeshConfig: def __init__(self, num_vertices: int, num_edges: int, num_cells: int): @@ -39,10 +50,7 @@ def get_num_edges(self): def get_num_cells(self): return self._num_cells - def get_index(self, dim:Dimension, marker:HorizontalIndexMarker) -> int: - if dim.kind != DimensionKind.HORIZONTAL: - raise ValueError("only defined for {} dimension kind ", DimensionKind.HORIZONTAL) - return 0 + diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py new file mode 100644 index 000000000..72ef0cc41 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py @@ -0,0 +1,145 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +from enum import Enum +from typing import Tuple + +import numpy as np +from functional.common import Field, Dimension, DimensionKind +from functional.iterator.embedded import np_as_located_field, NeighborTableOffsetProvider + +from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig, HorizontalMarkerIndex +from icon4py.common.dimension import EdgeDim, KDim, CellDim, VertexDim + + +class MeshConfig: + def __init__(self, horizontalMesh: HorizontalMeshConfig): + self._num_k_levels = 65 + self._n_shift_total = 0 + self._horizontal = horizontalMesh + + @property + def num_k_levels(self): + return self._num_k_levels + + @property + def n_shift_total(self): + return self._n_shift_total + + @property + def get_num_vertices(self): + return self._horizontal._num_vertices + + def get_num_edges(self): + return self._horizontal._num_edges + + def get_num_cells(self): + return self._horizontal._num_cells + +def builder(func): + def wrapper(self, *args, **kwargs): + func(self, *args, **kwargs) + return self + return wrapper + +class IconGrid: + def __init__(self): + self.config:MeshConfig = None + self.start_indices = {} + self.end_indices = {} + self.connectivities={} + + @builder + def with_config(self, config: MeshConfig): + self.config = config + + @builder + def with_start_end_indices(self, dim:Dimension, start_indices:np.ndarray, end_indices: np.ndarray): + self.start_indices[dim] = start_indices + self.end_indices[dim] = end_indices + + @builder + def with_connectivity(self, **connectivity): + self.connectivities.update(**connectivity) + + def k_levels(self): + return self.config.num_k_levels + + + + def get_indices_from_to(self, dim:Dimension, start_marker:int, end_marker: int) -> Tuple[int, int]: + if dim.kind != DimensionKind.HORIZONTAL: + raise ValueError("only defined for {} dimension kind ", DimensionKind.HORIZONTAL) + #print(self.start_indices[dim]) + #print(self.end_indices[dim]) + return self.start_indices[dim][start_marker], self.end_indices[dim][end_marker] + + def get_c2e_connectivity(self): + table = self.connectivities["c2e"] + return NeighborTableOffsetProvider(table, CellDim, EdgeDim, table.shape[1]) + + def get_e2c_connectivity(self): + table = self.connectivities["e2c"] + return NeighborTableOffsetProvider(table, EdgeDim, CellDim, table.shape[1]) + + def get_e2v_connectivity(self): + table = self.connectivities["e2v"] + return NeighborTableOffsetProvider(table, EdgeDim, VertexDim, table.shape[1]) + + def get_c2e2c_connectivity(self): + table = self.connectivities["c2e2c"] + return NeighborTableOffsetProvider(table, CellDim, CellDim, table.shape[1]) + + def get_c2e2c0_connectivity(self): + table = self.connectivities["c2e2c0"] + return NeighborTableOffsetProvider(table, CellDim, CellDim, table.shape[1]) + + +class VerticalModelParams: + def __init__(self, vct_a: np.ndarray, rayleigh_damping_height: float = 12500.0): + """ + Contains physical parameters defined on the grid. + + Args: + vct_a: + rayleigh_damping_height height of rayleigh damping in [m] mo_nonhydro_nml + """ + self.rayleigh_damping_height = rayleigh_damping_height + self.vct_a = vct_a + # TODO klevels in ICON are inverted! + self.index_of_damping_height = np.argmax( + self.vct_a >= self.rayleigh_damping_height + ) + + def get_index_of_damping_layer(self): + return self.index_of_damping_height + + def get_physical_heights(self) -> Field[[KDim], float]: + return np_as_located_field(KDim)(self.vct_a) + + def init_nabla2_factor_in_upper_damping_zone( + self, k_size: int + ) -> Field[[KDim], float]: + # this assumes n_shift == 0 + buffer = np.zeros(k_size) + buffer[2 : self.index_of_damping_height] = ( + 1.0 + / 12.0 + * ( + self.vct_a[2 : self.index_of_damping_height] + - self.vct_a[self.index_of_damping_height + 1] + ) + / (self.vct_a[2] - self.vct_a[self.index_of_damping_height + 1]) + ) + return np_as_located_field(KDim)(buffer) + + diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_mesh.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_mesh.py deleted file mode 100644 index 01a2ffdb0..000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_mesh.py +++ /dev/null @@ -1,124 +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 -from enum import Enum - -import numpy as np -from functional.common import Field -from functional.iterator.embedded import np_as_located_field - -from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig -from icon4py.common.dimension import EdgeDim, KDim - - - - - - -class MeshConfig: - def __init__(self, horizontalMesh: HorizontalMeshConfig): - self.num_k_levels = 65 - self.n_shift_total = 0 - self.horizontal = horizontalMesh - - def get_num_k_levels(self): - return self.num_k_levels - - def get_n_shift(self): - return self.n_shift_total - - def get_num_vertices(self): - return self.horizontal._num_vertices - - def get_num_edges(self): - return self.horizontal._num_edges - - def get_num_cells(self): - return self.horizontal._num_cells - - - -class IconMesh: - def __init__(self, config: MeshConfig): - self.config = config - # TODO [ml] calculate the content of this stuff? read in? - self._tangent_orientation = np_as_located_field(EdgeDim, KDim)( - np.zeros(self.config.get_num_edges(), self.config.get_num_k_levels()) - ) - self._triangle_edge_inverse_length = np_as_located_field(EdgeDim)( - np.zeros(self.config.get_num_edges()) - ) - # normal to triangle edge projected to the location of the vertices - self._primal_normal_vert = ( - np_as_located_field(EdgeDim)(np.zeros(self.config.get_num_edges())), - np_as_located_field(EdgeDim)(np.zeros(self.config.get_num_edges())), - ) - - # tangent to triangle edge, projected to the location of the vertices - self._dual_normal_vert_x = ( - np_as_located_field(EdgeDim)(np.zeros(self.config.get_num_edges())), - np_as_located_field(EdgeDim)(np.zeros(self.config.get_num_edges())), - ) - - # TODO [ml] geometry - def tangent_orientation(self): - return self._tangent_orientation - - def inv_primal_edge_length(self): - return self._triangle_edge_inverse_length - - def primal_normal_vert(self): - return self._primal_normal_vert - - def dual_normal_vert_x(self): - return self._dual_normal_vert_x - - -class VerticalModelParams: - def __init__(self, vct_a: np.ndarray, rayleigh_damping_height: float = 12500.0): - """ - Contains physical parameters defined on the grid. - - Args: - vct_a: TODO read from vertical_coord_tables! - rayleigh_damping_height height of rayleigh damping in [m] mo_nonhydro_nml - """ - self.rayleigh_damping_height = rayleigh_damping_height - self.vct_a = vct_a - # TODO klevels in ICON are inverted! - self.index_of_damping_height = np.argmax( - self.vct_a >= self.rayleigh_damping_height - ) - - def get_index_of_damping_layer(self): - return self.index_of_damping_height - - def get_physical_heights(self) -> Field[[KDim], float]: - return np_as_located_field(KDim)(self.vct_a) - - def init_nabla2_factor_in_upper_damping_zone( - self, k_size: int - ) -> Field[[KDim], float]: - # this assumes n_shift == 0 - buffer = np.zeros(k_size) - buffer[2 : self.index_of_damping_height] = ( - 1.0 - / 12.0 - * ( - self.vct_a[2 : self.index_of_damping_height] - - self.vct_a[self.index_of_damping_height + 1] - ) - / (self.vct_a[2] - self.vct_a[self.index_of_damping_height + 1]) - ) - return np_as_located_field(KDim)(buffer) - - diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 260826d57..654d8a4ed 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -30,7 +30,7 @@ from icon4py.atm_dyn_iconam.prognostic import PrognosticState from icon4py.common.dimension import KDim, VertexDim from icon4py.testutils.serialbox_utils import ( - read_from_call_diffusion_init_ser_data, SerializedDataProvider, + IconSerialDataProvider, ) from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import random_field, zero_field @@ -122,11 +122,6 @@ def test_set_zero_vertex_k(): assert np.allclose(0.0, f) -@pytest.mark.xfail -def test_diffusion_run(): - pytest.fail("not implemented yet") - - def test_diffusion_coefficients_with_hdiff_efdt_ratio(): config: DiffusionConfig = DiffusionConfig.create_with_defaults() config.hdiff_efdt_ratio = 1.0 @@ -190,18 +185,20 @@ def test_smagorinski_factor_diffusion_type_5(): def test_diffusion_init(): data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") - physical_heights_name = "vct_a" - serializer = SerializedDataProvider("icon_diffusion", data_path) + serializer = IconSerialDataProvider("icon_diffusion_init", data_path) serializer.print_info() - meta, fields = serializer.get_fields(metadata=["nlev", "linit", "date"], fields=[physical_heights_name]) + first_run_date = "2021-06-20T12:00:10.000" + savepoint = serializer.from_savepoint(linit=False, date=first_run_date) + vct_a = savepoint.physical_height_field() + meta = savepoint.get_metadata("nlev", "linit", "date") print(f"meta {meta}") assert meta["nlev"] == 65 assert meta["linit"] is False - assert meta["date"] == "2021-06-20T12:00:10.000" + assert meta["date"] == first_run_date config = DiffusionConfig.create_with_defaults() additional_parameters = DiffusionParams(config) - vct_a = np_as_located_field(KDim)(fields[physical_heights_name]) + diffusion = Diffusion(config, additional_parameters, vct_a) # assert static local fields are initialized and correct: assert ( @@ -236,31 +233,32 @@ def test_diffusion_init(): assert np.allclose(expected_enh_smag_fac, np.asarray(diffusion.enh_smag_fac)) +@pytest.mark.xfail def test_diffusion_run(): data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") - physical_heights_name = "vct_a" - serializer = SerializedDataProvider("icon_diffusion", data_path) - serializer.print_info() - meta, fields = serializer.get_fields(metadata=["nlev", "linit", "dtime"], - fields=[physical_heights_name, "tangent_orientation"]) + data_provider = IconSerialDataProvider("icon_diffusion_init", data_path) + data_provider.print_info() + sp = data_provider.from_savepoint(linit=False, date="2021-06-20T12:00:10.000") + config = DiffusionConfig.create_with_defaults() additional_parameters = DiffusionParams(config) - vct_a = np_as_located_field(KDim)(fields[physical_heights_name]) - diffusion = Diffusion(config, additional_parameters, vct_a) + vct_a = sp.physical_height_field() + diffusion = Diffusion(config, additional_parameters, vct_a) diagnostic_state = DiagnosticState() prognostic_state = PrognosticState() interpolation_state = InterpolationState() metric_state = MetricState() - dtime = meta["dtime"] - orientation = fields["tangent_orientation"] - - # inverse_primal_edge_lengths: Field[[EdgeDim], float], - inverse_vertical_vertex_lengths = fields["inv_vert_vert_length"] - primal_normal_vert: CartesianVectorTuple = (fields["primal_normal_vert_x"], fields["primal_normal_vert_y"]) - dual_normal_vert: CartesianVectorTuple = (fields["dual_normal_vert_x"], fields["dual_normal_vert_y"]) - + dtime = sp.get_metadata("dtime") + orientation = sp.tangent_orientation() + inverse_primal_edge_lengths= sp.inverse_primal_edge_lengths() + inverse_vertical_vertex_lengths = sp.inv_vert_vert_length() + inverse_dual_edge_length = sp.inv_dual_edge_length() + primal_normal_vert: CartesianVectorTuple = (sp.primal_normal_vert_x(), sp.primal_normal_vert_y()) + dual_normal_vert: CartesianVectorTuple = (sp.dual_normal_vert_x(), sp.dual_normal_vert_y()) + edge_areas = sp.edge_areas() + cell_areas = sp.cell_areas() diffusion.run(diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, @@ -269,7 +267,14 @@ def test_diffusion_run(): dtime=dtime, tangent_orientation=orientation, inverse_primal_edge_lengths=inverse_primal_edge_lengths, + inv_dual_edge_length=inverse_dual_edge_length, inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, primal_normal_vert=primal_normal_vert, - dual_normal_vert=dual_normal_vert + dual_normal_vert=dual_normal_vert, + edge_areas=edge_areas, + cell_areas=cell_areas ) + + pytest.fail("not implemented yet") + + diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py new file mode 100644 index 000000000..d82890659 --- /dev/null +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -0,0 +1,72 @@ +import os + +import numpy as np +from _pytest.fixtures import fixture + +from icon4py.atm_dyn_iconam.horizontal import HorizontalMarkerIndex, HorizontalMeshConfig, HorizontalMeshParams +from icon4py.atm_dyn_iconam.icon_grid import IconGrid, MeshConfig +from icon4py.common.dimension import VertexDim, EdgeDim, CellDim +from icon4py.testutils.serialbox_utils import IconSerialDataProvider + + +@fixture +def with_grid(): + data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") + sp = IconSerialDataProvider("icon_diffusion_init", data_path).from_savepoint(linit=True, + date="2021-06-20T12:00:10.000") + nproma, nlev, num_v, num_c, num_e = sp.get_metadata("nproma", "nlev", "num_vert", "num_cells", + "num_edges") + cell_starts = sp.cells_start_index() + cell_ends = sp.cells_end_index() + vertex_starts = sp.vertex_start_index() + vertex_ends = sp.vertex_end_index() + edge_starts = sp.edge_start_index() + edge_ends = sp.edge_end_index() + + config = MeshConfig(HorizontalMeshConfig(num_vertices=num_v, num_cells=num_c, num_edges=num_e)) + c2e2c = np.squeeze(sp.c2e2c(), axis=1) + c2e2c0 = np.column_stack((c2e2c, (np.asarray(range(c2e2c.shape[0]))))) + grid = IconGrid().with_config(config) \ + .with_start_end_indices(VertexDim, vertex_starts, vertex_ends) \ + .with_start_end_indices(EdgeDim, edge_starts, edge_ends)\ + .with_start_end_indices(CellDim, cell_starts, cell_ends)\ + .with_connectivity(c2e=sp.c2e()).with_connectivity(e2c=sp.e2c())\ + .with_connectivity(c2e2c=c2e2c) \ + .with_connectivity(e2v=sp.e2v())\ + .with_connectivity(c2e2c0=c2e2c0)\ + + return grid + +def test_horizontal_grid_cell_indices(with_grid): + assert with_grid.get_indices_from_to(CellDim, 3, 3) == (20897, 20896) # halo +1 + #assert grid.get_indices_from_to(CellDim, 4, 4) == (1, 20896) #halo instead is (0,20896) why + assert with_grid.get_indices_from_to(CellDim, 8, 8) == (4105, 20896) # interior + assert with_grid.get_indices_from_to(CellDim, 9, 9) == (1, 850) #lb+1 + assert with_grid.get_indices_from_to(CellDim, 10, 10) == (851, 1688) + assert with_grid.get_indices_from_to(CellDim, 11, 11) == (1689, 2511) #lb+2 + assert with_grid.get_indices_from_to(CellDim, 12,12) ==(2512, 3316) #lb+3 + assert with_grid.get_indices_from_to(CellDim, HorizontalMarkerIndex.START_PROG_CELL.value , HorizontalMarkerIndex.START_PROG_CELL.value) == (3317, 4104) #nudging + + +def test_horizontal_edge_indices(with_grid): + assert with_grid.get_indices_from_to(EdgeDim, 0, 0) == (31559, 31558) + assert with_grid.get_indices_from_to(EdgeDim, 3, 3) == (31559, 31558) + assert with_grid.get_indices_from_to(EdgeDim, 4, 4) == (31559, 31558) + # assert with_grid.get_indices_from_to(EdgeDim, 5, 5) == (1, 31558) #halo + assert with_grid.get_indices_from_to(EdgeDim, 23, 23) == (5388, 6176) # nudging +1 + assert with_grid.get_indices_from_to(EdgeDim, 22,22) == (4990, 5387) #nudging + assert with_grid.get_indices_from_to(EdgeDim, 21, 21) == (4185, 4989) #lb +7 + assert with_grid.get_indices_from_to(EdgeDim, 20, 20) == (3778, 4184) # lb +6 + assert with_grid.get_indices_from_to(EdgeDim, 19, 19) == (2955, 3777) # lb +5 + assert with_grid.get_indices_from_to(EdgeDim, 18, 18) == (2539, 2954) # lb +4 + assert with_grid.get_indices_from_to(EdgeDim, 17, 17) == (1701, 2538) # lb +3 + assert with_grid.get_indices_from_to(EdgeDim, 16, 16) == (1279, 1700) # lb +2 + assert with_grid.get_indices_from_to(EdgeDim, 15, 15) == (429, 1278) # lb +1 + assert with_grid.get_indices_from_to(EdgeDim, 14, 14) == (1, 428) # lb +0 + + +def test_horizontal_vertex_indices(with_grid): + assert with_grid.get_indices_from_to(VertexDim, 0, 0) == (10664, 10664) + + + diff --git a/atm_dyn_iconam/tests/test_vertical.py b/atm_dyn_iconam/tests/test_vertical.py index 225028d09..ab4c8f292 100644 --- a/atm_dyn_iconam/tests/test_vertical.py +++ b/atm_dyn_iconam/tests/test_vertical.py @@ -16,7 +16,7 @@ import numpy as np import pytest -from icon4py.atm_dyn_iconam.icon_mesh import VerticalModelParams +from icon4py.atm_dyn_iconam.icon_grid import VerticalModelParams @pytest.mark.parametrize( diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 233c81463..d2a83f243 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -14,6 +14,11 @@ import os from typing import List +from functional.common import Dimension +from functional.iterator.embedded import np_as_located_field +from serialbox import Savepoint + +from icon4py.common.dimension import KDim, EdgeDim, ECVDim, CellDim, VertexDim try: import serialbox as ser @@ -26,7 +31,94 @@ import serialbox as ser -class SerializedDataProvider: + + +class IconDiffusionSavepoint: + def __init__(self, sp: Savepoint, ser: ser.Serializer): + self.savepoint = sp + self.serializer = ser + + def print_meta_info(self): + print(self.savepoint.metainfo) + + def physical_height_field(self): + return self._get_field("vct_a", KDim) + + def _get_field(self, name, *dimensions): + return np_as_located_field(dimensions)(self.serializer.read(name, self.savepoint)) + + def get_metadata(self, *names): + metadata = self.savepoint.metainfo.to_dict() + return {n:metadata[n] for n in names if n in metadata} + + def tangent_orientation(self): + return self._get_field("tangent_orientation", EdgeDim) + + def inverse_primal_edge_lengths(self): + return self._get_field("inverse_primal_edge_lengths", EdgeDim) + + def inv_vert_vert_length(self): + return self._get_field("inv_vert_vert_length", EdgeDim) + + def primal_normal_vert_x(self): + return self._get_field("primal_normal_vert_x", ECVDim) + + def primal_normal_vert_y(self): + return self._get_field("primal_normal_vert_y", ECVDim) + + def dual_normal_vert_y(self): + return self._get_field("dual_normal_vert_y", ECVDim) + + def dual_normal_vert_x(self): + return self._get_field("dual_normal_vert_x", ECVDim) + + def cell_areas(self): + return self._get_field("cell_areas", CellDim) + + def edge_areas(self): + return self._get_field("cell_areas", EdgeDim) + + def inv_dual_edge_length(self): + return self._get_field("inv_dual_edge_length", EdgeDim) + + def cells_start_index(self): + return self.serializer.read("cells_start_index", self.savepoint) + + def cells_end_index(self): + return self.serializer.read("c_end_index", self.savepoint) + + def vertex_start_index(self): + return self.serializer.read("v_start_index", self.savepoint) + + def vertex_end_index(self): + return self.serializer.read("v_end_index", self.savepoint) + + def edge_start_index(self): + return self.serializer.read("e_start_index", self.savepoint) + + def edge_end_index(self): + return self.serializer.read("e_end_index", self.savepoint) + + def refin_ctrl(self, dim: Dimension): + if dim == CellDim: + return self.serializer.read("c_refin_ctl", self.savepoint) + elif dim == EdgeDim: + return self.serializer.read("e_refin_ctl", self.savepoint) + elif dim == VertexDim: + return self.serializer.read("v_refin_ctl", self.savepoint) + else: + return None + def c2e(self): + return self.serializer.read("c2e", self.savepoint) + + def c2e2c(self): + return self.serializer.read("c2e2c", self.savepoint) + + def e2c(self): + return self.serializer.read("e2c", self.savepoint) + def e2v(self): + return self.serializer.read("e2v", self.savepoint) +class IconSerialDataProvider: def __init__(self, fname_prefix, path="."): self.rank = 0 @@ -44,6 +136,17 @@ def print_info(self): print(f"SAVEPOINTS: {self.serializer.savepoint_list()}") print(f"FIELDNAMES: {self.serializer.fieldnames()}") + def from_savepoint(self, linit:bool, date:str) -> IconDiffusionSavepoint: + savepoint = ( + self.serializer.savepoint["call-diffusion-init"] + .linit[linit] + .date[date] + .as_savepoint() + ) + return IconDiffusionSavepoint(savepoint, self.serializer) + + + def get_fields(self, metadata: List[str], fields: List[str]): savepoint = ( self.serializer.savepoint["call-diffusion-init"] @@ -71,9 +174,7 @@ def get_fields(self, metadata: List[str], fields: List[str]): [print(f"metadata {f} not present in savepoint") for f in meta_absent] return meta_present, fields_present -def read_from_call_diffusion_init_ser_data( - path, fname_prefix,metadata: List[str], fields: List[str] -): - serializer = SerializedDataProvider(fname_prefix, path) - serializer.print_info() - return serializer.get_fields(metadata, fields) + + + + From bfb319e92be631f9bdf85dfd59da1abd6eb2ffc7 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 6 Dec 2022 10:49:16 +0100 Subject: [PATCH 020/263] -icon grid: calcuate start / end indices - add domains to diffusion.run --- .../src/icon4py/atm_dyn_iconam/constants.py | 24 +- .../src/icon4py/atm_dyn_iconam/diagnostic.py | 29 +- .../src/icon4py/atm_dyn_iconam/diffusion.py | 431 ++++++++++++------ .../fused_mo_nh_diffusion_stencil_02_03.py | 22 +- .../fused_mo_nh_diffusion_stencil_04_05_06.py | 60 ++- ...sed_mo_nh_diffusion_stencil_07_08_09_10.py | 62 ++- .../fused_mo_nh_diffusion_stencil_11_12.py | 31 +- .../fused_mo_nh_diffusion_stencil_13_14.py | 16 +- .../src/icon4py/atm_dyn_iconam/horizontal.py | 121 +++-- .../src/icon4py/atm_dyn_iconam/icon_grid.py | 49 +- .../atm_dyn_iconam/interpolation_state.py | 44 +- .../icon4py/atm_dyn_iconam/metric_state.py | 51 ++- atm_dyn_iconam/tests/test_diffusion.py | 55 ++- atm_dyn_iconam/tests/test_icon_grid.py | 289 ++++++++++-- .../src/icon4py/testutils/serialbox_utils.py | 56 +-- 15 files changed, 984 insertions(+), 356 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/constants.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/constants.py index ced31729a..15cc0e379 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/constants.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/constants.py @@ -1,6 +1,24 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + from typing import Annotated -GAS_CONSTANT_DRY_AIR: Annotated[float, "gas constant for dry air [J/K/kg], called 'rd' in ICON (mo_physical_constants.f90), https://glossary.ametsoc.org/wiki/Gas_constant"] = 287.04 -CPD: Annotated[float, "specific heat at constant pressure [J/K/kg]"] = 1004.64 -GAS_CONSTANT_WATER_VAPOR: Annotated[float, "gas constant for water vapor [J/K/kg], rv in Icon"] = 461.51 +GAS_CONSTANT_DRY_AIR: Annotated[ + float, + "gas constant for dry air [J/K/kg], called 'rd' in ICON (mo_physical_constants.f90), https://glossary.ametsoc.org/wiki/Gas_constant", +] = 287.04 +CPD: Annotated[float, "specific heat at constant pressure [J/K/kg]"] = 1004.64 +GAS_CONSTANT_WATER_VAPOR: Annotated[ + float, "gas constant for water vapor [J/K/kg], rv in Icon" +] = 461.51 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diagnostic.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diagnostic.py index e0634d202..78ed3e9cf 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diagnostic.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diagnostic.py @@ -1,14 +1,33 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + from dataclasses import dataclass from functional.common import Field -from icon4py.common.dimension import KDim, CellDim +from icon4py.common.dimension import CellDim, KDim @dataclass class DiagnosticState: # fields for 3D elements in turbdiff - hdef_ic: Field[[CellDim, KDim], float] #! divergence at half levels(nproma,nlevp1,nblks_c) [1/s] - div_ic = Field[[CellDim, KDim], float] #! horizontal wind field deformation (nproma,nlevp1,nblks_c) [1/s^2] - dwdx = None #zonal gradient of vertical wind speed (nproma,nlevp1,nblks_c) [1/s] - dwdy = None #meridional gradient of vertical wind speed (nproma,nlevp1,nblks_c) + hdef_ic: Field[ + [CellDim, KDim], float + ] # ! divergence at half levels(nproma,nlevp1,nblks_c) [1/s] + div_ic = Field[ + [CellDim, KDim], float + ] # ! horizontal wind field deformation (nproma,nlevp1,nblks_c) [1/s^2] + dwdx = ( + None # zonal gradient of vertical wind speed (nproma,nlevp1,nblks_c) [1/s] + ) + dwdy = None # meridional gradient of vertical wind speed (nproma,nlevp1,nblks_c) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index ace763af4..b8d1b8e68 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -16,37 +16,60 @@ from typing import Final import numpy as np -from functional.common import Dimension from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, broadcast, maximum, minimum from functional.iterator.embedded import np_as_located_field -from icon4py.atm_dyn_iconam.constants import GAS_CONSTANT_DRY_AIR, CPD +from icon4py.atm_dyn_iconam.constants import CPD, GAS_CONSTANT_DRY_AIR from icon4py.atm_dyn_iconam.diagnostic import DiagnosticState -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import \ - fused_mo_nh_diffusion_stencil_02_03 -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import \ - fused_mo_nh_diffusion_stencil_04_05_06 -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_07_08_09_10 import \ - fused_mo_nh_diffusion_stencil_07_08_09_10 -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_11_12 import \ - fused_mo_nh_diffusion_stencil_11_12 -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_13_14 import \ - fused_mo_nh_diffusion_stencil_13_14 -from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig -from icon4py.atm_dyn_iconam.icon_grid import MeshConfig, VerticalModelParams, IconGrid +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import ( + fused_mo_nh_diffusion_stencil_02_03, +) +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import ( + fused_mo_nh_diffusion_stencil_04_05_06, +) +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_07_08_09_10 import ( + fused_mo_nh_diffusion_stencil_07_08_09_10, +) +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_11_12 import ( + fused_mo_nh_diffusion_stencil_11_12, +) +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_13_14 import ( + fused_mo_nh_diffusion_stencil_13_14, +) +from icon4py.atm_dyn_iconam.horizontal import ( + HorizontalMarkerIndex, + HorizontalMeshConfig, +) +from icon4py.atm_dyn_iconam.icon_grid import ( + IconGrid, + MeshConfig, + VerticalModelParams, +) from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState from icon4py.atm_dyn_iconam.metric_state import MetricState from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import ( - _mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_01, + _mo_nh_diffusion_stencil_01, + mo_nh_diffusion_stencil_01, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_16 import ( + mo_nh_diffusion_stencil_16, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_16 import mo_nh_diffusion_stencil_16 from icon4py.atm_dyn_iconam.prognostic import PrognosticState -from icon4py.common.dimension import ECVDim, EdgeDim, KDim, Koff, VertexDim, CellDim, \ - C2E2CDim, C2E2CODim +from icon4py.common.dimension import ( + C2E2CDim, + C2E2CODim, + CellDim, + ECVDim, + EdgeDim, + KDim, + Koff, + VertexDim, +) + DiffusionTupleVT = namedtuple("DiffusionParamVT", "v t") CartesianVectorTuple = namedtuple("CartesianVectorTuple", "x y") @@ -71,15 +94,17 @@ def _setup_smag_limit(diff_multfac_vn: Field[[KDim], float]) -> Field[[KDim], fl @field_operator -def _scale_k(field: Field[[KDim], float], factor:float) -> Field[[KDim], float]: +def _scale_k(field: Field[[KDim], float], factor: float) -> Field[[KDim], float]: return field * factor + @program -def scale_k(field: Field[[KDim], float], factor:float, scaled_field: Field[[KDim],float]): +def scale_k( + field: Field[[KDim], float], factor: float, scaled_field: Field[[KDim], float] +): _scale_k(field, factor, out=scaled_field) - # TODO [ml] rename! @field_operator def _mo_nh_diffusion_stencil_01_scale_dtime( @@ -276,7 +301,7 @@ def create_with_defaults(cls): grid = IconGrid().with_config(MeshConfig(horizontalMesh=horizontal)) verticalParams = VerticalModelParams( - rayleigh_damping_height=12500, vct_a=np.zeros(grid.get_num_k_levels()) + rayleigh_damping_height=12500, vct_a=np.zeros(grid.k_levels()) ) return cls(grid=grid, vertical_params=verticalParams) @@ -290,8 +315,10 @@ def __init__(self, grid: IconGrid, vertical_params: VerticalModelParams): self.lhdiff_temp = True # ! diffusion on the temperature field self.lhdiff_w = True # ! diffusion on the vertical wind field self.lhdiff_rcf = True # namelist, remove if always true - self.itype_vn_diffu = 1 # ! reconstruction method used for Smagorinsky diffusion - self.l_smag_d = False #namelist lsmag_d, if `true`, compute 3D Smagorinsky diffusion coefficient. + self.itype_vn_diffu = ( + 1 # ! reconstruction method used for Smagorinsky diffusion + ) + self.l_smag_d = False # namelist lsmag_d, if `true`, compute 3D Smagorinsky diffusion coefficient. self.itype_t_diffu = 2 # ! discretization of temperature diffusion self.hdiff_efdt_ratio = 24.0 # ! ratio of e-folding time to time step @@ -309,21 +336,22 @@ def __init__(self, grid: IconGrid, vertical_params: VerticalModelParams): # namelist grid_nml -> TODO [ml] should go to grid config? self.l_limited_area = True - #name list: interpol_nml + # name list: interpol_nml self.nudge_max_coeff = 0.075 def substep_as_float(self): return float(self.ndyn_substeps) - - class DiffusionParams: """Calculates derived quantities depending on the diffusion config.""" def __init__(self, config: DiffusionConfig): # TODO [ml] logging for case KX == 0 # TODO [ml] generrally calculation for x_dom (jg) > 2..n_dom, why is jg special + self.boundary_diffusion_start_index_edges = ( + 5 # mo_nh_diffusion.start_bdydiff_e - 1 = 5 -1 + ) self.K2: Final[float] = ( 1.0 / (config.hdiff_efdt_ratio * 8.0) @@ -377,15 +405,17 @@ def diffusion_type_5_smagorinski_factor(config: DiffusionConfig): return factor, heights - -def mo_nh_diffusion_stencil_15_numpy(mask_hdiff: Field[[CellDim, KDim], int], - zd_vertidx: Field[[CellDim, C2E2CDim, KDim], int], - zd_diffcoef: Field[[CellDim, KDim], float], - geofac_n2s: Field[[CellDim, C2E2CODim], float], - vcoef: Field[[C2E2CDim, KDim], float], - theta_v: Field[[CellDim, KDim], float], - z_temp: Field[[CellDim, KDim], float] - ): +def mo_nh_diffusion_stencil_15_numpy( + mask_hdiff: Field[[CellDim, KDim], int], + zd_vertidx: Field[[CellDim, C2E2CDim, KDim], int], + zd_diffcoef: Field[[CellDim, KDim], float], + geofac_n2s: Field[[CellDim, C2E2CODim], float], + vcoef: Field[[C2E2CDim, KDim], float], + theta_v: Field[[CellDim, KDim], float], + z_temp: Field[[CellDim, KDim], float], + domain, + offset_provider, +): pass @@ -404,16 +434,20 @@ def __init__( TODO [ml]: initial run: linit = .TRUE.: smag_offset and diff_multfac_vn are defined differently. """ - self.params = params self.config = config self.grid = config.grid self.rd_o_cvd: float = GAS_CONSTANT_DRY_AIR / (CPD - GAS_CONSTANT_DRY_AIR) self.nudgezone_diff = 0.04 / (config.nudge_max_coeff + sys.float_info.epsilon) - self.bdy_diff = 0.015/(config.nudge_max_coeff + sys.float_info.epsilon) - self.fac_bdydiff_v = math.sqrt(config.substep_as_float())/config.lateral_boundary_denominator.v if config.lhdiff_rcf else 1. /config.lateral_boundary_denominator.v - self.thresh_tdiff = - 5.0 # threshold temperature deviation from neighboring grid points hat activates extra diffusion against runaway cooling - + self.bdy_diff = 0.015 / (config.nudge_max_coeff + sys.float_info.epsilon) + self.fac_bdydiff_v = ( + math.sqrt(config.substep_as_float()) / config.lateral_boundary_denominator.v + if config.lhdiff_rcf + else 1.0 / config.lateral_boundary_denominator.v + ) + self.thresh_tdiff = ( + -5.0 + ) # threshold temperature deviation from neighboring grid points hat activates extra diffusion against runaway cooling # different for init call: smag_offset = 0 self.smag_offset: float = 0.25 * params.K4 * config.substep_as_float() @@ -423,11 +457,9 @@ def __init__( # different for initial run!, through diff_multfac_vn self.diff_multfac_vn = np_as_located_field(KDim)( - np.zeros(config.grid.get_num_k_levels()) - ) - self.smag_limit = np_as_located_field(KDim)( - np.zeros(config.grid.get_num_k_levels()) + np.zeros(config.grid.k_levels()) ) + self.smag_limit = np_as_located_field(KDim)(np.zeros(config.grid.k_levels())) init_diffusion_local_fields( params.K4, @@ -438,7 +470,7 @@ def __init__( ) self.enh_smag_fac = np_as_located_field(KDim)( - np.zeros(config.grid.get_num_k_levels()) + np.zeros(config.grid.k_levels(), float) ) enhanced_smagorinski_factor( *params.smagorinski_factor, @@ -450,22 +482,28 @@ def __init__( self.diff_multfac_n2w = ( config.vertical_params.init_nabla2_factor_in_upper_damping_zone( - k_size=config.grid.get_num_k_levels() + k_size=config.grid.k_levels() ) ) self.diff_multfac_smag = np_as_located_field(KDim)( - np.zeros(config.grid.get_num_k_levels()) + np.zeros(config.grid.k_levels()) ) - shape_vk = (config.grid.get_num_vertices(), config.grid.get_num_k_levels()) - shape_ck = (config.grid.get_num_cells(), config.grid.get_num_k_levels()) + shape_vk = (config.grid.num_vertices(), config.grid.k_levels()) + shape_ck = (config.grid.num_cells(), config.grid.k_levels()) self.u_vert = np_as_located_field(VertexDim, KDim)(np.zeros(shape_vk, float)) self.v_vert = np_as_located_field(VertexDim, KDim)(np.zeros(shape_vk, float)) - shape_ek = (config.grid.get_num_edges(), config.grid.get_num_k_levels()) + shape_ek = (config.grid.num_edges(), config.grid.k_levels()) allocate_ek = np_as_located_field(EdgeDim, KDim)(np.zeros(shape_ek, float)) self.kh_smag_e = allocate_ek self.kh_smag_ec = allocate_ek self.z_nabla2_e = allocate_ek self.z_temp = np_as_located_field(CellDim, KDim)(np.zeros(shape_ck, float)) + self.vertical_index = np_as_located_field(KDim)( + np.arange(self.grid.k_levels() + 1) + ) + self.horizontal_cell_index = np_as_located_field(CellDim)( + np.arange((shape_ck[0])) + ) def run( self, @@ -473,7 +511,7 @@ def run( prognostic_state: PrognosticState, metric_state: MetricState, interpolation_state: InterpolationState, - dtime:float, + dtime: float, tangent_orientation: Field[[EdgeDim], float], inverse_primal_edge_lengths: Field[[EdgeDim], float], inv_dual_edge_length: Field[[EdgeDim], float], @@ -485,7 +523,7 @@ def run( Field[[ECVDim], float], Field[[ECVDim], float] ], edge_areas: Field[[EdgeDim], float], - cell_areas: Field[[CellDim], float] + cell_areas: Field[[CellDim], float], ): """ Run a diffusion step. @@ -501,8 +539,12 @@ def run( # Oa logging # 0b call timer start # 0c. dtime dependent stuff: enh_smag_factor, r_dtimensubsteps - r_dtimensubsteps = 1.0 / dtime if self.config.lhdiff_rcf else 1./ (dtime * self.config.substep_as_float) - klevels = self.config.grid.k_levels() + r_dtimensubsteps = ( + 1.0 / dtime + if self.config.lhdiff_rcf + else 1.0 / (dtime * self.config.substep_as_float()) + ) + klevels = self.grid.k_levels() # TODO is this needed? set_zero_v_k(self.u_vert) set_zero_v_k(self.v_vert) @@ -515,6 +557,15 @@ def run( interpolation_state.rbf_coeff_2, self.u_vert, self.v_vert, + domain={ + KDim: (0, klevels), + VertexDim: self.grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.local_boundary(VertexDim) + 1, + HorizontalMarkerIndex.halo(VertexDim) - 1, + ), + }, + offset_provider={"V2E": self.grid.get_v2e_offset_provider()}, ) # 2. HALO EXCHANGE -- CALL sync_patch_array_mult # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 @@ -538,108 +589,212 @@ def run( self.kh_smag_ec, self.z_nabla2_e, self.smag_offset, - dtime, domain={KDim:(0, klevels), EdgeDim:(), VertexDim:()} + dtime, + domain={ + KDim: (0, klevels), + EdgeDim: ( + self.grid.get_indices_from_to( + EdgeDim, + self.params.boundary_diffusion_start_index_edges, + HorizontalMarkerIndex.halo(EdgeDim) - 2, + ) + ), + }, + offset_provider={"E2C2V": self.grid.get_e2c2v_connectivity()}, ) - fused_mo_nh_diffusion_stencil_02_03(self.kh_smag_ec, - prognostic_state.normal_wind, - interpolation_state.e_bln_c_s, #:Field[[CellDim, C2EDim], float], - interpolation_state.geofac_div, # Field[[CellDim, C2EDim], float], - self.diff_multfac_smag, # Field[[KDim], float], - metric_state.wgtfac_c, # Field[[CellDim, KDim], float], - diagnostic_state.div_ic, #Field[[CellDim, KDim], float], - diagnostic_state.hdef_ic, #: Field[[CellDim, KDim], float]) - offset_provider={"C2E": self.grid.get_c2e_connectivity()}) + fused_mo_nh_diffusion_stencil_02_03( + self.kh_smag_ec, + prognostic_state.normal_wind, + interpolation_state.e_bln_c_s, + interpolation_state.geofac_div, + self.diff_multfac_smag, + metric_state.wgtfac_c, + diagnostic_state.div_ic, + diagnostic_state.hdef_ic, + domain={ + KDim: (0, klevels), + CellDim: ( + self.grid.get_indices_from_to( + HorizontalMarkerIndex._NUDGING_CELLS, + HorizontalMarkerIndex._HALO_CELLS, + ) + ), + }, + offset_provider={"C2E": self.grid.get_c2e_connectivity()}, + ) # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH # 5. CALL rbf_vec_interpol_vertex_wp - mo_intp_rbf_rbf_vec_interpol_vertex(self.z_nabla2_e, + mo_intp_rbf_rbf_vec_interpol_vertex( + self.z_nabla2_e, interpolation_state.rbf_coeff_1, interpolation_state.rbf_coeff_2, self.u_vert, - self.v_vert) + self.v_vert, + domain={ + KDim: (0, klevels), + VertexDim: self.grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.local_boundary(VertexDim) + 3, + HorizontalMarkerIndex.halo(VertexDim), + ), + }, + offset_provider={"V2E": self.grid.get_e2v_connectivity()}, + ) # 6. HALO EXCHANGE -- CALL sync_patch_array_mult # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 - #TODO where to get index field from - #horz_idx: Field[[EdgeDim], int32], - #start_2nd_nudge_line_idx_e: int32, - - fused_mo_nh_diffusion_stencil_04_05_06(u_vert=self.u_vert, - v_vert=self.v_vert, - primal_normal_vert_v1=primal_normal_vert.x, - primal_normal_vert_v2=primal_normal_vert.y, - z_nabla2_e=self.z_nabla2_e, - inv_vert_vert_length = inverse_vertical_vertex_lengths, - inv_primal_edge_length = inverse_primal_edge_lengths, - area_edges=edge_areas, - kh_smag_e=self.kh_smag_e, - diff_multfac_vn=self.diff_multfac_vn, - nudgecoeff_e=interpolation_state.nudgecoeff_e, - vn=prognostic_state.normal_wind, - horz_idx = None, - nudgezone_diff=self.nudgezone_diff, - fac_bdydiff_v= self.fac_bdydiff_v, - start_2nd_nudge_line_idx_e = None, - offset_provider={} - ) + # TODO where to get index field from + # horz_idx: Field[[EdgeDim], int32], + # start_2nd_nudge_line_idx_e: int32, + + fused_mo_nh_diffusion_stencil_04_05_06( + u_vert=self.u_vert, + v_vert=self.v_vert, + primal_normal_vert_v1=primal_normal_vert.x, + primal_normal_vert_v2=primal_normal_vert.y, + z_nabla2_e=self.z_nabla2_e, + inv_vert_vert_length=inverse_vertical_vertex_lengths, + inv_primal_edge_length=inverse_primal_edge_lengths, + area_edges=edge_areas, + kh_smag_e=self.kh_smag_e, + diff_multfac_vn=self.diff_multfac_vn, + nudgecoeff_e=interpolation_state.nudgecoeff_e, + vn=prognostic_state.normal_wind, + horz_idx=None, + nudgezone_diff=self.nudgezone_diff, + fac_bdydiff_v=self.fac_bdydiff_v, + start_2nd_nudge_line_idx_e=None, + domain={ + KDim: (0, klevels), + EdgeDim: self.grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.nudging(EdgeDim) + 1, + HorizontalMarkerIndex.halo(EdgeDim), + ), + }, + offset_provider={"E2C2V": self.grid.get_e2c2v_connectivity()}, + ) # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 - fused_mo_nh_diffusion_stencil_07_08_09_10(area = cell_areas, geofac_n2s=interpolation_state.geofac_n2s, - geofac_grg_x=interpolation_state.geofac_grg_x, - geofac_grg_y=interpolation_state.geofac_grg_y, - w_old=prognostic_state.vertical_wind, - w=prognostic_state.vertical_wind, - dwdx=diagnostic_state.dwdx, - dwdy=diagnostic_state.dwdy, - diff_multfac_w=self.diff_multfac_w, - diff_multfac_n2w=self.diff_multfac_n2w, - vert_idx= None, - horz_idx=None, - nrdmax=self.config.vertical_params.index_of_damping_height, - interior_idx=None, - halo_idx=None, - offset_provider={"C2E2CO": self.grid.get_c2e2c0_connectivity()} - ) + fused_mo_nh_diffusion_stencil_07_08_09_10( + area=cell_areas, + geofac_n2s=interpolation_state.geofac_n2s, + geofac_grg_x=interpolation_state.geofac_grg_x, + geofac_grg_y=interpolation_state.geofac_grg_y, + w_old=prognostic_state.vertical_wind, + w=prognostic_state.vertical_wind, + dwdx=diagnostic_state.dwdx, + dwdy=diagnostic_state.dwdy, + diff_multfac_w=self.diff_multfac_w, + diff_multfac_n2w=self.diff_multfac_n2w, + vert_idx=self.vertical_index, + horz_idx=self.horizontal_cell_index, + nrdmax=self.config.vertical_params.index_of_damping_height, + interior_idx=HorizontalMarkerIndex.interior(CellDim), + halo_idx=HorizontalMarkerIndex.halo(CellDim), + domain={ + KDim: (0, klevels), + CellDim: self.grid.get_indices_from_to( + CellDim, + # TODO: global mode is from NUDGING - 1 + HorizontalMarkerIndex.nudging(CellDim), + HorizontalMarkerIndex.halo(CellDim) + 1, + ), + }, + offset_provider={"C2E2CO": self.grid.get_c2e2c0_connectivity()}, + ) # 8. HALO EXCHANGE: CALL sync_patch_array # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 # TODO check: kh_smag_e is an out field, should not be calculated in init? - fused_mo_nh_diffusion_stencil_11_12(theta_v=prognostic_state.theta_v, - theta_ref_mc=metric_state.theta_ref_mc, - thresh_tdiff=self.thresh_tdiff, - kh_smag_e=self.kh_smag_e, - offset_provider={"E2C":self.grid.get_e2c_connectivity(), "C2E2C":self.grid.get_c2e2c_connectivity()} - ) - - - fused_mo_nh_diffusion_stencil_13_14(kh_smag_e=self.kh_smag_e, - inv_dual_edge_length=inv_dual_edge_length, - theta_v=prognostic_state.theta_v, - geofac_div=interpolation_state.geofac_div, - z_temp=self.z_temp, - offset_provider={"C2E":self.grid.get_c2e_connectivity(), "E2C":self.grid.get_e2c_connectivity()} - ) - - mo_nh_diffusion_stencil_15_numpy(mask_hdiff=metric_state.mask_hdiff, - zd_vertidx=metric_state.zd_vertidx_dsl, - vcoef=metric_state.zd_diffcoef, - zd_diffcoef=metric_state.zd_diffcoef, - geofac_n2s=interpolation_state.geofac_n2s, - theta_v=prognostic_state.theta_v, - z_temp=self.z_temp, - offset_provider={} - ) - - mo_nh_diffusion_stencil_16(z_temp = self.z_temp, - area=cell_areas, - theta_v=prognostic_state.theta_v, - exner=prognostic_state.exner_pressure, - rd_o_cvd=self.rd_o_cvd, - offset_provider={} - ) + fused_mo_nh_diffusion_stencil_11_12( + theta_v=prognostic_state.theta_v, + theta_ref_mc=metric_state.theta_ref_mc, + thresh_tdiff=self.thresh_tdiff, + kh_smag_e=self.kh_smag_e, + domain={ + KDim: (klevels - 2, klevels), + CellDim: self.grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.nudging(CellDim) - 1, + HorizontalMarkerIndex.halo(CellDim) + 1, + ), + EdgeDim: self.grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.nudging(EdgeDim) + 1, + HorizontalMarkerIndex.halo(EdgeDim), + ), + }, + offset_provider={ + "E2C": self.grid.get_e2c_connectivity(), + "C2E2C": self.grid.get_c2e2c_connectivity(), + }, + ) + + fused_mo_nh_diffusion_stencil_13_14( + kh_smag_e=self.kh_smag_e, + inv_dual_edge_length=inv_dual_edge_length, + theta_v=prognostic_state.theta_v, + geofac_div=interpolation_state.geofac_div, + z_temp=self.z_temp, + domain={ + KDim: (0, klevels), + CellDim: self.grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.nudging(CellDim), + HorizontalMarkerIndex.halo(CellDim), + ), + EdgeDim: self.grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.nudging(EdgeDim), + HorizontalMarkerIndex.halo(EdgeDim) - 1, + ), + }, + offset_provider={ + "C2E": self.grid.get_c2e_connectivity(), + "E2C": self.grid.get_e2c_connectivity(), + }, + ) + + mo_nh_diffusion_stencil_15_numpy( + mask_hdiff=metric_state.mask_hdiff, + zd_vertidx=metric_state.zd_vertidx_dsl, + vcoef=metric_state.zd_diffcoef, + zd_diffcoef=metric_state.zd_diffcoef, + geofac_n2s=interpolation_state.geofac_n2s, + theta_v=prognostic_state.theta_v, + z_temp=self.z_temp, + domain={ + KDim: (0, klevels), + CellDim: self.grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.nudging(CellDim).value, + HorizontalMarkerIndex.halo(CellDim), + ), + }, + offset_provider={}, + ) + + mo_nh_diffusion_stencil_16( + z_temp=self.z_temp, + area=cell_areas, + theta_v=prognostic_state.theta_v, + exner=prognostic_state.exner_pressure, + rd_o_cvd=self.rd_o_cvd, + domain={ + KDim: (0, klevels), + CellDim: self.grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.nudging(CellDim), + HorizontalMarkerIndex.halo(CellDim), + ), + }, + offset_provider={}, + ) # 10. HALO EXCHANGE sync_patch_array diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py index 8c111b5d4..2e7aa11c3 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -14,10 +14,14 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, neighbor_sum +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import ( + _mo_nh_diffusion_stencil_02, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import ( + _mo_nh_diffusion_stencil_03, +) from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim, Koff -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import _mo_nh_diffusion_stencil_02 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import _mo_nh_diffusion_stencil_03 @field_operator def _fused_mo_nh_diffusion_stencil_02_03( @@ -28,7 +32,9 @@ def _fused_mo_nh_diffusion_stencil_02_03( diff_multfac_smag: Field[[KDim], float], wgtfac_c: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - kh_c, div = _mo_nh_diffusion_stencil_02(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag) + kh_c, div = _mo_nh_diffusion_stencil_02( + kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag + ) div_ic, hdef_ic = _mo_nh_diffusion_stencil_03(div, kh_c, wgtfac_c) return div_ic, hdef_ic @@ -44,4 +50,12 @@ def fused_mo_nh_diffusion_stencil_02_03( div_ic: Field[[CellDim, KDim], float], hdef_ic: Field[[CellDim, KDim], float], ): - _fused_mo_nh_diffusion_stencil_02_03(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag, wgtfac_c, out=(div_ic, hdef_ic)) + _fused_mo_nh_diffusion_stencil_02_03( + kh_smag_ec, + vn, + e_bln_c_s, + geofac_div, + diff_multfac_smag, + wgtfac_c, + out=(div_ic, hdef_ic), + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py index 62f4316d5..91d24c868 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py @@ -12,8 +12,17 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, neighbor_sum, maximum, where, int32 +from functional.ffront.fbuiltins import Field, int32, maximum, neighbor_sum, where +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_04 import ( + _mo_nh_diffusion_stencil_04, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_05 import ( + _mo_nh_diffusion_stencil_05, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_06 import ( + _mo_nh_diffusion_stencil_06, +) from icon4py.common.dimension import ( E2C2V, E2ECV, @@ -23,9 +32,6 @@ VertexDim, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_04 import _mo_nh_diffusion_stencil_04 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_05 import _mo_nh_diffusion_stencil_05 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_06 import _mo_nh_diffusion_stencil_06 @field_operator def _fused_mo_nh_diffusion_stencil_04_05_06( @@ -47,12 +53,29 @@ def _fused_mo_nh_diffusion_stencil_04_05_06( start_2nd_nudge_line_idx_e: int32, ) -> Field[[EdgeDim, KDim], float]: - z_nabla4_e2 = _mo_nh_diffusion_stencil_04(u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2, z_nabla2_e, inv_vert_vert_length, inv_primal_edge_length) + z_nabla4_e2 = _mo_nh_diffusion_stencil_04( + u_vert, + v_vert, + primal_normal_vert_v1, + primal_normal_vert_v2, + z_nabla2_e, + inv_vert_vert_length, + inv_primal_edge_length, + ) vn = where( horz_idx >= start_2nd_nudge_line_idx_e, - _mo_nh_diffusion_stencil_05(area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, nudgecoeff_e, vn, nudgezone_diff), - _mo_nh_diffusion_stencil_06(z_nabla2_e, area_edge, vn, fac_bdydiff_v) + _mo_nh_diffusion_stencil_05( + area_edge, + kh_smag_e, + z_nabla2_e, + z_nabla4_e2, + diff_multfac_vn, + nudgecoeff_e, + vn, + nudgezone_diff, + ), + _mo_nh_diffusion_stencil_06(z_nabla2_e, area_edge, vn, fac_bdydiff_v), ) return vn @@ -77,7 +100,22 @@ def fused_mo_nh_diffusion_stencil_04_05_06( fac_bdydiff_v: float, start_2nd_nudge_line_idx_e: int32, ): - _fused_mo_nh_diffusion_stencil_04_05_06(u_vert, v_vert, primal_normal_vert_v1, - primal_normal_vert_v2, z_nabla2_e, inv_vert_vert_length, - inv_primal_edge_length, area_edge, kh_smag_e, diff_multfac_vn, nudgecoeff_e, vn, horz_idx, nudgezone_diff, fac_bdydiff_v, - start_2nd_nudge_line_idx_e, out=vn) + _fused_mo_nh_diffusion_stencil_04_05_06( + u_vert, + v_vert, + primal_normal_vert_v1, + primal_normal_vert_v2, + z_nabla2_e, + inv_vert_vert_length, + inv_primal_edge_length, + area_edge, + kh_smag_e, + diff_multfac_vn, + nudgecoeff_e, + vn, + horz_idx, + nudgezone_diff, + fac_bdydiff_v, + start_2nd_nudge_line_idx_e, + out=vn, + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py index a1e305ae5..796d0ff58 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -12,14 +12,28 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, neighbor_sum, where, int32, broadcast +from functional.ffront.fbuiltins import ( + Field, + broadcast, + int32, + neighbor_sum, + where, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_07 import ( + _mo_nh_diffusion_stencil_07, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_08 import ( + _mo_nh_diffusion_stencil_08, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_09 import ( + _mo_nh_diffusion_stencil_09, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_10 import ( + _mo_nh_diffusion_stencil_10, +) from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_07 import _mo_nh_diffusion_stencil_07 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_08 import _mo_nh_diffusion_stencil_08 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_09 import _mo_nh_diffusion_stencil_09 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_10 import _mo_nh_diffusion_stencil_10 @field_operator def _fused_mo_nh_diffusion_stencil_07_08_09_10( @@ -38,26 +52,35 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( nrdmax: int32, interior_idx: int32, halo_idx: int32, -) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: +) -> tuple[ + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], +]: - vert_idx = broadcast(vert_idx, (CellDim, KDim)) + vert_idx = broadcast(vert_idx, (CellDim, KDim)) dwdx, dwdy = where( vert_idx > int32(0), _mo_nh_diffusion_stencil_08(w_old, geofac_grg_x, geofac_grg_y), - (dwdx, dwdy) + (dwdx, dwdy), ) z_nabla2_c = _mo_nh_diffusion_stencil_07(w_old, geofac_n2s) w = where( (horz_idx >= interior_idx) & (horz_idx < halo_idx), - _mo_nh_diffusion_stencil_09(area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w), + _mo_nh_diffusion_stencil_09( + area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w + ), w_old, ) w = where( - (vert_idx > int32(0)) & (vert_idx < nrdmax) & (horz_idx >= interior_idx) & (horz_idx < halo_idx), + (vert_idx > int32(0)) + & (vert_idx < nrdmax) + & (horz_idx >= interior_idx) + & (horz_idx < halo_idx), _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, area, z_nabla2_c), w, ) @@ -83,4 +106,21 @@ def fused_mo_nh_diffusion_stencil_07_08_09_10( interior_idx: int32, halo_idx: int32, ): - _fused_mo_nh_diffusion_stencil_07_08_09_10(area, geofac_n2s, geofac_grg_x, geofac_grg_y, w_old, w, dwdx, dwdy, diff_multfac_w, diff_multfac_n2w, vert_idx, horz_idx, nrdmax, interior_idx, halo_idx, out = (w, dwdx, dwdy)) + _fused_mo_nh_diffusion_stencil_07_08_09_10( + area, + geofac_n2s, + geofac_grg_x, + geofac_grg_y, + w_old, + w, + dwdx, + dwdy, + diff_multfac_w, + diff_multfac_n2w, + vert_idx, + horz_idx, + nrdmax, + interior_idx, + halo_idx, + out=(w, dwdx, dwdy), + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py index 02fd00530..6709daba9 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py @@ -12,12 +12,30 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, where, neighbor_sum, max_over, maximum +from functional.ffront.fbuiltins import ( + Field, + max_over, + maximum, + neighbor_sum, + where, +) -from icon4py.common.dimension import E2C, E2CDim, C2E2C, C2E2CDim, CellDim, EdgeDim, KDim +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_11 import ( + _mo_nh_diffusion_stencil_11, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_12 import ( + _mo_nh_diffusion_stencil_12, +) +from icon4py.common.dimension import ( + C2E2C, + E2C, + C2E2CDim, + CellDim, + E2CDim, + EdgeDim, + KDim, +) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_11 import _mo_nh_diffusion_stencil_11 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_12 import _mo_nh_diffusion_stencil_12 @field_operator def _fused_mo_nh_diffusion_stencil_11_12( @@ -30,6 +48,7 @@ def _fused_mo_nh_diffusion_stencil_11_12( kh_smag_e = _mo_nh_diffusion_stencil_12(kh_smag_e, enh_diffu_3d) return kh_smag_e + @program def fused_mo_nh_diffusion_stencil_11_12( theta_v: Field[[CellDim, KDim], float], @@ -37,4 +56,6 @@ def fused_mo_nh_diffusion_stencil_11_12( thresh_tdiff: float, kh_smag_e: Field[[EdgeDim, KDim], float], ): - _fused_mo_nh_diffusion_stencil_11_12(theta_v, theta_ref_mc, thresh_tdiff, kh_smag_e, out=kh_smag_e) + _fused_mo_nh_diffusion_stencil_11_12( + theta_v, theta_ref_mc, thresh_tdiff, kh_smag_e, out=kh_smag_e + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py index 40faa10b5..28d8dbfb6 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -15,19 +15,20 @@ from functional.ffront.fbuiltins import Field, neighbor_sum from icon4py.common.dimension import ( - E2C, - E2CDim, - C2E2C, - C2E2CDim, C2CE, C2E, + C2E2C, + E2C, + C2E2CDim, C2EDim, CEDim, CellDim, + E2CDim, EdgeDim, KDim, ) + @field_operator def _mo_nh_diffusion_stencil_13( kh_smag_e: Field[[EdgeDim, KDim], float], @@ -37,6 +38,7 @@ def _mo_nh_diffusion_stencil_13( z_nabla2_e = kh_smag_e * inv_dual_edge_length * (theta_v(E2C[1]) - theta_v(E2C[0])) return z_nabla2_e + @field_operator def _mo_nh_diffusion_stencil_14( z_nabla2_e: Field[[EdgeDim, KDim], float], @@ -45,6 +47,7 @@ def _mo_nh_diffusion_stencil_14( z_temp = neighbor_sum(z_nabla2_e(C2E) * geofac_div(C2CE), axis=C2EDim) return z_temp + @field_operator def _fused_mo_nh_diffusion_stencil_13_14( kh_smag_e: Field[[EdgeDim, KDim], float], @@ -56,6 +59,7 @@ def _fused_mo_nh_diffusion_stencil_13_14( z_temp = _mo_nh_diffusion_stencil_14(z_nabla2_e, geofac_div) return z_temp + @program def fused_mo_nh_diffusion_stencil_13_14( kh_smag_e: Field[[EdgeDim, KDim], float], @@ -64,4 +68,6 @@ def fused_mo_nh_diffusion_stencil_13_14( geofac_div: Field[[CEDim], float], z_temp: Field[[CellDim, KDim], float], ): - _fused_mo_nh_diffusion_stencil_13_14(kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp) + _fused_mo_nh_diffusion_stencil_13_14( + kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py index b1b3a4afb..f893a7cbb 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py @@ -10,30 +10,101 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -from dataclasses import dataclass -from enum import Enum from typing import Final -from functional.common import Dimension, DimensionKind - -@dataclass(frozen=True) -class HorizontalMeshParams: - NUM_GHOST_ROWS:Final[int] = 2 - GRF_BOUNDARY_WIDTH_CELL:Final[int] = 12 - MIN_RL_CELL_INT:Final[int] = 4 - MIN_RL_CELL:Final[int] = MIN_RL_CELL_INT - 2 * NUM_GHOST_ROWS - MAX_RL_CELL:Final[int] = 13 - MIN_RL_VERTEX_INT:Final[int] = MIN_RL_CELL_INT - MIN_RL_VERTEX:Final[int] = MIN_RL_VERTEX_INT - (NUM_GHOST_ROWS + 1) - MAX_RL_VERTEX:Final[int] = MAX_RL_CELL - MIN_RL_EDGE_INT:Final[int] = 2 * MIN_RL_CELL_INT - MIN_RL_EDGE:Final[int] = MIN_RL_EDGE_INT - (2 * NUM_GHOST_ROWS + 1) - MAX_RL_EDGE:Final[int] = 2 * MAX_RL_CELL - GRF_BOUNDARY_WIDTH_EDGES: Final[int] = 9 - -class HorizontalMarkerIndex(Enum): - START_PROG_CELL = HorizontalMeshParams.GRF_BOUNDARY_WIDTH_CELL + 1 - END_PROG_CELL = HorizontalMeshParams.MIN_RL_CELL_INT +from functional.common import Dimension + +from icon4py.common import dimension + + +class HorizontalMarkerIndex: + NUM_GHOST_ROWS: Final[int] = 2 + # values from mo_impl_constants.f90 + _ICON_INDEX_OFFSET_CELLS: Final[int] = 8 + _GRF_BOUNDARY_WIDTH_CELL: Final[int] = 4 + _MIN_RL_CELL_INT: Final[int] = -4 + _MIN_RL_CELL: Final[int] = _MIN_RL_CELL_INT - 2 * NUM_GHOST_ROWS + _MAX_RL_CELL: Final[int] = 5 + + _ICON_INDEX_OFFSET_VERTEX: Final[int] = 7 + _MIN_RL_VERTEX_INT: Final[int] = _MIN_RL_CELL_INT + _MIN_RL_VERTEX: Final[int] = _MIN_RL_VERTEX_INT - (NUM_GHOST_ROWS + 1) + _MAX_RL_VERTEX: Final[int] = _MAX_RL_CELL + + _ICON_INDEX_OFFSET_EDGES: Final[int] = 13 + _GRF_BOUNDARY_WIDTH_EDGES: Final[int] = 9 + _MIN_RL_EDGE_INT: Final[int] = 2 * _MIN_RL_CELL_INT + _MIN_RL_EDGE: Final[int] = _MIN_RL_EDGE_INT - (2 * NUM_GHOST_ROWS + 1) + _MAX_RL_EDGE: Final[int] = 2 * _MAX_RL_CELL + + _LOCAL_BOUNDARY_EDGES: Final[int] = 1 + _ICON_INDEX_OFFSET_EDGES + _INTERIOR_EDGES: Final[int] = _ICON_INDEX_OFFSET_EDGES + _NUDGING_EDGES: Final[int] = _GRF_BOUNDARY_WIDTH_EDGES + _ICON_INDEX_OFFSET_EDGES + _HALO_EDGES: Final[int] = _MIN_RL_EDGE_INT + _ICON_INDEX_OFFSET_EDGES + _END_EDGES: Final[int] = 0 + + _LOCAL_BOUNDARY_CELLS: Final[int] = 1 + _ICON_INDEX_OFFSET_CELLS + _INTERIOR_CELLS: Final[int] = _ICON_INDEX_OFFSET_CELLS + _NUDGING_CELLS: Final[int] = _GRF_BOUNDARY_WIDTH_CELL + 1 + _ICON_INDEX_OFFSET_CELLS + _HALO_CELLS: Final[int] = _MIN_RL_CELL_INT + _ICON_INDEX_OFFSET_CELLS + _END_CELLS: Final[int] = 0 + + _LOCAL_BOUNDARY_VERTICES = 1 + _ICON_INDEX_OFFSET_VERTEX + _INTERIOR_VERTICES: Final[int] = _ICON_INDEX_OFFSET_VERTEX + _NUDGING_VERTICES: Final[int] = 0 + _HALO_VERTICES: Final[int] = _MIN_RL_VERTEX_INT + _ICON_INDEX_OFFSET_VERTEX + _END_VERTICES: Final[int] = 0 + + @classmethod + def local_boundary(cls, dim: Dimension) -> int: + match (dim): + case (dimension.CellDim): + return cls._LOCAL_BOUNDARY_CELLS + case (dimension.EdgeDim): + return cls._LOCAL_BOUNDARY_EDGES + case (dimension.VertexDim): + return cls._LOCAL_BOUNDARY_VERTICES + + @classmethod + def halo(cls, dim: Dimension) -> int: + match (dim): + case (dimension.CellDim): + return cls._HALO_CELLS + case (dimension.EdgeDim): + return cls._HALO_EDGES + case (dimension.VertexDim): + return cls._HALO_VERTICES + + @classmethod + def nudging(cls, dim: Dimension) -> int: + match (dim): + case (dimension.CellDim): + return cls._NUDGING_CELLS + case (dimension.EdgeDim): + return cls._NUDGING_EDGES + case (dimension.VertexDim): + return cls._NUDGING_VERTICES + + @classmethod + def interior(cls, dim: Dimension) -> int: + match (dim): + case (dimension.CellDim): + return cls._INTERIOR_CELLS + case (dimension.EdgeDim): + return cls._INTERIOR_EDGES + case (dimension.VertexDim): + return cls._INTERIOR_VERTICES + + @classmethod + def end(cls, dim: Dimension) -> int: + match (dim): + case (dimension.CellDim): + return cls._END_CELLS + case (dimension.EdgeDim): + return cls._END_EDGES + case (dimension.VertexDim): + return cls._END_VERTICES + class HorizontalMeshConfig: def __init__(self, num_vertices: int, num_edges: int, num_cells: int): @@ -49,9 +120,3 @@ def get_num_edges(self): def get_num_cells(self): return self._num_cells - - - - - - diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py index 72ef0cc41..a27dd5cc1 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py @@ -10,15 +10,17 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -from enum import Enum from typing import Tuple import numpy as np -from functional.common import Field, Dimension, DimensionKind -from functional.iterator.embedded import np_as_located_field, NeighborTableOffsetProvider +from functional.common import Dimension, DimensionKind, Field +from functional.iterator.embedded import ( + NeighborTableOffsetProvider, + np_as_located_field, +) -from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig, HorizontalMarkerIndex -from icon4py.common.dimension import EdgeDim, KDim, CellDim, VertexDim +from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig +from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim class MeshConfig: @@ -45,25 +47,30 @@ def get_num_edges(self): def get_num_cells(self): return self._horizontal._num_cells + def builder(func): def wrapper(self, *args, **kwargs): func(self, *args, **kwargs) return self + return wrapper + class IconGrid: def __init__(self): - self.config:MeshConfig = None + self.config: MeshConfig = None self.start_indices = {} self.end_indices = {} - self.connectivities={} + self.connectivities = {} @builder def with_config(self, config: MeshConfig): self.config = config @builder - def with_start_end_indices(self, dim:Dimension, start_indices:np.ndarray, end_indices: np.ndarray): + def with_start_end_indices( + self, dim: Dimension, start_indices: np.ndarray, end_indices: np.ndarray + ): self.start_indices[dim] = start_indices self.end_indices[dim] = end_indices @@ -74,13 +81,22 @@ def with_connectivity(self, **connectivity): def k_levels(self): return self.config.num_k_levels + def num_cells(self): + return self.config.get_num_cells() + def num_vertices(self): + return self.config.get_num_cells() - def get_indices_from_to(self, dim:Dimension, start_marker:int, end_marker: int) -> Tuple[int, int]: + def num_edges(self): + return self.config.get_num_edges() + + def get_indices_from_to( + self, dim: Dimension, start_marker: int, end_marker: int + ) -> Tuple[int, int]: if dim.kind != DimensionKind.HORIZONTAL: - raise ValueError("only defined for {} dimension kind ", DimensionKind.HORIZONTAL) - #print(self.start_indices[dim]) - #print(self.end_indices[dim]) + raise ValueError( + "only defined for {} dimension kind ", DimensionKind.HORIZONTAL + ) return self.start_indices[dim][start_marker], self.end_indices[dim][end_marker] def get_c2e_connectivity(self): @@ -103,6 +119,13 @@ def get_c2e2c0_connectivity(self): table = self.connectivities["c2e2c0"] return NeighborTableOffsetProvider(table, CellDim, CellDim, table.shape[1]) + def get_e2c2v_connectivity(self): + return self.get_e2v_connectivity() + + def get_v2e_offset_provider(self): + table = self.connectivities["v2e"] + return NeighborTableOffsetProvider(table, VertexDim, EdgeDim, table.shape[1]) + class VerticalModelParams: def __init__(self, vct_a: np.ndarray, rayleigh_damping_height: float = 12500.0): @@ -141,5 +164,3 @@ def init_nabla2_factor_in_upper_damping_zone( / (self.vct_a[2] - self.vct_a[self.index_of_damping_height + 1]) ) return np_as_located_field(KDim)(buffer) - - diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py index 9e94ea574..7a1e20d3e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py @@ -15,7 +15,15 @@ from functional.common import Field -from icon4py.common.dimension import V2EDim, VertexDim, CellDim, C2EDim, E2CDim, EdgeDim, C2E2CODim +from icon4py.common.dimension import ( + C2E2CODim, + C2EDim, + CellDim, + E2CDim, + EdgeDim, + V2EDim, + VertexDim, +) @dataclass @@ -27,15 +35,27 @@ def __init__(self): represents """ - c_lin_e: Field[[EdgeDim, E2CDim], float] # (nproma, 2, nblks_e) - e_bln_c_s: Field[[CellDim, C2EDim], float] # coefficent for bilinear interpolation from edge to cell () - rbf_coeff_1: Field[[VertexDim, V2EDim], float] # rbf_vec_coeff_v_1(nproma, rbf_vec_dim_v, nblks_v) - rbf_coeff_2: Field[[VertexDim, V2EDim], float] # rbf_vec_coeff_v_2(nproma, rbf_vec_dim_v, nblks_v) - - geofac_div: Field[[CellDim,C2EDim], float] # factor for divergence (nproma,cell_type,nblks_c) - - geofac_n2s: Field[[CellDim, C2E2CODim], float] # factor for nabla2-scalar (nproma,cell_type+1,nblks_c) - geofac_grg_x: Field[[CellDim, C2E2CODim], float] # factor for green gauss gradient (nproma,4,nblks_c,2) - geofac_grg_y: Field[[CellDim, C2E2CODim], float] # TODO combine to tuple - nudgecoeff_e: Field[[EdgeDim],float] #Nudgeing coeffients for edges + c_lin_e: Field[[EdgeDim, E2CDim], float] # (nproma, 2, nblks_e) + e_bln_c_s: Field[ + [CellDim, C2EDim], float + ] # coefficent for bilinear interpolation from edge to cell () + rbf_coeff_1: Field[ + [VertexDim, V2EDim], float + ] # rbf_vec_coeff_v_1(nproma, rbf_vec_dim_v, nblks_v) + rbf_coeff_2: Field[ + [VertexDim, V2EDim], float + ] # rbf_vec_coeff_v_2(nproma, rbf_vec_dim_v, nblks_v) + + geofac_div: Field[ + [CellDim, C2EDim], float + ] # factor for divergence (nproma,cell_type,nblks_c) + + geofac_n2s: Field[ + [CellDim, C2E2CODim], float + ] # factor for nabla2-scalar (nproma,cell_type+1,nblks_c) + geofac_grg_x: Field[ + [CellDim, C2E2CODim], float + ] # factor for green gauss gradient (nproma,4,nblks_c,2) + geofac_grg_y: Field[[CellDim, C2E2CODim], float] # TODO combine to tuple + nudgecoeff_e: Field[[EdgeDim], float] # Nudgeing coeffients for edges nudgecoeff_c: Field[[CellDim], float] # nudgeing coeffiecients for cells diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py index ecebde53f..94ce93d1c 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py @@ -1,24 +1,51 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + from dataclasses import dataclass from functional.common import Field -from icon4py.common.dimension import KDim, EdgeDim, CellDim, C2E2CDim +from icon4py.common.dimension import C2E2CDim, CellDim, EdgeDim, KDim @dataclass class MetricState: theta_ref_mc: Field[[CellDim, KDim], float] - enhfac_diffu: Field[[KDim], float] #Enhancement factor for nabla4 background diffusion TODO check dimension - wgtfac_e: Field[[EdgeDim, KDim], float] #weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_e) - wgtfac_c: Field[[CellDim, KDim], float] #weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_c) - wgtfacq1_e: Field[[EdgeDim, ], float] #weighting factor for quadratic interpolation to model top (nproma,3,nblks_e) - wgtfacq_e: Field[[EdgeDim, ], float] #weighting factor for quadratic interpolation to surface (nproma,3,nblks_e) - ddqz_z_full_e: Field[[EdgeDim, KDim], float] #functional determinant of the metrics [sqrt(gamma)] (nproma,nlev,nblks_e) + enhfac_diffu: Field[ + [KDim], float + ] # Enhancement factor for nabla4 background diffusion TODO check dimension + wgtfac_e: Field[ + [EdgeDim, KDim], float + ] # weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_e) + wgtfac_c: Field[ + [CellDim, KDim], float + ] # weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_c) + wgtfacq1_e: Field[ + [ + EdgeDim, + ], + float, + ] # weighting factor for quadratic interpolation to model top (nproma,3,nblks_e) + wgtfacq_e: Field[ + [ + EdgeDim, + ], + float, + ] # weighting factor for quadratic interpolation to surface (nproma,3,nblks_e) + ddqz_z_full_e: Field[ + [EdgeDim, KDim], float + ] # functional determinant of the metrics [sqrt(gamma)] (nproma,nlev,nblks_e) mask_hdiff: Field[[CellDim, KDim], int] - zd_vertidx_dsl: Field[[CellDim, C2E2CDim, KDim], int] + zd_vertidx_dsl: Field[[CellDim, C2E2CDim, KDim], int] zd_diffcoef: Field[[CellDim, KDim], float] zd_intcoef: Field[[CellDim, KDim], float] - - - - diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 654d8a4ed..f1420c5ac 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -14,24 +14,23 @@ import numpy as np import pytest -from functional.iterator.embedded import np_as_located_field from icon4py.atm_dyn_iconam.diagnostic import DiagnosticState from icon4py.atm_dyn_iconam.diffusion import ( + CartesianVectorTuple, Diffusion, DiffusionConfig, DiffusionParams, enhanced_smagorinski_factor, init_diffusion_local_fields, - set_zero_v_k, scale_k, CartesianVectorTuple, + scale_k, + set_zero_v_k, ) from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState from icon4py.atm_dyn_iconam.metric_state import MetricState from icon4py.atm_dyn_iconam.prognostic import PrognosticState from icon4py.common.dimension import KDim, VertexDim -from icon4py.testutils.serialbox_utils import ( - IconSerialDataProvider, -) +from icon4py.testutils.serialbox_utils import IconSerialDataProvider from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import random_field, zero_field @@ -41,9 +40,10 @@ def test_scale_k(): field = random_field(mesh, KDim) scaled_field = zero_field(mesh, KDim) factor = 2.0 - scale_k(field, factor, scaled_field, offset_provider = {}) + scale_k(field, factor, scaled_field, offset_provider={}) assert np.allclose(factor * np.asarray(field), scaled_field) + def smag_limit_numpy(shape, k4, substeps): return 0.125 - 4.0 * diff_multfac_vn_numpy(shape, k4, substeps) @@ -252,29 +252,34 @@ def test_diffusion_run(): dtime = sp.get_metadata("dtime") orientation = sp.tangent_orientation() - inverse_primal_edge_lengths= sp.inverse_primal_edge_lengths() + inverse_primal_edge_lengths = sp.inverse_primal_edge_lengths() inverse_vertical_vertex_lengths = sp.inv_vert_vert_length() inverse_dual_edge_length = sp.inv_dual_edge_length() - primal_normal_vert: CartesianVectorTuple = (sp.primal_normal_vert_x(), sp.primal_normal_vert_y()) - dual_normal_vert: CartesianVectorTuple = (sp.dual_normal_vert_x(), sp.dual_normal_vert_y()) + primal_normal_vert: CartesianVectorTuple = ( + sp.primal_normal_vert_x(), + sp.primal_normal_vert_y(), + ) + dual_normal_vert: CartesianVectorTuple = ( + sp.dual_normal_vert_x(), + sp.dual_normal_vert_y(), + ) edge_areas = sp.edge_areas() cell_areas = sp.cell_areas() - diffusion.run(diagnostic_state=diagnostic_state, - prognostic_state=prognostic_state, - metric_state=metric_state, - interpolation_state=interpolation_state, - dtime=dtime, - tangent_orientation=orientation, - inverse_primal_edge_lengths=inverse_primal_edge_lengths, - inv_dual_edge_length=inverse_dual_edge_length, - inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, - primal_normal_vert=primal_normal_vert, - dual_normal_vert=dual_normal_vert, - edge_areas=edge_areas, - cell_areas=cell_areas - ) + diffusion.run( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + metric_state=metric_state, + interpolation_state=interpolation_state, + dtime=dtime, + tangent_orientation=orientation, + inverse_primal_edge_lengths=inverse_primal_edge_lengths, + inv_dual_edge_length=inverse_dual_edge_length, + inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, + primal_normal_vert=primal_normal_vert, + dual_normal_vert=dual_normal_vert, + edge_areas=edge_areas, + cell_areas=cell_areas, + ) pytest.fail("not implemented yet") - - diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index d82890659..023062ad2 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -1,21 +1,39 @@ +# 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 os import numpy as np from _pytest.fixtures import fixture -from icon4py.atm_dyn_iconam.horizontal import HorizontalMarkerIndex, HorizontalMeshConfig, HorizontalMeshParams +from icon4py.atm_dyn_iconam.horizontal import ( + HorizontalMarkerIndex, + HorizontalMeshConfig, +) from icon4py.atm_dyn_iconam.icon_grid import IconGrid, MeshConfig -from icon4py.common.dimension import VertexDim, EdgeDim, CellDim +from icon4py.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.testutils.serialbox_utils import IconSerialDataProvider @fixture def with_grid(): data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") - sp = IconSerialDataProvider("icon_diffusion_init", data_path).from_savepoint(linit=True, - date="2021-06-20T12:00:10.000") - nproma, nlev, num_v, num_c, num_e = sp.get_metadata("nproma", "nlev", "num_vert", "num_cells", - "num_edges") + sp = IconSerialDataProvider("icon_diffusion_init", data_path).from_savepoint( + linit=True, date="2021-06-20T12:00:10.000" + ) + nproma, nlev, num_v, num_c, num_e = sp.get_metadata( + "nproma", "nlev", "num_vert", "num_cells", "num_edges" + ) cell_starts = sp.cells_start_index() cell_ends = sp.cells_end_index() vertex_starts = sp.vertex_start_index() @@ -23,50 +41,241 @@ def with_grid(): edge_starts = sp.edge_start_index() edge_ends = sp.edge_end_index() - config = MeshConfig(HorizontalMeshConfig(num_vertices=num_v, num_cells=num_c, num_edges=num_e)) + config = MeshConfig( + HorizontalMeshConfig(num_vertices=num_v, num_cells=num_c, num_edges=num_e) + ) c2e2c = np.squeeze(sp.c2e2c(), axis=1) c2e2c0 = np.column_stack((c2e2c, (np.asarray(range(c2e2c.shape[0]))))) - grid = IconGrid().with_config(config) \ - .with_start_end_indices(VertexDim, vertex_starts, vertex_ends) \ - .with_start_end_indices(EdgeDim, edge_starts, edge_ends)\ - .with_start_end_indices(CellDim, cell_starts, cell_ends)\ - .with_connectivity(c2e=sp.c2e()).with_connectivity(e2c=sp.e2c())\ - .with_connectivity(c2e2c=c2e2c) \ - .with_connectivity(e2v=sp.e2v())\ - .with_connectivity(c2e2c0=c2e2c0)\ - + grid = ( + IconGrid() + .with_config(config) + .with_start_end_indices(VertexDim, vertex_starts, vertex_ends) + .with_start_end_indices(EdgeDim, edge_starts, edge_ends) + .with_start_end_indices(CellDim, cell_starts, cell_ends) + .with_connectivity(c2e=sp.c2e()) + .with_connectivity(e2c=sp.e2c()) + .with_connectivity(c2e2c=c2e2c) + .with_connectivity(e2v=sp.e2v()) + .with_connectivity(c2e2c0=c2e2c0) + ) return grid + def test_horizontal_grid_cell_indices(with_grid): - assert with_grid.get_indices_from_to(CellDim, 3, 3) == (20897, 20896) # halo +1 - #assert grid.get_indices_from_to(CellDim, 4, 4) == (1, 20896) #halo instead is (0,20896) why - assert with_grid.get_indices_from_to(CellDim, 8, 8) == (4105, 20896) # interior - assert with_grid.get_indices_from_to(CellDim, 9, 9) == (1, 850) #lb+1 - assert with_grid.get_indices_from_to(CellDim, 10, 10) == (851, 1688) - assert with_grid.get_indices_from_to(CellDim, 11, 11) == (1689, 2511) #lb+2 - assert with_grid.get_indices_from_to(CellDim, 12,12) ==(2512, 3316) #lb+3 - assert with_grid.get_indices_from_to(CellDim, HorizontalMarkerIndex.START_PROG_CELL.value , HorizontalMarkerIndex.START_PROG_CELL.value) == (3317, 4104) #nudging + assert with_grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.halo(CellDim) - 1, + HorizontalMarkerIndex.halo(CellDim) - 1, + ) == ( + 20897, + 20896, + ) # halo +1 + assert with_grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.halo(CellDim), + HorizontalMarkerIndex.halo(CellDim), + ) == ( + 0, + 20896, + ) # halo in icon is (1,20896) why + assert with_grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.interior(CellDim), + HorizontalMarkerIndex.interior(CellDim), + ) == ( + 4105, + 20896, + ) # interior + assert with_grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.interior(CellDim) + 1, + HorizontalMarkerIndex._INTERIOR_CELLS + 1, + ) == ( + 1, + 850, + ) # lb+1 + assert with_grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.local_boundary(CellDim) + 1, + HorizontalMarkerIndex.local_boundary(CellDim) + 1, + ) == (851, 1688) + assert with_grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.local_boundary(CellDim) + 2, + HorizontalMarkerIndex.local_boundary(CellDim) + 2, + ) == ( + 1689, + 2511, + ) # lb+2 + assert with_grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.nudging(CellDim) - 1, + HorizontalMarkerIndex.nudging(CellDim) - 1, + ) == ( + 2512, + 3316, + ) # lb+3 + assert with_grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.nudging(CellDim), + HorizontalMarkerIndex.nudging(CellDim), + ) == ( + 3317, + 4104, + ) # nudging def test_horizontal_edge_indices(with_grid): - assert with_grid.get_indices_from_to(EdgeDim, 0, 0) == (31559, 31558) - assert with_grid.get_indices_from_to(EdgeDim, 3, 3) == (31559, 31558) - assert with_grid.get_indices_from_to(EdgeDim, 4, 4) == (31559, 31558) - # assert with_grid.get_indices_from_to(EdgeDim, 5, 5) == (1, 31558) #halo - assert with_grid.get_indices_from_to(EdgeDim, 23, 23) == (5388, 6176) # nudging +1 - assert with_grid.get_indices_from_to(EdgeDim, 22,22) == (4990, 5387) #nudging - assert with_grid.get_indices_from_to(EdgeDim, 21, 21) == (4185, 4989) #lb +7 - assert with_grid.get_indices_from_to(EdgeDim, 20, 20) == (3778, 4184) # lb +6 - assert with_grid.get_indices_from_to(EdgeDim, 19, 19) == (2955, 3777) # lb +5 - assert with_grid.get_indices_from_to(EdgeDim, 18, 18) == (2539, 2954) # lb +4 - assert with_grid.get_indices_from_to(EdgeDim, 17, 17) == (1701, 2538) # lb +3 - assert with_grid.get_indices_from_to(EdgeDim, 16, 16) == (1279, 1700) # lb +2 - assert with_grid.get_indices_from_to(EdgeDim, 15, 15) == (429, 1278) # lb +1 - assert with_grid.get_indices_from_to(EdgeDim, 14, 14) == (1, 428) # lb +0 + assert with_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.interior(EdgeDim), + HorizontalMarkerIndex.interior(EdgeDim), + ) == (6177, 31558) + assert with_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.halo(EdgeDim) - 2, + HorizontalMarkerIndex.halo(EdgeDim) - 2, + ) == (31559, 31558) + assert with_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.halo(EdgeDim) - 1, + HorizontalMarkerIndex.halo(EdgeDim) - 1, + ) == (31559, 31558) + assert with_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.halo(EdgeDim), + HorizontalMarkerIndex.halo(EdgeDim), + ) == ( + 0, + 31558, + ) # halo in icon is (1, 31558) + assert with_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.nudging(EdgeDim) + 1, + HorizontalMarkerIndex.nudging(EdgeDim) + 1, + ) == ( + 5388, + 6176, + ) # nudging +1 + assert with_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.nudging(EdgeDim), + HorizontalMarkerIndex.nudging(EdgeDim), + ) == ( + 4990, + 5387, + ) # nudging + assert with_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 7, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 7, + ) == ( + 4185, + 4989, + ) # lb +7 + assert with_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 6, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 6, + ) == ( + 3778, + 4184, + ) # lb +6 + assert with_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 5, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 5, + ) == ( + 2955, + 3777, + ) # lb +5 + assert with_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, + ) == ( + 2539, + 2954, + ) # lb +4 + assert with_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 3, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 3, + ) == ( + 1701, + 2538, + ) # lb +3 + assert with_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 2, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 2, + ) == ( + 1279, + 1700, + ) # lb +2 + assert with_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 1, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 1, + ) == ( + 429, + 1278, + ) # lb +1 + assert with_grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.local_boundary(EdgeDim), + HorizontalMarkerIndex.local_boundary(EdgeDim), + ) == ( + 1, + 428, + ) # lb +0 def test_horizontal_vertex_indices(with_grid): - assert with_grid.get_indices_from_to(VertexDim, 0, 0) == (10664, 10664) - + assert with_grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.end(VertexDim), + HorizontalMarkerIndex.end(VertexDim), + ) == (10664, 10663) + assert with_grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.halo(VertexDim), + HorizontalMarkerIndex.halo(VertexDim), + ) == (0, 10663) + assert with_grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.halo(VertexDim) - 1, + HorizontalMarkerIndex.halo(VertexDim) - 1, + ) == (10664, 10663) + assert with_grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.local_boundary(VertexDim), + HorizontalMarkerIndex.local_boundary(VertexDim), + ) == (1, 428) + assert with_grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.local_boundary(VertexDim) + 1, + HorizontalMarkerIndex.local_boundary(VertexDim) + 1, + ) == (429, 850) + assert with_grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.local_boundary(VertexDim) + 2, + HorizontalMarkerIndex.local_boundary(VertexDim) + 2, + ) == (851, 1266) + assert with_grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.local_boundary(VertexDim) + 3, + HorizontalMarkerIndex.local_boundary(VertexDim) + 3, + ) == (1267, 1673) + assert with_grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.local_boundary(VertexDim) + 4, + HorizontalMarkerIndex.local_boundary(VertexDim) + 4, + ) == (1674, 2071) + assert with_grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.interior(VertexDim), + HorizontalMarkerIndex.interior(VertexDim), + ) == (2072, 10663) diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index d2a83f243..4fa41919f 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -12,13 +12,13 @@ # SPDX-License-Identifier: GPL-3.0-or-later import os -from typing import List from functional.common import Dimension from functional.iterator.embedded import np_as_located_field from serialbox import Savepoint -from icon4py.common.dimension import KDim, EdgeDim, ECVDim, CellDim, VertexDim +from icon4py.common.dimension import CellDim, ECVDim, EdgeDim, KDim, VertexDim + try: import serialbox as ser @@ -31,8 +31,6 @@ import serialbox as ser - - class IconDiffusionSavepoint: def __init__(self, sp: Savepoint, ser: ser.Serializer): self.savepoint = sp @@ -45,11 +43,12 @@ def physical_height_field(self): return self._get_field("vct_a", KDim) def _get_field(self, name, *dimensions): - return np_as_located_field(dimensions)(self.serializer.read(name, self.savepoint)) + buffer = self.serializer.read(name, self.savepoint) + return np_as_located_field(*dimensions)(buffer) def get_metadata(self, *names): metadata = self.savepoint.metainfo.to_dict() - return {n:metadata[n] for n in names if n in metadata} + return {n: metadata[n] for n in names if n in metadata} def tangent_orientation(self): return self._get_field("tangent_orientation", EdgeDim) @@ -108,6 +107,7 @@ def refin_ctrl(self, dim: Dimension): return self.serializer.read("v_refin_ctl", self.savepoint) else: return None + def c2e(self): return self.serializer.read("c2e", self.savepoint) @@ -116,10 +116,12 @@ def c2e2c(self): def e2c(self): return self.serializer.read("e2c", self.savepoint) + def e2v(self): return self.serializer.read("e2v", self.savepoint) -class IconSerialDataProvider: + +class IconSerialDataProvider: def __init__(self, fname_prefix, path="."): self.rank = 0 self.serializer: ser.Serializer = None @@ -130,13 +132,15 @@ def __init__(self, fname_prefix, path="."): def _init_serializer(self): if not self.fname: print(" WARNING: no filename! closing serializer") - self.serializer = ser.Serializer(ser.OpenModeKind.Read, self.file_path, self.fname) + self.serializer = ser.Serializer( + ser.OpenModeKind.Read, self.file_path, self.fname + ) def print_info(self): print(f"SAVEPOINTS: {self.serializer.savepoint_list()}") print(f"FIELDNAMES: {self.serializer.fieldnames()}") - def from_savepoint(self, linit:bool, date:str) -> IconDiffusionSavepoint: + def from_savepoint(self, linit: bool, date: str) -> IconDiffusionSavepoint: savepoint = ( self.serializer.savepoint["call-diffusion-init"] .linit[linit] @@ -144,37 +148,3 @@ def from_savepoint(self, linit:bool, date:str) -> IconDiffusionSavepoint: .as_savepoint() ) return IconDiffusionSavepoint(savepoint, self.serializer) - - - - def get_fields(self, metadata: List[str], fields: List[str]): - savepoint = ( - self.serializer.savepoint["call-diffusion-init"] - .linit[False] - .date["2021-06-20T12:00:10.000"] - .as_savepoint() - ) - print(savepoint.metainfo) - meta_present = {} - meta_absent = [] - for md in metadata: - if md in savepoint.metainfo.to_dict(): - meta_present[md] = savepoint.metainfo[md] - else: - meta_absent.append(md) - - fields_present = {} - fields_absent = [] - for field_name in fields: - if field_name in self.serializer.fieldnames(): - fields_present[field_name] = self.serializer.read(field_name, savepoint) - else: - fields_absent.append(field_name) - [print(f"field {f} not present in savepoint") for f in fields_absent] - [print(f"metadata {f} not present in savepoint") for f in meta_absent] - return meta_present, fields_present - - - - - From 467a9012dd376c282917fdd3344091f93ccb8614 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 7 Dec 2022 10:01:46 +0100 Subject: [PATCH 021/263] - extract grid setup from test_icon_grid.py - clean up state data classes - initialize diffusion module in test --- .../src/icon4py/atm_dyn_iconam/diagnostic.py | 13 ++- .../src/icon4py/atm_dyn_iconam/diffusion.py | 31 +++-- ...sed_mo_nh_diffusion_stencil_07_08_09_10.py | 27 +++++ .../src/icon4py/atm_dyn_iconam/icon_grid.py | 6 +- .../atm_dyn_iconam/interpolation_state.py | 10 +- .../icon4py/atm_dyn_iconam/metric_state.py | 25 +--- .../src/icon4py/atm_dyn_iconam/prognostic.py | 1 - atm_dyn_iconam/tests/icon_grid_test_utils.py | 63 +++++++++++ atm_dyn_iconam/tests/test_diffusion.py | 66 ++++++++--- atm_dyn_iconam/tests/test_icon_grid.py | 49 +------- .../src/icon4py/testutils/serialbox_utils.py | 107 ++++++++++++++++-- 11 files changed, 269 insertions(+), 129 deletions(-) create mode 100644 atm_dyn_iconam/tests/icon_grid_test_utils.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diagnostic.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diagnostic.py index 78ed3e9cf..c3edbd21a 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diagnostic.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diagnostic.py @@ -24,10 +24,13 @@ class DiagnosticState: hdef_ic: Field[ [CellDim, KDim], float ] # ! divergence at half levels(nproma,nlevp1,nblks_c) [1/s] - div_ic = Field[ + div_ic: Field[ [CellDim, KDim], float ] # ! horizontal wind field deformation (nproma,nlevp1,nblks_c) [1/s^2] - dwdx = ( - None # zonal gradient of vertical wind speed (nproma,nlevp1,nblks_c) [1/s] - ) - dwdy = None # meridional gradient of vertical wind speed (nproma,nlevp1,nblks_c) + dwdx: Field[ + [CellDim, KDim], float + ] # zonal gradient of vertical wind speed (nproma,nlevp1,nblks_c) [1/s] + + dwdy: Field[ + [CellDim, KDim], float + ] # meridional gradient of vertical wind speed (nproma,nlevp1,nblks_c) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index b8d1b8e68..6bb3f250b 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -72,7 +72,7 @@ DiffusionTupleVT = namedtuple("DiffusionParamVT", "v t") -CartesianVectorTuple = namedtuple("CartesianVectorTuple", "x y") +VectorTuple = namedtuple("CartesianVectorTuple", "x y") # TODO [ml] initial RUN linit = TRUE @@ -105,7 +105,6 @@ def scale_k( _scale_k(field, factor, out=scaled_field) -# TODO [ml] rename! @field_operator def _mo_nh_diffusion_stencil_01_scale_dtime( enh_smag_fac: Field[[KDim], float], @@ -286,6 +285,7 @@ class DiffusionConfig: TODO: [ml] read from config TODO: [ml] handle dependencies on other namelists (below...) + TODO: move this to test: it sets up the mch_ch_r04b09_dsl experiment """ @classmethod @@ -514,14 +514,10 @@ def run( dtime: float, tangent_orientation: Field[[EdgeDim], float], inverse_primal_edge_lengths: Field[[EdgeDim], float], - inv_dual_edge_length: Field[[EdgeDim], float], + inverse_dual_edge_length: Field[[EdgeDim], float], inverse_vertical_vertex_lengths: Field[[EdgeDim], float], - primal_normal_vert: CartesianVectorTuple[ - Field[[ECVDim], float], Field[[ECVDim], float] - ], - dual_normal_vert: CartesianVectorTuple[ - Field[[ECVDim], float], Field[[ECVDim], float] - ], + primal_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], + dual_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], edge_areas: Field[[EdgeDim], float], cell_areas: Field[[CellDim], float], ): @@ -546,8 +542,8 @@ def run( ) klevels = self.grid.k_levels() # TODO is this needed? - set_zero_v_k(self.u_vert) - set_zero_v_k(self.v_vert) + set_zero_v_k(self.u_vert, offset_provider={}) + set_zero_v_k(self.v_vert, offset_provider={}) # 1. CALL rbf_vec_interpol_vertex @@ -567,10 +563,11 @@ def run( }, offset_provider={"V2E": self.grid.get_v2e_offset_provider()}, ) + # 2. HALO EXCHANGE -- CALL sync_patch_array_mult # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 # 0c. dtime dependent stuff: enh_smag_factor, ~~r_dtimensubsteps~~ - scale_k(self.enh_smag_fac, dtime, self.diff_multfac_smag) + scale_k(self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={}) mo_nh_diffusion_stencil_01( self.diff_multfac_smag, @@ -579,10 +576,10 @@ def run( inverse_vertical_vertex_lengths, self.u_vert, self.v_vert, - primal_normal_vert.x, - primal_normal_vert.y, - dual_normal_vert.x, - dual_normal_vert.y, + primal_normal_vert[0], + primal_normal_vert[1], + dual_normal_vert[0], + dual_normal_vert[1], prognostic_state.normal_wind, self.smag_limit, self.kh_smag_e, @@ -739,7 +736,7 @@ def run( fused_mo_nh_diffusion_stencil_13_14( kh_smag_e=self.kh_smag_e, - inv_dual_edge_length=inv_dual_edge_length, + inv_dual_edge_length=inverse_dual_edge_length, theta_v=prognostic_state.theta_v, geofac_div=interpolation_state.geofac_div, z_temp=self.z_temp, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py index 796d0ff58..1f9766257 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -65,6 +65,14 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( _mo_nh_diffusion_stencil_08(w_old, geofac_grg_x, geofac_grg_y), (dwdx, dwdy), ) + num_lev = 10 + + # _mo_nh_diffusion_stencil_08( + # w_old, + # geofac_grg_x, + # geofac_grg_y, + # out=(dwdx, dwdy), + # domain={KDim: (0, num_lev)}) z_nabla2_c = _mo_nh_diffusion_stencil_07(w_old, geofac_n2s) @@ -76,6 +84,16 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( w_old, ) + # _mo_nh_diffusion_stencil_07( + # area, + # z_nabla2_c, + # geofac_n2s, + # w_old, + # diff_multfac_w, + # out=w_old, + # domain={CellDim: (interior_idx, halo_idx)} + # ) + w = where( (vert_idx > int32(0)) & (vert_idx < nrdmax) @@ -85,6 +103,15 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( w, ) + # _mo_nh_diffusion_stencil_10( + # w_old, diff_multfac_n2w, + # area, + # z_nabla2_c, + # out=w_old, + # domain={CellDim: (interior_idx, halo_idx), KDim:(0, num_lev)} + # ) + + w = w_old return w, dwdx, dwdy diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py index a27dd5cc1..0c6dc2701 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py @@ -133,14 +133,14 @@ def __init__(self, vct_a: np.ndarray, rayleigh_damping_height: float = 12500.0): Contains physical parameters defined on the grid. Args: - vct_a: - rayleigh_damping_height height of rayleigh damping in [m] mo_nonhydro_nml + vct_a: field containing the physical heights of the k level + rayleigh_damping_height: height of rayleigh damping in [m] mo_nonhydro_nml """ self.rayleigh_damping_height = rayleigh_damping_height self.vct_a = vct_a # TODO klevels in ICON are inverted! self.index_of_damping_height = np.argmax( - self.vct_a >= self.rayleigh_damping_height + np.asarray(self.vct_a) >= self.rayleigh_damping_height ) def get_index_of_damping_layer(self): diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py index 7a1e20d3e..2d9a4035c 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py @@ -19,7 +19,6 @@ C2E2CODim, C2EDim, CellDim, - E2CDim, EdgeDim, V2EDim, VertexDim, @@ -28,14 +27,12 @@ @dataclass class InterpolationState: - def __init__(self): - self.g = None - """ - represents + represents the ICON interpolation state. + + TODO: keep? does this state make sense at all? """ - c_lin_e: Field[[EdgeDim, E2CDim], float] # (nproma, 2, nblks_e) e_bln_c_s: Field[ [CellDim, C2EDim], float ] # coefficent for bilinear interpolation from edge to cell () @@ -58,4 +55,3 @@ def __init__(self): ] # factor for green gauss gradient (nproma,4,nblks_c,2) geofac_grg_y: Field[[CellDim, C2E2CODim], float] # TODO combine to tuple nudgecoeff_e: Field[[EdgeDim], float] # Nudgeing coeffients for edges - nudgecoeff_c: Field[[CellDim], float] # nudgeing coeffiecients for cells diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py index 94ce93d1c..8ec577faf 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py @@ -15,37 +15,16 @@ from functional.common import Field -from icon4py.common.dimension import C2E2CDim, CellDim, EdgeDim, KDim +from icon4py.common.dimension import C2E2CDim, CellDim, KDim @dataclass class MetricState: theta_ref_mc: Field[[CellDim, KDim], float] - enhfac_diffu: Field[ - [KDim], float - ] # Enhancement factor for nabla4 background diffusion TODO check dimension - wgtfac_e: Field[ - [EdgeDim, KDim], float - ] # weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_e) wgtfac_c: Field[ [CellDim, KDim], float ] # weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_c) - wgtfacq1_e: Field[ - [ - EdgeDim, - ], - float, - ] # weighting factor for quadratic interpolation to model top (nproma,3,nblks_e) - wgtfacq_e: Field[ - [ - EdgeDim, - ], - float, - ] # weighting factor for quadratic interpolation to surface (nproma,3,nblks_e) - ddqz_z_full_e: Field[ - [EdgeDim, KDim], float - ] # functional determinant of the metrics [sqrt(gamma)] (nproma,nlev,nblks_e) mask_hdiff: Field[[CellDim, KDim], int] - zd_vertidx_dsl: Field[[CellDim, C2E2CDim, KDim], int] + zd_vertidx: Field[[CellDim, C2E2CDim, KDim], int] zd_diffcoef: Field[[CellDim, KDim], float] zd_intcoef: Field[[CellDim, KDim], float] diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/prognostic.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/prognostic.py index c7f9e1b95..5495e5004 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/prognostic.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/prognostic.py @@ -27,6 +27,5 @@ class PrognosticState: vertical_wind: Field[[CellDim, KDim], float] # w(nproma, nlevp1, nblks_c) [m/s] normal_wind: Field[[EdgeDim, KDim], float] # vn(nproma, nlev, nblks_e) [m/s] - density: Field[[CellDim, KDim], float] # rho(nproma, nlev, nblks_c) [kg/m^3] exner_pressure: Field[[CellDim, KDim], float] # exner(nrpoma, nlev, nblks_c) theta_v: Field[[CellDim, KDim], float] # (nproma, nlev, nlbks_c) [K] diff --git a/atm_dyn_iconam/tests/icon_grid_test_utils.py b/atm_dyn_iconam/tests/icon_grid_test_utils.py new file mode 100644 index 000000000..90999532e --- /dev/null +++ b/atm_dyn_iconam/tests/icon_grid_test_utils.py @@ -0,0 +1,63 @@ +# 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 os + +import numpy as np +from _pytest.fixtures import fixture + +from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig +from icon4py.atm_dyn_iconam.icon_grid import IconGrid, MeshConfig +from icon4py.common.dimension import CellDim, EdgeDim, VertexDim +from icon4py.testutils.serialbox_utils import IconSerialDataProvider + + +@fixture +def with_grid(): + data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") + sp = IconSerialDataProvider("icon_diffusion_init", data_path, True).from_savepoint( + linit=True, date="2021-06-20T12:00:10.000" + ) + + sp_meta = sp.get_metadata("nproma", "nlev", "num_vert", "num_cells", "num_edges") + + cell_starts = sp.cells_start_index() + cell_ends = sp.cells_end_index() + vertex_starts = sp.vertex_start_index() + vertex_ends = sp.vertex_end_index() + edge_starts = sp.edge_start_index() + edge_ends = sp.edge_end_index() + + config = MeshConfig( + HorizontalMeshConfig( + num_vertices=sp_meta["num_vert"], + num_cells=sp_meta["num_cells"], + num_edges=sp_meta["num_edges"], + ) + ) + c2e2c = sp.c2e2c() + c2e2c0 = np.column_stack((c2e2c, (np.asarray(range(c2e2c.shape[0]))))) + grid = ( + IconGrid() + .with_config(config) + .with_start_end_indices(VertexDim, vertex_starts, vertex_ends) + .with_start_end_indices(EdgeDim, edge_starts, edge_ends) + .with_start_end_indices(CellDim, cell_starts, cell_ends) + .with_connectivity(c2e=sp.c2e()) + .with_connectivity(e2c=sp.e2c()) + .with_connectivity(c2e2c=c2e2c) + .with_connectivity(e2v=sp.e2v()) + .with_connectivity(c2e2c0=c2e2c0) + .with_connectivity(v2e=sp.v2e()) + ) + return grid diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index f1420c5ac..26851e670 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -14,18 +14,20 @@ import numpy as np import pytest +from icon_grid_test_utils import with_grid from icon4py.atm_dyn_iconam.diagnostic import DiagnosticState from icon4py.atm_dyn_iconam.diffusion import ( - CartesianVectorTuple, Diffusion, DiffusionConfig, DiffusionParams, - enhanced_smagorinski_factor, + VectorTuple, + _en_smag_fac_for_zero_nshift, init_diffusion_local_fields, scale_k, set_zero_v_k, ) +from icon4py.atm_dyn_iconam.icon_grid import VerticalModelParams from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState from icon4py.atm_dyn_iconam.metric_state import MetricState from icon4py.atm_dyn_iconam.prognostic import PrognosticState @@ -110,7 +112,9 @@ def test_enhanced_smagorinski_factor(): result = zero_field(mesh, KDim) fac = (0.67, 0.5, 1.3, 0.8) z = (0.1, 0.2, 0.3, 0.4) - enhanced_smagorinski_factor(*fac, *z, a_vec, result, offset_provider={"Koff": KDim}) + _en_smag_fac_for_zero_nshift( + *fac, *z, a_vec, out=result, offset_provider={"Koff": KDim} + ) enhanced_smag_fac_np = enhanced_smagorinski_factor_np(fac, z, np.asarray(a_vec)) assert np.allclose(enhanced_smag_fac_np, np.asarray(result[:-1])) @@ -192,7 +196,6 @@ def test_diffusion_init(): vct_a = savepoint.physical_height_field() meta = savepoint.get_metadata("nlev", "linit", "date") - print(f"meta {meta}") assert meta["nlev"] == 65 assert meta["linit"] is False assert meta["date"] == first_run_date @@ -234,32 +237,65 @@ def test_diffusion_init(): @pytest.mark.xfail -def test_diffusion_run(): +def test_diffusion_run(with_grid): data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") data_provider = IconSerialDataProvider("icon_diffusion_init", data_path) data_provider.print_info() sp = data_provider.from_savepoint(linit=False, date="2021-06-20T12:00:10.000") + vct_a = sp.physical_height_field() + + config = DiffusionConfig( + with_grid, + vertical_params=VerticalModelParams( + vct_a=vct_a, rayleigh_damping_height=12500.0 + ), + ) - config = DiffusionConfig.create_with_defaults() additional_parameters = DiffusionParams(config) - vct_a = sp.physical_height_field() diffusion = Diffusion(config, additional_parameters, vct_a) - diagnostic_state = DiagnosticState() - prognostic_state = PrognosticState() - interpolation_state = InterpolationState() - metric_state = MetricState() - dtime = sp.get_metadata("dtime") + + diagnostic_state = DiagnosticState( + hdef_ic=sp.hdef_ic(), div_ic=sp.div_ic(), dwdx=sp.dwdx(), dwdy=sp.dwdy() + ) + prognostic_state = PrognosticState( + vertical_wind=sp.vn(), + normal_wind=sp.w(), + exner_pressure=sp.exner(), + theta_v=sp.theta_v(), + ) + grg = sp.geofac_grg() + + rbf_coef = sp.rbf_coeff() + interpolation_state = InterpolationState( + e_bln_c_s=sp.e_bln_c_s(), + rbf_coeff_1=rbf_coef[0], + rbf_coeff_2=sp.rbf_coeff[1], + geofac_div=sp.geofac_div(), + geofac_n2s=sp.geofac_n2s(), + geofac_grg_x=grg[0], + geofac_grg_y=grg[1], + nudgecoeff_e=sp.nudgecoeff_e(), + ) + metric_state = MetricState( + mask_hdiff=sp.theta_ref_mc(), + theta_ref_mc=sp.theta_ref_mc(), + wgtfac_c=sp.wgtfac_c(), + zd_intcoef=sp.zd_intcoef(), + zd_vertidx=sp.zd_vertidx(), + zd_diffcoef=sp.zd_diffcoef(), + ) + dtime = sp.get_metadata("dtime").get("dtime") orientation = sp.tangent_orientation() inverse_primal_edge_lengths = sp.inverse_primal_edge_lengths() inverse_vertical_vertex_lengths = sp.inv_vert_vert_length() inverse_dual_edge_length = sp.inv_dual_edge_length() - primal_normal_vert: CartesianVectorTuple = ( + primal_normal_vert: VectorTuple = ( sp.primal_normal_vert_x(), sp.primal_normal_vert_y(), ) - dual_normal_vert: CartesianVectorTuple = ( + dual_normal_vert: VectorTuple = ( sp.dual_normal_vert_x(), sp.dual_normal_vert_y(), ) @@ -274,7 +310,7 @@ def test_diffusion_run(): dtime=dtime, tangent_orientation=orientation, inverse_primal_edge_lengths=inverse_primal_edge_lengths, - inv_dual_edge_length=inverse_dual_edge_length, + inverse_dual_edge_length=inverse_dual_edge_length, inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, primal_normal_vert=primal_normal_vert, dual_normal_vert=dual_normal_vert, diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index 023062ad2..8f8fd7c78 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -10,55 +10,10 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +from icon_grid_test_utils import with_grid -import os - -import numpy as np -from _pytest.fixtures import fixture - -from icon4py.atm_dyn_iconam.horizontal import ( - HorizontalMarkerIndex, - HorizontalMeshConfig, -) -from icon4py.atm_dyn_iconam.icon_grid import IconGrid, MeshConfig +from icon4py.atm_dyn_iconam.horizontal import HorizontalMarkerIndex from icon4py.common.dimension import CellDim, EdgeDim, VertexDim -from icon4py.testutils.serialbox_utils import IconSerialDataProvider - - -@fixture -def with_grid(): - data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") - sp = IconSerialDataProvider("icon_diffusion_init", data_path).from_savepoint( - linit=True, date="2021-06-20T12:00:10.000" - ) - nproma, nlev, num_v, num_c, num_e = sp.get_metadata( - "nproma", "nlev", "num_vert", "num_cells", "num_edges" - ) - cell_starts = sp.cells_start_index() - cell_ends = sp.cells_end_index() - vertex_starts = sp.vertex_start_index() - vertex_ends = sp.vertex_end_index() - edge_starts = sp.edge_start_index() - edge_ends = sp.edge_end_index() - - config = MeshConfig( - HorizontalMeshConfig(num_vertices=num_v, num_cells=num_c, num_edges=num_e) - ) - c2e2c = np.squeeze(sp.c2e2c(), axis=1) - c2e2c0 = np.column_stack((c2e2c, (np.asarray(range(c2e2c.shape[0]))))) - grid = ( - IconGrid() - .with_config(config) - .with_start_end_indices(VertexDim, vertex_starts, vertex_ends) - .with_start_end_indices(EdgeDim, edge_starts, edge_ends) - .with_start_end_indices(CellDim, cell_starts, cell_ends) - .with_connectivity(c2e=sp.c2e()) - .with_connectivity(e2c=sp.e2c()) - .with_connectivity(c2e2c=c2e2c) - .with_connectivity(e2v=sp.e2v()) - .with_connectivity(c2e2c0=c2e2c0) - ) - return grid def test_horizontal_grid_cell_indices(with_grid): diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 4fa41919f..a3008ba9d 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -13,11 +13,22 @@ import os +import numpy as np from functional.common import Dimension from functional.iterator.embedded import np_as_located_field from serialbox import Savepoint -from icon4py.common.dimension import CellDim, ECVDim, EdgeDim, KDim, VertexDim +from icon4py.common.dimension import ( + C2E2CDim, + C2E2CODim, + C2EDim, + CellDim, + E2C2VDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) try: @@ -43,7 +54,8 @@ def physical_height_field(self): return self._get_field("vct_a", KDim) def _get_field(self, name, *dimensions): - buffer = self.serializer.read(name, self.savepoint) + buffer = np.squeeze(self.serializer.read(name, self.savepoint)) + print(f"{name} {buffer.shape}") return np_as_located_field(*dimensions)(buffer) def get_metadata(self, *names): @@ -54,22 +66,22 @@ def tangent_orientation(self): return self._get_field("tangent_orientation", EdgeDim) def inverse_primal_edge_lengths(self): - return self._get_field("inverse_primal_edge_lengths", EdgeDim) + return self._get_field("inv_primal_edge_length", EdgeDim) def inv_vert_vert_length(self): return self._get_field("inv_vert_vert_length", EdgeDim) def primal_normal_vert_x(self): - return self._get_field("primal_normal_vert_x", ECVDim) + return self._get_field("primal_normal_vert_x", VertexDim, E2C2VDim) def primal_normal_vert_y(self): - return self._get_field("primal_normal_vert_y", ECVDim) + return self._get_field("primal_normal_vert_y", VertexDim, E2C2VDim) def dual_normal_vert_y(self): - return self._get_field("dual_normal_vert_y", ECVDim) + return self._get_field("dual_normal_vert_y", VertexDim, E2C2VDim) def dual_normal_vert_x(self): - return self._get_field("dual_normal_vert_x", ECVDim) + return self._get_field("dual_normal_vert_x", VertexDim, E2C2VDim) def cell_areas(self): return self._get_field("cell_areas", CellDim) @@ -81,7 +93,7 @@ def inv_dual_edge_length(self): return self._get_field("inv_dual_edge_length", EdgeDim) def cells_start_index(self): - return self.serializer.read("cells_start_index", self.savepoint) + return self.serializer.read("c_start_index", self.savepoint) def cells_end_index(self): return self.serializer.read("c_end_index", self.savepoint) @@ -120,21 +132,94 @@ def e2c(self): def e2v(self): return self.serializer.read("e2v", self.savepoint) + def hdef_ic(self): + return self._get_field("hdef_ic", CellDim, KDim) + + def div_ic(self): + return self._get_field("div_ic", CellDim, KDim) + + def dwdx(self): + return self._get_field("dwdx", CellDim, KDim) + + def dwdy(self): + return self._get_field("dwdy", CellDim, KDim) + + def vn(self): + return self._get_field("vn", CellDim, KDim) + + def theta_v(self): + return self._get_field("theta_v", CellDim, KDim) + + def w(self): + return self._get_field("w", CellDim, KDim) + + def exner(self): + return self._get_field("exner", CellDim, KDim) + + def theta_ref_mc(self): + return self._get_field("theta_ref_mc", CellDim, KDim) + + def wgtfac_c(self): + return self._get_field("wgtfac_c", CellDim, KDim) + + def wgtfac_e(self): + return self._get_field("wgtfac_e", EdgeDim, KDim) + + def mask_diff(self): + return self._get_field("mask_diff", CellDim, KDim) + + def zd_diffcoef(self): + return self._get_field("zd_diffcoef", CellDim, KDim) + + def zd_intcoef(self): + return self._get_field("vcoef", CellDim, C2E2CDim, KDim) + + def e_bln_c_s(self): + return self._get_field("e_bln_c_s", CellDim, C2EDim) + + def geofac_div(self): + return self._get_field("geofac_div", CellDim, C2EDim) + + def geofac_n2s(self): + return self._get_field("geofac_n2s", CellDim, C2E2CODim) + + def geofac_grg(self): + grg = self.serializer.read("geofac_grg", self.savepoint) + return grg[:, :, :, 0], grg[:, :, :, 1] + + def nudgecoeff_e(self): + return self._get_field("nudgecoeff_e", EdgeDim) + + def zd_vertidx(self): + return self._get_field("zd_vertidx", CellDim, C2E2CDim, KDim) + + def rbf_coeff(self): + rbf_coef = self.serializer.read("rbf_vec_coeff_v", self.savepoint) + return ( + np_as_located_field(VertexDim, V2EDim)(rbf_coef[:, :, 0]), + np_as_located_field(VertexDim, V2EDim)(rbf_coef)[:, :, 1], + ) + + def v2e(self): + return self.serializer.read("v2e", self.savepoint) + class IconSerialDataProvider: - def __init__(self, fname_prefix, path="."): + def __init__(self, fname_prefix, path=".", do_print=False): self.rank = 0 self.serializer: ser.Serializer = None self.file_path: str = path self.fname = f"{fname_prefix}_rank{str(self.rank)}" - self._init_serializer() + self._init_serializer(do_print) - def _init_serializer(self): + def _init_serializer(self, do_print: bool): if not self.fname: print(" WARNING: no filename! closing serializer") self.serializer = ser.Serializer( ser.OpenModeKind.Read, self.file_path, self.fname ) + if do_print: + self.print_info() def print_info(self): print(f"SAVEPOINTS: {self.serializer.savepoint_list()}") From 21fb182623c630c3f9579e5b3ba24d373b14fd25 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Wed, 7 Dec 2022 15:19:32 +0100 Subject: [PATCH 022/263] Removed duplicate stencil definition from fused stencil 13 and 14 --- ...used_mo_nh_diffusion_stencil_07_08_09_10.py | 2 +- .../fused_mo_nh_diffusion_stencil_13_14.py | 18 ++---------------- 2 files changed, 3 insertions(+), 17 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py index a1e305ae5..a97a75d21 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -40,7 +40,7 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( halo_idx: int32, ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - vert_idx = broadcast(vert_idx, (CellDim, KDim)) + vert_idx = broadcast(vert_idx, (CellDim, KDim)) dwdx, dwdy = where( vert_idx > int32(0), diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py index 40faa10b5..8d2092008 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -28,22 +28,8 @@ KDim, ) -@field_operator -def _mo_nh_diffusion_stencil_13( - kh_smag_e: Field[[EdgeDim, KDim], float], - inv_dual_edge_length: Field[[EdgeDim], float], - theta_v: Field[[CellDim, KDim], float], -) -> Field[[EdgeDim, KDim], float]: - z_nabla2_e = kh_smag_e * inv_dual_edge_length * (theta_v(E2C[1]) - theta_v(E2C[0])) - return z_nabla2_e - -@field_operator -def _mo_nh_diffusion_stencil_14( - z_nabla2_e: Field[[EdgeDim, KDim], float], - geofac_div: Field[[CEDim], float], -) -> Field[[CellDim, KDim], float]: - z_temp = neighbor_sum(z_nabla2_e(C2E) * geofac_div(C2CE), axis=C2EDim) - return z_temp +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_13 import _mo_nh_diffusion_stencil_13 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_14 import _mo_nh_diffusion_stencil_14 @field_operator def _fused_mo_nh_diffusion_stencil_13_14( From 90bc6844f48f8358b7a2af67104928bc0094bde6 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 8 Dec 2022 14:18:47 +0100 Subject: [PATCH 023/263] - fix level offsets - add test for nrdmax calculation validating icon data --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 289 ++++++++++-------- ...sed_mo_nh_diffusion_stencil_07_08_09_10.py | 33 +- .../src/icon4py/atm_dyn_iconam/horizontal.py | 1 + .../src/icon4py/atm_dyn_iconam/icon_grid.py | 20 +- atm_dyn_iconam/tests/icon_grid_test_utils.py | 45 ++- atm_dyn_iconam/tests/test_diffusion.py | 103 ++++--- atm_dyn_iconam/tests/test_icon_grid.py | 132 ++++---- atm_dyn_iconam/tests/test_vertical.py | 22 +- .../src/icon4py/testutils/serialbox_utils.py | 98 ++++-- 9 files changed, 422 insertions(+), 321 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index 6bb3f250b..b16d8f8c6 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -19,23 +19,24 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, broadcast, maximum, minimum from functional.iterator.embedded import np_as_located_field +from functional.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.constants import CPD, GAS_CONSTANT_DRY_AIR from icon4py.atm_dyn_iconam.diagnostic import DiagnosticState from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import ( - fused_mo_nh_diffusion_stencil_02_03, + _fused_mo_nh_diffusion_stencil_02_03, ) from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import ( - fused_mo_nh_diffusion_stencil_04_05_06, + _fused_mo_nh_diffusion_stencil_04_05_06, ) from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_07_08_09_10 import ( - fused_mo_nh_diffusion_stencil_07_08_09_10, + _fused_mo_nh_diffusion_stencil_07_08_09_10, ) from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_11_12 import ( fused_mo_nh_diffusion_stencil_11_12, ) from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_13_14 import ( - fused_mo_nh_diffusion_stencil_13_14, + _fused_mo_nh_diffusion_stencil_13_14, ) from icon4py.atm_dyn_iconam.horizontal import ( HorizontalMarkerIndex, @@ -49,14 +50,13 @@ from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState from icon4py.atm_dyn_iconam.metric_state import MetricState from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( - mo_intp_rbf_rbf_vec_interpol_vertex, + _mo_intp_rbf_rbf_vec_interpol_vertex, ) from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import ( _mo_nh_diffusion_stencil_01, - mo_nh_diffusion_stencil_01, ) from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_16 import ( - mo_nh_diffusion_stencil_16, + _mo_nh_diffusion_stencil_16, ) from icon4py.atm_dyn_iconam.prognostic import PrognosticState from icon4py.common.dimension import ( @@ -272,6 +272,35 @@ def _set_zero_v_k() -> Field[[VertexDim, KDim], float]: def set_zero_v_k(field: Field[[VertexDim, KDim], float]): _set_zero_v_k(out=field) +def init_nabla2_factor_in_upper_damping_zone( + k_size: int, + nrdmax:int, + nshift: int, + physical_heights: np.ndarray +) -> Field[[KDim], float]: + """ + calculated diff_multfac_n2w + + numpy version gt4py does not allow non-constant indexing into fields + + Args + k_size: number of vertical levels + nrdmax: index of the level where rayleigh dampint starts + nshift: + physcial_heights: vector of physical heights [m] of the height levels + """ + buffer = np.zeros(k_size) + buffer[2: nrdmax + 1] = ( + 1.0 + / 12.0 + * (( + physical_heights[2 + nshift: nrdmax + 1 + nshift] + - physical_heights[nshift + nrdmax + 1] + ) + / (physical_heights[2] - physical_heights[nshift + nrdmax + 1]))**4 + ) + return np_as_located_field(KDim)(buffer) + class DiffusionConfig: """contains necessary parameter to configure a diffusion run. @@ -281,48 +310,39 @@ class DiffusionConfig: currently we use the MCH r04b09_dsl experiment as constants here. These should be read from config and the default from mo_diffusion_nml.f90 set as defaults. - - TODO: [ml] read from config - TODO: [ml] handle dependencies on other namelists (below...) - TODO: move this to test: it sets up the mch_ch_r04b09_dsl experiment + TODO: [ml] handle dependencies on other namelists (see below...) """ - - @classmethod - def create_with_defaults(cls): - """ - Create DiffusionConfig. - - initialize with values from exp.mch_ch_r04b09_dsl namelist - """ - horizontal = HorizontalMeshConfig( - num_vertices=50000, num_cells=50000, num_edges=50000 - ) - - grid = IconGrid().with_config(MeshConfig(horizontalMesh=horizontal)) - verticalParams = VerticalModelParams( - rayleigh_damping_height=12500, vct_a=np.zeros(grid.k_levels()) - ) - return cls(grid=grid, vertical_params=verticalParams) - - def __init__(self, grid: IconGrid, vertical_params: VerticalModelParams): - # TODO [ml]: external stuff grid related, p_patch, + def __init__(self, + grid: IconGrid, + vertical_params: VerticalModelParams, + diffusion_type:int = 5, + apply_to_horizontal_wind: bool = True, + apply_to_vertical_wind: bool = True, + apply_to_temperature: bool = True, + reconstruction_type_smag: int = 1, + compute_3d_smag_coeff:bool = False, + temperature_discretization: int=2, + horizontal_efdt_ratio: float = 24.0, + smag_scaling_factor: float = 0.025 + ): + # TODO [ml]: move external stuff out: grid related stuff, other than diffusion namelists (see below self.grid = grid self.vertical_params = vertical_params # from namelist diffusion_nml - self.diffusion_type = 5 # hdiff_order ! order of nabla operator for diffusion - self.lhdiff_vn = True # ! diffusion on the horizontal wind field - self.lhdiff_temp = True # ! diffusion on the temperature field - self.lhdiff_w = True # ! diffusion on the vertical wind field + self.diffusion_type = diffusion_type # hdiff_order ! order of nabla operator for diffusion + self.lhdiff_vn = apply_to_horizontal_wind # ! diffusion on the horizontal wind field + self.lhdiff_temp = apply_to_temperature # ! diffusion on the temperature field + self.lhdiff_w = apply_to_vertical_wind # ! diffusion on the vertical wind field self.lhdiff_rcf = True # namelist, remove if always true self.itype_vn_diffu = ( - 1 # ! reconstruction method used for Smagorinsky diffusion + reconstruction_type_smag # ! reconstruction method used for Smagorinsky diffusion ) - self.l_smag_d = False # namelist lsmag_d, if `true`, compute 3D Smagorinsky diffusion coefficient. + self.l_smag_d = compute_3d_smag_coeff # namelist lsmag_d, if `true`, compute 3D Smagorinsky diffusion coefficient. - self.itype_t_diffu = 2 # ! discretization of temperature diffusion - self.hdiff_efdt_ratio = 24.0 # ! ratio of e-folding time to time step - self.hdiff_smag_fac = 0.025 # ! scaling factor for Smagorinsky diffusion + self.itype_t_diffu = temperature_discretization # ! discretization of temperature diffusion + self.hdiff_efdt_ratio = horizontal_efdt_ratio # ! ratio of e-folding time to time step + self.hdiff_smag_fac = smag_scaling_factor # ! scaling factor for Smagorinsky diffusion # from other namelists # from parent namelist nonhydrostatic_nml @@ -333,7 +353,7 @@ def __init__(self, grid: IconGrid, vertical_params: VerticalModelParams): # denom_diffu_v = 150 ! denominator for lateral boundary diffusion of velocity self.lateral_boundary_denominator = DiffusionTupleVT(v=200.0, t=135.0) - # namelist grid_nml -> TODO [ml] should go to grid config? + # namelist grid_nml self.l_limited_area = True # name list: interpol_nml @@ -426,7 +446,7 @@ def __init__( self, config: DiffusionConfig, params: DiffusionParams, - a_vect: Field[[KDim], float], + vct_a: Field[[KDim], float], ): """ Initialize Diffusion granule. @@ -475,15 +495,15 @@ def __init__( enhanced_smagorinski_factor( *params.smagorinski_factor, *params.smagorinski_height, - a_vect, + vct_a, self.enh_smag_fac, offset_provider={"Koff": KDim}, ) - self.diff_multfac_n2w = ( - config.vertical_params.init_nabla2_factor_in_upper_damping_zone( - k_size=config.grid.k_levels() - ) + self.diff_multfac_n2w = init_nabla2_factor_in_upper_damping_zone( + k_size=config.grid.k_levels(), nshift = 0, physical_heights=np.asarray(vct_a), + nrdmax=self.config.vertical_params.index_of_damping_height + ) self.diff_multfac_smag = np_as_located_field(KDim)( np.zeros(config.grid.k_levels()) @@ -504,6 +524,9 @@ def __init__( self.horizontal_cell_index = np_as_located_field(CellDim)( np.arange((shape_ck[0])) ) + self.horizontal_edge_index = np_as_located_field(EdgeDim)(np.arange((shape_ek[0]))) + + def run( self, @@ -547,12 +570,11 @@ def run( # 1. CALL rbf_vec_interpol_vertex - mo_intp_rbf_rbf_vec_interpol_vertex( + _mo_intp_rbf_rbf_vec_interpol_vertex( prognostic_state.normal_wind, interpolation_state.rbf_coeff_1, interpolation_state.rbf_coeff_2, - self.u_vert, - self.v_vert, + out=(self.u_vert, self.v_vert), domain={ KDim: (0, klevels), VertexDim: self.grid.get_indices_from_to( @@ -562,6 +584,7 @@ def run( ), }, offset_provider={"V2E": self.grid.get_v2e_offset_provider()}, + ) # 2. HALO EXCHANGE -- CALL sync_patch_array_mult @@ -569,7 +592,7 @@ def run( # 0c. dtime dependent stuff: enh_smag_factor, ~~r_dtimensubsteps~~ scale_k(self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={}) - mo_nh_diffusion_stencil_01( + _mo_nh_diffusion_stencil_01( self.diff_multfac_smag, tangent_orientation, inverse_primal_edge_lengths, @@ -582,11 +605,7 @@ def run( dual_normal_vert[1], prognostic_state.normal_wind, self.smag_limit, - self.kh_smag_e, - self.kh_smag_ec, - self.z_nabla2_e, self.smag_offset, - dtime, domain={ KDim: (0, klevels), EdgeDim: ( @@ -597,74 +616,80 @@ def run( ) ), }, + out=(self.kh_smag_e, self.kh_smag_ec, self.z_nabla2_e), offset_provider={"E2C2V": self.grid.get_e2c2v_connectivity()}, - ) + backend = gtfn_cpu.run_gtfn + ) - fused_mo_nh_diffusion_stencil_02_03( + _fused_mo_nh_diffusion_stencil_02_03( self.kh_smag_ec, prognostic_state.normal_wind, interpolation_state.e_bln_c_s, interpolation_state.geofac_div, self.diff_multfac_smag, metric_state.wgtfac_c, - diagnostic_state.div_ic, - diagnostic_state.hdef_ic, + out=( + diagnostic_state.div_ic, + diagnostic_state.hdef_ic, + ), domain={ KDim: (0, klevels), CellDim: ( - self.grid.get_indices_from_to( - HorizontalMarkerIndex._NUDGING_CELLS, - HorizontalMarkerIndex._HALO_CELLS, + self.grid.get_indices_from_to(CellDim, + HorizontalMarkerIndex.nudging(CellDim), + HorizontalMarkerIndex.halo(CellDim) ) ), }, offset_provider={"C2E": self.grid.get_c2e_connectivity()}, + backend=gtfn_cpu.run_gtfn ) # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH # 5. CALL rbf_vec_interpol_vertex_wp - mo_intp_rbf_rbf_vec_interpol_vertex( + _mo_intp_rbf_rbf_vec_interpol_vertex( self.z_nabla2_e, interpolation_state.rbf_coeff_1, interpolation_state.rbf_coeff_2, - self.u_vert, - self.v_vert, - domain={ - KDim: (0, klevels), - VertexDim: self.grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 3, - HorizontalMarkerIndex.halo(VertexDim), - ), - }, + out=(self.u_vert, self.v_vert), + # domain={ + # KDim: (0, klevels), + # VertexDim: self.grid.get_indices_from_to( + # VertexDim, + # HorizontalMarkerIndex.local_boundary(VertexDim) + 3, + # HorizontalMarkerIndex.halo(VertexDim), + # ), + # }, offset_provider={"V2E": self.grid.get_e2v_connectivity()}, ) # 6. HALO EXCHANGE -- CALL sync_patch_array_mult # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 - # TODO where to get index field from - # horz_idx: Field[[EdgeDim], int32], - # start_2nd_nudge_line_idx_e: int32, - - fused_mo_nh_diffusion_stencil_04_05_06( - u_vert=self.u_vert, - v_vert=self.v_vert, - primal_normal_vert_v1=primal_normal_vert.x, - primal_normal_vert_v2=primal_normal_vert.y, - z_nabla2_e=self.z_nabla2_e, - inv_vert_vert_length=inverse_vertical_vertex_lengths, - inv_primal_edge_length=inverse_primal_edge_lengths, - area_edges=edge_areas, - kh_smag_e=self.kh_smag_e, - diff_multfac_vn=self.diff_multfac_vn, - nudgecoeff_e=interpolation_state.nudgecoeff_e, - vn=prognostic_state.normal_wind, - horz_idx=None, - nudgezone_diff=self.nudgezone_diff, - fac_bdydiff_v=self.fac_bdydiff_v, - start_2nd_nudge_line_idx_e=None, + + start_2nd_nudge_line = self.grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.nudging(EdgeDim)-1, + HorizontalMarkerIndex.nudging(EdgeDim)-1)[0] + _fused_mo_nh_diffusion_stencil_04_05_06( + self.u_vert, + self.v_vert, + primal_normal_vert.x, + primal_normal_vert.y, + self.z_nabla2_e, + inverse_vertical_vertex_lengths, + inverse_primal_edge_lengths, + edge_areas, + self.kh_smag_e, + self.diff_multfac_vn, + interpolation_state.nudgecoeff_e, + prognostic_state.normal_wind, + self.horizontal_edge_index, + self.nudgezone_diff, + self.fac_bdydiff_v, + start_2nd_nudge_line, + out=prognostic_state.normal_wind, domain={ KDim: (0, klevels), EdgeDim: self.grid.get_indices_from_to( @@ -677,23 +702,28 @@ def run( ) # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 - - fused_mo_nh_diffusion_stencil_07_08_09_10( - area=cell_areas, - geofac_n2s=interpolation_state.geofac_n2s, - geofac_grg_x=interpolation_state.geofac_grg_x, - geofac_grg_y=interpolation_state.geofac_grg_y, - w_old=prognostic_state.vertical_wind, - w=prognostic_state.vertical_wind, - dwdx=diagnostic_state.dwdx, - dwdy=diagnostic_state.dwdy, - diff_multfac_w=self.diff_multfac_w, - diff_multfac_n2w=self.diff_multfac_n2w, - vert_idx=self.vertical_index, - horz_idx=self.horizontal_cell_index, - nrdmax=self.config.vertical_params.index_of_damping_height, - interior_idx=HorizontalMarkerIndex.interior(CellDim), - halo_idx=HorizontalMarkerIndex.halo(CellDim), + interior_start_index, halo_endindex = self.grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.interior(CellDim), + HorizontalMarkerIndex.halo(CellDim)) + + _fused_mo_nh_diffusion_stencil_07_08_09_10( + cell_areas, + interpolation_state.geofac_n2s, + interpolation_state.geofac_grg_x, + interpolation_state.geofac_grg_y, + prognostic_state.vertical_wind, + prognostic_state.vertical_wind, + diagnostic_state.dwdx, + diagnostic_state.dwdy, + self.diff_multfac_w, + self.diff_multfac_n2w, + self.vertical_index, + self.horizontal_cell_index, + self.config.vertical_params.index_of_damping_height, + interior_start_index, + halo_endindex, + out=(prognostic_state.vertical_wind, diagnostic_state.dwdx, diagnostic_state.dwdy), domain={ KDim: (0, klevels), CellDim: self.grid.get_indices_from_to( @@ -711,10 +741,10 @@ def run( # TODO check: kh_smag_e is an out field, should not be calculated in init? fused_mo_nh_diffusion_stencil_11_12( - theta_v=prognostic_state.theta_v, - theta_ref_mc=metric_state.theta_ref_mc, - thresh_tdiff=self.thresh_tdiff, - kh_smag_e=self.kh_smag_e, + prognostic_state.theta_v, + metric_state.theta_ref_mc, + self.thresh_tdiff, + out=self.kh_smag_e, domain={ KDim: (klevels - 2, klevels), CellDim: self.grid.get_indices_from_to( @@ -734,12 +764,12 @@ def run( }, ) - fused_mo_nh_diffusion_stencil_13_14( - kh_smag_e=self.kh_smag_e, - inv_dual_edge_length=inverse_dual_edge_length, - theta_v=prognostic_state.theta_v, - geofac_div=interpolation_state.geofac_div, - z_temp=self.z_temp, + _fused_mo_nh_diffusion_stencil_13_14( + self.kh_smag_e, + inverse_dual_edge_length, + prognostic_state.theta_v, + interpolation_state.geofac_div, + out = self.z_temp, domain={ KDim: (0, klevels), CellDim: self.grid.get_indices_from_to( @@ -761,7 +791,7 @@ def run( mo_nh_diffusion_stencil_15_numpy( mask_hdiff=metric_state.mask_hdiff, - zd_vertidx=metric_state.zd_vertidx_dsl, + zd_vertidx=metric_state.zd_vertidx, vcoef=metric_state.zd_diffcoef, zd_diffcoef=metric_state.zd_diffcoef, geofac_n2s=interpolation_state.geofac_n2s, @@ -771,19 +801,20 @@ def run( KDim: (0, klevels), CellDim: self.grid.get_indices_from_to( CellDim, - HorizontalMarkerIndex.nudging(CellDim).value, + HorizontalMarkerIndex.nudging(CellDim), HorizontalMarkerIndex.halo(CellDim), ), }, offset_provider={}, ) - mo_nh_diffusion_stencil_16( - z_temp=self.z_temp, - area=cell_areas, - theta_v=prognostic_state.theta_v, - exner=prognostic_state.exner_pressure, - rd_o_cvd=self.rd_o_cvd, + _mo_nh_diffusion_stencil_16( + self.z_temp, + cell_areas, + prognostic_state.theta_v, + prognostic_state.exner_pressure, + self.rd_o_cvd, + out=(prognostic_state.theta_v, prognostic_state.exner_pressure), domain={ KDim: (0, klevels), CellDim: self.grid.get_indices_from_to( diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py index 1f9766257..ed3f19c1b 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -33,9 +33,10 @@ _mo_nh_diffusion_stencil_10, ) from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim +from functional.program_processors.runners import gtfn_cpu -@field_operator +@field_operator(backend=gtfn_cpu.run_gtfn) def _fused_mo_nh_diffusion_stencil_07_08_09_10( area: Field[[CellDim], float], geofac_n2s: Field[[CellDim, C2E2CODim], float], @@ -65,14 +66,6 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( _mo_nh_diffusion_stencil_08(w_old, geofac_grg_x, geofac_grg_y), (dwdx, dwdy), ) - num_lev = 10 - - # _mo_nh_diffusion_stencil_08( - # w_old, - # geofac_grg_x, - # geofac_grg_y, - # out=(dwdx, dwdy), - # domain={KDim: (0, num_lev)}) z_nabla2_c = _mo_nh_diffusion_stencil_07(w_old, geofac_n2s) @@ -84,16 +77,6 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( w_old, ) - # _mo_nh_diffusion_stencil_07( - # area, - # z_nabla2_c, - # geofac_n2s, - # w_old, - # diff_multfac_w, - # out=w_old, - # domain={CellDim: (interior_idx, halo_idx)} - # ) - w = where( (vert_idx > int32(0)) & (vert_idx < nrdmax) @@ -102,20 +85,10 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, area, z_nabla2_c), w, ) - - # _mo_nh_diffusion_stencil_10( - # w_old, diff_multfac_n2w, - # area, - # z_nabla2_c, - # out=w_old, - # domain={CellDim: (interior_idx, halo_idx), KDim:(0, num_lev)} - # ) - - w = w_old return w, dwdx, dwdy -@program + def fused_mo_nh_diffusion_stencil_07_08_09_10( area: Field[[CellDim], float], geofac_n2s: Field[[CellDim, C2E2CODim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py index f893a7cbb..120d38890 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py @@ -66,6 +66,7 @@ def local_boundary(cls, dim: Dimension) -> int: return cls._LOCAL_BOUNDARY_VERTICES @classmethod + #TODO should be renamed "halo" in dusk means "after this the halo starts", so is essentially local def halo(cls, dim: Dimension) -> int: match (dim): case (dimension.CellDim): diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py index 0c6dc2701..0eac7acb5 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py @@ -138,29 +138,15 @@ def __init__(self, vct_a: np.ndarray, rayleigh_damping_height: float = 12500.0): """ self.rayleigh_damping_height = rayleigh_damping_height self.vct_a = vct_a - # TODO klevels in ICON are inverted! + # TODO klevels in ICON are inverted! TODO test against ICON LOG.exp... self.index_of_damping_height = np.argmax( - np.asarray(self.vct_a) >= self.rayleigh_damping_height + np.where(np.asarray(self.vct_a) >= self.rayleigh_damping_height) ) + def get_index_of_damping_layer(self): return self.index_of_damping_height def get_physical_heights(self) -> Field[[KDim], float]: return np_as_located_field(KDim)(self.vct_a) - def init_nabla2_factor_in_upper_damping_zone( - self, k_size: int - ) -> Field[[KDim], float]: - # this assumes n_shift == 0 - buffer = np.zeros(k_size) - buffer[2 : self.index_of_damping_height] = ( - 1.0 - / 12.0 - * ( - self.vct_a[2 : self.index_of_damping_height] - - self.vct_a[self.index_of_damping_height + 1] - ) - / (self.vct_a[2] - self.vct_a[self.index_of_damping_height + 1]) - ) - return np_as_located_field(KDim)(buffer) diff --git a/atm_dyn_iconam/tests/icon_grid_test_utils.py b/atm_dyn_iconam/tests/icon_grid_test_utils.py index 90999532e..c186452f8 100644 --- a/atm_dyn_iconam/tests/icon_grid_test_utils.py +++ b/atm_dyn_iconam/tests/icon_grid_test_utils.py @@ -16,16 +16,17 @@ import numpy as np from _pytest.fixtures import fixture +from icon4py.atm_dyn_iconam.diffusion import DiffusionConfig from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig -from icon4py.atm_dyn_iconam.icon_grid import IconGrid, MeshConfig +from icon4py.atm_dyn_iconam.icon_grid import IconGrid, MeshConfig, VerticalModelParams from icon4py.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.testutils.serialbox_utils import IconSerialDataProvider @fixture -def with_grid(): +def with_icon_grid(): data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") - sp = IconSerialDataProvider("icon_diffusion_init", data_path, True).from_savepoint( + sp = IconSerialDataProvider("icon_diffusion_init", data_path, True).from_savepoint_init( linit=True, date="2021-06-20T12:00:10.000" ) @@ -38,13 +39,22 @@ def with_grid(): edge_starts = sp.edge_start_index() edge_ends = sp.edge_end_index() + # config = MeshConfig( + # HorizontalMeshConfig( + # num_vertices=sp_meta["num_vert"], + # num_cells=sp_meta["num_cells"], + # num_edges=sp_meta["num_edges"], + # ) + # ) + config = MeshConfig( HorizontalMeshConfig( - num_vertices=sp_meta["num_vert"], - num_cells=sp_meta["num_cells"], - num_edges=sp_meta["num_edges"], + num_vertices=sp_meta["nproma"], + num_cells=sp_meta["nproma"], + num_edges=sp_meta["nproma"], ) ) + c2e2c = sp.c2e2c() c2e2c0 = np.column_stack((c2e2c, (np.asarray(range(c2e2c.shape[0]))))) grid = ( @@ -59,5 +69,28 @@ def with_grid(): .with_connectivity(e2v=sp.e2v()) .with_connectivity(c2e2c0=c2e2c0) .with_connectivity(v2e=sp.v2e()) + .with_connectivity(e2v=sp.e2v()) ) return grid + +@fixture +def with_r04b09_diffusion_config()->DiffusionConfig: + """ + Create DiffusionConfig. + + that uses the parameters of MCH.CH_r04b09_dsl experiment + """ + data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") + sp = IconSerialDataProvider("icon_diffusion_init", data_path, True).from_savepoint_init( + linit=True, date="2021-06-20T12:00:10.000" + ) + nproma = sp.get_metadata("nproma")["nproma"] + horizontalConfig = HorizontalMeshConfig( + num_vertices=nproma, num_cells=nproma, num_edges=nproma + ) + + grid = IconGrid().with_config(MeshConfig(horizontalMesh=horizontalConfig)) + verticalParams = VerticalModelParams( + rayleigh_damping_height=12500, vct_a=sp.vct_a() + ) + return DiffusionConfig(grid=grid, vertical_params=verticalParams) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 26851e670..16d7de8e3 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -14,7 +14,10 @@ import numpy as np import pytest -from icon_grid_test_utils import with_grid + +from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import \ + mo_intp_rbf_rbf_vec_interpol_vertex +from icon_grid_test_utils import with_icon_grid, with_r04b09_diffusion_config from icon4py.atm_dyn_iconam.diagnostic import DiagnosticState from icon4py.atm_dyn_iconam.diffusion import ( @@ -126,8 +129,8 @@ def test_set_zero_vertex_k(): assert np.allclose(0.0, f) -def test_diffusion_coefficients_with_hdiff_efdt_ratio(): - config: DiffusionConfig = DiffusionConfig.create_with_defaults() +def test_diffusion_coefficients_with_hdiff_efdt_ratio(with_r04b09_diffusion_config): + config = with_r04b09_diffusion_config config.hdiff_efdt_ratio = 1.0 params = DiffusionParams(config) @@ -138,8 +141,8 @@ def test_diffusion_coefficients_with_hdiff_efdt_ratio(): assert params.K4W == pytest.approx(0.125 / 4.0, abs=1e-12) -def test_diffusion_coefficients_without_hdiff_efdt_ratio(): - config: DiffusionConfig = DiffusionConfig.create_with_defaults() +def test_diffusion_coefficients_without_hdiff_efdt_ratio(with_r04b09_diffusion_config): + config = with_r04b09_diffusion_config config.hdiff_efdt_ratio = 0.0 params = DiffusionParams(config) @@ -150,8 +153,8 @@ def test_diffusion_coefficients_without_hdiff_efdt_ratio(): assert params.K4W == 0.0 -def test_smagorinski_factor_for_diffusion_type_4(): - config: DiffusionConfig = DiffusionConfig.create_with_defaults() +def test_smagorinski_factor_for_diffusion_type_4(with_r04b09_diffusion_config): + config = with_r04b09_diffusion_config config.hdiff_smag_fac = 0.15 config.diffusion_type = 4 @@ -161,8 +164,8 @@ def test_smagorinski_factor_for_diffusion_type_4(): assert params.smagorinski_height is None -def test_smagorinski_heights_diffusion_type_5_are_consistent(): - config: DiffusionConfig = DiffusionConfig.create_with_defaults() +def test_smagorinski_heights_diffusion_type_5_are_consistent(with_r04b09_diffusion_config): + config = with_r04b09_diffusion_config config.hdiff_smag_fac = 0.15 config.diffusion_type = 5 @@ -176,30 +179,27 @@ def test_smagorinski_heights_diffusion_type_5_are_consistent(): assert params.smagorinski_height[2] != params.smagorinski_height[3] -def test_smagorinski_factor_diffusion_type_5(): - config: DiffusionConfig = DiffusionConfig.create_with_defaults() - config.hdiff_smag_fac = 0.15 - config.diffusion_type = 5 - +def test_smagorinski_factor_diffusion_type_5(with_r04b09_diffusion_config): + config = with_r04b09_diffusion_config params = DiffusionParams(config) assert len(params.smagorinski_factor) == len(params.smagorinski_height) assert len(params.smagorinski_factor) == 4 assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor))) -def test_diffusion_init(): +def test_diffusion_init(with_r04b09_diffusion_config): + config = with_r04b09_diffusion_config data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") serializer = IconSerialDataProvider("icon_diffusion_init", data_path) serializer.print_info() first_run_date = "2021-06-20T12:00:10.000" - savepoint = serializer.from_savepoint(linit=False, date=first_run_date) - vct_a = savepoint.physical_height_field() + savepoint = serializer.from_savepoint_init(linit=False, date=first_run_date) + vct_a = savepoint.vct_a() meta = savepoint.get_metadata("nlev", "linit", "date") assert meta["nlev"] == 65 assert meta["linit"] is False assert meta["date"] == first_run_date - config = DiffusionConfig.create_with_defaults() additional_parameters = DiffusionParams(config) diffusion = Diffusion(config, additional_parameters, vct_a) @@ -214,11 +214,14 @@ def test_diffusion_init(): assert np.allclose(0.0, np.asarray(diffusion.v_vert)) assert np.allclose(0.0, np.asarray(diffusion.u_vert)) - assert np.allclose(0.0, np.asarray(diffusion.diff_multfac_n2w)) + #TODO test against ICON serialized + #assert np.allclose(0.0, np.asarray(diffusion.diff_multfac_n2w)) assert np.allclose(0.0, np.asarray(diffusion.kh_smag_ec)) assert np.allclose(0.0, np.asarray(diffusion.kh_smag_e)) shape_k = np.asarray(diffusion.diff_multfac_vn.shape) + + expected_smag_limit = smag_limit_numpy( shape_k, additional_parameters.K4, config.substep_as_float() ) @@ -237,15 +240,15 @@ def test_diffusion_init(): @pytest.mark.xfail -def test_diffusion_run(with_grid): +def test_diffusion_run(with_icon_grid): data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") data_provider = IconSerialDataProvider("icon_diffusion_init", data_path) data_provider.print_info() - sp = data_provider.from_savepoint(linit=False, date="2021-06-20T12:00:10.000") - vct_a = sp.physical_height_field() + sp = data_provider.from_savepoint_init(linit=False, date="2021-06-20T12:00:10.000") + vct_a = sp.vct_a() config = DiffusionConfig( - with_grid, + with_icon_grid, vertical_params=VerticalModelParams( vct_a=vct_a, rayleigh_damping_height=12500.0 ), @@ -259,24 +262,24 @@ def test_diffusion_run(with_grid): hdef_ic=sp.hdef_ic(), div_ic=sp.div_ic(), dwdx=sp.dwdx(), dwdy=sp.dwdy() ) prognostic_state = PrognosticState( - vertical_wind=sp.vn(), - normal_wind=sp.w(), + vertical_wind=sp.w(), + normal_wind=sp.vn(), exner_pressure=sp.exner(), theta_v=sp.theta_v(), ) grg = sp.geofac_grg() - rbf_coef = sp.rbf_coeff() interpolation_state = InterpolationState( e_bln_c_s=sp.e_bln_c_s(), - rbf_coeff_1=rbf_coef[0], - rbf_coeff_2=sp.rbf_coeff[1], + rbf_coeff_1=sp.rbf_vec_coeff_v1(), + rbf_coeff_2=sp.rbf_vec_coeff_v2(), geofac_div=sp.geofac_div(), geofac_n2s=sp.geofac_n2s(), geofac_grg_x=grg[0], geofac_grg_y=grg[1], nudgecoeff_e=sp.nudgecoeff_e(), ) + metric_state = MetricState( mask_hdiff=sp.theta_ref_mc(), theta_ref_mc=sp.theta_ref_mc(), @@ -302,20 +305,32 @@ def test_diffusion_run(with_grid): edge_areas = sp.edge_areas() cell_areas = sp.cell_areas() - diffusion.run( - diagnostic_state=diagnostic_state, - prognostic_state=prognostic_state, - metric_state=metric_state, - interpolation_state=interpolation_state, - dtime=dtime, - tangent_orientation=orientation, - inverse_primal_edge_lengths=inverse_primal_edge_lengths, - inverse_dual_edge_length=inverse_dual_edge_length, - inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, - primal_normal_vert=primal_normal_vert, - dual_normal_vert=dual_normal_vert, - edge_areas=edge_areas, - cell_areas=cell_areas, - ) - pytest.fail("not implemented yet") + # diffusion.run( + # diagnostic_state=diagnostic_state, + # prognostic_state=prognostic_state, + # metric_state=metric_state, + # interpolation_state=interpolation_state, + # dtime=dtime, + # tangent_orientation=orientation, + # inverse_primal_edge_lengths=inverse_primal_edge_lengths, + # inverse_dual_edge_length=inverse_dual_edge_length, + # inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, + # primal_normal_vert=primal_normal_vert, + # dual_normal_vert=dual_normal_vert, + # edge_areas=edge_areas, + # cell_areas=cell_areas, + # ) + + exit_savepoint = data_provider.from_save_point_exit(linit=False, date="2021-06-20T12:00:10.000") + icon_result_exner = exit_savepoint.exner() + icon_result_vn = exit_savepoint.vn() + icon_result_w = exit_savepoint.w() + icon_result_theta_w = exit_savepoint.theta_v() + + assert np.allclose(icon_result_w, np.asarray(prognostic_state.vertical_wind)) + assert np.allclose(np.asarray(icon_result_vn), np.asarray(prognostic_state.normal_wind)) + assert np.allclose(np.asarray(icon_result_theta_w), np.asarray(prognostic_state.theta_v)) + assert np.allclose(np.asarray(icon_result_exner), np.asarray(prognostic_state.exner_pressure)) + + pytest.fail("not finished yet") diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index 8f8fd7c78..262d63128 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -10,227 +10,227 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -from icon_grid_test_utils import with_grid +from icon_grid_test_utils import with_icon_grid from icon4py.atm_dyn_iconam.horizontal import HorizontalMarkerIndex from icon4py.common.dimension import CellDim, EdgeDim, VertexDim -def test_horizontal_grid_cell_indices(with_grid): - assert with_grid.get_indices_from_to( +def test_horizontal_grid_cell_indices(with_icon_grid): + assert with_icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.halo(CellDim) - 1, HorizontalMarkerIndex.halo(CellDim) - 1, ) == ( - 20897, + 20896, 20896, ) # halo +1 - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.halo(CellDim), HorizontalMarkerIndex.halo(CellDim), ) == ( - 0, + -1, 20896, ) # halo in icon is (1,20896) why - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.interior(CellDim), HorizontalMarkerIndex.interior(CellDim), ) == ( - 4105, + 4104, 20896, ) # interior - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, HorizontalMarkerIndex._INTERIOR_CELLS + 1, ) == ( - 1, + 0, 850, ) # lb+1 - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.local_boundary(CellDim) + 1, HorizontalMarkerIndex.local_boundary(CellDim) + 1, - ) == (851, 1688) - assert with_grid.get_indices_from_to( + ) == (850, 1688) + assert with_icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.local_boundary(CellDim) + 2, HorizontalMarkerIndex.local_boundary(CellDim) + 2, ) == ( - 1689, + 1688, 2511, ) # lb+2 - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.nudging(CellDim) - 1, HorizontalMarkerIndex.nudging(CellDim) - 1, ) == ( - 2512, + 2511, 3316, ) # lb+3 - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.nudging(CellDim), HorizontalMarkerIndex.nudging(CellDim), ) == ( - 3317, + 3316, 4104, ) # nudging -def test_horizontal_edge_indices(with_grid): - assert with_grid.get_indices_from_to( +def test_horizontal_edge_indices(with_icon_grid): + assert with_icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), HorizontalMarkerIndex.interior(EdgeDim), - ) == (6177, 31558) - assert with_grid.get_indices_from_to( + ) == (6176, 31558) + assert with_icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.halo(EdgeDim) - 2, HorizontalMarkerIndex.halo(EdgeDim) - 2, - ) == (31559, 31558) - assert with_grid.get_indices_from_to( + ) == (31558, 31558) + assert with_icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.halo(EdgeDim) - 1, HorizontalMarkerIndex.halo(EdgeDim) - 1, - ) == (31559, 31558) - assert with_grid.get_indices_from_to( + ) == (31558, 31558) + assert with_icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.halo(EdgeDim), HorizontalMarkerIndex.halo(EdgeDim), ) == ( - 0, + -1, 31558, ) # halo in icon is (1, 31558) - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, HorizontalMarkerIndex.nudging(EdgeDim) + 1, ) == ( - 5388, + 5387, 6176, ) # nudging +1 - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), HorizontalMarkerIndex.nudging(EdgeDim), ) == ( - 4990, + 4989, 5387, ) # nudging - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local_boundary(EdgeDim) + 7, HorizontalMarkerIndex.local_boundary(EdgeDim) + 7, ) == ( - 4185, + 4184, 4989, ) # lb +7 - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local_boundary(EdgeDim) + 6, HorizontalMarkerIndex.local_boundary(EdgeDim) + 6, ) == ( - 3778, + 3777, 4184, ) # lb +6 - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local_boundary(EdgeDim) + 5, HorizontalMarkerIndex.local_boundary(EdgeDim) + 5, ) == ( - 2955, + 2954, 3777, ) # lb +5 - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, ) == ( - 2539, + 2538, 2954, ) # lb +4 - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local_boundary(EdgeDim) + 3, HorizontalMarkerIndex.local_boundary(EdgeDim) + 3, ) == ( - 1701, + 1700, 2538, ) # lb +3 - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local_boundary(EdgeDim) + 2, HorizontalMarkerIndex.local_boundary(EdgeDim) + 2, ) == ( - 1279, + 1278, 1700, ) # lb +2 - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local_boundary(EdgeDim) + 1, HorizontalMarkerIndex.local_boundary(EdgeDim) + 1, ) == ( - 429, + 428, 1278, ) # lb +1 - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local_boundary(EdgeDim), HorizontalMarkerIndex.local_boundary(EdgeDim), ) == ( - 1, + 0, 428, ) # lb +0 -def test_horizontal_vertex_indices(with_grid): - assert with_grid.get_indices_from_to( +def test_horizontal_vertex_indices(with_icon_grid): + assert with_icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.end(VertexDim), HorizontalMarkerIndex.end(VertexDim), - ) == (10664, 10663) - assert with_grid.get_indices_from_to( + ) == (10663, 10663) + assert with_icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.halo(VertexDim), HorizontalMarkerIndex.halo(VertexDim), - ) == (0, 10663) - assert with_grid.get_indices_from_to( + ) == (-1, 10663) + assert with_icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.halo(VertexDim) - 1, HorizontalMarkerIndex.halo(VertexDim) - 1, - ) == (10664, 10663) + ) == (10663, 10663) - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim), HorizontalMarkerIndex.local_boundary(VertexDim), - ) == (1, 428) - assert with_grid.get_indices_from_to( + ) == (0, 428) + assert with_icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 1, HorizontalMarkerIndex.local_boundary(VertexDim) + 1, - ) == (429, 850) - assert with_grid.get_indices_from_to( + ) == (428, 850) + assert with_icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 2, HorizontalMarkerIndex.local_boundary(VertexDim) + 2, - ) == (851, 1266) - assert with_grid.get_indices_from_to( + ) == (850, 1266) + assert with_icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 3, HorizontalMarkerIndex.local_boundary(VertexDim) + 3, - ) == (1267, 1673) - assert with_grid.get_indices_from_to( + ) == (1266, 1673) + assert with_icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 4, HorizontalMarkerIndex.local_boundary(VertexDim) + 4, - ) == (1674, 2071) + ) == (1673, 2071) - assert with_grid.get_indices_from_to( + assert with_icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.interior(VertexDim), HorizontalMarkerIndex.interior(VertexDim), - ) == (2072, 10663) + ) == (2071, 10663) diff --git a/atm_dyn_iconam/tests/test_vertical.py b/atm_dyn_iconam/tests/test_vertical.py index ab4c8f292..80fd3b3ba 100644 --- a/atm_dyn_iconam/tests/test_vertical.py +++ b/atm_dyn_iconam/tests/test_vertical.py @@ -12,19 +12,35 @@ # SPDX-License-Identifier: GPL-3.0-or-later import math +import os import numpy as np import pytest from icon4py.atm_dyn_iconam.icon_grid import VerticalModelParams +from icon4py.testutils.serialbox_utils import IconSerialDataProvider +from icon_grid_test_utils import with_icon_grid @pytest.mark.parametrize( "max_h,damping,delta", - [(60000, 34000, 612), (12000, 10000, 100), (109000, 45000, 123)], + [(60000, 34000, 612), (12000, 10000, 100), (109050, 45000, 123)], ) -# TODO [ml] klevels run from num_lev (ground) to 1 (top most) def test_nrdmax_calculation(max_h, damping, delta): vct_a = np.arange(0, max_h, delta) + vct_a = vct_a[::-1] vertical_params = VerticalModelParams(rayleigh_damping_height=damping, vct_a=vct_a) - assert vertical_params.get_index_of_damping_layer() == math.ceil(damping / delta) + assert vertical_params.get_index_of_damping_layer() == vct_a.shape[0] - math.ceil(damping / delta) -1 + +def test_nrdmax_calculation_from_icon_input(with_icon_grid): + data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") + sp = IconSerialDataProvider("icon_diffusion_init", data_path, True).from_savepoint_init( + linit=True, date="2021-06-20T12:00:10.000" + ) + a = sp.vct_a() + damping_height = 12500 + vertical_params = VerticalModelParams(rayleigh_damping_height=damping_height, vct_a=a) + assert 9 == vertical_params.get_index_of_damping_layer() + a_array = np.asarray(a) + assert a_array[9] > damping_height + assert a_array[10] < damping_height diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index a3008ba9d..421326839 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -41,18 +41,13 @@ ) import serialbox as ser - -class IconDiffusionSavepoint: +class IconDiffustionSavepoint: def __init__(self, sp: Savepoint, ser: ser.Serializer): self.savepoint = sp self.serializer = ser - def print_meta_info(self): print(self.savepoint.metainfo) - def physical_height_field(self): - return self._get_field("vct_a", KDim) - def _get_field(self, name, *dimensions): buffer = np.squeeze(self.serializer.read(name, self.savepoint)) print(f"{name} {buffer.shape}") @@ -62,6 +57,12 @@ def get_metadata(self, *names): metadata = self.savepoint.metainfo.to_dict() return {n: metadata[n] for n in names if n in metadata} +class IconDiffusionInitSavepoint(IconDiffustionSavepoint): + + def vct_a(self): + return self._get_field("vct_a", KDim) + + def tangent_orientation(self): return self._get_field("tangent_orientation", EdgeDim) @@ -93,21 +94,30 @@ def inv_dual_edge_length(self): return self._get_field("inv_dual_edge_length", EdgeDim) def cells_start_index(self): - return self.serializer.read("c_start_index", self.savepoint) + #subtract 1 for python 0 based indexing + return self.serializer.read("c_start_index", self.savepoint) - 1 def cells_end_index(self): + # don't need to subtract 1, because FORTRAN slices are inclusive [from:to] so the being + # one off accounts for being exclusive [from:to) return self.serializer.read("c_end_index", self.savepoint) def vertex_start_index(self): - return self.serializer.read("v_start_index", self.savepoint) + + return self.serializer.read("v_start_index", self.savepoint) - 1 def vertex_end_index(self): + # don't need to subtract 1, because FORTRAN slices are inclusive [from:to] so the being + # one off accounts for being exclusive [from:to) return self.serializer.read("v_end_index", self.savepoint) def edge_start_index(self): - return self.serializer.read("e_start_index", self.savepoint) + # subtract 1 for python 0 based indexing + return self.serializer.read("e_start_index", self.savepoint) - 1 def edge_end_index(self): + # don't need to subtract 1, because FORTRAN slices are inclusive [from:to] so the being + # one off accounts for being exclusive [from:to) return self.serializer.read("e_end_index", self.savepoint) def refin_ctrl(self, dim: Dimension): @@ -121,16 +131,24 @@ def refin_ctrl(self, dim: Dimension): return None def c2e(self): - return self.serializer.read("c2e", self.savepoint) + #subtract 1 to account for python being 0 based + return self.serializer.read("c2e", self.savepoint)-1 def c2e2c(self): - return self.serializer.read("c2e2c", self.savepoint) + # subtract 1 to account for python being 0 based + return self.serializer.read("c2e2c", self.savepoint)-1 def e2c(self): - return self.serializer.read("e2c", self.savepoint) + # subtract 1 to account for python being 0 based + return self.serializer.read("e2c", self.savepoint)-1 def e2v(self): - return self.serializer.read("e2v", self.savepoint) + # subtract 1 to account for python being 0 based + return self.serializer.read("e2v", self.savepoint)-1 + + def v2e(self): + # subtract 1 to account for python being 0 based + return self.serializer.read("v2e", self.savepoint)-1 def hdef_ic(self): return self._get_field("hdef_ic", CellDim, KDim) @@ -145,7 +163,7 @@ def dwdy(self): return self._get_field("dwdy", CellDim, KDim) def vn(self): - return self._get_field("vn", CellDim, KDim) + return self._get_field("vn", EdgeDim, KDim) def theta_v(self): return self._get_field("theta_v", CellDim, KDim) @@ -184,8 +202,8 @@ def geofac_n2s(self): return self._get_field("geofac_n2s", CellDim, C2E2CODim) def geofac_grg(self): - grg = self.serializer.read("geofac_grg", self.savepoint) - return grg[:, :, :, 0], grg[:, :, :, 1] + grg = np.squeeze(self.serializer.read("geofac_grg", self.savepoint)) + return np_as_located_field(CellDim, C2E2CODim)(grg[:, :, 0]), np_as_located_field(CellDim, C2E2CODim)(grg[:, :,1]) def nudgecoeff_e(self): return self._get_field("nudgecoeff_e", EdgeDim) @@ -193,16 +211,27 @@ def nudgecoeff_e(self): def zd_vertidx(self): return self._get_field("zd_vertidx", CellDim, C2E2CDim, KDim) - def rbf_coeff(self): - rbf_coef = self.serializer.read("rbf_vec_coeff_v", self.savepoint) - return ( - np_as_located_field(VertexDim, V2EDim)(rbf_coef[:, :, 0]), - np_as_located_field(VertexDim, V2EDim)(rbf_coef)[:, :, 1], - ) + def rbf_vec_coeff_v1(self): + return self._get_field("rbf_vec_coeff_v1", VertexDim, V2EDim) + + def rbf_vec_coeff_v2(self): + return self._get_field("rbf_vec_coeff_v2", VertexDim, V2EDim) + - def v2e(self): - return self.serializer.read("v2e", self.savepoint) +class IconDiffusionExitSavepoint(IconDiffustionSavepoint): + + def vn(self): + return self._get_field("x_vn", EdgeDim, KDim) + + def theta_v(self): + return self._get_field("x_theta_v", CellDim, KDim) + + def w(self): + return self._get_field("x_w", CellDim, KDim) + + def exner(self): + return self._get_field("x_exner", CellDim, KDim) class IconSerialDataProvider: def __init__(self, fname_prefix, path=".", do_print=False): @@ -225,11 +254,28 @@ def print_info(self): print(f"SAVEPOINTS: {self.serializer.savepoint_list()}") print(f"FIELDNAMES: {self.serializer.fieldnames()}") - def from_savepoint(self, linit: bool, date: str) -> IconDiffusionSavepoint: + def from_savepoint_init(self, linit: bool, date: str) -> IconDiffusionInitSavepoint: savepoint = ( self.serializer.savepoint["call-diffusion-init"] .linit[linit] .date[date] .as_savepoint() ) - return IconDiffusionSavepoint(savepoint, self.serializer) + return IconDiffusionInitSavepoint(savepoint, self.serializer) + + def from_save_point_exit(self, linit: bool, date: str) -> IconDiffusionExitSavepoint: + savepoint = ( + self.serializer.savepoint["call-diffusion-exit"] + .linit[linit] + .date[date] + .as_savepoint() + ) + return IconDiffusionExitSavepoint(savepoint, self.serializer) + + + + + + + + From f627d15c855ae82f680312c82499f9224f86156e Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 8 Dec 2022 17:34:16 +0100 Subject: [PATCH 024/263] add verification for fields calculated in init --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 67 ++++++----- .../src/icon4py/atm_dyn_iconam/horizontal.py | 9 +- .../src/icon4py/atm_dyn_iconam/icon_grid.py | 22 ++-- atm_dyn_iconam/tests/icon_grid_test_utils.py | 14 +-- atm_dyn_iconam/tests/test_diffusion.py | 107 +++++++++++++----- .../src/icon4py/testutils/serialbox_utils.py | 30 +++++ 6 files changed, 169 insertions(+), 80 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index fefbe5097..4a931dc2d 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -38,13 +38,8 @@ from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_13_14 import ( _fused_mo_nh_diffusion_stencil_13_14, ) -from icon4py.atm_dyn_iconam.horizontal import ( - HorizontalMarkerIndex, -) -from icon4py.atm_dyn_iconam.icon_grid import ( - IconGrid, - VerticalModelParams, -) +from icon4py.atm_dyn_iconam.horizontal import HorizontalMarkerIndex +from icon4py.atm_dyn_iconam.icon_grid import IconGrid, VerticalModelParams from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState from icon4py.atm_dyn_iconam.metric_state import MetricState from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( @@ -69,8 +64,8 @@ ) -DiffusionTupleVT = namedtuple("DiffusionParamVT", "v t") -VectorTuple = namedtuple("CartesianVectorTuple", "x y") +TupleVT = namedtuple("TupleVT", "v t") +VectorTuple = namedtuple("VectorTuple", "x y") # TODO [ml] initial RUN linit = TRUE @@ -286,15 +281,15 @@ def init_nabla2_factor_in_upper_damping_zone( physcial_heights: vector of physical heights [m] of the height levels """ buffer = np.zeros(k_size) - buffer[2 : nrdmax + 1] = ( + buffer[1 : nrdmax + 1] = ( 1.0 / 12.0 * ( ( - physical_heights[2 + nshift : nrdmax + 1 + nshift] + physical_heights[1 + nshift : nrdmax + 1 + nshift] - physical_heights[nshift + nrdmax + 1] ) - / (physical_heights[2] - physical_heights[nshift + nrdmax + 1]) + / (physical_heights[1] - physical_heights[nshift + nrdmax + 1]) ) ** 4 ) @@ -306,7 +301,7 @@ class DiffusionConfig: - encapsulates namelist parameters and derived parameters (for now) - currently we use the MCH r04b09_dsl experiment as constants here. These should + currently we use the MCH r04b09_dsl experiment as defaults here. These should be read from config and the default from mo_diffusion_nml.f90 set as defaults. TODO: [ml] read from config @@ -325,6 +320,7 @@ def __init__( compute_3d_smag_coeff: bool = False, temperature_discretization: int = 2, horizontal_efdt_ratio: float = 24.0, + horizontal_w_efdt_ratio: float = 15.0, smag_scaling_factor: float = 0.025, ): # TODO [ml]: move external stuff out: grid related stuff, other than diffusion namelists (see below @@ -349,6 +345,7 @@ def __init__( self.hdiff_efdt_ratio = ( horizontal_efdt_ratio # ! ratio of e-folding time to time step ) + self.hdiff_w_efdt_ratio = horizontal_w_efdt_ratio # ratio of e-folding time to time step for w diffusion (NH only) self.hdiff_smag_fac = ( smag_scaling_factor # ! scaling factor for Smagorinsky diffusion ) @@ -359,8 +356,8 @@ def __init__( self.ndyn_substeps = 5 # namelist gridref_nml - # denom_diffu_v = 150 ! denominator for lateral boundary diffusion of velocity - self.lateral_boundary_denominator = DiffusionTupleVT(v=200.0, t=135.0) + # default is v=200.0, t=135.0 + self.lateral_boundary_denominator = TupleVT(v=150.0, t=135.0) # namelist grid_nml self.l_limited_area = True @@ -388,8 +385,14 @@ def __init__(self, config: DiffusionConfig): else 0.0 ) self.K4: Final[float] = self.K2 / 8.0 - self.K8: Final[float] = self.K2 / 64.0 - self.K4W: Final[float] = self.K2 / 4.0 + self.K6: Final[float] = self.K2 / 64.0 + + self.K4W: Final[float] = ( + 1.0 / (config.hdiff_w_efdt_ratio * 36.0) + if config.hdiff_w_efdt_ratio > 0 + else 0.0 + ) + ( self.smagorinski_factor, self.smagorinski_height, @@ -435,6 +438,7 @@ def diffusion_type_5_smagorinski_factor(config: DiffusionConfig): def mo_nh_diffusion_stencil_15_numpy( + c2e2c, mask_hdiff: Field[[CellDim, KDim], int], zd_vertidx: Field[[CellDim, C2E2CDim, KDim], int], zd_diffcoef: Field[[CellDim, KDim], float], @@ -445,7 +449,12 @@ def mo_nh_diffusion_stencil_15_numpy( domain, offset_provider, ): - pass + + z_temp = np.sum( + geofac_n2s * vcoef * theta_v[c2e2c[zd_vertidx]] + + (1.0 - vcoef) + + theta_v[c2e2c[zd_vertidx + 1]] + ) class Diffusion: @@ -485,10 +494,8 @@ def __init__( ) # different for initial run!, through diff_multfac_vn - self.diff_multfac_vn = np_as_located_field(KDim)( - np.zeros(config.grid.k_levels()) - ) - self.smag_limit = np_as_located_field(KDim)(np.zeros(config.grid.k_levels())) + self.diff_multfac_vn = np_as_located_field(KDim)(np.zeros(config.grid.n_lev())) + self.smag_limit = np_as_located_field(KDim)(np.zeros(config.grid.n_lev())) init_diffusion_local_fields( params.K4, @@ -499,7 +506,7 @@ def __init__( ) self.enh_smag_fac = np_as_located_field(KDim)( - np.zeros(config.grid.k_levels(), float) + np.zeros(config.grid.n_lev(), float) ) enhanced_smagorinski_factor( *params.smagorinski_factor, @@ -510,26 +517,26 @@ def __init__( ) self.diff_multfac_n2w = init_nabla2_factor_in_upper_damping_zone( - k_size=config.grid.k_levels(), + k_size=config.grid.n_lev(), nshift=0, physical_heights=np.asarray(vct_a), nrdmax=self.config.vertical_params.index_of_damping_height, ) self.diff_multfac_smag = np_as_located_field(KDim)( - np.zeros(config.grid.k_levels()) + np.zeros(config.grid.n_lev()) ) - shape_vk = (config.grid.num_vertices(), config.grid.k_levels()) - shape_ck = (config.grid.num_cells(), config.grid.k_levels()) + shape_vk = (config.grid.num_vertices(), config.grid.n_lev()) + shape_ck = (config.grid.num_cells(), config.grid.n_lev()) self.u_vert = np_as_located_field(VertexDim, KDim)(np.zeros(shape_vk, float)) self.v_vert = np_as_located_field(VertexDim, KDim)(np.zeros(shape_vk, float)) - shape_ek = (config.grid.num_edges(), config.grid.k_levels()) + shape_ek = (config.grid.num_edges(), config.grid.n_lev()) allocate_ek = np_as_located_field(EdgeDim, KDim)(np.zeros(shape_ek, float)) self.kh_smag_e = allocate_ek self.kh_smag_ec = allocate_ek self.z_nabla2_e = allocate_ek self.z_temp = np_as_located_field(CellDim, KDim)(np.zeros(shape_ck, float)) self.vertical_index = np_as_located_field(KDim)( - np.arange(self.grid.k_levels() + 1) + np.arange(self.grid.n_lev() + 1) ) self.horizontal_cell_index = np_as_located_field(CellDim)( np.arange((shape_ck[0])) @@ -573,7 +580,7 @@ def run( if self.config.lhdiff_rcf else 1.0 / (dtime * self.config.substep_as_float()) ) - klevels = self.grid.k_levels() + klevels = self.grid.n_lev() # TODO is this needed? set_zero_v_k(self.u_vert, offset_provider={}) set_zero_v_k(self.v_vert, offset_provider={}) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py index a413b01d3..ab18f397e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py @@ -113,11 +113,14 @@ def __init__(self, num_vertices: int, num_edges: int, num_cells: int): self._num_edges = num_edges self._num_cells = num_cells - def get_num_vertices(self): + @property + def num_vertices(self): return self._num_vertices - def get_num_edges(self): + @property + def num_edges(self): return self._num_edges - def get_num_cells(self): + @property + def num_cells(self): return self._num_cells diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py index 9ba6c3f30..1e6fadd6c 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py @@ -38,14 +38,16 @@ def n_shift_total(self): return self._n_shift_total @property - def get_num_vertices(self): - return self._horizontal._num_vertices + def num_vertices(self): + return self._horizontal.num_vertices - def get_num_edges(self): - return self._horizontal._num_edges + @property + def num_edges(self): + return self._horizontal.num_edges - def get_num_cells(self): - return self._horizontal._num_cells + @property + def num_cells(self): + return self._horizontal.num_cells def builder(func): @@ -78,17 +80,17 @@ def with_start_end_indices( def with_connectivity(self, **connectivity): self.connectivities.update(**connectivity) - def k_levels(self): + def n_lev(self): return self.config.num_k_levels def num_cells(self): - return self.config.get_num_cells() + return self.config.num_cells def num_vertices(self): - return self.config.get_num_cells() + return self.config.num_vertices def num_edges(self): - return self.config.get_num_edges() + return self.config.num_edges def get_indices_from_to( self, dim: Dimension, start_marker: int, end_marker: int diff --git a/atm_dyn_iconam/tests/icon_grid_test_utils.py b/atm_dyn_iconam/tests/icon_grid_test_utils.py index 0f342c88f..f3aa67283 100644 --- a/atm_dyn_iconam/tests/icon_grid_test_utils.py +++ b/atm_dyn_iconam/tests/icon_grid_test_utils.py @@ -43,19 +43,11 @@ def with_icon_grid(): edge_starts = sp.edge_start_index() edge_ends = sp.edge_end_index() - # config = MeshConfig( - # HorizontalMeshConfig( - # num_vertices=sp_meta["num_vert"], - # num_cells=sp_meta["num_cells"], - # num_edges=sp_meta["num_edges"], - # ) - # ) - config = MeshConfig( HorizontalMeshConfig( - num_vertices=sp_meta["nproma"], - num_cells=sp_meta["nproma"], - num_edges=sp_meta["nproma"], + num_vertices=sp_meta["nproma"], # or rather "num_vert" + num_cells=sp_meta["nproma"], # or rather "num_cells" + num_edges=sp_meta["nproma"], # or rather "num_edges" ) ) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 64a7b4172..eaf785b0a 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -30,12 +30,12 @@ from icon4py.atm_dyn_iconam.icon_grid import VerticalModelParams from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState from icon4py.atm_dyn_iconam.metric_state import MetricState -from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( - mo_intp_rbf_rbf_vec_interpol_vertex, -) from icon4py.atm_dyn_iconam.prognostic import PrognosticState from icon4py.common.dimension import KDim, VertexDim -from icon4py.testutils.serialbox_utils import IconSerialDataProvider +from icon4py.testutils.serialbox_utils import ( + IconDiffusionInitSavepoint, + IconSerialDataProvider, +) from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import random_field, zero_field @@ -132,24 +132,26 @@ def test_set_zero_vertex_k(): def test_diffusion_coefficients_with_hdiff_efdt_ratio(with_r04b09_diffusion_config): config = with_r04b09_diffusion_config config.hdiff_efdt_ratio = 1.0 + config.hdiff_w_efdt_ratio = 2.0 params = DiffusionParams(config) assert params.K2 == pytest.approx(0.125, abs=1e-12) assert params.K4 == pytest.approx(0.125 / 8.0, abs=1e-12) - assert params.K8 == pytest.approx(0.125 / 64.0, abs=1e-12) - assert params.K4W == pytest.approx(0.125 / 4.0, abs=1e-12) + assert params.K6 == pytest.approx(0.125 / 64.0, abs=1e-12) + assert params.K4W == pytest.approx(1.0 / 72.0, abs=1e-12) def test_diffusion_coefficients_without_hdiff_efdt_ratio(with_r04b09_diffusion_config): config = with_r04b09_diffusion_config config.hdiff_efdt_ratio = 0.0 + config.hdiff_w_efdt_ratio = 0.0 params = DiffusionParams(config) assert params.K2 == 0.0 assert params.K4 == 0.0 - assert params.K8 == 0.0 + assert params.K6 == 0.0 assert params.K4W == 0.0 @@ -216,8 +218,6 @@ def test_diffusion_init(with_r04b09_diffusion_config): assert np.allclose(0.0, np.asarray(diffusion.v_vert)) assert np.allclose(0.0, np.asarray(diffusion.u_vert)) - # TODO test against ICON serialized - # assert np.allclose(0.0, np.asarray(diffusion.diff_multfac_n2w)) assert np.allclose(0.0, np.asarray(diffusion.kh_smag_ec)) assert np.allclose(0.0, np.asarray(diffusion.kh_smag_e)) @@ -240,6 +240,63 @@ def test_diffusion_init(with_r04b09_diffusion_config): assert np.allclose(expected_enh_smag_fac, np.asarray(diffusion.enh_smag_fac)) +def verify_init_values_against_savepoint( + savepoint: IconDiffusionInitSavepoint, diffusion: Diffusion +): + dtime = savepoint.get_metadata("dtime")["dtime"] + assert savepoint.nudgezone_diff() == diffusion.nudgezone_diff + assert savepoint.bdy_diff() == diffusion.bdy_diff + assert savepoint.fac_bdydiff_v() == diffusion.fac_bdydiff_v + assert savepoint.smag_offset() == diffusion.smag_offset + assert savepoint.diff_multfac_w() == diffusion.diff_multfac_w + + # this is done in diffusion.run(...) because it depends on the dtime + scale_k( + diffusion.enh_smag_fac, dtime, diffusion.diff_multfac_smag, offset_provider={} + ) + assert np.allclose(savepoint.diff_multfac_smag(), diffusion.diff_multfac_smag) + + assert np.allclose(savepoint.smag_limit(), diffusion.smag_limit) + assert np.allclose( + savepoint.diff_multfac_n2w(), np.asarray(diffusion.diff_multfac_n2w) + ) + assert np.allclose(savepoint.diff_multfac_vn(), diffusion.diff_multfac_vn) + + +def test_verify_diffusion_init_against_first_regular_savepoint( + with_r04b09_diffusion_config, +): + config = with_r04b09_diffusion_config + data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") + serializer = IconSerialDataProvider("icon_diffusion_init", data_path) + serializer.print_info() + first_run_date = "2021-06-20T12:00:10.000" + savepoint = serializer.from_savepoint_init(linit=False, date=first_run_date) + vct_a = savepoint.vct_a() + + additional_parameters = DiffusionParams(config) + diffusion = Diffusion(config, additional_parameters, vct_a) + + verify_init_values_against_savepoint(savepoint, diffusion) + + +def test_verify_diffusion_init_against_other_regular_savepoint( + with_r04b09_diffusion_config, +): + config = with_r04b09_diffusion_config + data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") + serializer = IconSerialDataProvider("icon_diffusion_init", data_path) + serializer.print_info() + run_date = "2021-06-20T12:00:50.000" + savepoint = serializer.from_savepoint_init(linit=False, date=run_date) + vct_a = savepoint.vct_a() + + additional_parameters = DiffusionParams(config) + diffusion = Diffusion(config, additional_parameters, vct_a) + + verify_init_values_against_savepoint(savepoint, diffusion) + + @pytest.mark.xfail def test_diffusion_run(with_icon_grid): data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") @@ -306,21 +363,21 @@ def test_diffusion_run(with_icon_grid): edge_areas = sp.edge_areas() cell_areas = sp.cell_areas() - # diffusion.run( - # diagnostic_state=diagnostic_state, - # prognostic_state=prognostic_state, - # metric_state=metric_state, - # interpolation_state=interpolation_state, - # dtime=dtime, - # tangent_orientation=orientation, - # inverse_primal_edge_lengths=inverse_primal_edge_lengths, - # inverse_dual_edge_length=inverse_dual_edge_length, - # inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, - # primal_normal_vert=primal_normal_vert, - # dual_normal_vert=dual_normal_vert, - # edge_areas=edge_areas, - # cell_areas=cell_areas, - # ) + diffusion.run( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + metric_state=metric_state, + interpolation_state=interpolation_state, + dtime=dtime, + tangent_orientation=orientation, + inverse_primal_edge_lengths=inverse_primal_edge_lengths, + inverse_dual_edge_length=inverse_dual_edge_length, + inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, + primal_normal_vert=primal_normal_vert, + dual_normal_vert=dual_normal_vert, + edge_areas=edge_areas, + cell_areas=cell_areas, + ) exit_savepoint = data_provider.from_save_point_exit( linit=False, date="2021-06-20T12:00:10.000" @@ -340,5 +397,3 @@ def test_diffusion_run(with_icon_grid): assert np.allclose( np.asarray(icon_result_exner), np.asarray(prognostic_state.exner_pressure) ) - - pytest.fail("not finished yet") diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 942abf39a..a299a88e1 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -220,6 +220,36 @@ def rbf_vec_coeff_v1(self): def rbf_vec_coeff_v2(self): return self._get_field("rbf_vec_coeff_v2", VertexDim, V2EDim) + def diff_multfac_smag(self): + return np.squeeze(self.serializer.read("diff_multfac_smag", self.savepoint)) + + def smag_limit(self): + return np.squeeze(self.serializer.read("smag_limit", self.savepoint)) + + def diff_multfac_n2w(self): + return np.squeeze(self.serializer.read("diff_multfac_n2w", self.savepoint)) + + def nrdmax(self): + return self.serializer.read("nrdmax", self.savepoint)[0] + + def nudgezone_diff(self) -> int: + return self.serializer.read("nudgezone_diff", self.savepoint)[0] + + def bdy_diff(self) -> int: + return self.serializer.read("bdy_diff", self.savepoint)[0] + + def fac_bdydiff_v(self) -> int: + return self.serializer.read("fac_bdydiff_v", self.savepoint)[0] + + def smag_offset(self): + return self.serializer.read("smag_offset", self.savepoint)[0] + + def diff_multfac_w(self): + return self.serializer.read("diff_multfac_w", self.savepoint)[0] + + def diff_multfac_vn(self): + return self.serializer.read("diff_multfac_vn", self.savepoint) + class IconDiffusionExitSavepoint(IconDiffustionSavepoint): def vn(self): From bf832cb64d95a2a2c76668a377c909a6bc57a927 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 9 Dec 2022 11:21:50 +0100 Subject: [PATCH 025/263] fuse diffusion initialization stencils --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 178 ++++++++---------- atm_dyn_iconam/tests/test_diffusion.py | 58 +++--- 2 files changed, 110 insertions(+), 126 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index 4a931dc2d..73aca6ab9 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -137,66 +137,9 @@ def _mo_nh_diffusion_stencil_01_scale_dtime( ) -@program -def mo_nh_diffusion_stencil_01_scaled_dtime( - enh_smag_fac: Field[[KDim], float], - tangent_orientation: Field[[EdgeDim], float], - inv_primal_edge_length: Field[[EdgeDim], float], - inv_vert_vert_length: Field[[EdgeDim], float], - u_vert: Field[[VertexDim, KDim], float], - v_vert: Field[[VertexDim, KDim], float], - primal_normal_vert_x: Field[[ECVDim], float], - primal_normal_vert_y: Field[[ECVDim], float], - dual_normal_vert_x: Field[[ECVDim], float], - dual_normal_vert_y: Field[[ECVDim], float], - vn: Field[[EdgeDim, KDim], float], - smag_limit: Field[[KDim], float], - kh_smag_e: Field[[EdgeDim, KDim], float], - kh_smag_ec: Field[[EdgeDim, KDim], float], - z_nabla2_e: Field[[EdgeDim, KDim], float], - smag_offset: float, - dtime: float, -): - _mo_nh_diffusion_stencil_01_scale_dtime( - enh_smag_fac, - tangent_orientation, - inv_primal_edge_length, - inv_vert_vert_length, - u_vert, - v_vert, - primal_normal_vert_x, - primal_normal_vert_y, - dual_normal_vert_x, - dual_normal_vert_y, - vn, - smag_limit, - smag_offset, - dtime, - out=(kh_smag_e, kh_smag_ec, z_nabla2_e), - ) - - -@field_operator -def _init_diffusion_local_fields( - k4: float, dyn_substeps: float -) -> tuple[Field[[KDim], float], Field[[KDim], float]]: - diff_multfac_vn = _setup_runtime_diff_multfac_vn(k4, dyn_substeps) - smag_limit = _setup_smag_limit(diff_multfac_vn) - return diff_multfac_vn, smag_limit - - -@program -def init_diffusion_local_fields( - k4: float, - dyn_substeps: float, - diff_multfac_vn: Field[[KDim], float], - smag_limit: Field[[KDim], float], -): - _init_diffusion_local_fields(k4, dyn_substeps, out=(diff_multfac_vn, smag_limit)) - - @field_operator def _en_smag_fac_for_zero_nshift( + vect_a: Field[[KDim], float], hdiff_smag_fac: float, hdiff_smag_fac2: float, hdiff_smag_fac3: float, @@ -205,7 +148,6 @@ def _en_smag_fac_for_zero_nshift( hdiff_smag_z2: float, hdiff_smag_z3: float, hdiff_smag_z4: float, - vect_a: Field[[KDim], float], ) -> Field[[KDim], float]: dz21 = hdiff_smag_z2 - hdiff_smag_z alin = (hdiff_smag_fac2 - hdiff_smag_fac) / dz21 @@ -224,8 +166,40 @@ def _en_smag_fac_for_zero_nshift( return enh_smag_fac +@field_operator +def _init_diffusion_local_fields( + k4: float, + dyn_substeps: float, + hdiff_smag_fac: float, + hdiff_smag_fac2: float, + hdiff_smag_fac3: float, + hdiff_smag_fac4: float, + hdiff_smag_z: float, + hdiff_smag_z2: float, + hdiff_smag_z3: float, + hdiff_smag_z4: float, + vect_a: Field[[KDim], float], +) -> tuple[Field[[KDim], float], Field[[KDim], float], Field[[KDim], float]]: + diff_multfac_vn = _setup_runtime_diff_multfac_vn(k4, dyn_substeps) + smag_limit = _setup_smag_limit(diff_multfac_vn) + enh_smag_fac = _en_smag_fac_for_zero_nshift( + vect_a, + hdiff_smag_fac, + hdiff_smag_fac2, + hdiff_smag_fac3, + hdiff_smag_fac4, + hdiff_smag_z, + hdiff_smag_z2, + hdiff_smag_z3, + hdiff_smag_z4, + ) + return diff_multfac_vn, smag_limit, enh_smag_fac + + @program -def enhanced_smagorinski_factor( +def init_diffusion_local_fields( + k4: float, + dyn_substeps: float, hdiff_smag_fac: float, hdiff_smag_fac2: float, hdiff_smag_fac3: float, @@ -235,9 +209,13 @@ def enhanced_smagorinski_factor( hdiff_smag_z3: float, hdiff_smag_z4: float, vect_a: Field[[KDim], float], + diff_multfac_vn: Field[[KDim], float], + smag_limit: Field[[KDim], float], enh_smag_fac: Field[[KDim], float], ): - _en_smag_fac_for_zero_nshift( + _init_diffusion_local_fields( + k4, + dyn_substeps, hdiff_smag_fac, hdiff_smag_fac2, hdiff_smag_fac3, @@ -247,15 +225,10 @@ def enhanced_smagorinski_factor( hdiff_smag_z3, hdiff_smag_z4, vect_a, - out=enh_smag_fac, + out=(diff_multfac_vn, smag_limit, enh_smag_fac), ) -@field_operator -def _set_zero_k() -> Field[[KDim], float]: - return broadcast(0.0, (KDim,)) - - @field_operator def _set_zero_v_k() -> Field[[VertexDim, KDim], float]: return broadcast(0.0, (VertexDim, KDim)) @@ -373,8 +346,6 @@ class DiffusionParams: """Calculates derived quantities depending on the diffusion config.""" def __init__(self, config: DiffusionConfig): - # TODO [ml] logging for case KX == 0 - # TODO [ml] generrally calculation for x_dom (jg) > 2..n_dom, why is jg special self.boundary_diffusion_start_index_edges = ( 5 # mo_nh_diffusion.start_bdydiff_e - 1 = 5 -1 ) @@ -411,7 +382,7 @@ def determine_smagorinski_factor(self, config: DiffusionConfig): ( smagorinski_factor, smagorinski_height, - ) = self.diffusion_type_5_smagorinski_factor(config) + ) = self._diffusion_type_5_smagorinski_factor(config) case 4: # according to mo_nh_diffusion.f90 this isn't used anywhere the factor is only # used for diffusion_type (3,5) but the defaults are only defined for iequations=3 @@ -427,7 +398,7 @@ def determine_smagorinski_factor(self, config: DiffusionConfig): return smagorinski_factor, smagorinski_height @staticmethod - def diffusion_type_5_smagorinski_factor(config: DiffusionConfig): + def _diffusion_type_5_smagorinski_factor(config: DiffusionConfig): # initial values from mo_diffusion_nml.f90 magic_sqrt = math.sqrt(1600.0 * (1600 + 50000.0)) magic_fac2_value = 2e-6 * (1600.0 + 25000.0 + magic_sqrt) @@ -472,18 +443,20 @@ def __init__( TODO [ml]: initial run: linit = .TRUE.: smag_offset and diff_multfac_vn are defined differently. """ - self.params = params - self.config = config + self.params: DiffusionParams = params + self.config: DiffusionConfig = config self.grid = config.grid self.rd_o_cvd: float = GAS_CONSTANT_DRY_AIR / (CPD - GAS_CONSTANT_DRY_AIR) - self.nudgezone_diff = 0.04 / (config.nudge_max_coeff + sys.float_info.epsilon) - self.bdy_diff = 0.015 / (config.nudge_max_coeff + sys.float_info.epsilon) - self.fac_bdydiff_v = ( + self.nudgezone_diff: float = 0.04 / ( + config.nudge_max_coeff + sys.float_info.epsilon + ) + self.bdy_diff: float = 0.015 / (config.nudge_max_coeff + sys.float_info.epsilon) + self.fac_bdydiff_v: float = ( math.sqrt(config.substep_as_float()) / config.lateral_boundary_denominator.v if config.lhdiff_rcf else 1.0 / config.lateral_boundary_denominator.v ) - self.thresh_tdiff = ( + self.thresh_tdiff: float = ( -5.0 ) # threshold temperature deviation from neighboring grid points hat activates extra diffusion against runaway cooling @@ -493,38 +466,12 @@ def __init__( 1.0 / 48.0, params.K4W * config.substep_as_float() ) - # different for initial run!, through diff_multfac_vn + # TODO different for initial run!, through diff_multfac_vn self.diff_multfac_vn = np_as_located_field(KDim)(np.zeros(config.grid.n_lev())) self.smag_limit = np_as_located_field(KDim)(np.zeros(config.grid.n_lev())) - - init_diffusion_local_fields( - params.K4, - config.substep_as_float(), - self.diff_multfac_vn, - self.smag_limit, - offset_provider={}, - ) - self.enh_smag_fac = np_as_located_field(KDim)( np.zeros(config.grid.n_lev(), float) ) - enhanced_smagorinski_factor( - *params.smagorinski_factor, - *params.smagorinski_height, - vct_a, - self.enh_smag_fac, - offset_provider={"Koff": KDim}, - ) - - self.diff_multfac_n2w = init_nabla2_factor_in_upper_damping_zone( - k_size=config.grid.n_lev(), - nshift=0, - physical_heights=np.asarray(vct_a), - nrdmax=self.config.vertical_params.index_of_damping_height, - ) - self.diff_multfac_smag = np_as_located_field(KDim)( - np.zeros(config.grid.n_lev()) - ) shape_vk = (config.grid.num_vertices(), config.grid.n_lev()) shape_ck = (config.grid.num_cells(), config.grid.n_lev()) self.u_vert = np_as_located_field(VertexDim, KDim)(np.zeros(shape_vk, float)) @@ -545,6 +492,29 @@ def __init__( np.arange((shape_ek[0])) ) + self.diff_multfac_smag = np_as_located_field(KDim)( + np.zeros(config.grid.n_lev()) + ) + + init_diffusion_local_fields( + params.K4, + config.substep_as_float(), + *params.smagorinski_factor, + *params.smagorinski_height, + vct_a, + self.diff_multfac_vn, + self.smag_limit, + self.enh_smag_fac, + offset_provider={"Koff": KDim}, + ) + + self.diff_multfac_n2w = init_nabla2_factor_in_upper_damping_zone( + k_size=config.grid.n_lev(), + nshift=0, + physical_heights=np.asarray(vct_a), + nrdmax=self.config.vertical_params.index_of_damping_height, + ) + def run( self, diagnostic_state: DiagnosticState, diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index eaf785b0a..3f15dc33d 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -31,7 +31,7 @@ from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState from icon4py.atm_dyn_iconam.metric_state import MetricState from icon4py.atm_dyn_iconam.prognostic import PrognosticState -from icon4py.common.dimension import KDim, VertexDim +from icon4py.common.dimension import KDim, Koff, VertexDim from icon4py.testutils.serialbox_utils import ( IconDiffusionInitSavepoint, IconSerialDataProvider, @@ -56,18 +56,33 @@ def smag_limit_numpy(shape, k4, substeps): def test_init_diff_multifac_vn_const(): mesh = SimpleMesh() diff_multfac_vn = zero_field(mesh, KDim) - smag_offset = zero_field(mesh, KDim) + smag_limit = zero_field(mesh, KDim) k4 = 1.0 substeps = 5.0 shape = np.asarray(diff_multfac_vn).shape - expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps) + enh_smag_fac = zero_field(mesh, KDim) + a_vec = random_field(mesh, KDim, low=1.0, high=10.0) + fac = (0.67, 0.5, 1.3, 0.8) + z = (0.1, 0.2, 0.3, 0.4) + expected_smag_limit = smag_limit_numpy(shape, k4, substeps) + expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps) + enhanced_smag_fac_np = enhanced_smagorinski_factor_numpy(fac, z, np.asarray(a_vec)) init_diffusion_local_fields( - k4, substeps, diff_multfac_vn, smag_offset, offset_provider={} + k4, + substeps, + *fac, + *z, + a_vec, + diff_multfac_vn, + smag_limit, + enh_smag_fac, + offset_provider={"Koff": KDim}, ) assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) - assert np.allclose(expected_smag_limit, smag_offset) + assert np.allclose(expected_smag_limit, smag_limit) + assert np.allclose(enhanced_smag_fac_np, np.asarray(enh_smag_fac[:-1])) def test_init_diff_multifac_vn_k4_substeps(): @@ -76,13 +91,25 @@ def test_init_diff_multifac_vn_k4_substeps(): smag_limit = zero_field(mesh, KDim) k4 = 0.003 substeps = 1.0 + enh_smag_fac = zero_field(mesh, KDim) + a_vec = random_field(mesh, KDim, low=1.0, high=10.0) + fac = (0.67, 0.5, 1.3, 0.8) + z = (0.1, 0.2, 0.3, 0.4) + shape = np.asarray(diff_multfac_vn).shape expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps) - expected_smag_limit = smag_limit_numpy(shape, k4, substeps) init_diffusion_local_fields( - k4, substeps, diff_multfac_vn, smag_limit, offset_provider={} + k4, + substeps, + *fac, + *z, + a_vec, + diff_multfac_vn, + smag_limit, + enh_smag_fac, + offset_provider={"Koff": KDim}, ) assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) assert np.allclose(expected_smag_limit, smag_limit) @@ -93,7 +120,7 @@ def diff_multfac_vn_numpy(shape, k4, substeps): return factor * np.ones(shape) -def enhanced_smagorinski_factor_np(factor_in, heigths_in, a_vec): +def enhanced_smagorinski_factor_numpy(factor_in, heigths_in, a_vec): alin = (factor_in[1] - factor_in[0]) / (heigths_in[1] - heigths_in[0]) df32 = factor_in[2] - factor_in[1] df42 = factor_in[3] - factor_in[1] @@ -109,19 +136,6 @@ def enhanced_smagorinski_factor_np(factor_in, heigths_in, a_vec): return factor_in[0] + dzlin * alin + dzqdr * (aqdr + dzqdr * bqdr) -def test_enhanced_smagorinski_factor(): - mesh = SimpleMesh() - a_vec = random_field(mesh, KDim, low=1.0, high=10.0) - result = zero_field(mesh, KDim) - fac = (0.67, 0.5, 1.3, 0.8) - z = (0.1, 0.2, 0.3, 0.4) - _en_smag_fac_for_zero_nshift( - *fac, *z, a_vec, out=result, offset_provider={"Koff": KDim} - ) - enhanced_smag_fac_np = enhanced_smagorinski_factor_np(fac, z, np.asarray(a_vec)) - assert np.allclose(enhanced_smag_fac_np, np.asarray(result[:-1])) - - def test_set_zero_vertex_k(): mesh = SimpleMesh() f = random_field(mesh, VertexDim, KDim) @@ -232,7 +246,7 @@ def test_diffusion_init(with_r04b09_diffusion_config): shape_k, additional_parameters.K4, config.substep_as_float() ) assert np.allclose(expected_diff_multfac_vn, np.asarray(diffusion.diff_multfac_vn)) - expected_enh_smag_fac = enhanced_smagorinski_factor_np( + expected_enh_smag_fac = enhanced_smagorinski_factor_numpy( additional_parameters.smagorinski_factor, additional_parameters.smagorinski_height, vct_a, From 6fa2bc2db9d7a1b1f2aa42d7276e8a42fe38956f Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 9 Dec 2022 15:01:34 +0100 Subject: [PATCH 026/263] refactor allocation of fields: TODO duplicated zero_field function --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 56 +++++++++---------- .../src/icon4py/atm_dyn_iconam/icon_grid.py | 52 ++++++++++++----- .../src/icon4py/atm_dyn_iconam/utils.py | 8 +++ atm_dyn_iconam/tests/icon_grid_test_utils.py | 29 +++++----- 4 files changed, 90 insertions(+), 55 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index 73aca6ab9..20a9c7e0c 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -16,6 +16,7 @@ from typing import Final import numpy as np +from functional.common import Dimension from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, broadcast, maximum, minimum from functional.iterator.embedded import np_as_located_field @@ -52,6 +53,7 @@ _mo_nh_diffusion_stencil_16, ) from icon4py.atm_dyn_iconam.prognostic import PrognosticState +from icon4py.atm_dyn_iconam.utils import zero_field from icon4py.common.dimension import ( C2E2CDim, C2E2CODim, @@ -467,34 +469,7 @@ def __init__( ) # TODO different for initial run!, through diff_multfac_vn - self.diff_multfac_vn = np_as_located_field(KDim)(np.zeros(config.grid.n_lev())) - self.smag_limit = np_as_located_field(KDim)(np.zeros(config.grid.n_lev())) - self.enh_smag_fac = np_as_located_field(KDim)( - np.zeros(config.grid.n_lev(), float) - ) - shape_vk = (config.grid.num_vertices(), config.grid.n_lev()) - shape_ck = (config.grid.num_cells(), config.grid.n_lev()) - self.u_vert = np_as_located_field(VertexDim, KDim)(np.zeros(shape_vk, float)) - self.v_vert = np_as_located_field(VertexDim, KDim)(np.zeros(shape_vk, float)) - shape_ek = (config.grid.num_edges(), config.grid.n_lev()) - allocate_ek = np_as_located_field(EdgeDim, KDim)(np.zeros(shape_ek, float)) - self.kh_smag_e = allocate_ek - self.kh_smag_ec = allocate_ek - self.z_nabla2_e = allocate_ek - self.z_temp = np_as_located_field(CellDim, KDim)(np.zeros(shape_ck, float)) - self.vertical_index = np_as_located_field(KDim)( - np.arange(self.grid.n_lev() + 1) - ) - self.horizontal_cell_index = np_as_located_field(CellDim)( - np.arange((shape_ck[0])) - ) - self.horizontal_edge_index = np_as_located_field(EdgeDim)( - np.arange((shape_ek[0])) - ) - - self.diff_multfac_smag = np_as_located_field(KDim)( - np.zeros(config.grid.n_lev()) - ) + self._allocate_local_fields() init_diffusion_local_fields( params.K4, @@ -515,6 +490,31 @@ def __init__( nrdmax=self.config.vertical_params.index_of_damping_height, ) + def _allocate_local_fields(self): + def _allocate(*dims: Dimension): + return zero_field(self.grid, *dims) + + def _index_field(dim:Dimension, size=None): + size = size if size else self.grid.size[dim] + return np_as_located_field(dim)( + np.arange(size) + ) + + self.diff_multfac_vn = _allocate(KDim) + self.smag_limit = _allocate(KDim) + self.enh_smag_fac = _allocate(KDim) + self.u_vert = _allocate(VertexDim, KDim) + self.v_vert = _allocate(VertexDim, KDim) + self.kh_smag_e = _allocate(EdgeDim, KDim) + self.kh_smag_ec = _allocate(EdgeDim, KDim) + self.z_nabla2_e = _allocate(EdgeDim, KDim) + self.z_temp = _allocate(EdgeDim, KDim) + self.diff_multfac_smag = _allocate(KDim) + self.vertical_index = _index_field(KDim, self.grid.n_lev() + 1) + self.horizontal_cell_index = _index_field(CellDim) + self.horizontal_edge_index = _index_field(EdgeDim) + + def run( self, diagnostic_state: DiagnosticState, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py index 1e6fadd6c..344ee5c14 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py @@ -10,7 +10,8 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -from typing import Tuple +from dataclasses import dataclass +from typing import Tuple, Dict import numpy as np from functional.common import Dimension, DimensionKind, Field @@ -20,18 +21,27 @@ ) from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig -from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim +from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim, E2CDim + + +class VerticalMeshConfig: + def __init__(self, num_lev: int): + self._num_lev = num_lev + + @property + def num_lev(self)->int: + return self._num_lev class MeshConfig: - def __init__(self, horizontalMesh: HorizontalMeshConfig): - self._num_k_levels = 65 + def __init__(self, horizontal_config: HorizontalMeshConfig, vertical_config: VerticalMeshConfig): + self._vertical = vertical_config self._n_shift_total = 0 - self._horizontal = horizontalMesh + self._horizontal = horizontal_config @property def num_k_levels(self): - return self._num_k_levels + return self._vertical.num_lev @property def n_shift_total(self): @@ -63,11 +73,24 @@ def __init__(self): self.config: MeshConfig = None self.start_indices = {} self.end_indices = {} - self.connectivities = {} + self.connectivities: Dict[Dimension, np.ndarray] = {} + self.size: Dict[Dimension, int] = {} + + + def _update_size(self, config: MeshConfig): + self.size[VertexDim]= config.num_vertices + self.size[CellDim] = config.num_cells + self.size[EdgeDim] = config.num_edges + self.size[KDim] = config.num_k_levels + + @builder def with_config(self, config: MeshConfig): self.config = config + self._update_size(config) + + @builder def with_start_end_indices( @@ -77,17 +100,19 @@ def with_start_end_indices( self.end_indices[dim] = end_indices @builder - def with_connectivity(self, **connectivity): - self.connectivities.update(**connectivity) + def with_connectivities(self, connectivity:Dict[Dimension, np.ndarray]): + self.connectivities.update({d.value.lower(): k for d, k in connectivity.items()}) + self.size.update({d:t.shape[1] for d, t in connectivity.items()}) + def n_lev(self): - return self.config.num_k_levels + return self.config.num_k_levels if self.config else 0 def num_cells(self): - return self.config.num_cells + return self.config.num_cells if self.config else 0 def num_vertices(self): - return self.config.num_vertices + return self.config.num_vertices if self.config else 0 def num_edges(self): return self.config.num_edges @@ -130,7 +155,7 @@ def get_v2e_offset_provider(self): class VerticalModelParams: - def __init__(self, vct_a: np.ndarray, rayleigh_damping_height: float = 12500.0): + def __init__(self, vct_a: np.ndarray, rayleigh_damping_height: float): """ Contains physical parameters defined on the grid. @@ -140,7 +165,6 @@ def __init__(self, vct_a: np.ndarray, rayleigh_damping_height: float = 12500.0): """ self.rayleigh_damping_height = rayleigh_damping_height self.vct_a = vct_a - # TODO klevels in ICON are inverted! TODO test against ICON LOG.exp... self.index_of_damping_height = np.argmax( np.where(np.asarray(self.vct_a) >= self.rayleigh_damping_height) ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py new file mode 100644 index 000000000..b8fdf6ca7 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py @@ -0,0 +1,8 @@ +import numpy as np +from functional.common import Dimension +from functional.iterator.embedded import np_as_located_field + + +def zero_field(mesh, *dims: Dimension, dtype=float): + shapex = tuple(map(lambda x: mesh.size[x], dims)) + return np_as_located_field(*dims)(np.zeros(shapex, dtype=dtype)) diff --git a/atm_dyn_iconam/tests/icon_grid_test_utils.py b/atm_dyn_iconam/tests/icon_grid_test_utils.py index f3aa67283..7b8c50706 100644 --- a/atm_dyn_iconam/tests/icon_grid_test_utils.py +++ b/atm_dyn_iconam/tests/icon_grid_test_utils.py @@ -21,12 +21,15 @@ from icon4py.atm_dyn_iconam.icon_grid import ( IconGrid, MeshConfig, - VerticalModelParams, + VerticalModelParams, VerticalMeshConfig, ) -from icon4py.common.dimension import CellDim, EdgeDim, VertexDim +from icon4py.common.dimension import CellDim, EdgeDim, VertexDim, E2CDim, C2EDim, C2E2CDim, \ + C2E2CODim, E2VDim, V2EDim from icon4py.testutils.serialbox_utils import IconSerialDataProvider + + @fixture def with_icon_grid(): data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") @@ -48,7 +51,7 @@ def with_icon_grid(): num_vertices=sp_meta["nproma"], # or rather "num_vert" num_cells=sp_meta["nproma"], # or rather "num_cells" num_edges=sp_meta["nproma"], # or rather "num_edges" - ) + ), VerticalMeshConfig(num_lev=sp_meta["nlev"]) ) c2e2c = sp.c2e2c() @@ -59,14 +62,9 @@ def with_icon_grid(): .with_start_end_indices(VertexDim, vertex_starts, vertex_ends) .with_start_end_indices(EdgeDim, edge_starts, edge_ends) .with_start_end_indices(CellDim, cell_starts, cell_ends) - .with_connectivity(c2e=sp.c2e()) - .with_connectivity(e2c=sp.e2c()) - .with_connectivity(c2e2c=c2e2c) - .with_connectivity(e2v=sp.e2v()) - .with_connectivity(c2e2c0=c2e2c0) - .with_connectivity(v2e=sp.v2e()) - .with_connectivity(e2v=sp.e2v()) - ) + .with_connectivities({C2EDim:sp.c2e(), E2CDim:sp.e2c(), C2E2CDim:c2e2c, C2E2CODim:c2e2c0}) + .with_connectivities({E2VDim:sp.e2v(), V2EDim:sp.v2e()}) + ) return grid @@ -82,11 +80,16 @@ def with_r04b09_diffusion_config() -> DiffusionConfig: "icon_diffusion_init", data_path, True ).from_savepoint_init(linit=True, date="2021-06-20T12:00:10.000") nproma = sp.get_metadata("nproma")["nproma"] - horizontalConfig = HorizontalMeshConfig( + num_lev = sp.get_metadata("nlev")["nlev"] + horizontal_config = HorizontalMeshConfig( num_vertices=nproma, num_cells=nproma, num_edges=nproma ) + vertical_config = VerticalMeshConfig(num_lev = num_lev) + + grid = IconGrid().with_config(MeshConfig( + horizontal_config=horizontal_config, + vertical_config = vertical_config)) - grid = IconGrid().with_config(MeshConfig(horizontalMesh=horizontalConfig)) verticalParams = VerticalModelParams( rayleigh_damping_height=12500, vct_a=sp.vct_a() ) From 7f6c6c3670bd71571fc2fd38894c98b560d54284 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 9 Dec 2022 16:57:41 +0100 Subject: [PATCH 027/263] - cleanup diffusion config (1) --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 83 +++++++++++-------- .../src/icon4py/atm_dyn_iconam/icon_grid.py | 9 +- atm_dyn_iconam/tests/icon_grid_test_utils.py | 12 ++- atm_dyn_iconam/tests/test_diffusion.py | 4 +- 4 files changed, 71 insertions(+), 37 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index 20a9c7e0c..d78f4c1e0 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -287,58 +287,75 @@ def __init__( self, grid: IconGrid, vertical_params: VerticalModelParams, - diffusion_type: int = 5, - apply_to_horizontal_wind: bool = True, - apply_to_vertical_wind: bool = True, - apply_to_temperature: bool = True, - reconstruction_type_smag: int = 1, - compute_3d_smag_coeff: bool = False, - temperature_discretization: int = 2, - horizontal_efdt_ratio: float = 24.0, - horizontal_w_efdt_ratio: float = 15.0, - smag_scaling_factor: float = 0.025, + diffusion_type: int = 5, # TODO: use enum + hdiff_w = True, + type_vn_diffu: int = 1, + smag_3d: bool = False, + type_t_diffu: int = 2, + hdiff_efdt_ratio: float = 36.0, + hdiff_w_efdt_ratio: float = 15.0, + smag_scaling_fac: float = 0.015, + zdiffu_t = True ): # TODO [ml]: move external stuff out: grid related stuff, other than diffusion namelists (see below self.grid = grid self.vertical_params = vertical_params - # from namelist diffusion_nml + + # parameters from namelist diffusion_nml self.diffusion_type = ( - diffusion_type # hdiff_order ! order of nabla operator for diffusion - ) - self.lhdiff_vn = ( - apply_to_horizontal_wind # ! diffusion on the horizontal wind field + diffusion_type # hdiff_order : order of nabla operator for diffusion ) - self.lhdiff_temp = apply_to_temperature # ! diffusion on the temperature field - self.lhdiff_w = apply_to_vertical_wind # ! diffusion on the vertical wind field - self.lhdiff_rcf = True # namelist, remove if always true - self.itype_vn_diffu = reconstruction_type_smag # ! reconstruction method used for Smagorinsky diffusion - self.l_smag_d = compute_3d_smag_coeff # namelist lsmag_d, if `true`, compute 3D Smagorinsky diffusion coefficient. + self.apply_to_vertical_wind = hdiff_w # lhdiff_w, diffusion on the vertical wind field + self.apply_to_horizontal_wind = True # lhdiff_vn, diffusion on horizonal wind field, ONLY used in mo_nh_stepping.f90 + self.apply_to_temperature = True # lhdiff_temp, apply horizontal diffusion to temperature + self.compute_3d_smag_coeff = smag_3d # lsmag_3d, if `true`, compute 3D Smagorinsky diffusion coefficient. + self.itype_vn_diffu = type_vn_diffu # itype_vn_diffu, options for discretizing the Smagorinsky momentum diffusion self.itype_t_diffu = ( - temperature_discretization # ! discretization of temperature diffusion + type_t_diffu # itype_t_diffu, options for discretizing the Smagorinsky temperature diffusion ) self.hdiff_efdt_ratio = ( - horizontal_efdt_ratio # ! ratio of e-folding time to time step + hdiff_efdt_ratio # hdiff_efdt_ratio, ratio of e-folding time to (2*)time step ) - self.hdiff_w_efdt_ratio = horizontal_w_efdt_ratio # ratio of e-folding time to time step for w diffusion (NH only) - self.hdiff_smag_fac = ( - smag_scaling_factor # ! scaling factor for Smagorinsky diffusion + self.hdiff_w_efdt_ratio = hdiff_w_efdt_ratio # hdiff_w_efdt_ratio, ratio of e-folding time to time step for w diffusion (NH only) + self.hdiff_smag_factor = ( + smag_scaling_fac # hdiff_smag_fac, scaling factor for Smagorinsky diffusion at height hdiff_smag_z and below ) + self.zdiffu_t = zdiffu_t #l_zdiffu_t, apply truly horizontal temperature diffusion, from parent namelist nonhydrostatic_nml, but is only used in diffusion, and in mo_vertical_grid.prepare_zdiffu. + + # from other namelists - # from parent namelist nonhydrostatic_nml - self.l_zdiffu_t = True # ! l_zdiffu_t: specifies computation of Smagorinsky temperature diffusion + #from parent namelist nonhydrostatic_nml self.ndyn_substeps = 5 + self.lhdiff_rcf = True # namelist gridref_nml # default is v=200.0, t=135.0 self.lateral_boundary_denominator = TupleVT(v=150.0, t=135.0) - # namelist grid_nml - self.l_limited_area = True - # name list: interpol_nml self.nudge_max_coeff = 0.075 + self._validate() + + def _validate(self): + if self.diffusion_type != 5: + raise NotImplementedError("only diffusion type 5 : `Smagorinsky diffusion with fourth-order background diffusion` is implemented") + + if not self.grid.limited_area: + raise NotImplementedError("only limited area mode is implemented") + + if self.diffusion_type < 0: + self.apply_to_temperature = False + self.apply_to_horizontal_wind = False + self.apply_to_vertical_wind = False + else: + self.apply_to_temperature = True + self.apply_to_horizontal_wind = True + + if not self.zdiffu_t: + raise NotImplementedError("zdiffu_t = False is not implemented (leaves out stencil_15)") + def substep_as_float(self): return float(self.ndyn_substeps) @@ -389,11 +406,11 @@ def determine_smagorinski_factor(self, config: DiffusionConfig): # according to mo_nh_diffusion.f90 this isn't used anywhere the factor is only # used for diffusion_type (3,5) but the defaults are only defined for iequations=3 smagorinski_factor = ( - config.hdiff_smag_fac if config.hdiff_smag_fac else 0.15, + config.hdiff_smag_factor if config.hdiff_smag_factor else 0.15, ) smagorinski_height = None case _: - print("not implemented") + raise NotImplementedError("Only implemented for diffusion type 4 and 5") smagorinski_factor = None smagorinski_height = None pass @@ -405,7 +422,7 @@ def _diffusion_type_5_smagorinski_factor(config: DiffusionConfig): magic_sqrt = math.sqrt(1600.0 * (1600 + 50000.0)) magic_fac2_value = 2e-6 * (1600.0 + 25000.0 + magic_sqrt) magic_z2 = 1600.0 + 50000.0 + magic_sqrt - factor = (config.hdiff_smag_fac, magic_fac2_value, 0.0, 1.0) + factor = (config.hdiff_smag_factor, magic_fac2_value, 0.0, 1.0) heights = (32500.0, magic_z2, 50000.0, 90000.0) return factor, heights diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py index 344ee5c14..a3219e0be 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py @@ -34,11 +34,16 @@ def num_lev(self)->int: class MeshConfig: - def __init__(self, horizontal_config: HorizontalMeshConfig, vertical_config: VerticalMeshConfig): + def __init__(self, horizontal_config: HorizontalMeshConfig, vertical_config: VerticalMeshConfig, limited_area = True): self._vertical = vertical_config self._n_shift_total = 0 + self._limited_area = limited_area self._horizontal = horizontal_config + @property + def limited_area(self): + return self._limited_area + @property def num_k_levels(self): return self._vertical.num_lev @@ -105,6 +110,8 @@ def with_connectivities(self, connectivity:Dict[Dimension, np.ndarray]): self.size.update({d:t.shape[1] for d, t in connectivity.items()}) + def limited_area(self): + return self.config.limited_area def n_lev(self): return self.config.num_k_levels if self.config else 0 diff --git a/atm_dyn_iconam/tests/icon_grid_test_utils.py b/atm_dyn_iconam/tests/icon_grid_test_utils.py index 7b8c50706..a90563a89 100644 --- a/atm_dyn_iconam/tests/icon_grid_test_utils.py +++ b/atm_dyn_iconam/tests/icon_grid_test_utils.py @@ -93,4 +93,14 @@ def with_r04b09_diffusion_config() -> DiffusionConfig: verticalParams = VerticalModelParams( rayleigh_damping_height=12500, vct_a=sp.vct_a() ) - return DiffusionConfig(grid=grid, vertical_params=verticalParams) + + return DiffusionConfig(grid=grid, + vertical_params=verticalParams, + diffusion_type=5, + hdiff_w=True, + type_t_diffu=2, + type_vn_diffu=1, + hdiff_efdt_ratio=24.0, + hdiff_w_efdt_ratio=15.0, + smag_scaling_fac=0.025, + zdiffu_t=True) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 3f15dc33d..bebb2f77b 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -171,7 +171,7 @@ def test_diffusion_coefficients_without_hdiff_efdt_ratio(with_r04b09_diffusion_c def test_smagorinski_factor_for_diffusion_type_4(with_r04b09_diffusion_config): config = with_r04b09_diffusion_config - config.hdiff_smag_fac = 0.15 + config.hdiff_smag_factor = 0.15 config.diffusion_type = 4 params = DiffusionParams(config) @@ -184,7 +184,7 @@ def test_smagorinski_heights_diffusion_type_5_are_consistent( with_r04b09_diffusion_config, ): config = with_r04b09_diffusion_config - config.hdiff_smag_fac = 0.15 + config.hdiff_smag_factor = 0.15 config.diffusion_type = 5 params = DiffusionParams(config) From 39439c2e2feb35804bf4717b10b8521a1433b8f7 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 9 Dec 2022 20:05:20 +0100 Subject: [PATCH 028/263] - add linit=true case --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 251 +++++++++++++----- .../src/icon4py/atm_dyn_iconam/icon_grid.py | 32 +-- .../src/icon4py/atm_dyn_iconam/utils.py | 17 +- atm_dyn_iconam/tests/icon_grid_test_utils.py | 57 ++-- atm_dyn_iconam/tests/test_diffusion.py | 167 ++++++++---- 5 files changed, 357 insertions(+), 167 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index d78f4c1e0..64e06d1ae 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -70,8 +70,9 @@ VectorTuple = namedtuple("VectorTuple", "x y") -# TODO [ml] initial RUN linit = TRUE -# def _setup_initial_diff_multfac_vn +@field_operator +def _setup_smag_limit(diff_multfac_vn: Field[[KDim], float]) -> Field[[KDim], float]: + return 0.125 - 4.0 * diff_multfac_vn @field_operator @@ -84,8 +85,17 @@ def _setup_runtime_diff_multfac_vn( @field_operator -def _setup_smag_limit(diff_multfac_vn: Field[[KDim], float]) -> Field[[KDim], float]: - return 0.125 - 4.0 * diff_multfac_vn +def _setup_initial_diff_multfac_vn( + k4: float, hdiff_efdt_ratio: float +) -> Field[[KDim], float]: + return broadcast(k4 / 3.0 * hdiff_efdt_ratio, (KDim,)) + + +@field_operator +def setup_fields_for_initial_step(k4: float, hdiff_efdt_ratio: float): + diff_multfac_vn = _setup_initial_diff_multfac_vn(k4, hdiff_efdt_ratio) + smag_limit = _setup_smag_limit(diff_multfac_vn) + return diff_multfac_vn, smag_limit @field_operator @@ -195,7 +205,11 @@ def _init_diffusion_local_fields( hdiff_smag_z3, hdiff_smag_z4, ) - return diff_multfac_vn, smag_limit, enh_smag_fac + return ( + diff_multfac_vn, + smag_limit, + enh_smag_fac, + ) @program @@ -227,7 +241,11 @@ def init_diffusion_local_fields( hdiff_smag_z3, hdiff_smag_z4, vect_a, - out=(diff_multfac_vn, smag_limit, enh_smag_fac), + out=( + diff_multfac_vn, + smag_limit, + enh_smag_fac, + ), ) @@ -287,15 +305,15 @@ def __init__( self, grid: IconGrid, vertical_params: VerticalModelParams, - diffusion_type: int = 5, # TODO: use enum - hdiff_w = True, + diffusion_type: int = 5, # TODO: use enum + hdiff_w=True, type_vn_diffu: int = 1, smag_3d: bool = False, type_t_diffu: int = 2, hdiff_efdt_ratio: float = 36.0, hdiff_w_efdt_ratio: float = 15.0, smag_scaling_fac: float = 0.015, - zdiffu_t = True + zdiffu_t=True, ): # TODO [ml]: move external stuff out: grid related stuff, other than diffusion namelists (see below self.grid = grid @@ -305,28 +323,25 @@ def __init__( self.diffusion_type = ( diffusion_type # hdiff_order : order of nabla operator for diffusion ) - self.apply_to_vertical_wind = hdiff_w # lhdiff_w, diffusion on the vertical wind field + self.apply_to_vertical_wind = ( + hdiff_w # lhdiff_w, diffusion on the vertical wind field + ) self.apply_to_horizontal_wind = True # lhdiff_vn, diffusion on horizonal wind field, ONLY used in mo_nh_stepping.f90 - self.apply_to_temperature = True # lhdiff_temp, apply horizontal diffusion to temperature + self.apply_to_temperature = ( + True # lhdiff_temp, apply horizontal diffusion to temperature + ) self.compute_3d_smag_coeff = smag_3d # lsmag_3d, if `true`, compute 3D Smagorinsky diffusion coefficient. self.itype_vn_diffu = type_vn_diffu # itype_vn_diffu, options for discretizing the Smagorinsky momentum diffusion - self.itype_t_diffu = ( - type_t_diffu # itype_t_diffu, options for discretizing the Smagorinsky temperature diffusion - ) - self.hdiff_efdt_ratio = ( - hdiff_efdt_ratio # hdiff_efdt_ratio, ratio of e-folding time to (2*)time step - ) + self.itype_t_diffu = type_t_diffu # itype_t_diffu, options for discretizing the Smagorinsky temperature diffusion + self.hdiff_efdt_ratio = hdiff_efdt_ratio # hdiff_efdt_ratio, ratio of e-folding time to (2*)time step self.hdiff_w_efdt_ratio = hdiff_w_efdt_ratio # hdiff_w_efdt_ratio, ratio of e-folding time to time step for w diffusion (NH only) - self.hdiff_smag_factor = ( - smag_scaling_fac # hdiff_smag_fac, scaling factor for Smagorinsky diffusion at height hdiff_smag_z and below - ) - - self.zdiffu_t = zdiffu_t #l_zdiffu_t, apply truly horizontal temperature diffusion, from parent namelist nonhydrostatic_nml, but is only used in diffusion, and in mo_vertical_grid.prepare_zdiffu. + self.hdiff_smag_factor = smag_scaling_fac # hdiff_smag_fac, scaling factor for Smagorinsky diffusion at height hdiff_smag_z and below + self.zdiffu_t = zdiffu_t # l_zdiffu_t, apply truly horizontal temperature diffusion, from parent namelist nonhydrostatic_nml, but is only used in diffusion, and in mo_vertical_grid.prepare_zdiffu. # from other namelists - #from parent namelist nonhydrostatic_nml + # from parent namelist nonhydrostatic_nml self.ndyn_substeps = 5 self.lhdiff_rcf = True @@ -340,7 +355,9 @@ def __init__( def _validate(self): if self.diffusion_type != 5: - raise NotImplementedError("only diffusion type 5 : `Smagorinsky diffusion with fourth-order background diffusion` is implemented") + raise NotImplementedError( + "only diffusion type 5 : `Smagorinsky diffusion with fourth-order background diffusion` is implemented" + ) if not self.grid.limited_area: raise NotImplementedError("only limited area mode is implemented") @@ -354,8 +371,9 @@ def _validate(self): self.apply_to_horizontal_wind = True if not self.zdiffu_t: - raise NotImplementedError("zdiffu_t = False is not implemented (leaves out stencil_15)") - + raise NotImplementedError( + "zdiffu_t = False is not implemented (leaves out stencil_15)" + ) def substep_as_float(self): return float(self.ndyn_substeps) @@ -459,8 +477,7 @@ def __init__( """ Initialize Diffusion granule. - TODO [ml]: initial run: linit = .TRUE.: smag_offset and diff_multfac_vn are defined - differently. + calculates all local fields that are used in diffusion within the time loop """ self.params: DiffusionParams = params self.config: DiffusionConfig = config @@ -479,23 +496,22 @@ def __init__( -5.0 ) # threshold temperature deviation from neighboring grid points hat activates extra diffusion against runaway cooling - # different for init call: smag_offset = 0 - self.smag_offset: float = 0.25 * params.K4 * config.substep_as_float() + self._smag_offset: float = 0.25 * params.K4 * config.substep_as_float() self.diff_multfac_w: float = min( 1.0 / 48.0, params.K4W * config.substep_as_float() ) - # TODO different for initial run!, through diff_multfac_vn self._allocate_local_fields() + # TODO different for initial run!, through diff_multfac_vn init_diffusion_local_fields( params.K4, config.substep_as_float(), *params.smagorinski_factor, *params.smagorinski_height, vct_a, - self.diff_multfac_vn, - self.smag_limit, + self._diff_multfac_vn, + self._smag_limit, self.enh_smag_fac, offset_provider={"Koff": KDim}, ) @@ -511,14 +527,13 @@ def _allocate_local_fields(self): def _allocate(*dims: Dimension): return zero_field(self.grid, *dims) - def _index_field(dim:Dimension, size=None): + def _index_field(dim: Dimension, size=None): size = size if size else self.grid.size[dim] - return np_as_located_field(dim)( - np.arange(size) - ) + return np_as_located_field(dim)(np.arange(size)) + + self._diff_multfac_vn = _allocate(KDim) - self.diff_multfac_vn = _allocate(KDim) - self.smag_limit = _allocate(KDim) + self._smag_limit = _allocate(KDim) self.enh_smag_fac = _allocate(KDim) self.u_vert = _allocate(VertexDim, KDim) self.v_vert = _allocate(VertexDim, KDim) @@ -531,8 +546,61 @@ def _index_field(dim:Dimension, size=None): self.horizontal_cell_index = _index_field(CellDim) self.horizontal_edge_index = _index_field(EdgeDim) + def initial_step( + self, + diagnostic_state: DiagnosticState, + prognostic_state: PrognosticState, + metric_state: MetricState, + interpolation_state: InterpolationState, + dtime: float, + tangent_orientation: Field[[EdgeDim], float], + inverse_primal_edge_lengths: Field[[EdgeDim], float], + inverse_dual_edge_length: Field[[EdgeDim], float], + inverse_vertical_vertex_lengths: Field[[EdgeDim], float], + primal_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], + dual_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], + edge_areas: Field[[EdgeDim], float], + cell_areas: Field[[CellDim], float], + ): + """ + Calculate initial diffusion step. + + In ICON at the start of the simulation diffusion is run with a parameter linit = True: + + 'For real-data runs, perform an extra diffusion call before the first time + step because no other filtering of the interpolated velocity field is done' + + This run uses special values for diff_multfac_vn, smag_limit and smag_offset + + """ + diff_multfac_vn = zero_field(self.grid, KDim) + smag_limit = zero_field(self.grid, KDim) - def run( + setup_fields_for_initial_step( + self.params.K4, + self.config.hdiff_efdt_ratio, + out=(diff_multfac_vn, smag_limit), + ) + self._do_diffusion_step( + diagnostic_state, + prognostic_state, + metric_state, + interpolation_state, + dtime, + tangent_orientation, + inverse_primal_edge_lengths, + inverse_dual_edge_length, + inverse_vertical_vertex_lengths, + primal_normal_vert, + dual_normal_vert, + edge_areas, + cell_areas, + diff_multfac_vn, + smag_limit, + 0.0, + ) + + def time_step( self, diagnostic_state: DiagnosticState, prognostic_state: PrognosticState, @@ -547,40 +615,92 @@ def run( dual_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], edge_areas: Field[[EdgeDim], float], cell_areas: Field[[CellDim], float], + ): + """ + Do one diffusion step within regular time loop. + + runs a diffusion step for the parameter linit=False, within regular time loop. + """ + self._do_diffusion_step( + diagnostic_state, + prognostic_state, + metric_state, + interpolation_state, + dtime, + tangent_orientation, + inverse_primal_edge_lengths, + inverse_dual_edge_length, + inverse_vertical_vertex_lengths, + primal_normal_vert, + dual_normal_vert, + edge_areas, + cell_areas, + self._diff_multfac_vn, + self._smag_limit, + self._smag_offset, + ) + + def _do_diffusion_step( + self, + diagnostic_state: DiagnosticState, + prognostic_state: PrognosticState, + metric_state: MetricState, + interpolation_state: InterpolationState, + dtime: float, + tangent_orientation: Field[[EdgeDim], float], + inverse_primal_edge_lengths: Field[[EdgeDim], float], + inverse_dual_edge_length: Field[[EdgeDim], float], + inverse_vertical_vertex_lengths: Field[[EdgeDim], float], + primal_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], + dual_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], + edge_areas: Field[[EdgeDim], float], + cell_areas: Field[[CellDim], float], + diff_multfac_vn: Field[[KDim], float], + smag_limit: Field[[KDim], float], + smag_offset: float, ): """ Run a diffusion step. - inputs are - - output fields: - - input fields that have changed from one time step to another: - - simulation parameters: dtime - timestep + Args: + diagnostic_state: output argument, data class that contains diagnostic variables + prognostic_state: output argument, data class that contains prognostic variables + metric_state: + interpolation_state: + dtime: the time step, + tangent_orientation: + inverse_primal_edge_lengths: + inverse_dual_edge_length: + inverse_vertical_vertex_lengths: + primal_normal_vert: + dual_normal_vert: + edge_areas: + cell_areas: + diff_multfac_vn: + smag_limit: + smag_offset: + """ # ------- # OUTLINE # ------- - # Oa logging - # 0b call timer start - # 0c. dtime dependent stuff: enh_smag_factor, r_dtimensubsteps - r_dtimensubsteps = ( - 1.0 / dtime - if self.config.lhdiff_rcf - else 1.0 / (dtime * self.config.substep_as_float()) - ) - klevels = self.grid.n_lev() - # TODO is this needed? + # Oa TODO: logging + # 0b TODO: call timer start + # + # 0c. dtime dependent stuff: enh_smag_factor, + scale_k(self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={}) + + # TODO: is this needed, if not remove set_zero_v_k(self.u_vert, offset_provider={}) set_zero_v_k(self.v_vert, offset_provider={}) # 1. CALL rbf_vec_interpol_vertex - _mo_intp_rbf_rbf_vec_interpol_vertex( prognostic_state.normal_wind, interpolation_state.rbf_coeff_1, interpolation_state.rbf_coeff_2, out=(self.u_vert, self.v_vert), domain={ - KDim: (0, klevels), VertexDim: self.grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 1, @@ -592,9 +712,6 @@ def run( # 2. HALO EXCHANGE -- CALL sync_patch_array_mult # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 - # 0c. dtime dependent stuff: enh_smag_factor, ~~r_dtimensubsteps~~ - scale_k(self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={}) - _mo_nh_diffusion_stencil_01( self.diff_multfac_smag, tangent_orientation, @@ -607,10 +724,9 @@ def run( dual_normal_vert[0], dual_normal_vert[1], prognostic_state.normal_wind, - self.smag_limit, - self.smag_offset, + smag_limit, + smag_offset, domain={ - KDim: (0, klevels), EdgeDim: ( self.grid.get_indices_from_to( EdgeDim, @@ -636,7 +752,6 @@ def run( diagnostic_state.hdef_ic, ), domain={ - KDim: (0, klevels), CellDim: ( self.grid.get_indices_from_to( CellDim, @@ -658,7 +773,6 @@ def run( interpolation_state.rbf_coeff_2, out=(self.u_vert, self.v_vert), domain={ - KDim: (0, klevels), VertexDim: self.grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 3, @@ -677,6 +791,7 @@ def run( HorizontalMarkerIndex.nudging(EdgeDim) - 1, HorizontalMarkerIndex.nudging(EdgeDim) - 1, )[0] + _fused_mo_nh_diffusion_stencil_04_05_06( self.u_vert, self.v_vert, @@ -687,7 +802,7 @@ def run( inverse_primal_edge_lengths, edge_areas, self.kh_smag_e, - self.diff_multfac_vn, + diff_multfac_vn, interpolation_state.nudgecoeff_e, prognostic_state.normal_wind, self.horizontal_edge_index, @@ -696,7 +811,6 @@ def run( start_2nd_nudge_line, out=prognostic_state.normal_wind, domain={ - KDim: (0, klevels), EdgeDim: self.grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, @@ -735,7 +849,6 @@ def run( diagnostic_state.dwdy, ), domain={ - KDim: (0, klevels), CellDim: self.grid.get_indices_from_to( CellDim, # TODO: global mode is from NUDGING - 1 @@ -750,6 +863,7 @@ def run( # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 # TODO check: kh_smag_e is an out field, should not be calculated in init? + klevels = self.grid.n_lev() fused_mo_nh_diffusion_stencil_11_12( prognostic_state.theta_v, metric_state.theta_ref_mc, @@ -781,7 +895,6 @@ def run( interpolation_state.geofac_div, out=self.z_temp, domain={ - KDim: (0, klevels), CellDim: self.grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.nudging(CellDim), @@ -808,7 +921,6 @@ def run( theta_v=prognostic_state.theta_v, z_temp=self.z_temp, domain={ - KDim: (0, klevels), CellDim: self.grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.nudging(CellDim), @@ -826,7 +938,6 @@ def run( self.rd_o_cvd, out=(prognostic_state.theta_v, prognostic_state.exner_pressure), domain={ - KDim: (0, klevels), CellDim: self.grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.nudging(CellDim), diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py index a3219e0be..0a79a69ba 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py @@ -10,8 +10,8 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -from dataclasses import dataclass -from typing import Tuple, Dict + +from typing import Dict, Tuple import numpy as np from functional.common import Dimension, DimensionKind, Field @@ -21,7 +21,7 @@ ) from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig -from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim, E2CDim +from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim class VerticalMeshConfig: @@ -29,12 +29,17 @@ def __init__(self, num_lev: int): self._num_lev = num_lev @property - def num_lev(self)->int: + def num_lev(self) -> int: return self._num_lev class MeshConfig: - def __init__(self, horizontal_config: HorizontalMeshConfig, vertical_config: VerticalMeshConfig, limited_area = True): + def __init__( + self, + horizontal_config: HorizontalMeshConfig, + vertical_config: VerticalMeshConfig, + limited_area=True, + ): self._vertical = vertical_config self._n_shift_total = 0 self._limited_area = limited_area @@ -81,22 +86,17 @@ def __init__(self): self.connectivities: Dict[Dimension, np.ndarray] = {} self.size: Dict[Dimension, int] = {} - def _update_size(self, config: MeshConfig): - self.size[VertexDim]= config.num_vertices + self.size[VertexDim] = config.num_vertices self.size[CellDim] = config.num_cells self.size[EdgeDim] = config.num_edges self.size[KDim] = config.num_k_levels - - @builder def with_config(self, config: MeshConfig): self.config = config self._update_size(config) - - @builder def with_start_end_indices( self, dim: Dimension, start_indices: np.ndarray, end_indices: np.ndarray @@ -105,13 +105,15 @@ def with_start_end_indices( self.end_indices[dim] = end_indices @builder - def with_connectivities(self, connectivity:Dict[Dimension, np.ndarray]): - self.connectivities.update({d.value.lower(): k for d, k in connectivity.items()}) - self.size.update({d:t.shape[1] for d, t in connectivity.items()}) - + def with_connectivities(self, connectivity: Dict[Dimension, np.ndarray]): + self.connectivities.update( + {d.value.lower(): k for d, k in connectivity.items()} + ) + self.size.update({d: t.shape[1] for d, t in connectivity.items()}) def limited_area(self): return self.config.limited_area + def n_lev(self): return self.config.num_k_levels if self.config else 0 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py index b8fdf6ca7..56447c01c 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py @@ -1,8 +1,21 @@ +# 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 functional.common import Dimension from functional.iterator.embedded import np_as_located_field - +#TODO fix duplication: duplicated from test testutils/utils.py def zero_field(mesh, *dims: Dimension, dtype=float): - shapex = tuple(map(lambda x: mesh.size[x], dims)) + shapex = tuple(map(lambda x: mesh.size[x], dims)) return np_as_located_field(*dims)(np.zeros(shapex, dtype=dtype)) diff --git a/atm_dyn_iconam/tests/icon_grid_test_utils.py b/atm_dyn_iconam/tests/icon_grid_test_utils.py index a90563a89..3ec41e820 100644 --- a/atm_dyn_iconam/tests/icon_grid_test_utils.py +++ b/atm_dyn_iconam/tests/icon_grid_test_utils.py @@ -21,15 +21,23 @@ from icon4py.atm_dyn_iconam.icon_grid import ( IconGrid, MeshConfig, - VerticalModelParams, VerticalMeshConfig, + VerticalMeshConfig, + VerticalModelParams, +) +from icon4py.common.dimension import ( + C2E2CDim, + C2E2CODim, + C2EDim, + CellDim, + E2CDim, + E2VDim, + EdgeDim, + V2EDim, + VertexDim, ) -from icon4py.common.dimension import CellDim, EdgeDim, VertexDim, E2CDim, C2EDim, C2E2CDim, \ - C2E2CODim, E2VDim, V2EDim from icon4py.testutils.serialbox_utils import IconSerialDataProvider - - @fixture def with_icon_grid(): data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") @@ -51,7 +59,8 @@ def with_icon_grid(): num_vertices=sp_meta["nproma"], # or rather "num_vert" num_cells=sp_meta["nproma"], # or rather "num_cells" num_edges=sp_meta["nproma"], # or rather "num_edges" - ), VerticalMeshConfig(num_lev=sp_meta["nlev"]) + ), + VerticalMeshConfig(num_lev=sp_meta["nlev"]), ) c2e2c = sp.c2e2c() @@ -62,9 +71,11 @@ def with_icon_grid(): .with_start_end_indices(VertexDim, vertex_starts, vertex_ends) .with_start_end_indices(EdgeDim, edge_starts, edge_ends) .with_start_end_indices(CellDim, cell_starts, cell_ends) - .with_connectivities({C2EDim:sp.c2e(), E2CDim:sp.e2c(), C2E2CDim:c2e2c, C2E2CODim:c2e2c0}) - .with_connectivities({E2VDim:sp.e2v(), V2EDim:sp.v2e()}) + .with_connectivities( + {C2EDim: sp.c2e(), E2CDim: sp.e2c(), C2E2CDim: c2e2c, C2E2CODim: c2e2c0} ) + .with_connectivities({E2VDim: sp.e2v(), V2EDim: sp.v2e()}) + ) return grid @@ -84,23 +95,25 @@ def with_r04b09_diffusion_config() -> DiffusionConfig: horizontal_config = HorizontalMeshConfig( num_vertices=nproma, num_cells=nproma, num_edges=nproma ) - vertical_config = VerticalMeshConfig(num_lev = num_lev) + vertical_config = VerticalMeshConfig(num_lev=num_lev) - grid = IconGrid().with_config(MeshConfig( - horizontal_config=horizontal_config, - vertical_config = vertical_config)) + grid = IconGrid().with_config( + MeshConfig(horizontal_config=horizontal_config, vertical_config=vertical_config) + ) verticalParams = VerticalModelParams( rayleigh_damping_height=12500, vct_a=sp.vct_a() ) - return DiffusionConfig(grid=grid, - vertical_params=verticalParams, - diffusion_type=5, - hdiff_w=True, - type_t_diffu=2, - type_vn_diffu=1, - hdiff_efdt_ratio=24.0, - hdiff_w_efdt_ratio=15.0, - smag_scaling_fac=0.025, - zdiffu_t=True) + return DiffusionConfig( + grid=grid, + vertical_params=verticalParams, + diffusion_type=5, + hdiff_w=True, + type_t_diffu=2, + type_vn_diffu=1, + hdiff_efdt_ratio=24.0, + hdiff_w_efdt_ratio=15.0, + smag_scaling_fac=0.025, + zdiffu_t=True, + ) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index bebb2f77b..794bcaacc 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -16,6 +16,7 @@ import pytest from icon_grid_test_utils import with_icon_grid, with_r04b09_diffusion_config +from icon4py.atm_dyn_iconam import utils from icon4py.atm_dyn_iconam.diagnostic import DiagnosticState from icon4py.atm_dyn_iconam.diffusion import ( Diffusion, @@ -23,15 +24,17 @@ DiffusionParams, VectorTuple, _en_smag_fac_for_zero_nshift, - init_diffusion_local_fields, + _setup_runtime_diff_multfac_vn, + _setup_smag_limit, scale_k, set_zero_v_k, + setup_fields_for_initial_step, ) from icon4py.atm_dyn_iconam.icon_grid import VerticalModelParams from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState from icon4py.atm_dyn_iconam.metric_state import MetricState from icon4py.atm_dyn_iconam.prognostic import PrognosticState -from icon4py.common.dimension import KDim, Koff, VertexDim +from icon4py.common.dimension import KDim, VertexDim from icon4py.testutils.serialbox_utils import ( IconDiffusionInitSavepoint, IconSerialDataProvider, @@ -49,72 +52,90 @@ def test_scale_k(): assert np.allclose(factor * np.asarray(field), scaled_field) -def smag_limit_numpy(shape, k4, substeps): - return 0.125 - 4.0 * diff_multfac_vn_numpy(shape, k4, substeps) +def initial_diff_multfac_vn_numpy(shape, k4, hdiff_efdt_ratio): + return k4 * hdiff_efdt_ratio / 3.0 * np.ones(shape) -def test_init_diff_multifac_vn_const(): +def smag_limit_numpy(func, *args): + return 0.125 - 4.0 * func(*args) + + +def test_diff_multfac_vn_and_smag_limit_for_initial_step(): + mesh = SimpleMesh() + diff_multfac_vn_init = zero_field(mesh, KDim) + smag_limit_init = zero_field(mesh, KDim) + k4 = 1.0 + efdt_ratio = 24.0 + shape = np.asarray(diff_multfac_vn_init).shape + + expected_diff_multfac_vn_init = initial_diff_multfac_vn_numpy(shape, k4, efdt_ratio) + expected_smag_limit_init = smag_limit_numpy( + initial_diff_multfac_vn_numpy, shape, k4, efdt_ratio + ) + + setup_fields_for_initial_step( + k4, efdt_ratio, out=(diff_multfac_vn_init, smag_limit_init), offset_provider={} + ) + + assert np.allclose(expected_diff_multfac_vn_init, diff_multfac_vn_init) + assert np.allclose(expected_smag_limit_init, smag_limit_init) + + +def test_diff_multfac_vn_smag_limit_for_time_step_with_const_value(): mesh = SimpleMesh() diff_multfac_vn = zero_field(mesh, KDim) smag_limit = zero_field(mesh, KDim) k4 = 1.0 substeps = 5.0 + efdt_ratio = 24.0 shape = np.asarray(diff_multfac_vn).shape - enh_smag_fac = zero_field(mesh, KDim) - a_vec = random_field(mesh, KDim, low=1.0, high=10.0) - fac = (0.67, 0.5, 1.3, 0.8) - z = (0.1, 0.2, 0.3, 0.4) - expected_smag_limit = smag_limit_numpy(shape, k4, substeps) expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps) - enhanced_smag_fac_np = enhanced_smagorinski_factor_numpy(fac, z, np.asarray(a_vec)) + expected_smag_limit = smag_limit_numpy(diff_multfac_vn_numpy, shape, k4, substeps) - init_diffusion_local_fields( - k4, - substeps, - *fac, - *z, - a_vec, - diff_multfac_vn, - smag_limit, - enh_smag_fac, - offset_provider={"Koff": KDim}, + _setup_runtime_diff_multfac_vn( + k4, efdt_ratio, out=diff_multfac_vn, offset_provider={} ) + _setup_smag_limit(diff_multfac_vn, out=smag_limit, offset_provider={}) + assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) assert np.allclose(expected_smag_limit, smag_limit) - assert np.allclose(enhanced_smag_fac_np, np.asarray(enh_smag_fac[:-1])) -def test_init_diff_multifac_vn_k4_substeps(): +def test_diff_multfac_vn_smag_limit_for_loop_run_with_k4_substeps(): mesh = SimpleMesh() diff_multfac_vn = zero_field(mesh, KDim) smag_limit = zero_field(mesh, KDim) k4 = 0.003 substeps = 1.0 - enh_smag_fac = zero_field(mesh, KDim) - a_vec = random_field(mesh, KDim, low=1.0, high=10.0) - fac = (0.67, 0.5, 1.3, 0.8) - z = (0.1, 0.2, 0.3, 0.4) shape = np.asarray(diff_multfac_vn).shape expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps) - expected_smag_limit = smag_limit_numpy(shape, k4, substeps) - - init_diffusion_local_fields( - k4, - substeps, - *fac, - *z, - a_vec, - diff_multfac_vn, - smag_limit, - enh_smag_fac, - offset_provider={"Koff": KDim}, + expected_smag_limit = smag_limit_numpy(diff_multfac_vn_numpy, shape, k4, substeps) + _setup_runtime_diff_multfac_vn( + k4, substeps, out=diff_multfac_vn, offset_provider={} ) + _setup_smag_limit(diff_multfac_vn, out=smag_limit, offset_provider={}) + assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) assert np.allclose(expected_smag_limit, smag_limit) +def test_init_enh_smag_fac(): + mesh = SimpleMesh() + enh_smag_fac = zero_field(mesh, KDim) + a_vec = random_field(mesh, KDim, low=1.0, high=10.0) + fac = (0.67, 0.5, 1.3, 0.8) + z = (0.1, 0.2, 0.3, 0.4) + + enhanced_smag_fac_np = enhanced_smagorinski_factor_numpy(fac, z, np.asarray(a_vec)) + + _en_smag_fac_for_zero_nshift( + a_vec, *fac, *z, out=enh_smag_fac, offset_provider={"Koff": KDim} + ) + assert np.allclose(enhanced_smag_fac_np, np.asarray(enh_smag_fac[:-1])) + + def diff_multfac_vn_numpy(shape, k4, substeps): factor = min(1.0 / 128.0, k4 * substeps / 3.0) return factor * np.ones(shape) @@ -221,11 +242,7 @@ def test_diffusion_init(with_r04b09_diffusion_config): additional_parameters = DiffusionParams(config) diffusion = Diffusion(config, additional_parameters, vct_a) - # assert static local fields are initialized and correct: - assert ( - diffusion.smag_offset - == 0.25 * additional_parameters.K4 * config.substep_as_float() - ) + assert diffusion.diff_multfac_w == min( 1.0 / 48.0, additional_parameters.K4W * config.substep_as_float() ) @@ -235,17 +252,24 @@ def test_diffusion_init(with_r04b09_diffusion_config): assert np.allclose(0.0, np.asarray(diffusion.kh_smag_ec)) assert np.allclose(0.0, np.asarray(diffusion.kh_smag_e)) - shape_k = np.asarray(diffusion.diff_multfac_vn.shape) - + shape_k = (config.grid.n_lev(),) expected_smag_limit = smag_limit_numpy( - shape_k, additional_parameters.K4, config.substep_as_float() + diff_multfac_vn_numpy, + shape_k, + additional_parameters.K4, + config.substep_as_float(), + ) + + assert ( + diffusion._smag_offset + == 0.25 * additional_parameters.K4 * config.substep_as_float() ) - assert np.allclose(expected_smag_limit, np.asarray(diffusion.smag_limit)) + assert np.allclose(expected_smag_limit, diffusion._smag_limit) expected_diff_multfac_vn = diff_multfac_vn_numpy( shape_k, additional_parameters.K4, config.substep_as_float() ) - assert np.allclose(expected_diff_multfac_vn, np.asarray(diffusion.diff_multfac_vn)) + assert np.allclose(expected_diff_multfac_vn, diffusion._diff_multfac_vn) expected_enh_smag_fac = enhanced_smagorinski_factor_numpy( additional_parameters.smagorinski_factor, additional_parameters.smagorinski_height, @@ -254,14 +278,15 @@ def test_diffusion_init(with_r04b09_diffusion_config): assert np.allclose(expected_enh_smag_fac, np.asarray(diffusion.enh_smag_fac)) -def verify_init_values_against_savepoint( +def _verify_init_values_against_savepoint( savepoint: IconDiffusionInitSavepoint, diffusion: Diffusion ): dtime = savepoint.get_metadata("dtime")["dtime"] + assert savepoint.nudgezone_diff() == diffusion.nudgezone_diff assert savepoint.bdy_diff() == diffusion.bdy_diff assert savepoint.fac_bdydiff_v() == diffusion.fac_bdydiff_v - assert savepoint.smag_offset() == diffusion.smag_offset + assert savepoint.smag_offset() == diffusion._smag_offset assert savepoint.diff_multfac_w() == diffusion.diff_multfac_w # this is done in diffusion.run(...) because it depends on the dtime @@ -270,11 +295,37 @@ def verify_init_values_against_savepoint( ) assert np.allclose(savepoint.diff_multfac_smag(), diffusion.diff_multfac_smag) - assert np.allclose(savepoint.smag_limit(), diffusion.smag_limit) - assert np.allclose( - savepoint.diff_multfac_n2w(), np.asarray(diffusion.diff_multfac_n2w) + assert np.allclose(savepoint.smag_limit(), diffusion._smag_limit) + np.allclose(savepoint.diff_multfac_n2w(), np.asarray(diffusion.diff_multfac_n2w)) + assert np.allclose(savepoint.diff_multfac_vn(), diffusion._diff_multfac_vn) + + +def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( + with_r04b09_diffusion_config, +): + config = with_r04b09_diffusion_config + + params = DiffusionParams(config) + data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") + serializer = IconSerialDataProvider("icon_diffusion_init", data_path) + serializer.print_info() + first_run_date = "2021-06-20T12:00:10.000" + savepoint = serializer.from_savepoint_init(linit=True, date=first_run_date) + expected_diff_multfac_vn = savepoint.diff_multfac_vn() + expected_smag_limit = savepoint.smag_limit() + exptected_smag_offset = savepoint.smag_offset() + + diff_multfac_vn = utils.zero_field(config.grid, KDim) + smag_limit = utils.zero_field(config.grid, KDim) + setup_fields_for_initial_step( + params.K4, + config.hdiff_efdt_ratio, + out=(diff_multfac_vn, smag_limit), + offset_provider={}, ) - assert np.allclose(savepoint.diff_multfac_vn(), diffusion.diff_multfac_vn) + assert np.allclose(expected_smag_limit, smag_limit) + assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) + assert exptected_smag_offset == 0.0 def test_verify_diffusion_init_against_first_regular_savepoint( @@ -291,7 +342,7 @@ def test_verify_diffusion_init_against_first_regular_savepoint( additional_parameters = DiffusionParams(config) diffusion = Diffusion(config, additional_parameters, vct_a) - verify_init_values_against_savepoint(savepoint, diffusion) + _verify_init_values_against_savepoint(savepoint, diffusion) def test_verify_diffusion_init_against_other_regular_savepoint( @@ -308,7 +359,7 @@ def test_verify_diffusion_init_against_other_regular_savepoint( additional_parameters = DiffusionParams(config) diffusion = Diffusion(config, additional_parameters, vct_a) - verify_init_values_against_savepoint(savepoint, diffusion) + _verify_init_values_against_savepoint(savepoint, diffusion) @pytest.mark.xfail @@ -377,7 +428,7 @@ def test_diffusion_run(with_icon_grid): edge_areas = sp.edge_areas() cell_areas = sp.cell_areas() - diffusion.run( + diffusion.time_step( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, metric_state=metric_state, From aaf26fd664bd115191822f778783c68267e57e17 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 13 Dec 2022 14:55:25 +0100 Subject: [PATCH 029/263] -m switch to gt4py branch https://github.com/tehrengruber/gt4py/tree/reduce_foast_lowering_complexity - fix icon4py unit tests to run again --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 2 ++ .../fused_mo_nh_diffusion_stencil_13_14.py | 6 +++--- .../atm_dyn_iconam/mo_nh_diffusion_stencil_02.py | 3 ++- .../atm_dyn_iconam/mo_nh_diffusion_stencil_14.py | 16 ++++------------ .../src/icon4py/atm_dyn_iconam/utils.py | 3 ++- ...rpolation_scalar_cells2verts_scalar_ri_dsl.py | 2 +- .../test_mo_intp_rbf_rbf_vec_interpol_vertex.py | 2 +- .../test_mo_math_divrot_rot_vertex_ri_dsl.py | 4 +--- ...o_math_gradients_grad_green_gauss_cell_dsl.py | 5 ++++- .../tests/test_mo_nh_diffusion_stencil_02.py | 1 + .../tests/test_mo_nh_diffusion_stencil_07.py | 5 ++++- .../tests/test_mo_nh_diffusion_stencil_08.py | 5 ++++- .../tests/test_mo_nh_diffusion_stencil_09.py | 5 ++++- .../tests/test_mo_nh_diffusion_stencil_14.py | 7 ++++--- .../tests/test_mo_solve_nonhydro_stencil_19.py | 4 +--- .../tests/test_mo_solve_nonhydro_stencil_25.py | 5 ++++- .../tests/test_mo_solve_nonhydro_stencil_30.py | 2 ++ .../tests/test_mo_solve_nonhydro_stencil_31.py | 5 ++++- .../tests/test_mo_solve_nonhydro_stencil_39.py | 6 +++++- .../tests/test_mo_solve_nonhydro_stencil_40.py | 6 +++++- .../tests/test_mo_solve_nonhydro_stencil_41.py | 2 +- .../test_mo_velocity_advection_stencil_01.py | 5 ++++- .../test_mo_velocity_advection_stencil_08.py | 2 +- .../test_mo_velocity_advection_stencil_09.py | 2 +- .../test_mo_velocity_advection_stencil_17.py | 2 +- .../test_mo_velocity_advection_stencil_18.py | 5 ++++- .../test_mo_velocity_advection_stencil_19.py | 1 + .../test_mo_velocity_advection_stencil_20.py | 2 ++ base-requirements-dev.txt | 3 ++- 29 files changed, 75 insertions(+), 43 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index 64e06d1ae..0df504805 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -21,6 +21,7 @@ from functional.ffront.fbuiltins import Field, broadcast, maximum, minimum from functional.iterator.embedded import np_as_located_field from functional.program_processors.runners import gtfn_cpu +from functional.program_processors.runners.gtfn_cpu import run_gtfn from icon4py.atm_dyn_iconam.constants import CPD, GAS_CONSTANT_DRY_AIR from icon4py.atm_dyn_iconam.diagnostic import DiagnosticState @@ -84,6 +85,7 @@ def _setup_runtime_diff_multfac_vn( return broadcast(minimum(con, dyn), (KDim,)) +# @field_operator(backend=run_gtfn) @field_operator def _setup_initial_diff_multfac_vn( k4: float, hdiff_efdt_ratio: float diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py index 0941bbd77..6ebe54825 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -20,7 +20,7 @@ from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_14 import ( _mo_nh_diffusion_stencil_14, ) -from icon4py.common.dimension import CEDim, CellDim, EdgeDim, KDim +from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim @field_operator @@ -28,7 +28,7 @@ def _fused_mo_nh_diffusion_stencil_13_14( kh_smag_e: Field[[EdgeDim, KDim], float], inv_dual_edge_length: Field[[EdgeDim], float], theta_v: Field[[CellDim, KDim], float], - geofac_div: Field[[CEDim], float], + geofac_div: Field[[CellDim, C2EDim], float], ) -> Field[[CellDim, KDim], float]: z_nabla2_e = _mo_nh_diffusion_stencil_13(kh_smag_e, inv_dual_edge_length, theta_v) z_temp = _mo_nh_diffusion_stencil_14(z_nabla2_e, geofac_div) @@ -40,7 +40,7 @@ def fused_mo_nh_diffusion_stencil_13_14( kh_smag_e: Field[[EdgeDim, KDim], float], inv_dual_edge_length: Field[[EdgeDim], float], theta_v: Field[[CellDim, KDim], float], - geofac_div: Field[[CEDim], float], + geofac_div: Field[[CellDim, C2EDim], float], z_temp: Field[[CellDim, KDim], float], ): _fused_mo_nh_diffusion_stencil_13_14( diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_02.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_02.py index c7a8baa85..bd5a454fc 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_02.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_02.py @@ -13,11 +13,12 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, neighbor_sum +from functional.program_processors.runners.gtfn_cpu import run_gtfn from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim -@field_operator +@field_operator(backend=run_gtfn) def _mo_nh_diffusion_stencil_02( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_14.py index 54ad87213..4fbe63e4d 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_14.py @@ -14,30 +14,22 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import ( - C2CE, - C2E, - C2EDim, - CEDim, - CellDim, - EdgeDim, - KDim, -) +from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim @field_operator def _mo_nh_diffusion_stencil_14( z_nabla2_e: Field[[EdgeDim, KDim], float], - geofac_div: Field[[CEDim], float], + geofac_div: Field[[CellDim, C2EDim], float], ) -> Field[[CellDim, KDim], float]: - z_temp = neighbor_sum(z_nabla2_e(C2E) * geofac_div(C2CE), axis=C2EDim) + z_temp = neighbor_sum(z_nabla2_e(C2E) * geofac_div, axis=C2EDim) return z_temp @program def mo_nh_diffusion_stencil_14( z_nabla2_e: Field[[EdgeDim, KDim], float], - geofac_div: Field[[CEDim], float], + geofac_div: Field[[CellDim, C2EDim], float], z_temp: Field[[CellDim, KDim], float], ): _mo_nh_diffusion_stencil_14(z_nabla2_e, geofac_div, out=z_temp) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py index 56447c01c..593fcd40b 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py @@ -15,7 +15,8 @@ from functional.common import Dimension from functional.iterator.embedded import np_as_located_field -#TODO fix duplication: duplicated from test testutils/utils.py + +# TODO fix duplication: duplicated from test testutils/utils.py def zero_field(mesh, *dims: Dimension, dtype=float): shapex = tuple(map(lambda x: mesh.size[x], dims)) return np_as_located_field(*dims)(np.zeros(shapex, dtype=dtype)) diff --git a/atm_dyn_iconam/tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py b/atm_dyn_iconam/tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py index 26fca2a7d..8bf252fdc 100644 --- a/atm_dyn_iconam/tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py +++ b/atm_dyn_iconam/tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py @@ -43,7 +43,7 @@ def test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl(): p_cell_in, c_intp, p_vert_out, - offset_provider={"V2C": mesh.get_v2c_offset_provider()}, + offset_provider={"V2C": mesh.get_v2c_offset_provider(), "V2CDim": V2CDim}, ) assert np.allclose(p_vert_out, ref) diff --git a/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index 7290ac375..a52651ffe 100644 --- a/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -51,7 +51,7 @@ def test_mo_intp_rbf_rbf_vec_interpol_vertex(): ptr_coeff_2, p_v_out, p_u_out, - offset_provider={"V2E": mesh.get_v2e_offset_provider()}, + offset_provider={"V2E": mesh.get_v2e_offset_provider(), "V2EDim": V2EDim}, ) assert np.allclose(p_v_out, p_v_out_ref) diff --git a/atm_dyn_iconam/tests/test_mo_math_divrot_rot_vertex_ri_dsl.py b/atm_dyn_iconam/tests/test_mo_math_divrot_rot_vertex_ri_dsl.py index 65b5ed02e..881683025 100644 --- a/atm_dyn_iconam/tests/test_mo_math_divrot_rot_vertex_ri_dsl.py +++ b/atm_dyn_iconam/tests/test_mo_math_divrot_rot_vertex_ri_dsl.py @@ -46,9 +46,7 @@ def test_mo_math_divrot_rot_vertex_ri_dsl_numpy(): vec_e, geofac_rot, rot_vec, - offset_provider={ - "V2E": mesh.get_v2e_offset_provider(), - }, + offset_provider={"V2E": mesh.get_v2e_offset_provider(), "V2EDim": V2EDim}, ) assert np.allclose(rot_vec, ref) diff --git a/atm_dyn_iconam/tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py b/atm_dyn_iconam/tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py index 2a8d8fbd7..d0938b7d3 100644 --- a/atm_dyn_iconam/tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py +++ b/atm_dyn_iconam/tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py @@ -74,7 +74,10 @@ def test_mo_math_gradients_grad_green_gauss_cell_dsl_numpy(): p_ccpr2, geofac_grg_x, geofac_grg_y, - offset_provider={"C2E2CO": mesh.get_c2e2cO_offset_provider()}, + offset_provider={ + "C2E2CO": mesh.get_c2e2cO_offset_provider(), + "C2E2CODim": C2E2CODim, + }, ) assert np.allclose(p_grad_1_u_ref, p_grad_1_u) assert np.allclose(p_grad_1_v_ref, p_grad_1_v) diff --git a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_02.py b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_02.py index 0dd68565a..4721e1b9d 100644 --- a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_02.py +++ b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_02.py @@ -73,6 +73,7 @@ def test_mo_nh_diffusion_stencil_02(): div, offset_provider={ "C2E": mesh.get_c2e_offset_provider(), + "C2EDim": C2EDim, }, ) assert np.allclose(kh_c, kh_c_ref) diff --git a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_07.py b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_07.py index fe0830d17..cc960c107 100644 --- a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_07.py +++ b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_07.py @@ -43,6 +43,9 @@ def test_mo_nh_diffusion_stencil_07(): w, geofac_n2s, z_nabla2_c, - offset_provider={"C2E2CO": mesh.get_c2e2cO_offset_provider()}, + offset_provider={ + "C2E2CO": mesh.get_c2e2cO_offset_provider(), + "C2E2CODim": C2E2CODim, + }, ) assert np.allclose(z_nabla2_c, ref) diff --git a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_08.py b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_08.py index 8d2605c8c..d5b7f8533 100644 --- a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_08.py +++ b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_08.py @@ -53,7 +53,10 @@ def test_mo_nh_diffusion_stencil_08(): geofac_grg_y, dwdx, dwdy, - offset_provider={"C2E2CO": mesh.get_c2e2cO_offset_provider()}, + offset_provider={ + "C2E2CO": mesh.get_c2e2cO_offset_provider(), + "C2E2CODim": C2E2CODim, + }, ) assert np.allclose(dwdx, dwdx_ref) diff --git a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_09.py b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_09.py index 77abc6367..62329b5f4 100644 --- a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_09.py +++ b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_09.py @@ -61,7 +61,10 @@ def test_mo_nh_diffusion_stencil_09(): geofac_n2s, w, diff_multfac_w, - offset_provider={"C2E2CO": mesh.get_c2e2cO_offset_provider()}, + offset_provider={ + "C2E2CO": mesh.get_c2e2cO_offset_provider(), + "C2E2CODim": C2E2CODim, + }, ) assert np.allclose(w, ref) diff --git a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_14.py b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_14.py index 38977d33c..afb99aa76 100644 --- a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_14.py +++ b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_14.py @@ -19,7 +19,7 @@ ) from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import as_1D_sparse_field, random_field, zero_field +from icon4py.testutils.utils import random_field, zero_field def mo_nh_diffusion_stencil_14_numpy( @@ -35,7 +35,6 @@ def test_mo_nh_diffusion_stencil_14(): z_nabla2_e = random_field(mesh, EdgeDim, KDim) geofac_div = random_field(mesh, CellDim, C2EDim) - geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) out = zero_field(mesh, CellDim, KDim) @@ -44,11 +43,13 @@ def test_mo_nh_diffusion_stencil_14(): ) mo_nh_diffusion_stencil_14( z_nabla2_e, - geofac_div_new, + geofac_div, out, offset_provider={ "C2E": mesh.get_c2e_offset_provider(), "C2CE": StridedNeighborOffsetProvider(CellDim, CEDim, mesh.n_c2e), + "C2EDim": C2EDim, + "CEDim": CEDim, }, ) assert np.allclose(out, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_19.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_19.py index fa7e92ae1..b0591692c 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_19.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_19.py @@ -67,8 +67,6 @@ def test_mo_solve_nonhydro_stencil_19(): c_lin_e, z_dexner_dz_c_1, z_gradh_exner, - offset_provider={ - "E2C": mesh.get_e2c_offset_provider(), - }, + offset_provider={"E2C": mesh.get_e2c_offset_provider(), "E2CDim": E2CDim}, ) assert np.allclose(z_gradh_exner, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_25.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_25.py index 5b1deb566..7780b7c9f 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_25.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_25.py @@ -43,6 +43,9 @@ def test_mo_solve_nonhydro_stencil_25(): geofac_grdiv, z_graddiv_vn, z_graddiv2_vn, - offset_provider={"E2C2EO": mesh.get_e2c2eO_offset_provider()}, + offset_provider={ + "E2C2EO": mesh.get_e2c2eO_offset_provider(), + "E2C2EODim": E2C2EODim, + }, ) assert np.allclose(z_graddiv2_vn, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_30.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_30.py index 6f19b08f1..93744b799 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_30.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_30.py @@ -68,6 +68,8 @@ def test_mo_solve_nonhydro_stencil_30_vt(): offset_provider={ "E2C2EO": mesh.get_e2c2eO_offset_provider(), "E2C2E": mesh.get_e2c2e_offset_provider(), + "E2C2EODim": E2C2EODim, + "E2C2EDim": E2C2EDim, }, ) assert np.allclose(z_vn_avg_ref, z_vn_avg) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_31.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_31.py index 59c342604..b6b8f257d 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_31.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_31.py @@ -43,6 +43,9 @@ def test_mo_solve_nonhydro_stencil_31(): e_flx_avg, vn, z_vn_avg, - offset_provider={"E2C2EO": mesh.get_e2c2eO_offset_provider()}, + offset_provider={ + "E2C2EO": mesh.get_e2c2eO_offset_provider(), + "E2C2EODim": E2C2EODim, + }, ) assert np.allclose(z_vn_avg, ref) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_39.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_39.py index de638b97a..ffa484cbb 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_39.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_39.py @@ -56,7 +56,11 @@ def test_mo_solve_nonhydro_stencil_39(): z_w_concorr_me, wgtfac_c, w_concorr_c, - offset_provider={"Koff": KDim, "C2E": mesh.get_c2e_offset_provider()}, + offset_provider={ + "Koff": KDim, + "C2E": mesh.get_c2e_offset_provider(), + "C2EDim": C2EDim, + }, ) assert np.allclose(w_concorr_c[:, 1:], w_concorr_c_ref[:, 1:]) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_40.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_40.py index bd7d528cf..e8f61e8cd 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_40.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_40.py @@ -63,7 +63,11 @@ def test_mo_solve_nonhydro_stencil_40(): z_w_concorr_me, wgtfacq_c, w_concorr_c, - offset_provider={"Koff": KDim, "C2E": mesh.get_c2e_offset_provider()}, + offset_provider={ + "Koff": KDim, + "C2E": mesh.get_c2e_offset_provider(), + "C2EDim": C2EDim, + }, ) assert np.allclose(w_concorr_c[:, 3:], w_concorr_c_ref[:, 3:]) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_41.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_41.py index ebdf0260c..294416fb7 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_41.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_41.py @@ -51,7 +51,7 @@ def test_mo_solve_nonhydro_stencil_41_z_flxdiv_theta(): z_theta_v_fl_e, z_flxdiv_mass, z_flxdiv_theta, - offset_provider={"C2E": mesh.get_c2e_offset_provider()}, + offset_provider={"C2E": mesh.get_c2e_offset_provider(), "C2EDim": C2EDim}, ) assert np.allclose(z_flxdiv_mass, z_flxdiv_mass_ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_01.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_01.py index 9c4faffd0..c03cfed72 100644 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_01.py +++ b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_01.py @@ -45,6 +45,9 @@ def test_mo_velocity_advection_stencil_01(): vn, rbf_vec_coeff, vt, - offset_provider={"E2C2E": mesh.get_e2c2e_offset_provider()}, + offset_provider={ + "E2C2E": mesh.get_e2c2e_offset_provider(), + "E2C2EDim": E2C2EDim, + }, ) assert np.allclose(vt, ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_08.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_08.py index 993c7b67d..6e52203b4 100644 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_08.py +++ b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_08.py @@ -45,6 +45,6 @@ def test_mo_velocity_advection_stencil_08(): z_kin_hor_e, e_bln_c_s, z_ekinh, - offset_provider={"C2E": mesh.get_c2e_offset_provider()}, + offset_provider={"C2E": mesh.get_c2e_offset_provider(), "C2EDim": C2EDim}, ) assert np.allclose(z_ekinh, ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_09.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_09.py index 59c8552c7..ac42453ba 100644 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_09.py +++ b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_09.py @@ -45,6 +45,6 @@ def test_mo_velocity_advection_stencil_09(): z_w_concorr_me, e_bln_c_s, z_w_concorr_mc, - offset_provider={"C2E": mesh.get_c2e_offset_provider()}, + offset_provider={"C2E": mesh.get_c2e_offset_provider(), "C2EDim": C2EDim}, ) assert np.allclose(z_w_concorr_mc, ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_17.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_17.py index a44a07410..f0b93546d 100644 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_17.py +++ b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_17.py @@ -43,6 +43,6 @@ def test_mo_velocity_advection_stencil_17(): e_bln_c_s, z_v_grad_w, ddt_w_adv, - offset_provider={"C2E": mesh.get_c2e_offset_provider()}, + offset_provider={"C2E": mesh.get_c2e_offset_provider(), "C2EDim": C2EDim}, ) assert np.allclose(ddt_w_adv, ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_18.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_18.py index d9c29162f..9df20eb20 100644 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_18.py +++ b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_18.py @@ -105,7 +105,10 @@ def test_mo_velocity_advection_stencil_18(): scalfac_exdiff, cfl_w_limit, dtime, - offset_provider={"C2E2CO": mesh.get_c2e2cO_offset_provider()}, + offset_provider={ + "C2E2CO": mesh.get_c2e2cO_offset_provider(), + "C2E2CODim": C2E2CODim, + }, ) assert np.allclose(ddt_w_adv, ref) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_19.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_19.py index d43b4058c..bb7fd3caa 100644 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_19.py +++ b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_19.py @@ -121,6 +121,7 @@ def test_mo_velocity_advection_stencil_19(): offset_provider={ "E2V": mesh.get_e2v_offset_provider(), "E2C": mesh.get_e2c_offset_provider(), + "E2CDim": E2CDim, "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, mesh.n_e2c), "Koff": KDim, }, diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_20.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_20.py index e8075489f..4c830b9c0 100644 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_20.py +++ b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_20.py @@ -140,7 +140,9 @@ def test_mo_velocity_advection_stencil_20(): offset_provider={ "Koff": KDim, "E2C": mesh.get_e2c_offset_provider(), + "E2CDim": E2CDim, "E2C2EO": mesh.get_e2c2eO_offset_provider(), + "E2C2EODim": E2C2EODim, "E2V": mesh.get_e2v_offset_provider(), }, ) diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 2c793f104..b76b2a58b 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -1,5 +1,6 @@ # VCS --e git+https://github.com/GridTools/gt4py.git@functional#egg=gt4py-functional +#-e git+https://github.com/GridTools/gt4py.git@functional#egg=gt4py-functional +-e git+https://github.com/tehrengruber/gt4py.git@reduce_foast_lowering_complexity#egg=gt4py-functional # PyPI flake8>=3.8 From 46eb0533cf75e578056320d212c1636e236624a7 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 15 Dec 2022 21:27:07 +0100 Subject: [PATCH 030/263] gt4py program fro diffusion step (WIP) --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 565 ++++++++++-------- .../atm_dyn_iconam/diffusion_program.py | 293 +++++++++ .../icon4py/atm_dyn_iconam/diffusion_utils.py | 32 + .../fused_mo_nh_diffusion_stencil_02_03.py | 1 + .../src/icon4py/atm_dyn_iconam/horizontal.py | 6 +- .../src/icon4py/atm_dyn_iconam/icon_grid.py | 35 +- .../icon4py/atm_dyn_iconam/metric_state.py | 2 +- .../mo_intp_rbf_rbf_vec_interpol_vertex.py | 7 +- atm_dyn_iconam/tests/icon_grid_test_utils.py | 6 +- atm_dyn_iconam/tests/test_diffusion.py | 33 +- atm_dyn_iconam/tests/test_icon_grid.py | 30 +- .../src/icon4py/testutils/serialbox_utils.py | 51 +- 12 files changed, 771 insertions(+), 290 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_utils.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index 0df504805..75a3bb04b 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -13,18 +13,20 @@ import math import sys from collections import namedtuple -from typing import Final +from typing import Final, Tuple import numpy as np from functional.common import Dimension from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, broadcast, maximum, minimum -from functional.iterator.embedded import np_as_located_field +from functional.ffront.fbuiltins import Field, broadcast, maximum, minimum, int32 +from functional.iterator.embedded import np_as_located_field, StridedNeighborOffsetProvider from functional.program_processors.runners import gtfn_cpu from functional.program_processors.runners.gtfn_cpu import run_gtfn from icon4py.atm_dyn_iconam.constants import CPD, GAS_CONSTANT_DRY_AIR from icon4py.atm_dyn_iconam.diagnostic import DiagnosticState +from icon4py.atm_dyn_iconam.diffusion_program import diffusion_run +from icon4py.atm_dyn_iconam.diffusion_utils import set_zero_v_k, scale_k from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import ( _fused_mo_nh_diffusion_stencil_02_03, ) @@ -35,7 +37,7 @@ _fused_mo_nh_diffusion_stencil_07_08_09_10, ) from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_11_12 import ( - fused_mo_nh_diffusion_stencil_11_12, + _fused_mo_nh_diffusion_stencil_11_12, ) from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_13_14 import ( _fused_mo_nh_diffusion_stencil_13_14, @@ -45,7 +47,7 @@ from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState from icon4py.atm_dyn_iconam.metric_state import MetricState from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( - _mo_intp_rbf_rbf_vec_interpol_vertex, + _mo_intp_rbf_rbf_vec_interpol_vertex, mo_intp_rbf_rbf_vec_interpol_vertex, ) from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import ( _mo_nh_diffusion_stencil_01, @@ -63,14 +65,12 @@ EdgeDim, KDim, Koff, - VertexDim, + VertexDim, V2EDim, C2EDim, E2ECV, ) TupleVT = namedtuple("TupleVT", "v t") VectorTuple = namedtuple("VectorTuple", "x y") - - @field_operator def _setup_smag_limit(diff_multfac_vn: Field[[KDim], float]) -> Field[[KDim], float]: return 0.125 - 4.0 * diff_multfac_vn @@ -99,58 +99,6 @@ def setup_fields_for_initial_step(k4: float, hdiff_efdt_ratio: float): smag_limit = _setup_smag_limit(diff_multfac_vn) return diff_multfac_vn, smag_limit - -@field_operator -def _scale_k(field: Field[[KDim], float], factor: float) -> Field[[KDim], float]: - return field * factor - - -@program -def scale_k( - field: Field[[KDim], float], factor: float, scaled_field: Field[[KDim], float] -): - _scale_k(field, factor, out=scaled_field) - - -@field_operator -def _mo_nh_diffusion_stencil_01_scale_dtime( - enh_smag_fac: Field[[KDim], float], - tangent_orientation: Field[[EdgeDim], float], - inv_primal_edge_length: Field[[EdgeDim], float], - inv_vert_vert_length: Field[[EdgeDim], float], - u_vert: Field[[VertexDim, KDim], float], - v_vert: Field[[VertexDim, KDim], float], - primal_normal_vert_x: Field[[ECVDim], float], - primal_normal_vert_y: Field[[ECVDim], float], - dual_normal_vert_x: Field[[ECVDim], float], - dual_normal_vert_y: Field[[ECVDim], float], - vn: Field[[EdgeDim, KDim], float], - smag_limit: Field[[KDim], float], - smag_offset: float, - dtime: float, -) -> tuple[ - Field[[EdgeDim, KDim], float], - Field[[EdgeDim, KDim], float], - Field[[EdgeDim, KDim], float], -]: - diff_multfac_smag = _scale_k(enh_smag_fac, dtime) - return _mo_nh_diffusion_stencil_01( - diff_multfac_smag, - tangent_orientation, - inv_primal_edge_length, - inv_vert_vert_length, - u_vert, - v_vert, - primal_normal_vert_x, - primal_normal_vert_y, - dual_normal_vert_x, - dual_normal_vert_y, - vn, - smag_limit, - smag_offset, - ) - - @field_operator def _en_smag_fac_for_zero_nshift( vect_a: Field[[KDim], float], @@ -249,18 +197,6 @@ def init_diffusion_local_fields( enh_smag_fac, ), ) - - -@field_operator -def _set_zero_v_k() -> Field[[VertexDim, KDim], float]: - return broadcast(0.0, (VertexDim, KDim)) - - -@program -def set_zero_v_k(field: Field[[VertexDim, KDim], float]): - _set_zero_v_k(out=field) - - def init_nabla2_factor_in_upper_damping_zone( k_size: int, nrdmax: int, nshift: int, physical_heights: np.ndarray ) -> Field[[KDim], float]: @@ -291,6 +227,7 @@ def init_nabla2_factor_in_upper_damping_zone( return np_as_located_field(KDim)(buffer) + class DiffusionConfig: """contains necessary parameter to configure a diffusion run. @@ -385,7 +322,7 @@ class DiffusionParams: """Calculates derived quantities depending on the diffusion config.""" def __init__(self, config: DiffusionConfig): - self.boundary_diffusion_start_index_edges = ( + self.boundary_diffusion_start_index_edges = int32( 5 # mo_nh_diffusion.start_bdydiff_e - 1 = 5 -1 ) @@ -450,21 +387,30 @@ def _diffusion_type_5_smagorinski_factor(config: DiffusionConfig): def mo_nh_diffusion_stencil_15_numpy( c2e2c, mask_hdiff: Field[[CellDim, KDim], int], - zd_vertidx: Field[[CellDim, C2E2CDim, KDim], int], + zd_vertidx: Field[[C2E2CDim, KDim], int], zd_diffcoef: Field[[CellDim, KDim], float], geofac_n2s: Field[[CellDim, C2E2CODim], float], vcoef: Field[[C2E2CDim, KDim], float], theta_v: Field[[CellDim, KDim], float], z_temp: Field[[CellDim, KDim], float], - domain, - offset_provider, ): + geofac_n2s = np.asarray(geofac_n2s) + geofac_n2s_nbh = geofac_n2s[:,1:] + geofac_n2s_c = geofac_n2s[:, 0] + mask_hdiff = np.asarray(mask_hdiff) + zd_vertidx = np.asarray(zd_vertidx) + zd_diffcoef = np.asarray(zd_diffcoef) + vcoef = np.asarray(vcoef) + second = (1 - vcoef) * theta_v[c2e2c][zd_vertidx + 1] + + z_temp = np.asarray(z_temp) + vertidx_ = vcoef * theta_v[c2e2c][zd_vertidx] + first = geofac_n2s_nbh * vertidx_ + sum = np.sum(first + second, axis=0) + + z_temp = np.where(mask_hdiff, z_temp + zd_diffcoef* (theta_v* geofac_n2s_c + sum), z_temp) + return z_temp - z_temp = np.sum( - geofac_n2s * vcoef * theta_v[c2e2c[zd_vertidx]] - + (1.0 - vcoef) - + theta_v[c2e2c[zd_vertidx + 1]] - ) class Diffusion: @@ -531,7 +477,7 @@ def _allocate(*dims: Dimension): def _index_field(dim: Dimension, size=None): size = size if size else self.grid.size[dim] - return np_as_located_field(dim)(np.arange(size)) + return np_as_located_field(dim)(np.arange(size, dtype=int32)) self._diff_multfac_vn = _allocate(KDim) @@ -542,7 +488,7 @@ def _index_field(dim: Dimension, size=None): self.kh_smag_e = _allocate(EdgeDim, KDim) self.kh_smag_ec = _allocate(EdgeDim, KDim) self.z_nabla2_e = _allocate(EdgeDim, KDim) - self.z_temp = _allocate(EdgeDim, KDim) + self.z_temp = _allocate(CellDim, KDim) self.diff_multfac_smag = _allocate(KDim) self.vertical_index = _index_field(KDim, self.grid.n_lev() + 1) self.horizontal_cell_index = _index_field(CellDim) @@ -623,23 +569,153 @@ def time_step( runs a diffusion step for the parameter linit=False, within regular time loop. """ - self._do_diffusion_step( - diagnostic_state, - prognostic_state, - metric_state, - interpolation_state, - dtime, + # self._do_diffusion_step( + # diagnostic_state, + # prognostic_state, + # metric_state, + # interpolation_state, + # dtime, + # tangent_orientation, + # inverse_primal_edge_lengths, + # inverse_dual_edge_length, + # inverse_vertical_vertex_lengths, + # primal_normal_vert, + # dual_normal_vert, + # edge_areas, + # cell_areas, + # self._diff_multfac_vn, + # self._smag_limit, + # self._smag_offset, + # ) + + + cell_start_nudging_minus1, cell_end_local_plus1 = self.grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.nudging(CellDim) - 1, + HorizontalMarkerIndex.local(CellDim) + 1, + ) + + cell_start_interior, cell_end_halo = self.grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.interior(CellDim), + HorizontalMarkerIndex.local(CellDim), + ) + + cell_start_nudging, cell_end_local = self.grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.nudging(CellDim), + HorizontalMarkerIndex.local(CellDim), + ) + + edge_start_nudging_plus_one, edge_end_local = self.grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.nudging(EdgeDim) + 1, + HorizontalMarkerIndex.local(EdgeDim), + ) + + edge_start_nudging_minus1, edge_end_local_minus2 = self.grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.nudging(EdgeDim) - 1, + HorizontalMarkerIndex.local(EdgeDim) - 2, + ) + + vertex_start_local_boundary_plus3, vertex_end_local = self.grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.local_boundary(VertexDim) + 3, + HorizontalMarkerIndex.local(VertexDim), + ) + vertex_start_local_boundary_plus1, vertex_end_local_minus1 = self.grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.local_boundary(VertexDim) + 1, + HorizontalMarkerIndex.local(VertexDim) - 1, + ) + + diffusion_run( + diagnostic_state.div_ic, + diagnostic_state.hdef_ic, + diagnostic_state.dwdx, + diagnostic_state.dwdy, + prognostic_state.vertical_wind, + prognostic_state.normal_wind, + prognostic_state.exner_pressure, + prognostic_state.theta_v, + metric_state.theta_ref_mc, + metric_state.wgtfac_c, + metric_state.mask_hdiff, + metric_state.zd_vertidx, + metric_state.zd_diffcoef, + metric_state.zd_intcoef, + interpolation_state.e_bln_c_s, + interpolation_state.rbf_coeff_1, + interpolation_state.rbf_coeff_2, + interpolation_state.geofac_div, + interpolation_state.geofac_grg_x, + interpolation_state.geofac_grg_y, + interpolation_state.nudgecoeff_e, + interpolation_state.geofac_n2s, tangent_orientation, inverse_primal_edge_lengths, inverse_dual_edge_length, inverse_vertical_vertex_lengths, - primal_normal_vert, - dual_normal_vert, + primal_normal_vert[0], + primal_normal_vert[1], + dual_normal_vert[0], + dual_normal_vert[1], edge_areas, cell_areas, self._diff_multfac_vn, + dtime, + self.rd_o_cvd, + self.thresh_tdiff, self._smag_limit, + self.u_vert, + self.v_vert, + self.enh_smag_fac, + self.kh_smag_e, + self.kh_smag_ec, + self.z_nabla2_e, + self.z_temp, + self.diff_multfac_smag, + self.diff_multfac_n2w, self._smag_offset, + self. nudgezone_diff, + self.fac_bdydiff_v, + self.diff_multfac_w, + self.vertical_index, + self.horizontal_cell_index, + self.horizontal_edge_index, + cell_start_nudging_minus1, + cell_start_interior, + cell_start_nudging, + cell_end_local_plus1, + cell_end_local, + edge_start_nudging_plus_one, + edge_start_nudging_minus1, + edge_end_local, + edge_end_local_minus2, + vertex_start_local_boundary_plus3, + vertex_start_local_boundary_plus1, + vertex_end_local, + vertex_end_local_minus1, + self.config.vertical_params.index_of_damping_height, + self.grid.n_lev(), + self.params.boundary_diffusion_start_index_edges, + + offset_provider = { + "V2E": self.grid.get_v2e_connectivity(), + "V2EDim": V2EDim, + "E2C2V": self.grid.get_e2c2v_connectivity(), + "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, 4), + "C2E": self.grid.get_c2e_connectivity(), + "C2EDim": C2EDim, + "E2C": self.grid.get_e2c_connectivity(), + "C2E2C": self.grid.get_c2e2c_connectivity(), + "C2E2CDim": C2E2CDim, + "ECVDim": ECVDim, + "C2E2CODim": C2E2CODim, + "C2E2CO": self.grid.get_c2e2co_connectivity(), + "Koff": KDim + } ) def _do_diffusion_step( @@ -653,8 +729,8 @@ def _do_diffusion_step( inverse_primal_edge_lengths: Field[[EdgeDim], float], inverse_dual_edge_length: Field[[EdgeDim], float], inverse_vertical_vertex_lengths: Field[[EdgeDim], float], - primal_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], - dual_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], + primal_normal_vert: Tuple[Field[[ECVDim], float], Field[[ECVDim], float]], + dual_normal_vert: Tuple[Field[[ECVDim], float], Field[[ECVDim], float]], edge_areas: Field[[EdgeDim], float], cell_areas: Field[[CellDim], float], diff_multfac_vn: Field[[KDim], float], @@ -686,34 +762,89 @@ def _do_diffusion_step( # ------- # OUTLINE # ------- + klevels = self.grid.n_lev() + k_start_end_minus2 = klevels - 2 + + cell_start_nudging_minus1, cell_end_local_plus1 = self.grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.nudging(CellDim) - 1, + HorizontalMarkerIndex.local(CellDim) + 1, + ) + + cell_start_interior, cell_end_halo = self.grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.interior(CellDim), + HorizontalMarkerIndex.local(CellDim), + ) + + cell_start_nudging, cell_end_local = self.grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.nudging(CellDim), + HorizontalMarkerIndex.local(CellDim), + ) + + edge_start_nudging_plus_one, edge_end_local = self.grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.nudging(EdgeDim) + 1, + HorizontalMarkerIndex.local(EdgeDim), + ) + + edge_start_nudging_minus1, edge_end_local_minus2 = self.grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.nudging(EdgeDim) - 1, + HorizontalMarkerIndex.local(EdgeDim) - 2, + ) + + vertex_start_local_boundary_plus3, vertex_end_local = self.grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.local_boundary(VertexDim) + 3, + HorizontalMarkerIndex.local(VertexDim), + ) + vertex_start_local_boundary_plus1, vertex_end_local_minus1 = self.grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.local_boundary(VertexDim) + 1, + HorizontalMarkerIndex.local(VertexDim) - 1, + ) + # Oa TODO: logging # 0b TODO: call timer start # # 0c. dtime dependent stuff: enh_smag_factor, scale_k(self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={}) - # TODO: is this needed, if not remove + + + # TODO: is this needed?, if not remove set_zero_v_k(self.u_vert, offset_provider={}) set_zero_v_k(self.v_vert, offset_provider={}) - # 1. CALL rbf_vec_interpol_vertex - _mo_intp_rbf_rbf_vec_interpol_vertex( + # # 1. CALL rbf_vec_interpol_vertex + + v_start = vertex_start_local_boundary_plus1 + v_end = vertex_end_local_minus1 + k_start = 0 + k_end = klevels + mo_intp_rbf_rbf_vec_interpol_vertex( prognostic_state.normal_wind, interpolation_state.rbf_coeff_1, interpolation_state.rbf_coeff_2, - out=(self.u_vert, self.v_vert), - domain={ - VertexDim: self.grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 1, - HorizontalMarkerIndex.halo(VertexDim) - 1, - ), - }, - offset_provider={"V2E": self.grid.get_v2e_offset_provider()}, + self.u_vert, + self.v_vert, + v_start, + v_end, + k_start, + k_end, + offset_provider={"V2E": self.grid.get_v2e_connectivity(), "V2EDim":V2EDim}, ) # 2. HALO EXCHANGE -- CALL sync_patch_array_mult # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 + + + e_start = self.params.boundary_diffusion_start_index_edges + e_end = edge_end_local_minus2 + k_start = 0 + k_end = klevels _mo_nh_diffusion_stencil_01( self.diff_multfac_smag, tangent_orientation, @@ -728,20 +859,16 @@ def _do_diffusion_step( prognostic_state.normal_wind, smag_limit, smag_offset, - domain={ - EdgeDim: ( - self.grid.get_indices_from_to( - EdgeDim, - self.params.boundary_diffusion_start_index_edges, - HorizontalMarkerIndex.halo(EdgeDim) - 2, - ) - ), - }, out=(self.kh_smag_e, self.kh_smag_ec, self.z_nabla2_e), - offset_provider={"E2C2V": self.grid.get_e2c2v_connectivity()}, - backend=gtfn_cpu.run_gtfn, + offset_provider={"E2C2V": self.grid.get_e2c2v_connectivity(), + "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, 4)}, ) + + c_start = cell_start_nudging, + c_end = cell_end_local, + k_start = 0 + k_end = klevels _fused_mo_nh_diffusion_stencil_02_03( self.kh_smag_ec, prognostic_state.normal_wind, @@ -753,52 +880,47 @@ def _do_diffusion_step( diagnostic_state.div_ic, diagnostic_state.hdef_ic, ), - domain={ - CellDim: ( - self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.halo(CellDim), - ) - ), - }, - offset_provider={"C2E": self.grid.get_c2e_connectivity()}, - backend=gtfn_cpu.run_gtfn, - ) - # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH + offset_provider={"C2E": self.grid.get_c2e_connectivity(), "C2EDim":C2EDim}, + ) + # + # # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH + # + # # 5. CALL rbf_vec_interpol_vertex_wp - # 5. CALL rbf_vec_interpol_vertex_wp - _mo_intp_rbf_rbf_vec_interpol_vertex( + k_start = 0 + k_end = klevels + v_start = vertex_start_local_boundary_plus3 + v_end = vertex_end_local + mo_intp_rbf_rbf_vec_interpol_vertex( self.z_nabla2_e, interpolation_state.rbf_coeff_1, interpolation_state.rbf_coeff_2, - out=(self.u_vert, self.v_vert), - domain={ - VertexDim: self.grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 3, - HorizontalMarkerIndex.halo(VertexDim), - ), - }, - offset_provider={"V2E": self.grid.get_e2v_connectivity()}, + self.u_vert, + self.v_vert, + v_start, + v_end, + k_start, k_end, + offset_provider={"V2E": self.grid.get_v2e_connectivity(), + "V2EDim":V2EDim}, ) - # 6. HALO EXCHANGE -- CALL sync_patch_array_mult + # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult + # + # # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 + # # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 + # - # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 - # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 - start_2nd_nudge_line = self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim) - 1, - HorizontalMarkerIndex.nudging(EdgeDim) - 1, - )[0] + e_start = edge_start_nudging_plus_one + e_end = edge_end_local + k_start = 0 + k_end = klevels _fused_mo_nh_diffusion_stencil_04_05_06( self.u_vert, self.v_vert, - primal_normal_vert.x, - primal_normal_vert.y, + primal_normal_vert[0], + primal_normal_vert[1], self.z_nabla2_e, inverse_vertical_vertex_lengths, inverse_primal_edge_lengths, @@ -810,25 +932,19 @@ def _do_diffusion_step( self.horizontal_edge_index, self.nudgezone_diff, self.fac_bdydiff_v, - start_2nd_nudge_line, + edge_start_nudging_minus1, out=prognostic_state.normal_wind, - domain={ - EdgeDim: self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim) + 1, - HorizontalMarkerIndex.halo(EdgeDim), - ), - }, - offset_provider={"E2C2V": self.grid.get_e2c2v_connectivity()}, - ) - # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, - # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 - interior_start_index, halo_endindex = self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.interior(CellDim), - HorizontalMarkerIndex.halo(CellDim), + offset_provider={"E2C2V": self.grid.get_e2c2v_connectivity(), + "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, 4), + "ECVDim":ECVDim,} ) + # # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, + # # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 + c_start = cell_start_nudging + c_end = cell_end_local_plus1 + k_start = 0 + k_end = klevels _fused_mo_nh_diffusion_stencil_07_08_09_10( cell_areas, interpolation_state.geofac_n2s, @@ -843,109 +959,90 @@ def _do_diffusion_step( self.vertical_index, self.horizontal_cell_index, self.config.vertical_params.index_of_damping_height, - interior_start_index, - halo_endindex, + cell_start_interior, + cell_end_halo, out=( prognostic_state.vertical_wind, diagnostic_state.dwdx, diagnostic_state.dwdy, ), - domain={ - CellDim: self.grid.get_indices_from_to( - CellDim, - # TODO: global mode is from NUDGING - 1 - HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.halo(CellDim) + 1, - ), - }, - offset_provider={"C2E2CO": self.grid.get_c2e2c0_connectivity()}, + + offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity(), "C2E2CODim":C2E2CODim}, ) - # 8. HALO EXCHANGE: CALL sync_patch_array - # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, - # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 + # # 8. HALO EXCHANGE: CALL sync_patch_array + # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, + # # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 + # + # # TODO check: kh_smag_e is an out field, should not be calculated in init? + # - # TODO check: kh_smag_e is an out field, should not be calculated in init? - klevels = self.grid.n_lev() - fused_mo_nh_diffusion_stencil_11_12( + e_start = edge_start_nudging_plus_one + e_end = edge_end_local + c_start= cell_start_nudging_minus1 + c_end = cell_end_local_plus1 + k_start = k_start_end_minus2 + k_end = klevels + + _fused_mo_nh_diffusion_stencil_11_12( prognostic_state.theta_v, metric_state.theta_ref_mc, self.thresh_tdiff, + self.kh_smag_e, out=self.kh_smag_e, - domain={ - KDim: (klevels - 2, klevels), - CellDim: self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim) - 1, - HorizontalMarkerIndex.halo(CellDim) + 1, - ), - EdgeDim: self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim) + 1, - HorizontalMarkerIndex.halo(EdgeDim), - ), - }, + offset_provider={ "E2C": self.grid.get_e2c_connectivity(), "C2E2C": self.grid.get_c2e2c_connectivity(), + "C2E2CDim":C2E2CDim }, ) + + + c_start = cell_start_nudging, + c_end = cell_end_local _fused_mo_nh_diffusion_stencil_13_14( self.kh_smag_e, inverse_dual_edge_length, prognostic_state.theta_v, interpolation_state.geofac_div, out=self.z_temp, - domain={ - CellDim: self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.halo(CellDim), - ), - EdgeDim: self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim), - HorizontalMarkerIndex.halo(EdgeDim) - 1, - ), - }, offset_provider={ "C2E": self.grid.get_c2e_connectivity(), "E2C": self.grid.get_e2c_connectivity(), + "C2EDim":C2EDim, }, ) - mo_nh_diffusion_stencil_15_numpy( - mask_hdiff=metric_state.mask_hdiff, - zd_vertidx=metric_state.zd_vertidx, - vcoef=metric_state.zd_diffcoef, - zd_diffcoef=metric_state.zd_diffcoef, - geofac_n2s=interpolation_state.geofac_n2s, - theta_v=prognostic_state.theta_v, - z_temp=self.z_temp, - domain={ - CellDim: self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.halo(CellDim), - ), - }, - offset_provider={}, - ) - + # mo_nh_diffusion_stencil_15_numpy( + # mask_hdiff=metric_state.mask_hdiff, + # zd_vertidx=metric_state.zd_vertidx, + # vcoef=metric_state.zd_diffcoef, + # zd_diffcoef=metric_state.zd_diffcoef, + # geofac_n2s=interpolation_state.geofac_n2s, + # theta_v=prognostic_state.theta_v, + # z_temp=self.z_temp, + # domain={ + # CellDim: self.grid.get_indices_from_to( + # CellDim, + # HorizontalMarkerIndex.nudging(CellDim), + # HorizontalMarkerIndex.halo(CellDim), + # ), + # }, + # offset_provider={}, + # ) + + + c_start = cell_start_nudging + c_end = cell_end_local _mo_nh_diffusion_stencil_16( self.z_temp, cell_areas, prognostic_state.theta_v, prognostic_state.exner_pressure, self.rd_o_cvd, - out=(prognostic_state.theta_v, prognostic_state.exner_pressure), - domain={ - CellDim: self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.halo(CellDim), - ), - }, offset_provider={}, ) # 10. HALO EXCHANGE sync_patch_array + + diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py new file mode 100644 index 000000000..c4d7f111b --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py @@ -0,0 +1,293 @@ +from typing import Tuple + +from functional.common import Field +from functional.ffront.fbuiltins import int32 +from functional.ffront.decorator import program + +from functional.program_processors.runners import gtfn_cpu + +from icon4py.atm_dyn_iconam.diffusion_utils import _scale_k, _set_zero_v_k +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import \ + _fused_mo_nh_diffusion_stencil_02_03 +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import \ + _fused_mo_nh_diffusion_stencil_04_05_06 +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_07_08_09_10 import \ + _fused_mo_nh_diffusion_stencil_07_08_09_10 +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_11_12 import \ + _fused_mo_nh_diffusion_stencil_11_12 +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_13_14 import \ + _fused_mo_nh_diffusion_stencil_13_14 +from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import \ + _mo_intp_rbf_rbf_vec_interpol_vertex +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import _mo_nh_diffusion_stencil_01 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_16 import _mo_nh_diffusion_stencil_16 +from icon4py.common.dimension import CellDim, KDim, EdgeDim, C2E2CDim, C2EDim, V2EDim, VertexDim, \ + C2E2CODim, ECVDim + + +@program(backend=gtfn_cpu.run_gtfn) +def diffusion_run( + diagnostic_hdef_ic: Field[[CellDim, KDim], float], + diagnostic_div_ic: Field[[CellDim, KDim], float], + diagnostic_dwdx: Field[[CellDim, KDim], float], + diagnostic_dwdy: Field[[CellDim, KDim], float], + prognostic_vertical_wind: Field[[CellDim, KDim], float], + prognostic_normal_wind: Field[[EdgeDim, KDim], float], + prognostic_exner_pressure: Field[[CellDim, KDim], float], + prognostic_theta_v: Field[[CellDim, KDim], float], + metric_theta_ref_mc: Field[[CellDim, KDim], float], + metric_wgtfac_c: Field[[CellDim, KDim], float], + metric_mask_hdiff: Field[[CellDim, KDim], int], + metric_zd_vertidx: Field[[CellDim, C2E2CDim, KDim], int], + metric_zd_diffcoef: Field[[CellDim, KDim], float], + metric_zd_intcoef: Field[[CellDim, C2E2CDim, KDim], float], + interpolation_e_bln_c_s: Field[[CellDim, C2EDim], float], + interpolation_rbf_coeff_1: Field[[VertexDim, V2EDim], float], + interpolation_rbf_coeff_2: Field[[VertexDim, V2EDim], float], + interpolation_geofac_div: Field[[CellDim, C2EDim], float], + interpolation_geofac_grg_x: Field[[CellDim, C2E2CODim], float] , + interpolation_geofac_grg_y: Field[[CellDim, C2E2CODim], float], + interpolation_nudgecoeff_e: Field[[EdgeDim], float], + interpolation_geofac_n2s:Field[[CellDim, C2E2CODim], float], + tangent_orientation: Field[[EdgeDim], float], + inverse_primal_edge_lengths: Field[[EdgeDim], float], + inverse_dual_edge_length: Field[[EdgeDim], float], + inverse_vertical_vertex_lengths: Field[[EdgeDim], float], + primal_normal_vert_1: Field[[ECVDim], float], + primal_normal_vert_2: Field[[ECVDim], float], + dual_normal_vert_1: Field[[ECVDim], float], + dual_normal_vert_2: Field[[ECVDim], float], + edge_areas: Field[[EdgeDim], float], + cell_areas: Field[[CellDim], float], + diff_multfac_vn: Field[[KDim], float], + dtime: float, + rd_o_cvd:float, + local_thresh_tdiff: float, + local_smag_limit: Field[[KDim], float], + local_u_vert: Field[[VertexDim, KDim], float], + local_v_vert: Field[[VertexDim, KDim], float], + local_enh_smag_fac:Field[[KDim], float], + local_kh_smag_e:Field[[EdgeDim, KDim], float], + local_kh_smag_ec:Field[[EdgeDim, KDim], float], + local_z_nabla2_e:Field[[EdgeDim, KDim], float], + local_z_temp :Field[[CellDim, KDim], float], + local_diff_multfac_smag:Field[[KDim], float], + local_diff_multfac_n2w: Field[[KDim], float], + local_smag_offset: float, + local_nudgezone_diff: float, + local_fac_bdydiff_v: float, + local_diff_multfac_w: float, + + local_vertical_index:Field[[KDim], int32], #local_grid.n_lev() + 1 + local_horizontal_cell_index:Field[[CellDim], int32], + local_horizontal_edge_index: Field[[EdgeDim], int32], + cell_startindex_nudging_minus1: int32, + cell_startindex_interior:int32, + cell_startindex_nudging:int32, + cell_endindex_local_plus1: int32, + cell_endindex_local:int32, + edge_startindex_nudging_plus1:int32, + edge_startindex_nudging_minus1:int32, + edge_endindex_local:int32, + edge_endindex_local_minus2:int32, + vertex_startindex_lb_plus3:int32, + vertex_startindex_lb_plus1:int32, + vertex_endindex_local:int32, + vertex_endindex_local_minus1: int32, + index_of_damping_height: int32, + nlev:int, + boundary_diffusion_start_index_edges: int32 +): + _scale_k(local_enh_smag_fac, dtime, out=local_diff_multfac_smag) + + # TODO: is this needed?, if not remove + _set_zero_v_k(out=local_u_vert) + _set_zero_v_k(out=local_v_vert) + + # # 1. CALL rbf_vec_interpol_vertex + _mo_intp_rbf_rbf_vec_interpol_vertex( + prognostic_normal_wind, + interpolation_rbf_coeff_1, + interpolation_rbf_coeff_2, + out=(local_u_vert,local_v_vert), + domain={ + VertexDim: (vertex_startindex_lb_plus1, vertex_endindex_local_minus1), + KDim: (0,nlev) + } + ) + + # 2. HALO EXCHANGE -- CALL sync_patch_array_mult + # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 + + _mo_nh_diffusion_stencil_01( + local_diff_multfac_smag, + tangent_orientation, + inverse_primal_edge_lengths, + inverse_vertical_vertex_lengths, + local_u_vert, + local_v_vert, + primal_normal_vert_1, + primal_normal_vert_2, + dual_normal_vert_1, + dual_normal_vert_2, + prognostic_normal_wind, + local_smag_limit, + local_smag_offset, + out=(local_kh_smag_e, local_kh_smag_ec, local_z_nabla2_e), + domain={ + EdgeDim:(boundary_diffusion_start_index_edges, edge_endindex_local_minus2), + KDim: (0,nlev) + } + + ) + + _fused_mo_nh_diffusion_stencil_02_03( + local_kh_smag_ec, + prognostic_normal_wind, + interpolation_e_bln_c_s, + interpolation_geofac_div, + local_diff_multfac_smag, + metric_wgtfac_c, + out=( + diagnostic_div_ic, + diagnostic_hdef_ic, + ), + domain={CellDim:(cell_startindex_nudging, cell_endindex_local), + KDim: (0, nlev)} + + ) + # + # # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH + # + # # 5. CALL rbf_vec_interpol_vertex_wp + _mo_intp_rbf_rbf_vec_interpol_vertex( + local_z_nabla2_e, + interpolation_rbf_coeff_1, + interpolation_rbf_coeff_2, + out = (local_u_vert,local_v_vert), + domain={VertexDim:(vertex_startindex_lb_plus3, vertex_endindex_local), + KDim: (0, nlev) + } + + ) + # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult + # + # # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 + # # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 + # + _fused_mo_nh_diffusion_stencil_04_05_06( + local_u_vert, + local_v_vert, + primal_normal_vert_1, + primal_normal_vert_2, + local_z_nabla2_e, + inverse_vertical_vertex_lengths, + inverse_primal_edge_lengths, + edge_areas, + local_kh_smag_e, + diff_multfac_vn, + interpolation_nudgecoeff_e, + prognostic_normal_wind, + local_horizontal_edge_index, + local_nudgezone_diff, + local_fac_bdydiff_v, + edge_startindex_nudging_minus1, + out=prognostic_normal_wind, + domain={ + EdgeDim: (edge_startindex_nudging_plus1, edge_endindex_local), + KDim: (0, nlev) + } + + ) + # # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, + # # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 + + _fused_mo_nh_diffusion_stencil_07_08_09_10( + cell_areas, + interpolation_geofac_n2s, + interpolation_geofac_grg_x, + interpolation_geofac_grg_y, + prognostic_vertical_wind, + prognostic_vertical_wind, + diagnostic_dwdx, + diagnostic_dwdy, + local_diff_multfac_w, + local_diff_multfac_n2w, + local_vertical_index, + local_horizontal_cell_index, + index_of_damping_height, + cell_startindex_interior, + cell_endindex_local, + out=( + prognostic_vertical_wind, + diagnostic_dwdx, + diagnostic_dwdy, + ), + domain={ + CellDim: (cell_startindex_nudging, cell_endindex_local_plus1), + KDim: (0, nlev) + } + ) + # # 8. HALO EXCHANGE: CALL sync_patch_array + # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, + # # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 + # + # # TODO check: kh_smag_e is an out field, should not be calculated in init? + # + + _fused_mo_nh_diffusion_stencil_11_12( + prognostic_theta_v, + metric_theta_ref_mc, + local_thresh_tdiff, + local_kh_smag_e, + out=local_kh_smag_e, + domain={ + #CellDim: (cell_startindex_nudging_minus1, cell_endindex_local_plus1), + EdgeDim: (edge_startindex_nudging_plus1, edge_endindex_local), + KDim: (nlev-2, nlev) + } + + ) + _fused_mo_nh_diffusion_stencil_13_14( + local_kh_smag_e, + inverse_dual_edge_length, + prognostic_theta_v, + interpolation_geofac_div, + out=local_z_temp, + domain={ + CellDim:(cell_startindex_nudging, cell_endindex_local), + KDim: (0, nlev) + } + + ) + + # mo_nh_diffusion_stencil_15_numpy( + # mask_hdiff=metric_state.mask_hdiff, + # zd_vertidx=metric_state.zd_vertidx, + # vcoef=metric_state.zd_diffcoef, + # zd_diffcoef=metric_state.zd_diffcoef, + # geofac_n2s=interpolation_state.geofac_n2s, + # theta_v=prognostic_state.theta_v, + # z_temp=self.z_temp, + # domain={ + # CellDim: self.grid.get_indices_from_to( + # CellDim, + # HorizontalMarkerIndex.nudging(CellDim), + # HorizontalMarkerIndex.local(CellDim), + # ), + # }, + # + # ) + + _mo_nh_diffusion_stencil_16( + local_z_temp, + cell_areas, + prognostic_theta_v, + prognostic_exner_pressure, + rd_o_cvd, + out=(prognostic_theta_v,prognostic_exner_pressure), + domain = {CellDim: (cell_startindex_nudging, cell_endindex_local), + KDim: (0, nlev)} + ) + # 10. HALO EXCHANGE sync_patch_array + diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_utils.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_utils.py new file mode 100644 index 000000000..73c4a64f5 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_utils.py @@ -0,0 +1,32 @@ +from functional.common import Field +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import broadcast + + +from icon4py.common.dimension import KDim, VertexDim + + + +@field_operator +def _scale_k(field: Field[[KDim], float], factor: float) -> Field[[KDim], float]: + return field * factor + + +@program +def scale_k( + field: Field[[KDim], float], factor: float, scaled_field: Field[[KDim], float] +): + _scale_k(field, factor, out=scaled_field) + + + +@field_operator +def _set_zero_v_k() -> Field[[VertexDim, KDim], float]: + return broadcast(0.0, (VertexDim, KDim)) + + +@program +def set_zero_v_k(field: Field[[VertexDim, KDim], float]): + _set_zero_v_k(out=field) + + diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py index 6b32d5a67..da9b3f494 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -13,6 +13,7 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field +from functional.program_processors.runners.gtfn_cpu import run_gtfn from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import ( _mo_nh_diffusion_stencil_02, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py index ab18f397e..00e5ff27c 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py @@ -66,8 +66,10 @@ def local_boundary(cls, dim: Dimension) -> int: return cls._LOCAL_BOUNDARY_VERTICES @classmethod - # TODO should be renamed "halo" in dusk means "after this the halo starts", so is essentially local - def halo(cls, dim: Dimension) -> int: + def local(cls, dim: Dimension) -> int: + """ + + """ match (dim): case (dimension.CellDim): return cls._HALO_CELLS diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py index 0a79a69ba..ae08bc384 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py @@ -14,6 +14,7 @@ from typing import Dict, Tuple import numpy as np +from functional.ffront.fbuiltins import int32, int64 from functional.common import Dimension, DimensionKind, Field from functional.iterator.embedded import ( NeighborTableOffsetProvider, @@ -21,7 +22,8 @@ ) from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig -from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim +from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim, C2E2CODim, C2EDim, C2E2CDim, \ + E2VDim, E2CDim class VerticalMeshConfig: @@ -78,12 +80,14 @@ def wrapper(self, *args, **kwargs): return wrapper + + class IconGrid: def __init__(self): self.config: MeshConfig = None self.start_indices = {} self.end_indices = {} - self.connectivities: Dict[Dimension, np.ndarray] = {} + self.connectivities: Dict[str, np.ndarray] = {} self.size: Dict[Dimension, int] = {} def _update_size(self, config: MeshConfig): @@ -128,12 +132,21 @@ def num_edges(self): def get_indices_from_to( self, dim: Dimension, start_marker: int, end_marker: int - ) -> Tuple[int, int]: + ) -> Tuple[int32, int32]: + """ + Use to specifzy domains of a field for field_operator. + + For a given dimension, returns the start and end index if a + horizontal region in a field given by the markers. + + field operators will then run from start of the region given by the + start_marker to the end of the region given by the end_marker + """ if dim.kind != DimensionKind.HORIZONTAL: raise ValueError( "only defined for {} dimension kind ", DimensionKind.HORIZONTAL ) - return self.start_indices[dim][start_marker], self.end_indices[dim][end_marker] + return self.start_indices[dim][start_marker],self.end_indices[dim][end_marker] def get_c2e_connectivity(self): table = self.connectivities["c2e"] @@ -151,14 +164,18 @@ def get_c2e2c_connectivity(self): table = self.connectivities["c2e2c"] return NeighborTableOffsetProvider(table, CellDim, CellDim, table.shape[1]) - def get_c2e2c0_connectivity(self): - table = self.connectivities["c2e2c0"] + def get_c2e2co_connectivity(self): + table = self.connectivities["c2e2co"] return NeighborTableOffsetProvider(table, CellDim, CellDim, table.shape[1]) def get_e2c2v_connectivity(self): return self.get_e2v_connectivity() - def get_v2e_offset_provider(self): + def get_e2c2v_size(self): + self.connectivities["e2v"].shape[1] + + + def get_v2e_connectivity(self): table = self.connectivities["v2e"] return NeighborTableOffsetProvider(table, VertexDim, EdgeDim, table.shape[1]) @@ -174,8 +191,8 @@ def __init__(self, vct_a: np.ndarray, rayleigh_damping_height: float): """ self.rayleigh_damping_height = rayleigh_damping_height self.vct_a = vct_a - self.index_of_damping_height = np.argmax( - np.where(np.asarray(self.vct_a) >= self.rayleigh_damping_height) + self.index_of_damping_height = int32(np.argmax( + np.where(np.asarray(self.vct_a) >= self.rayleigh_damping_height)) ) def get_index_of_damping_layer(self): diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py index 8ec577faf..06d66ca82 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py @@ -27,4 +27,4 @@ class MetricState: mask_hdiff: Field[[CellDim, KDim], int] zd_vertidx: Field[[CellDim, C2E2CDim, KDim], int] zd_diffcoef: Field[[CellDim, KDim], float] - zd_intcoef: Field[[CellDim, KDim], float] + zd_intcoef: Field[[CellDim,C2E2CDim, KDim], float] diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py index c8668b357..9913f48f4 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -10,6 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +from typing import Tuple from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, neighbor_sum @@ -35,7 +36,11 @@ def mo_intp_rbf_rbf_vec_interpol_vertex( ptr_coeff_2: Field[[VertexDim, V2EDim], float], p_u_out: Field[[VertexDim, KDim], float], p_v_out: Field[[VertexDim, KDim], float], + v_start: int, + v_end: int, + k_start: int, + k_end:int, ): _mo_intp_rbf_rbf_vec_interpol_vertex( - p_e_in, ptr_coeff_1, ptr_coeff_2, out=(p_u_out, p_v_out) + p_e_in, ptr_coeff_1, ptr_coeff_2, out=(p_u_out, p_v_out), domain={VertexDim: (v_start, v_end), KDim:(k_start, k_end)} ) diff --git a/atm_dyn_iconam/tests/icon_grid_test_utils.py b/atm_dyn_iconam/tests/icon_grid_test_utils.py index 3ec41e820..6705dc20b 100644 --- a/atm_dyn_iconam/tests/icon_grid_test_utils.py +++ b/atm_dyn_iconam/tests/icon_grid_test_utils.py @@ -14,7 +14,7 @@ import os import numpy as np -from _pytest.fixtures import fixture +import pytest from icon4py.atm_dyn_iconam.diffusion import DiffusionConfig from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig @@ -38,7 +38,7 @@ from icon4py.testutils.serialbox_utils import IconSerialDataProvider -@fixture +@pytest.fixture def with_icon_grid(): data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") sp = IconSerialDataProvider( @@ -79,7 +79,7 @@ def with_icon_grid(): return grid -@fixture +@pytest.fixture def with_r04b09_diffusion_config() -> DiffusionConfig: """ Create DiffusionConfig. diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 794bcaacc..da0a7297d 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -15,7 +15,7 @@ import numpy as np import pytest from icon_grid_test_utils import with_icon_grid, with_r04b09_diffusion_config - +from functional.ffront.fbuiltins import int32 from icon4py.atm_dyn_iconam import utils from icon4py.atm_dyn_iconam.diagnostic import DiagnosticState from icon4py.atm_dyn_iconam.diffusion import ( @@ -28,20 +28,34 @@ _setup_smag_limit, scale_k, set_zero_v_k, - setup_fields_for_initial_step, + setup_fields_for_initial_step, mo_nh_diffusion_stencil_15_numpy, ) from icon4py.atm_dyn_iconam.icon_grid import VerticalModelParams from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState from icon4py.atm_dyn_iconam.metric_state import MetricState from icon4py.atm_dyn_iconam.prognostic import PrognosticState -from icon4py.common.dimension import KDim, VertexDim +from icon4py.common.dimension import KDim, VertexDim, ECVDim, CellDim, C2E2CODim, C2E2CDim from icon4py.testutils.serialbox_utils import ( IconDiffusionInitSavepoint, IconSerialDataProvider, ) from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field +from icon4py.testutils.utils import random_field, zero_field, as_1D_sparse_field, constant_field, \ + random_mask + + +def test_mo_nh_diffusion_stencil_15(): + mesh = SimpleMesh() + mask_diff = constant_field(mesh, 1, CellDim, KDim, dtype=int32) + zd_vertidx = constant_field(mesh, 1, C2E2CDim, KDim, dtype=int32) + zd_diffcoef = random_field(mesh, CellDim, KDim) + geofac_n2s = random_field(mesh, CellDim, C2E2CODim) + vcoef = random_field(mesh, C2E2CDim, KDim) + theta_v = random_field(mesh, CellDim, KDim) + z_temp = random_field(mesh, CellDim, KDim) + # mo_nh_diffusion_stencil_15_numpy(mesh.c2e2c, + # mask_diff, zd_vertidx, zd_diffcoef, geofac_n2s, vcoef, theta_v, z_temp) def test_scale_k(): mesh = SimpleMesh() @@ -362,7 +376,6 @@ def test_verify_diffusion_init_against_other_regular_savepoint( _verify_init_values_against_savepoint(savepoint, diffusion) -@pytest.mark.xfail def test_diffusion_run(with_icon_grid): data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") data_provider = IconSerialDataProvider("icon_diffusion_init", data_path) @@ -404,7 +417,7 @@ def test_diffusion_run(with_icon_grid): ) metric_state = MetricState( - mask_hdiff=sp.theta_ref_mc(), + mask_hdiff=sp.mask_diff(), theta_ref_mc=sp.theta_ref_mc(), wgtfac_c=sp.wgtfac_c(), zd_intcoef=sp.zd_intcoef(), @@ -418,12 +431,12 @@ def test_diffusion_run(with_icon_grid): inverse_vertical_vertex_lengths = sp.inv_vert_vert_length() inverse_dual_edge_length = sp.inv_dual_edge_length() primal_normal_vert: VectorTuple = ( - sp.primal_normal_vert_x(), - sp.primal_normal_vert_y(), + as_1D_sparse_field(sp.primal_normal_vert_x(), ECVDim), + as_1D_sparse_field(sp.primal_normal_vert_y(), ECVDim), ) dual_normal_vert: VectorTuple = ( - sp.dual_normal_vert_x(), - sp.dual_normal_vert_y(), + as_1D_sparse_field(sp.dual_normal_vert_x(), ECVDim), + as_1D_sparse_field(sp.dual_normal_vert_y(), ECVDim), ) edge_areas = sp.edge_areas() cell_areas = sp.cell_areas() diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index 262d63128..cda0b8d31 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -10,6 +10,8 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +from typing import Tuple + from icon_grid_test_utils import with_icon_grid from icon4py.atm_dyn_iconam.horizontal import HorizontalMarkerIndex @@ -19,16 +21,16 @@ def test_horizontal_grid_cell_indices(with_icon_grid): assert with_icon_grid.get_indices_from_to( CellDim, - HorizontalMarkerIndex.halo(CellDim) - 1, - HorizontalMarkerIndex.halo(CellDim) - 1, + HorizontalMarkerIndex.local(CellDim) - 1, + HorizontalMarkerIndex.local(CellDim) - 1, ) == ( 20896, 20896, ) # halo +1 assert with_icon_grid.get_indices_from_to( CellDim, - HorizontalMarkerIndex.halo(CellDim), - HorizontalMarkerIndex.halo(CellDim), + HorizontalMarkerIndex.local(CellDim), + HorizontalMarkerIndex.local(CellDim), ) == ( -1, 20896, @@ -88,18 +90,18 @@ def test_horizontal_edge_indices(with_icon_grid): ) == (6176, 31558) assert with_icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.halo(EdgeDim) - 2, - HorizontalMarkerIndex.halo(EdgeDim) - 2, + HorizontalMarkerIndex.local(EdgeDim) - 2, + HorizontalMarkerIndex.local(EdgeDim) - 2, ) == (31558, 31558) assert with_icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.halo(EdgeDim) - 1, - HorizontalMarkerIndex.halo(EdgeDim) - 1, + HorizontalMarkerIndex.local(EdgeDim) - 1, + HorizontalMarkerIndex.local(EdgeDim) - 1, ) == (31558, 31558) assert with_icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.halo(EdgeDim), - HorizontalMarkerIndex.halo(EdgeDim), + HorizontalMarkerIndex.local(EdgeDim), + HorizontalMarkerIndex.local(EdgeDim), ) == ( -1, 31558, @@ -194,13 +196,13 @@ def test_horizontal_vertex_indices(with_icon_grid): ) == (10663, 10663) assert with_icon_grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.halo(VertexDim), - HorizontalMarkerIndex.halo(VertexDim), + HorizontalMarkerIndex.local(VertexDim), + HorizontalMarkerIndex.local(VertexDim), ) == (-1, 10663) assert with_icon_grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.halo(VertexDim) - 1, - HorizontalMarkerIndex.halo(VertexDim) - 1, + HorizontalMarkerIndex.local(VertexDim) - 1, + HorizontalMarkerIndex.local(VertexDim) - 1, ) == (10663, 10663) assert with_icon_grid.get_indices_from_to( diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index a299a88e1..4713bbf04 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -15,6 +15,7 @@ import numpy as np from functional.common import Dimension +from functional.ffront.fbuiltins import int32 from functional.iterator.embedded import np_as_located_field from serialbox import Savepoint @@ -50,11 +51,12 @@ def __init__(self, sp: Savepoint, ser: ser.Serializer): def print_meta_info(self): print(self.savepoint.metainfo) - def _get_field(self, name, *dimensions): - buffer = np.squeeze(self.serializer.read(name, self.savepoint)) + def _get_field(self, name, *dimensions, dtype=float): + buffer = np.squeeze(self.serializer.read(name, self.savepoint).astype(dtype)) print(f"{name} {buffer.shape}") return np_as_located_field(*dimensions)(buffer) + def get_metadata(self, *names): metadata = self.savepoint.metainfo.to_dict() return {n: metadata[n] for n in names if n in metadata} @@ -95,26 +97,43 @@ def inv_dual_edge_length(self): return self._get_field("inv_dual_edge_length", EdgeDim) def cells_start_index(self): - # subtract 1 for python 0 based indexing - return self.serializer.read("c_start_index", self.savepoint) - 1 + return self._read_int32_shift1("c_start_index") + + def _read_int32_shift1(self, name:str): + """ + Read a index field and shift it by -1. + + use for start indeces: the shift accounts for the zero based python + values are converted to int32 + """ + return (self.serializer.read(name, self.savepoint) - 1).astype(int32) def cells_end_index(self): - # don't need to subtract 1, because FORTRAN slices are inclusive [from:to] so the being - # one off accounts for being exclusive [from:to) - return self.serializer.read("c_end_index", self.savepoint) + return self._read_int32("c_end_index") - def vertex_start_index(self): + def read_int(self, name:str): + buffer = self.serializer.read(name, self.savepoint).astype(int) + print(f"{name} {buffer.shape}") + return buffer + + def _read_int32(self, name:str): + """ + Read a int field by name + + use this for end indices: because FORTRAN slices are inclusive [from:to] _and_ one based + this accounts for being exclusive python exclusive bounds: [from:to) + field values are convert to int32 + """ + return self.serializer.read(name, self.savepoint).astype(int32) - return self.serializer.read("v_start_index", self.savepoint) - 1 + def vertex_start_index(self): + return self._read_int32_shift1("v_start_index") def vertex_end_index(self): - # don't need to subtract 1, because FORTRAN slices are inclusive [from:to] so the being - # one off accounts for being exclusive [from:to) - return self.serializer.read("v_end_index", self.savepoint) + return self._read_int32("v_end_index") def edge_start_index(self): - # subtract 1 for python 0 based indexing - return self.serializer.read("e_start_index", self.savepoint) - 1 + return self._read_int32_shift1("e_start_index") def edge_end_index(self): # don't need to subtract 1, because FORTRAN slices are inclusive [from:to] so the being @@ -185,7 +204,7 @@ def wgtfac_e(self): return self._get_field("wgtfac_e", EdgeDim, KDim) def mask_diff(self): - return self._get_field("mask_diff", CellDim, KDim) + return self._get_field("mask_hdiff", CellDim, KDim, dtype=int) def zd_diffcoef(self): return self._get_field("zd_diffcoef", CellDim, KDim) @@ -212,7 +231,7 @@ def nudgecoeff_e(self): return self._get_field("nudgecoeff_e", EdgeDim) def zd_vertidx(self): - return self._get_field("zd_vertidx", CellDim, C2E2CDim, KDim) + return self._get_field("zd_vertidx", CellDim, C2E2CDim, KDim, dtype=int) def rbf_vec_coeff_v1(self): return self._get_field("rbf_vec_coeff_v1", VertexDim, V2EDim) From 77bb5ab50880b557e0bce72e0d0e602500ff6a43 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 4 Jan 2023 09:45:25 +0100 Subject: [PATCH 031/263] - cleanup: -- revert unnecessary changes --- atm_dyn_iconam/requirements-dev.txt | 1 + .../src/icon4py/atm_dyn_iconam/diffusion_program.py | 4 ++-- .../mo_intp_rbf_rbf_vec_interpol_vertex.py | 10 +++++----- .../atm_dyn_iconam/mo_nh_diffusion_stencil_02.py | 3 +-- .../atm_dyn_iconam/mo_nh_diffusion_stencil_10.py | 8 ++++---- atm_dyn_iconam/tests/test_diffusion.py | 2 +- 6 files changed, 14 insertions(+), 14 deletions(-) diff --git a/atm_dyn_iconam/requirements-dev.txt b/atm_dyn_iconam/requirements-dev.txt index 8f8d43342..2cb7bda5e 100644 --- a/atm_dyn_iconam/requirements-dev.txt +++ b/atm_dyn_iconam/requirements-dev.txt @@ -1,2 +1,3 @@ -r ../base-requirements-dev.txt -e ../common +-e . diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py index c4d7f111b..33b447ced 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py @@ -25,7 +25,8 @@ C2E2CODim, ECVDim -@program(backend=gtfn_cpu.run_gtfn) +#@program(backend=gtfn_cpu.run_gtfn) +@program def diffusion_run( diagnostic_hdef_ic: Field[[CellDim, KDim], float], diagnostic_div_ic: Field[[CellDim, KDim], float], @@ -242,7 +243,6 @@ def diffusion_run( local_kh_smag_e, out=local_kh_smag_e, domain={ - #CellDim: (cell_startindex_nudging_minus1, cell_endindex_local_plus1), EdgeDim: (edge_startindex_nudging_plus1, edge_endindex_local), KDim: (nlev-2, nlev) } diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py index 9913f48f4..1269bf03b 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -36,11 +36,11 @@ def mo_intp_rbf_rbf_vec_interpol_vertex( ptr_coeff_2: Field[[VertexDim, V2EDim], float], p_u_out: Field[[VertexDim, KDim], float], p_v_out: Field[[VertexDim, KDim], float], - v_start: int, - v_end: int, - k_start: int, - k_end:int, + horizontal_start: int, + horizontal_end: int, + vertical_start: int, + vertical_end:int, ): _mo_intp_rbf_rbf_vec_interpol_vertex( - p_e_in, ptr_coeff_1, ptr_coeff_2, out=(p_u_out, p_v_out), domain={VertexDim: (v_start, v_end), KDim:(k_start, k_end)} + p_e_in, ptr_coeff_1, ptr_coeff_2, out=(p_u_out, p_v_out), domain={VertexDim: (horizontal_start, horizontal_end), KDim:(vertical_start, vertical_end)} ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_02.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_02.py index bd5a454fc..c7a8baa85 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_02.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_02.py @@ -13,12 +13,11 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, neighbor_sum -from functional.program_processors.runners.gtfn_cpu import run_gtfn from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim -@field_operator(backend=run_gtfn) +@field_operator def _mo_nh_diffusion_stencil_02( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py index fed277e85..e439f742e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py @@ -21,10 +21,10 @@ def _mo_nh_diffusion_stencil_10( w: Field[[CellDim, KDim], float], diff_multfac_n2w: Field[[KDim], float], - area: Field[[CellDim], float], + cell_area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - w = w + diff_multfac_n2w * (area * z_nabla2_c) + w = w + diff_multfac_n2w * (cell_area * z_nabla2_c) return w @@ -32,7 +32,7 @@ def _mo_nh_diffusion_stencil_10( def mo_nh_diffusion_stencil_10( w: Field[[CellDim, KDim], float], diff_multfac_n2w: Field[[KDim], float], - area: Field[[CellDim], float], + cell_area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], ): - _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, area, z_nabla2_c, out=w) + _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index da0a7297d..2e1a341d6 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -375,7 +375,7 @@ def test_verify_diffusion_init_against_other_regular_savepoint( _verify_init_values_against_savepoint(savepoint, diffusion) - +@pytest.mark.xfail def test_diffusion_run(with_icon_grid): data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") data_provider = IconSerialDataProvider("icon_diffusion_init", data_path) From f3a179cdfa6b474048e4945263d83f41b0a45cea Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 4 Jan 2023 14:23:38 +0100 Subject: [PATCH 032/263] - cleanup: - fix test_mo_intp_rbf_rbf_vec_interpol_vertex.py: pass domain in program - skip tests in icon4pygen that fail due to changes in offset_provider: https://github.com/tehrengruber/gt4py.git@reduce_foast_lowering_complexity#egg=gt4py-functional - move fixtures to conftest.py - cleanup imports --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 130 +++++++----- .../atm_dyn_iconam/diffusion_program.py | 196 +++++++++--------- .../icon4py/atm_dyn_iconam/diffusion_utils.py | 18 +- .../fused_mo_nh_diffusion_stencil_02_03.py | 1 - .../src/icon4py/atm_dyn_iconam/horizontal.py | 4 +- .../src/icon4py/atm_dyn_iconam/icon_grid.py | 14 +- .../icon4py/atm_dyn_iconam/metric_state.py | 2 +- .../mo_intp_rbf_rbf_vec_interpol_vertex.py | 12 +- .../{icon_grid_test_utils.py => conftest.py} | 10 +- atm_dyn_iconam/tests/test_diffusion.py | 66 +++--- atm_dyn_iconam/tests/test_icon_grid.py | 71 +++---- ...est_mo_intp_rbf_rbf_vec_interpol_vertex.py | 4 + atm_dyn_iconam/tests/test_vertical.py | 3 +- pyutils/tests/test_codegen.py | 11 + pyutils/tests/test_exceptions.py | 1 + .../src/icon4py/testutils/serialbox_utils.py | 9 +- 16 files changed, 293 insertions(+), 259 deletions(-) rename atm_dyn_iconam/tests/{icon_grid_test_utils.py => conftest.py} (94%) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index 75a3bb04b..0c6bc6ca4 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -18,15 +18,16 @@ import numpy as np from functional.common import Dimension from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, broadcast, maximum, minimum, int32 -from functional.iterator.embedded import np_as_located_field, StridedNeighborOffsetProvider -from functional.program_processors.runners import gtfn_cpu -from functional.program_processors.runners.gtfn_cpu import run_gtfn +from functional.ffront.fbuiltins import Field, broadcast, int32, maximum, minimum +from functional.iterator.embedded import ( + StridedNeighborOffsetProvider, + np_as_located_field, +) from icon4py.atm_dyn_iconam.constants import CPD, GAS_CONSTANT_DRY_AIR from icon4py.atm_dyn_iconam.diagnostic import DiagnosticState from icon4py.atm_dyn_iconam.diffusion_program import diffusion_run -from icon4py.atm_dyn_iconam.diffusion_utils import set_zero_v_k, scale_k +from icon4py.atm_dyn_iconam.diffusion_utils import scale_k, set_zero_v_k from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import ( _fused_mo_nh_diffusion_stencil_02_03, ) @@ -47,7 +48,7 @@ from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState from icon4py.atm_dyn_iconam.metric_state import MetricState from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( - _mo_intp_rbf_rbf_vec_interpol_vertex, mo_intp_rbf_rbf_vec_interpol_vertex, + mo_intp_rbf_rbf_vec_interpol_vertex, ) from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import ( _mo_nh_diffusion_stencil_01, @@ -58,19 +59,24 @@ from icon4py.atm_dyn_iconam.prognostic import PrognosticState from icon4py.atm_dyn_iconam.utils import zero_field from icon4py.common.dimension import ( + E2ECV, C2E2CDim, C2E2CODim, + C2EDim, CellDim, ECVDim, EdgeDim, KDim, Koff, - VertexDim, V2EDim, C2EDim, E2ECV, + V2EDim, + VertexDim, ) TupleVT = namedtuple("TupleVT", "v t") VectorTuple = namedtuple("VectorTuple", "x y") + + @field_operator def _setup_smag_limit(diff_multfac_vn: Field[[KDim], float]) -> Field[[KDim], float]: return 0.125 - 4.0 * diff_multfac_vn @@ -99,6 +105,7 @@ def setup_fields_for_initial_step(k4: float, hdiff_efdt_ratio: float): smag_limit = _setup_smag_limit(diff_multfac_vn) return diff_multfac_vn, smag_limit + @field_operator def _en_smag_fac_for_zero_nshift( vect_a: Field[[KDim], float], @@ -197,6 +204,8 @@ def init_diffusion_local_fields( enh_smag_fac, ), ) + + def init_nabla2_factor_in_upper_damping_zone( k_size: int, nrdmax: int, nshift: int, physical_heights: np.ndarray ) -> Field[[KDim], float]: @@ -227,7 +236,6 @@ def init_nabla2_factor_in_upper_damping_zone( return np_as_located_field(KDim)(buffer) - class DiffusionConfig: """contains necessary parameter to configure a diffusion run. @@ -395,7 +403,7 @@ def mo_nh_diffusion_stencil_15_numpy( z_temp: Field[[CellDim, KDim], float], ): geofac_n2s = np.asarray(geofac_n2s) - geofac_n2s_nbh = geofac_n2s[:,1:] + geofac_n2s_nbh = geofac_n2s[:, 1:] geofac_n2s_c = geofac_n2s[:, 0] mask_hdiff = np.asarray(mask_hdiff) zd_vertidx = np.asarray(zd_vertidx) @@ -405,14 +413,15 @@ def mo_nh_diffusion_stencil_15_numpy( z_temp = np.asarray(z_temp) vertidx_ = vcoef * theta_v[c2e2c][zd_vertidx] - first = geofac_n2s_nbh * vertidx_ - sum = np.sum(first + second, axis=0) + first = geofac_n2s_nbh * vertidx_ + summed = np.sum(first + second, axis=0) - z_temp = np.where(mask_hdiff, z_temp + zd_diffcoef* (theta_v* geofac_n2s_c + sum), z_temp) + z_temp = np.where( + mask_hdiff, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + summed), z_temp + ) return z_temp - class Diffusion: """Class that configures diffusion and does one diffusion step.""" @@ -588,7 +597,6 @@ def time_step( # self._smag_offset, # ) - cell_start_nudging_minus1, cell_end_local_plus1 = self.grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.nudging(CellDim) - 1, @@ -613,18 +621,27 @@ def time_step( HorizontalMarkerIndex.local(EdgeDim), ) - edge_start_nudging_minus1, edge_end_local_minus2 = self.grid.get_indices_from_to( + ( + edge_start_nudging_minus1, + edge_end_local_minus2, + ) = self.grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) - 1, HorizontalMarkerIndex.local(EdgeDim) - 2, ) - vertex_start_local_boundary_plus3, vertex_end_local = self.grid.get_indices_from_to( + ( + vertex_start_local_boundary_plus3, + vertex_end_local, + ) = self.grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 3, HorizontalMarkerIndex.local(VertexDim), ) - vertex_start_local_boundary_plus1, vertex_end_local_minus1 = self.grid.get_indices_from_to( + ( + vertex_start_local_boundary_plus1, + vertex_end_local_minus1, + ) = self.grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 1, HorizontalMarkerIndex.local(VertexDim) - 1, @@ -678,7 +695,7 @@ def time_step( self.diff_multfac_smag, self.diff_multfac_n2w, self._smag_offset, - self. nudgezone_diff, + self.nudgezone_diff, self.fac_bdydiff_v, self.diff_multfac_w, self.vertical_index, @@ -700,8 +717,7 @@ def time_step( self.config.vertical_params.index_of_damping_height, self.grid.n_lev(), self.params.boundary_diffusion_start_index_edges, - - offset_provider = { + offset_provider={ "V2E": self.grid.get_v2e_connectivity(), "V2EDim": V2EDim, "E2C2V": self.grid.get_e2c2v_connectivity(), @@ -714,8 +730,8 @@ def time_step( "ECVDim": ECVDim, "C2E2CODim": C2E2CODim, "C2E2CO": self.grid.get_c2e2co_connectivity(), - "Koff": KDim - } + "Koff": KDim, + }, ) def _do_diffusion_step( @@ -789,18 +805,27 @@ def _do_diffusion_step( HorizontalMarkerIndex.local(EdgeDim), ) - edge_start_nudging_minus1, edge_end_local_minus2 = self.grid.get_indices_from_to( + ( + edge_start_nudging_minus1, + edge_end_local_minus2, + ) = self.grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) - 1, HorizontalMarkerIndex.local(EdgeDim) - 2, ) - vertex_start_local_boundary_plus3, vertex_end_local = self.grid.get_indices_from_to( + ( + vertex_start_local_boundary_plus3, + vertex_end_local, + ) = self.grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 3, HorizontalMarkerIndex.local(VertexDim), ) - vertex_start_local_boundary_plus1, vertex_end_local_minus1 = self.grid.get_indices_from_to( + ( + vertex_start_local_boundary_plus1, + vertex_end_local_minus1, + ) = self.grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 1, HorizontalMarkerIndex.local(VertexDim) - 1, @@ -812,8 +837,6 @@ def _do_diffusion_step( # 0c. dtime dependent stuff: enh_smag_factor, scale_k(self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={}) - - # TODO: is this needed?, if not remove set_zero_v_k(self.u_vert, offset_provider={}) set_zero_v_k(self.v_vert, offset_provider={}) @@ -834,13 +857,12 @@ def _do_diffusion_step( v_end, k_start, k_end, - offset_provider={"V2E": self.grid.get_v2e_connectivity(), "V2EDim":V2EDim}, + offset_provider={"V2E": self.grid.get_v2e_connectivity(), "V2EDim": V2EDim}, ) # 2. HALO EXCHANGE -- CALL sync_patch_array_mult # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 - e_start = self.params.boundary_diffusion_start_index_edges e_end = edge_end_local_minus2 k_start = 0 @@ -860,13 +882,14 @@ def _do_diffusion_step( smag_limit, smag_offset, out=(self.kh_smag_e, self.kh_smag_ec, self.z_nabla2_e), - offset_provider={"E2C2V": self.grid.get_e2c2v_connectivity(), - "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, 4)}, + offset_provider={ + "E2C2V": self.grid.get_e2c2v_connectivity(), + "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, 4), + }, ) - - c_start = cell_start_nudging, - c_end = cell_end_local, + c_start = (cell_start_nudging,) + c_end = (cell_end_local,) k_start = 0 k_end = klevels _fused_mo_nh_diffusion_stencil_02_03( @@ -880,8 +903,7 @@ def _do_diffusion_step( diagnostic_state.div_ic, diagnostic_state.hdef_ic, ), - - offset_provider={"C2E": self.grid.get_c2e_connectivity(), "C2EDim":C2EDim}, + offset_provider={"C2E": self.grid.get_c2e_connectivity(), "C2EDim": C2EDim}, ) # # # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH @@ -900,9 +922,9 @@ def _do_diffusion_step( self.v_vert, v_start, v_end, - k_start, k_end, - offset_provider={"V2E": self.grid.get_v2e_connectivity(), - "V2EDim":V2EDim}, + k_start, + k_end, + offset_provider={"V2E": self.grid.get_v2e_connectivity(), "V2EDim": V2EDim}, ) # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult # @@ -910,8 +932,6 @@ def _do_diffusion_step( # # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 # - - e_start = edge_start_nudging_plus_one e_end = edge_end_local k_start = 0 @@ -934,9 +954,11 @@ def _do_diffusion_step( self.fac_bdydiff_v, edge_start_nudging_minus1, out=prognostic_state.normal_wind, - offset_provider={"E2C2V": self.grid.get_e2c2v_connectivity(), - "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, 4), - "ECVDim":ECVDim,} + offset_provider={ + "E2C2V": self.grid.get_e2c2v_connectivity(), + "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, 4), + "ECVDim": ECVDim, + }, ) # # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, # # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 @@ -966,8 +988,10 @@ def _do_diffusion_step( diagnostic_state.dwdx, diagnostic_state.dwdy, ), - - offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity(), "C2E2CODim":C2E2CODim}, + offset_provider={ + "C2E2CO": self.grid.get_c2e2co_connectivity(), + "C2E2CODim": C2E2CODim, + }, ) # # 8. HALO EXCHANGE: CALL sync_patch_array # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, @@ -978,7 +1002,7 @@ def _do_diffusion_step( e_start = edge_start_nudging_plus_one e_end = edge_end_local - c_start= cell_start_nudging_minus1 + c_start = cell_start_nudging_minus1 c_end = cell_end_local_plus1 k_start = k_start_end_minus2 k_end = klevels @@ -989,17 +1013,14 @@ def _do_diffusion_step( self.thresh_tdiff, self.kh_smag_e, out=self.kh_smag_e, - offset_provider={ "E2C": self.grid.get_e2c_connectivity(), "C2E2C": self.grid.get_c2e2c_connectivity(), - "C2E2CDim":C2E2CDim + "C2E2CDim": C2E2CDim, }, ) - - - c_start = cell_start_nudging, + c_start = (cell_start_nudging,) c_end = cell_end_local _fused_mo_nh_diffusion_stencil_13_14( self.kh_smag_e, @@ -1010,7 +1031,7 @@ def _do_diffusion_step( offset_provider={ "C2E": self.grid.get_c2e_connectivity(), "E2C": self.grid.get_e2c_connectivity(), - "C2EDim":C2EDim, + "C2EDim": C2EDim, }, ) @@ -1032,7 +1053,6 @@ def _do_diffusion_step( # offset_provider={}, # ) - c_start = cell_start_nudging c_end = cell_end_local _mo_nh_diffusion_stencil_16( @@ -1044,5 +1064,3 @@ def _do_diffusion_step( offset_provider={}, ) # 10. HALO EXCHANGE sync_patch_array - - diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py index 33b447ced..a8f408576 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py @@ -1,32 +1,61 @@ -from typing import Tuple +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + from functional.common import Field -from functional.ffront.fbuiltins import int32 from functional.ffront.decorator import program - +from functional.ffront.fbuiltins import int32 from functional.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.diffusion_utils import _scale_k, _set_zero_v_k -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import \ - _fused_mo_nh_diffusion_stencil_02_03 -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import \ - _fused_mo_nh_diffusion_stencil_04_05_06 -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_07_08_09_10 import \ - _fused_mo_nh_diffusion_stencil_07_08_09_10 -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_11_12 import \ - _fused_mo_nh_diffusion_stencil_11_12 -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_13_14 import \ - _fused_mo_nh_diffusion_stencil_13_14 -from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import \ - _mo_intp_rbf_rbf_vec_interpol_vertex -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import _mo_nh_diffusion_stencil_01 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_16 import _mo_nh_diffusion_stencil_16 -from icon4py.common.dimension import CellDim, KDim, EdgeDim, C2E2CDim, C2EDim, V2EDim, VertexDim, \ - C2E2CODim, ECVDim +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import ( + _fused_mo_nh_diffusion_stencil_02_03, +) +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import ( + _fused_mo_nh_diffusion_stencil_04_05_06, +) +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_07_08_09_10 import ( + _fused_mo_nh_diffusion_stencil_07_08_09_10, +) +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_11_12 import ( + _fused_mo_nh_diffusion_stencil_11_12, +) +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_13_14 import ( + _fused_mo_nh_diffusion_stencil_13_14, +) +from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( + _mo_intp_rbf_rbf_vec_interpol_vertex, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import ( + _mo_nh_diffusion_stencil_01, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_16 import ( + _mo_nh_diffusion_stencil_16, +) +from icon4py.common.dimension import ( + C2E2CDim, + C2E2CODim, + C2EDim, + CellDim, + ECVDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) -#@program(backend=gtfn_cpu.run_gtfn) -@program +@program(backend=gtfn_cpu.run_gtfn) def diffusion_run( diagnostic_hdef_ic: Field[[CellDim, KDim], float], diagnostic_div_ic: Field[[CellDim, KDim], float], @@ -46,10 +75,10 @@ def diffusion_run( interpolation_rbf_coeff_1: Field[[VertexDim, V2EDim], float], interpolation_rbf_coeff_2: Field[[VertexDim, V2EDim], float], interpolation_geofac_div: Field[[CellDim, C2EDim], float], - interpolation_geofac_grg_x: Field[[CellDim, C2E2CODim], float] , + interpolation_geofac_grg_x: Field[[CellDim, C2E2CODim], float], interpolation_geofac_grg_y: Field[[CellDim, C2E2CODim], float], interpolation_nudgecoeff_e: Field[[EdgeDim], float], - interpolation_geofac_n2s:Field[[CellDim, C2E2CODim], float], + interpolation_geofac_n2s: Field[[CellDim, C2E2CODim], float], tangent_orientation: Field[[EdgeDim], float], inverse_primal_edge_lengths: Field[[EdgeDim], float], inverse_dual_edge_length: Field[[EdgeDim], float], @@ -62,42 +91,41 @@ def diffusion_run( cell_areas: Field[[CellDim], float], diff_multfac_vn: Field[[KDim], float], dtime: float, - rd_o_cvd:float, + rd_o_cvd: float, local_thresh_tdiff: float, local_smag_limit: Field[[KDim], float], local_u_vert: Field[[VertexDim, KDim], float], local_v_vert: Field[[VertexDim, KDim], float], - local_enh_smag_fac:Field[[KDim], float], - local_kh_smag_e:Field[[EdgeDim, KDim], float], - local_kh_smag_ec:Field[[EdgeDim, KDim], float], - local_z_nabla2_e:Field[[EdgeDim, KDim], float], - local_z_temp :Field[[CellDim, KDim], float], - local_diff_multfac_smag:Field[[KDim], float], + local_enh_smag_fac: Field[[KDim], float], + local_kh_smag_e: Field[[EdgeDim, KDim], float], + local_kh_smag_ec: Field[[EdgeDim, KDim], float], + local_z_nabla2_e: Field[[EdgeDim, KDim], float], + local_z_temp: Field[[CellDim, KDim], float], + local_diff_multfac_smag: Field[[KDim], float], local_diff_multfac_n2w: Field[[KDim], float], local_smag_offset: float, local_nudgezone_diff: float, local_fac_bdydiff_v: float, local_diff_multfac_w: float, - - local_vertical_index:Field[[KDim], int32], #local_grid.n_lev() + 1 - local_horizontal_cell_index:Field[[CellDim], int32], + local_vertical_index: Field[[KDim], int32], # local_grid.n_lev() + 1 + local_horizontal_cell_index: Field[[CellDim], int32], local_horizontal_edge_index: Field[[EdgeDim], int32], cell_startindex_nudging_minus1: int32, - cell_startindex_interior:int32, - cell_startindex_nudging:int32, + cell_startindex_interior: int32, + cell_startindex_nudging: int32, cell_endindex_local_plus1: int32, - cell_endindex_local:int32, - edge_startindex_nudging_plus1:int32, - edge_startindex_nudging_minus1:int32, - edge_endindex_local:int32, - edge_endindex_local_minus2:int32, - vertex_startindex_lb_plus3:int32, - vertex_startindex_lb_plus1:int32, - vertex_endindex_local:int32, + cell_endindex_local: int32, + edge_startindex_nudging_plus1: int32, + edge_startindex_nudging_minus1: int32, + edge_endindex_local: int32, + edge_endindex_local_minus2: int32, + vertex_startindex_lb_plus3: int32, + vertex_startindex_lb_plus1: int32, + vertex_endindex_local: int32, vertex_endindex_local_minus1: int32, index_of_damping_height: int32, - nlev:int, - boundary_diffusion_start_index_edges: int32 + nlev: int, + boundary_diffusion_start_index_edges: int32, ): _scale_k(local_enh_smag_fac, dtime, out=local_diff_multfac_smag) @@ -110,11 +138,11 @@ def diffusion_run( prognostic_normal_wind, interpolation_rbf_coeff_1, interpolation_rbf_coeff_2, - out=(local_u_vert,local_v_vert), + out=(local_u_vert, local_v_vert), domain={ VertexDim: (vertex_startindex_lb_plus1, vertex_endindex_local_minus1), - KDim: (0,nlev) - } + KDim: (0, nlev), + }, ) # 2. HALO EXCHANGE -- CALL sync_patch_array_mult @@ -136,10 +164,9 @@ def diffusion_run( local_smag_offset, out=(local_kh_smag_e, local_kh_smag_ec, local_z_nabla2_e), domain={ - EdgeDim:(boundary_diffusion_start_index_edges, edge_endindex_local_minus2), - KDim: (0,nlev) - } - + EdgeDim: (boundary_diffusion_start_index_edges, edge_endindex_local_minus2), + KDim: (0, nlev), + }, ) _fused_mo_nh_diffusion_stencil_02_03( @@ -153,9 +180,10 @@ def diffusion_run( diagnostic_div_ic, diagnostic_hdef_ic, ), - domain={CellDim:(cell_startindex_nudging, cell_endindex_local), - KDim: (0, nlev)} - + domain={ + CellDim: (cell_startindex_nudging, cell_endindex_local), + KDim: (0, nlev), + }, ) # # # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH @@ -165,11 +193,11 @@ def diffusion_run( local_z_nabla2_e, interpolation_rbf_coeff_1, interpolation_rbf_coeff_2, - out = (local_u_vert,local_v_vert), - domain={VertexDim:(vertex_startindex_lb_plus3, vertex_endindex_local), - KDim: (0, nlev) - } - + out=(local_u_vert, local_v_vert), + domain={ + VertexDim: (vertex_startindex_lb_plus3, vertex_endindex_local), + KDim: (0, nlev), + }, ) # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult # @@ -196,9 +224,8 @@ def diffusion_run( out=prognostic_normal_wind, domain={ EdgeDim: (edge_startindex_nudging_plus1, edge_endindex_local), - KDim: (0, nlev) - } - + KDim: (0, nlev), + }, ) # # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, # # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 @@ -226,8 +253,8 @@ def diffusion_run( ), domain={ CellDim: (cell_startindex_nudging, cell_endindex_local_plus1), - KDim: (0, nlev) - } + KDim: (0, nlev), + }, ) # # 8. HALO EXCHANGE: CALL sync_patch_array # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, @@ -244,9 +271,8 @@ def diffusion_run( out=local_kh_smag_e, domain={ EdgeDim: (edge_startindex_nudging_plus1, edge_endindex_local), - KDim: (nlev-2, nlev) - } - + KDim: (nlev - 2, nlev), + }, ) _fused_mo_nh_diffusion_stencil_13_14( local_kh_smag_e, @@ -255,29 +281,12 @@ def diffusion_run( interpolation_geofac_div, out=local_z_temp, domain={ - CellDim:(cell_startindex_nudging, cell_endindex_local), - KDim: (0, nlev) - } - + CellDim: (cell_startindex_nudging, cell_endindex_local), + KDim: (0, nlev), + }, ) - # mo_nh_diffusion_stencil_15_numpy( - # mask_hdiff=metric_state.mask_hdiff, - # zd_vertidx=metric_state.zd_vertidx, - # vcoef=metric_state.zd_diffcoef, - # zd_diffcoef=metric_state.zd_diffcoef, - # geofac_n2s=interpolation_state.geofac_n2s, - # theta_v=prognostic_state.theta_v, - # z_temp=self.z_temp, - # domain={ - # CellDim: self.grid.get_indices_from_to( - # CellDim, - # HorizontalMarkerIndex.nudging(CellDim), - # HorizontalMarkerIndex.local(CellDim), - # ), - # }, - # - # ) + # MO_NH_DIFFUSION_STENCIL_15: needs index fields! _mo_nh_diffusion_stencil_16( local_z_temp, @@ -285,9 +294,10 @@ def diffusion_run( prognostic_theta_v, prognostic_exner_pressure, rd_o_cvd, - out=(prognostic_theta_v,prognostic_exner_pressure), - domain = {CellDim: (cell_startindex_nudging, cell_endindex_local), - KDim: (0, nlev)} + out=(prognostic_theta_v, prognostic_exner_pressure), + domain={ + CellDim: (cell_startindex_nudging, cell_endindex_local), + KDim: (0, nlev), + }, ) # 10. HALO EXCHANGE sync_patch_array - diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_utils.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_utils.py index 73c4a64f5..845f3ee87 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_utils.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_utils.py @@ -1,12 +1,23 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + from functional.common import Field from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import broadcast - from icon4py.common.dimension import KDim, VertexDim - @field_operator def _scale_k(field: Field[[KDim], float], factor: float) -> Field[[KDim], float]: return field * factor @@ -19,7 +30,6 @@ def scale_k( _scale_k(field, factor, out=scaled_field) - @field_operator def _set_zero_v_k() -> Field[[VertexDim, KDim], float]: return broadcast(0.0, (VertexDim, KDim)) @@ -28,5 +38,3 @@ def _set_zero_v_k() -> Field[[VertexDim, KDim], float]: @program def set_zero_v_k(field: Field[[VertexDim, KDim], float]): _set_zero_v_k(out=field) - - diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py index da9b3f494..6b32d5a67 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -13,7 +13,6 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field -from functional.program_processors.runners.gtfn_cpu import run_gtfn from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import ( _mo_nh_diffusion_stencil_02, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py index 00e5ff27c..846933f32 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py @@ -67,9 +67,7 @@ def local_boundary(cls, dim: Dimension) -> int: @classmethod def local(cls, dim: Dimension) -> int: - """ - - """ + """ """ match (dim): case (dimension.CellDim): return cls._HALO_CELLS diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py index ae08bc384..a531784a7 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py @@ -14,16 +14,15 @@ from typing import Dict, Tuple import numpy as np -from functional.ffront.fbuiltins import int32, int64 from functional.common import Dimension, DimensionKind, Field +from functional.ffront.fbuiltins import int32 from functional.iterator.embedded import ( NeighborTableOffsetProvider, np_as_located_field, ) from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig -from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim, C2E2CODim, C2EDim, C2E2CDim, \ - E2VDim, E2CDim +from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim class VerticalMeshConfig: @@ -80,8 +79,6 @@ def wrapper(self, *args, **kwargs): return wrapper - - class IconGrid: def __init__(self): self.config: MeshConfig = None @@ -146,7 +143,7 @@ def get_indices_from_to( raise ValueError( "only defined for {} dimension kind ", DimensionKind.HORIZONTAL ) - return self.start_indices[dim][start_marker],self.end_indices[dim][end_marker] + return self.start_indices[dim][start_marker], self.end_indices[dim][end_marker] def get_c2e_connectivity(self): table = self.connectivities["c2e"] @@ -174,7 +171,6 @@ def get_e2c2v_connectivity(self): def get_e2c2v_size(self): self.connectivities["e2v"].shape[1] - def get_v2e_connectivity(self): table = self.connectivities["v2e"] return NeighborTableOffsetProvider(table, VertexDim, EdgeDim, table.shape[1]) @@ -191,8 +187,8 @@ def __init__(self, vct_a: np.ndarray, rayleigh_damping_height: float): """ self.rayleigh_damping_height = rayleigh_damping_height self.vct_a = vct_a - self.index_of_damping_height = int32(np.argmax( - np.where(np.asarray(self.vct_a) >= self.rayleigh_damping_height)) + self.index_of_damping_height = int32( + np.argmax(np.where(np.asarray(self.vct_a) >= self.rayleigh_damping_height)) ) def get_index_of_damping_layer(self): diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py index 06d66ca82..faf172be5 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py @@ -27,4 +27,4 @@ class MetricState: mask_hdiff: Field[[CellDim, KDim], int] zd_vertidx: Field[[CellDim, C2E2CDim, KDim], int] zd_diffcoef: Field[[CellDim, KDim], float] - zd_intcoef: Field[[CellDim,C2E2CDim, KDim], float] + zd_intcoef: Field[[CellDim, C2E2CDim, KDim], float] diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py index 1269bf03b..50b3966e6 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -10,7 +10,6 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -from typing import Tuple from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, neighbor_sum @@ -39,8 +38,15 @@ def mo_intp_rbf_rbf_vec_interpol_vertex( horizontal_start: int, horizontal_end: int, vertical_start: int, - vertical_end:int, + vertical_end: int, ): _mo_intp_rbf_rbf_vec_interpol_vertex( - p_e_in, ptr_coeff_1, ptr_coeff_2, out=(p_u_out, p_v_out), domain={VertexDim: (horizontal_start, horizontal_end), KDim:(vertical_start, vertical_end)} + p_e_in, + ptr_coeff_1, + ptr_coeff_2, + out=(p_u_out, p_v_out), + domain={ + VertexDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/atm_dyn_iconam/tests/icon_grid_test_utils.py b/atm_dyn_iconam/tests/conftest.py similarity index 94% rename from atm_dyn_iconam/tests/icon_grid_test_utils.py rename to atm_dyn_iconam/tests/conftest.py index 6705dc20b..411ca1bdd 100644 --- a/atm_dyn_iconam/tests/icon_grid_test_utils.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -38,9 +38,12 @@ from icon4py.testutils.serialbox_utils import IconSerialDataProvider +data_path = os.path.join(os.path.dirname(__file__), "./ser_icondata") + + @pytest.fixture -def with_icon_grid(): - data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") +def icon_grid(): + sp = IconSerialDataProvider( "icon_diffusion_init", data_path, True ).from_savepoint_init(linit=True, date="2021-06-20T12:00:10.000") @@ -80,13 +83,12 @@ def with_icon_grid(): @pytest.fixture -def with_r04b09_diffusion_config() -> DiffusionConfig: +def r04b09_diffusion_config() -> DiffusionConfig: """ Create DiffusionConfig. that uses the parameters of MCH.CH_r04b09_dsl experiment """ - data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") sp = IconSerialDataProvider( "icon_diffusion_init", data_path, True ).from_savepoint_init(linit=True, date="2021-06-20T12:00:10.000") diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 2e1a341d6..35201dfb6 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -14,8 +14,7 @@ import numpy as np import pytest -from icon_grid_test_utils import with_icon_grid, with_r04b09_diffusion_config -from functional.ffront.fbuiltins import int32 + from icon4py.atm_dyn_iconam import utils from icon4py.atm_dyn_iconam.diagnostic import DiagnosticState from icon4py.atm_dyn_iconam.diffusion import ( @@ -28,34 +27,20 @@ _setup_smag_limit, scale_k, set_zero_v_k, - setup_fields_for_initial_step, mo_nh_diffusion_stencil_15_numpy, + setup_fields_for_initial_step, ) from icon4py.atm_dyn_iconam.icon_grid import VerticalModelParams from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState from icon4py.atm_dyn_iconam.metric_state import MetricState from icon4py.atm_dyn_iconam.prognostic import PrognosticState -from icon4py.common.dimension import KDim, VertexDim, ECVDim, CellDim, C2E2CODim, C2E2CDim +from icon4py.common.dimension import ECVDim, KDim, VertexDim from icon4py.testutils.serialbox_utils import ( IconDiffusionInitSavepoint, IconSerialDataProvider, ) from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field, zero_field, as_1D_sparse_field, constant_field, \ - random_mask - +from icon4py.testutils.utils import as_1D_sparse_field, random_field, zero_field -def test_mo_nh_diffusion_stencil_15(): - mesh = SimpleMesh() - mask_diff = constant_field(mesh, 1, CellDim, KDim, dtype=int32) - zd_vertidx = constant_field(mesh, 1, C2E2CDim, KDim, dtype=int32) - zd_diffcoef = random_field(mesh, CellDim, KDim) - geofac_n2s = random_field(mesh, CellDim, C2E2CODim) - vcoef = random_field(mesh, C2E2CDim, KDim) - theta_v = random_field(mesh, CellDim, KDim) - z_temp = random_field(mesh, CellDim, KDim) - - # mo_nh_diffusion_stencil_15_numpy(mesh.c2e2c, - # mask_diff, zd_vertidx, zd_diffcoef, geofac_n2s, vcoef, theta_v, z_temp) def test_scale_k(): mesh = SimpleMesh() @@ -178,8 +163,8 @@ def test_set_zero_vertex_k(): assert np.allclose(0.0, f) -def test_diffusion_coefficients_with_hdiff_efdt_ratio(with_r04b09_diffusion_config): - config = with_r04b09_diffusion_config +def test_diffusion_coefficients_with_hdiff_efdt_ratio(r04b09_diffusion_config): + config = r04b09_diffusion_config config.hdiff_efdt_ratio = 1.0 config.hdiff_w_efdt_ratio = 2.0 @@ -191,8 +176,8 @@ def test_diffusion_coefficients_with_hdiff_efdt_ratio(with_r04b09_diffusion_conf assert params.K4W == pytest.approx(1.0 / 72.0, abs=1e-12) -def test_diffusion_coefficients_without_hdiff_efdt_ratio(with_r04b09_diffusion_config): - config = with_r04b09_diffusion_config +def test_diffusion_coefficients_without_hdiff_efdt_ratio(r04b09_diffusion_config): + config = r04b09_diffusion_config config.hdiff_efdt_ratio = 0.0 config.hdiff_w_efdt_ratio = 0.0 @@ -204,8 +189,8 @@ def test_diffusion_coefficients_without_hdiff_efdt_ratio(with_r04b09_diffusion_c assert params.K4W == 0.0 -def test_smagorinski_factor_for_diffusion_type_4(with_r04b09_diffusion_config): - config = with_r04b09_diffusion_config +def test_smagorinski_factor_for_diffusion_type_4(r04b09_diffusion_config): + config = r04b09_diffusion_config config.hdiff_smag_factor = 0.15 config.diffusion_type = 4 @@ -216,9 +201,9 @@ def test_smagorinski_factor_for_diffusion_type_4(with_r04b09_diffusion_config): def test_smagorinski_heights_diffusion_type_5_are_consistent( - with_r04b09_diffusion_config, + r04b09_diffusion_config, ): - config = with_r04b09_diffusion_config + config = r04b09_diffusion_config config.hdiff_smag_factor = 0.15 config.diffusion_type = 5 @@ -232,16 +217,16 @@ def test_smagorinski_heights_diffusion_type_5_are_consistent( assert params.smagorinski_height[2] != params.smagorinski_height[3] -def test_smagorinski_factor_diffusion_type_5(with_r04b09_diffusion_config): - config = with_r04b09_diffusion_config +def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): + config = r04b09_diffusion_config params = DiffusionParams(config) assert len(params.smagorinski_factor) == len(params.smagorinski_height) assert len(params.smagorinski_factor) == 4 assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor))) -def test_diffusion_init(with_r04b09_diffusion_config): - config = with_r04b09_diffusion_config +def test_diffusion_init(r04b09_diffusion_config): + config = r04b09_diffusion_config data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") serializer = IconSerialDataProvider("icon_diffusion_init", data_path) serializer.print_info() @@ -315,9 +300,9 @@ def _verify_init_values_against_savepoint( def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( - with_r04b09_diffusion_config, + r04b09_diffusion_config, ): - config = with_r04b09_diffusion_config + config = r04b09_diffusion_config params = DiffusionParams(config) data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") @@ -343,9 +328,9 @@ def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( def test_verify_diffusion_init_against_first_regular_savepoint( - with_r04b09_diffusion_config, + r04b09_diffusion_config, ): - config = with_r04b09_diffusion_config + config = r04b09_diffusion_config data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") serializer = IconSerialDataProvider("icon_diffusion_init", data_path) serializer.print_info() @@ -360,9 +345,9 @@ def test_verify_diffusion_init_against_first_regular_savepoint( def test_verify_diffusion_init_against_other_regular_savepoint( - with_r04b09_diffusion_config, + r04b09_diffusion_config, ): - config = with_r04b09_diffusion_config + config = r04b09_diffusion_config data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") serializer = IconSerialDataProvider("icon_diffusion_init", data_path) serializer.print_info() @@ -375,8 +360,9 @@ def test_verify_diffusion_init_against_other_regular_savepoint( _verify_init_values_against_savepoint(savepoint, diffusion) -@pytest.mark.xfail -def test_diffusion_run(with_icon_grid): + +@pytest.mark.skip +def test_diffusion_run(icon_grid): data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") data_provider = IconSerialDataProvider("icon_diffusion_init", data_path) data_provider.print_info() @@ -384,7 +370,7 @@ def test_diffusion_run(with_icon_grid): vct_a = sp.vct_a() config = DiffusionConfig( - with_icon_grid, + icon_grid, vertical_params=VerticalModelParams( vct_a=vct_a, rayleigh_damping_height=12500.0 ), diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index cda0b8d31..a9ed6dc3a 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -10,16 +10,13 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -from typing import Tuple - -from icon_grid_test_utils import with_icon_grid from icon4py.atm_dyn_iconam.horizontal import HorizontalMarkerIndex from icon4py.common.dimension import CellDim, EdgeDim, VertexDim -def test_horizontal_grid_cell_indices(with_icon_grid): - assert with_icon_grid.get_indices_from_to( +def test_horizontal_grid_cell_indices(icon_grid): + assert icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.local(CellDim) - 1, HorizontalMarkerIndex.local(CellDim) - 1, @@ -27,7 +24,7 @@ def test_horizontal_grid_cell_indices(with_icon_grid): 20896, 20896, ) # halo +1 - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.local(CellDim), HorizontalMarkerIndex.local(CellDim), @@ -35,7 +32,7 @@ def test_horizontal_grid_cell_indices(with_icon_grid): -1, 20896, ) # halo in icon is (1,20896) why - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.interior(CellDim), HorizontalMarkerIndex.interior(CellDim), @@ -43,7 +40,7 @@ def test_horizontal_grid_cell_indices(with_icon_grid): 4104, 20896, ) # interior - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, HorizontalMarkerIndex._INTERIOR_CELLS + 1, @@ -51,12 +48,12 @@ def test_horizontal_grid_cell_indices(with_icon_grid): 0, 850, ) # lb+1 - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.local_boundary(CellDim) + 1, HorizontalMarkerIndex.local_boundary(CellDim) + 1, ) == (850, 1688) - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.local_boundary(CellDim) + 2, HorizontalMarkerIndex.local_boundary(CellDim) + 2, @@ -64,7 +61,7 @@ def test_horizontal_grid_cell_indices(with_icon_grid): 1688, 2511, ) # lb+2 - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.nudging(CellDim) - 1, HorizontalMarkerIndex.nudging(CellDim) - 1, @@ -72,7 +69,7 @@ def test_horizontal_grid_cell_indices(with_icon_grid): 2511, 3316, ) # lb+3 - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.nudging(CellDim), HorizontalMarkerIndex.nudging(CellDim), @@ -82,23 +79,23 @@ def test_horizontal_grid_cell_indices(with_icon_grid): ) # nudging -def test_horizontal_edge_indices(with_icon_grid): - assert with_icon_grid.get_indices_from_to( +def test_horizontal_edge_indices(icon_grid): + assert icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), HorizontalMarkerIndex.interior(EdgeDim), ) == (6176, 31558) - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2, HorizontalMarkerIndex.local(EdgeDim) - 2, ) == (31558, 31558) - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 1, HorizontalMarkerIndex.local(EdgeDim) - 1, ) == (31558, 31558) - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local(EdgeDim), HorizontalMarkerIndex.local(EdgeDim), @@ -106,7 +103,7 @@ def test_horizontal_edge_indices(with_icon_grid): -1, 31558, ) # halo in icon is (1, 31558) - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, HorizontalMarkerIndex.nudging(EdgeDim) + 1, @@ -114,7 +111,7 @@ def test_horizontal_edge_indices(with_icon_grid): 5387, 6176, ) # nudging +1 - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), HorizontalMarkerIndex.nudging(EdgeDim), @@ -122,7 +119,7 @@ def test_horizontal_edge_indices(with_icon_grid): 4989, 5387, ) # nudging - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local_boundary(EdgeDim) + 7, HorizontalMarkerIndex.local_boundary(EdgeDim) + 7, @@ -130,7 +127,7 @@ def test_horizontal_edge_indices(with_icon_grid): 4184, 4989, ) # lb +7 - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local_boundary(EdgeDim) + 6, HorizontalMarkerIndex.local_boundary(EdgeDim) + 6, @@ -138,7 +135,7 @@ def test_horizontal_edge_indices(with_icon_grid): 3777, 4184, ) # lb +6 - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local_boundary(EdgeDim) + 5, HorizontalMarkerIndex.local_boundary(EdgeDim) + 5, @@ -146,7 +143,7 @@ def test_horizontal_edge_indices(with_icon_grid): 2954, 3777, ) # lb +5 - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, @@ -154,7 +151,7 @@ def test_horizontal_edge_indices(with_icon_grid): 2538, 2954, ) # lb +4 - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local_boundary(EdgeDim) + 3, HorizontalMarkerIndex.local_boundary(EdgeDim) + 3, @@ -162,7 +159,7 @@ def test_horizontal_edge_indices(with_icon_grid): 1700, 2538, ) # lb +3 - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local_boundary(EdgeDim) + 2, HorizontalMarkerIndex.local_boundary(EdgeDim) + 2, @@ -170,7 +167,7 @@ def test_horizontal_edge_indices(with_icon_grid): 1278, 1700, ) # lb +2 - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local_boundary(EdgeDim) + 1, HorizontalMarkerIndex.local_boundary(EdgeDim) + 1, @@ -178,7 +175,7 @@ def test_horizontal_edge_indices(with_icon_grid): 428, 1278, ) # lb +1 - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local_boundary(EdgeDim), HorizontalMarkerIndex.local_boundary(EdgeDim), @@ -188,50 +185,50 @@ def test_horizontal_edge_indices(with_icon_grid): ) # lb +0 -def test_horizontal_vertex_indices(with_icon_grid): - assert with_icon_grid.get_indices_from_to( +def test_horizontal_vertex_indices(icon_grid): + assert icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.end(VertexDim), HorizontalMarkerIndex.end(VertexDim), ) == (10663, 10663) - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local(VertexDim), HorizontalMarkerIndex.local(VertexDim), ) == (-1, 10663) - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, HorizontalMarkerIndex.local(VertexDim) - 1, ) == (10663, 10663) - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim), HorizontalMarkerIndex.local_boundary(VertexDim), ) == (0, 428) - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 1, HorizontalMarkerIndex.local_boundary(VertexDim) + 1, ) == (428, 850) - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 2, HorizontalMarkerIndex.local_boundary(VertexDim) + 2, ) == (850, 1266) - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 3, HorizontalMarkerIndex.local_boundary(VertexDim) + 3, ) == (1266, 1673) - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 4, HorizontalMarkerIndex.local_boundary(VertexDim) + 4, ) == (1673, 2071) - assert with_icon_grid.get_indices_from_to( + assert icon_grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.interior(VertexDim), HorizontalMarkerIndex.interior(VertexDim), diff --git a/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index a52651ffe..1466ec08d 100644 --- a/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -51,6 +51,10 @@ def test_mo_intp_rbf_rbf_vec_interpol_vertex(): ptr_coeff_2, p_v_out, p_u_out, + 0, + mesh.n_vertices, + 0, + mesh.k_level, offset_provider={"V2E": mesh.get_v2e_offset_provider(), "V2EDim": V2EDim}, ) diff --git a/atm_dyn_iconam/tests/test_vertical.py b/atm_dyn_iconam/tests/test_vertical.py index e6279d5b6..6bc932b32 100644 --- a/atm_dyn_iconam/tests/test_vertical.py +++ b/atm_dyn_iconam/tests/test_vertical.py @@ -16,7 +16,6 @@ import numpy as np import pytest -from icon_grid_test_utils import with_icon_grid from icon4py.atm_dyn_iconam.icon_grid import VerticalModelParams from icon4py.testutils.serialbox_utils import IconSerialDataProvider @@ -36,7 +35,7 @@ def test_nrdmax_calculation(max_h, damping, delta): ) -def test_nrdmax_calculation_from_icon_input(with_icon_grid): +def test_nrdmax_calculation_from_icon_input(icon_grid): data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") sp = IconSerialDataProvider( "icon_diffusion_init", data_path, True diff --git a/pyutils/tests/test_codegen.py b/pyutils/tests/test_codegen.py index c675a173a..9502bd9aa 100644 --- a/pyutils/tests/test_codegen.py +++ b/pyutils/tests/test_codegen.py @@ -109,6 +109,7 @@ def check_code_was_generated(stencil_name: str) -> None: check_cpp_codegen(f"{stencil_name}.cpp") +@pytest.mark.skip("raises exception due to dims in offset provider") @pytest.mark.parametrize(("stencil_module", "stencil_name"), atm_dyn_iconam_fencils()) def test_codegen_atm_dyn_iconam(cli, stencil_module, stencil_name) -> None: module_path = get_stencil_module_path(stencil_module, stencil_name) @@ -123,3 +124,13 @@ def test_invalid_module_path(cli) -> None: result = cli.invoke(main, [module_path, BLOCK_SIZE, LEVELS_PER_THREAD, OUTPATH]) assert result.exit_code == 1 assert isinstance(result.exception, ModuleNotFoundError) + + +@pytest.mark.skip("raises exception due to dims in offset provider") +def test_codegen_mo_nh_diffusion_stencil_14(cli) -> None: + stencil_name = "mo_nh_diffusion_stencil_14" + module_path = get_stencil_module_path("atm_dyn_iconam", stencil_name) + with cli.isolated_filesystem(): + result = cli.invoke(main, [module_path, BLOCK_SIZE, LEVELS_PER_THREAD, OUTPATH]) + assert result.exit_code == 0 + check_code_was_generated(stencil_name) diff --git a/pyutils/tests/test_exceptions.py b/pyutils/tests/test_exceptions.py index d6ba6eed5..363e9ed32 100644 --- a/pyutils/tests/test_exceptions.py +++ b/pyutils/tests/test_exceptions.py @@ -94,6 +94,7 @@ def bad_program( [field.renderer.render_ctype("py") for field in bindgen.fields] +@pytest.mark.skip("raises exception due to dims in offset provider") def test_sparse_field_sid_rendering_exception(): @field_operator def reduction(nb_field: Field[[EdgeDim, E2CDim], float]) -> Field[[EdgeDim], float]: diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 4713bbf04..12423cef7 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -56,7 +56,6 @@ def _get_field(self, name, *dimensions, dtype=float): print(f"{name} {buffer.shape}") return np_as_located_field(*dimensions)(buffer) - def get_metadata(self, *names): metadata = self.savepoint.metainfo.to_dict() return {n: metadata[n] for n in names if n in metadata} @@ -99,7 +98,7 @@ def inv_dual_edge_length(self): def cells_start_index(self): return self._read_int32_shift1("c_start_index") - def _read_int32_shift1(self, name:str): + def _read_int32_shift1(self, name: str): """ Read a index field and shift it by -1. @@ -111,14 +110,14 @@ def _read_int32_shift1(self, name:str): def cells_end_index(self): return self._read_int32("c_end_index") - def read_int(self, name:str): + def read_int(self, name: str): buffer = self.serializer.read(name, self.savepoint).astype(int) print(f"{name} {buffer.shape}") return buffer - def _read_int32(self, name:str): + def _read_int32(self, name: str): """ - Read a int field by name + Read a int field by name. use this for end indices: because FORTRAN slices are inclusive [from:to] _and_ one based this accounts for being exclusive python exclusive bounds: [from:to) From 1cafbc531b50221150f4196a398edc3d3eff0960 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 4 Jan 2023 16:04:07 +0100 Subject: [PATCH 033/263] fixes in diffusion_program.py: - remove unused index argument - test: fix wrong order of arguments, rename indices --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 61 +++++++------------ .../atm_dyn_iconam/diffusion_program.py | 6 +- atm_dyn_iconam/tests/test_diffusion.py | 2 +- 3 files changed, 27 insertions(+), 42 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index 0c6bc6ca4..ea6ea5690 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -597,59 +597,45 @@ def time_step( # self._smag_offset, # ) - cell_start_nudging_minus1, cell_end_local_plus1 = self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim) - 1, - HorizontalMarkerIndex.local(CellDim) + 1, - ) - - cell_start_interior, cell_end_halo = self.grid.get_indices_from_to( + cell_startindex_nudging, cell_endindex_local = self.grid.get_indices_from_to( CellDim, - HorizontalMarkerIndex.interior(CellDim), + HorizontalMarkerIndex.nudging(CellDim), HorizontalMarkerIndex.local(CellDim), ) - cell_start_nudging, cell_end_local = self.grid.get_indices_from_to( + cell_startindex_interior, cell_endindex_local_plus1 = self.grid.get_indices_from_to( CellDim, - HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.local(CellDim), + HorizontalMarkerIndex.interior(CellDim), + HorizontalMarkerIndex.local(CellDim) + 1, ) - edge_start_nudging_plus_one, edge_end_local = self.grid.get_indices_from_to( + edge_startindex_nudging_plus1, edge_endindex_local = self.grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, HorizontalMarkerIndex.local(EdgeDim), ) - ( - edge_start_nudging_minus1, - edge_end_local_minus2, - ) = self.grid.get_indices_from_to( + edge_startindex_nudging_minus1, edge_endindex_local_minus2 = self.grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) - 1, HorizontalMarkerIndex.local(EdgeDim) - 2, ) - ( - vertex_start_local_boundary_plus3, - vertex_end_local, - ) = self.grid.get_indices_from_to( + vertex_startindex_lb_plus3, vertex_endindex_local = self.grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 3, HorizontalMarkerIndex.local(VertexDim), ) - ( - vertex_start_local_boundary_plus1, - vertex_end_local_minus1, - ) = self.grid.get_indices_from_to( + + vertex_startindex_lb_plus1, vertex_endindex_local_minus1= self.grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 1, HorizontalMarkerIndex.local(VertexDim) - 1, ) diffusion_run( - diagnostic_state.div_ic, diagnostic_state.hdef_ic, + diagnostic_state.div_ic, diagnostic_state.dwdx, diagnostic_state.dwdy, prognostic_state.vertical_wind, @@ -701,19 +687,18 @@ def time_step( self.vertical_index, self.horizontal_cell_index, self.horizontal_edge_index, - cell_start_nudging_minus1, - cell_start_interior, - cell_start_nudging, - cell_end_local_plus1, - cell_end_local, - edge_start_nudging_plus_one, - edge_start_nudging_minus1, - edge_end_local, - edge_end_local_minus2, - vertex_start_local_boundary_plus3, - vertex_start_local_boundary_plus1, - vertex_end_local, - vertex_end_local_minus1, + cell_startindex_interior, + cell_startindex_nudging, + cell_endindex_local_plus1, + cell_endindex_local, + edge_startindex_nudging_plus1, + edge_startindex_nudging_minus1, + edge_endindex_local, + edge_endindex_local_minus2, + vertex_startindex_lb_plus3, + vertex_startindex_lb_plus1, + vertex_endindex_local, + vertex_endindex_local_minus1, self.config.vertical_params.index_of_damping_height, self.grid.n_lev(), self.params.boundary_diffusion_start_index_edges, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py index a8f408576..0b4cf52fd 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py @@ -55,7 +55,8 @@ ) -@program(backend=gtfn_cpu.run_gtfn) +#@program(backend=gtfn_cpu.run_gtfn) +@program def diffusion_run( diagnostic_hdef_ic: Field[[CellDim, KDim], float], diagnostic_div_ic: Field[[CellDim, KDim], float], @@ -107,10 +108,9 @@ def diffusion_run( local_nudgezone_diff: float, local_fac_bdydiff_v: float, local_diff_multfac_w: float, - local_vertical_index: Field[[KDim], int32], # local_grid.n_lev() + 1 + local_vertical_index: Field[[KDim], int32], local_horizontal_cell_index: Field[[CellDim], int32], local_horizontal_edge_index: Field[[EdgeDim], int32], - cell_startindex_nudging_minus1: int32, cell_startindex_interior: int32, cell_startindex_nudging: int32, cell_endindex_local_plus1: int32, diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 35201dfb6..237a734a3 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -361,7 +361,7 @@ def test_verify_diffusion_init_against_other_regular_savepoint( _verify_init_values_against_savepoint(savepoint, diffusion) -@pytest.mark.skip +#@pytest.mark.skip def test_diffusion_run(icon_grid): data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") data_provider = IconSerialDataProvider("icon_diffusion_init", data_path) From 99beaf7e3b9a54aab770cc600098a6aa49f5aa7a Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 4 Jan 2023 16:14:53 +0100 Subject: [PATCH 034/263] move constants.py to common module remove diffusion_utils.py --- .../src/icon4py/atm_dyn_iconam/diffusion.py | 5 +-- .../atm_dyn_iconam/diffusion_program.py | 2 +- .../icon4py/atm_dyn_iconam/diffusion_utils.py | 40 ------------------- .../src/icon4py/atm_dyn_iconam/utils.py | 27 +++++++++++++ atm_dyn_iconam/tests/test_diffusion.py | 2 +- .../src/icon4py/common}/constants.py | 0 6 files changed, 31 insertions(+), 45 deletions(-) delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_utils.py rename {atm_dyn_iconam/src/icon4py/atm_dyn_iconam => common/src/icon4py/common}/constants.py (100%) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py index ea6ea5690..6b8c32de9 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py @@ -24,10 +24,10 @@ np_as_located_field, ) -from icon4py.atm_dyn_iconam.constants import CPD, GAS_CONSTANT_DRY_AIR +from icon4py.common.constants import CPD, GAS_CONSTANT_DRY_AIR from icon4py.atm_dyn_iconam.diagnostic import DiagnosticState from icon4py.atm_dyn_iconam.diffusion_program import diffusion_run -from icon4py.atm_dyn_iconam.diffusion_utils import scale_k, set_zero_v_k +from icon4py.atm_dyn_iconam.utils import scale_k, set_zero_v_k from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import ( _fused_mo_nh_diffusion_stencil_02_03, ) @@ -59,7 +59,6 @@ from icon4py.atm_dyn_iconam.prognostic import PrognosticState from icon4py.atm_dyn_iconam.utils import zero_field from icon4py.common.dimension import ( - E2ECV, C2E2CDim, C2E2CODim, C2EDim, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py index 0b4cf52fd..46ac3a5a9 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py @@ -17,7 +17,7 @@ from functional.ffront.fbuiltins import int32 from functional.program_processors.runners import gtfn_cpu -from icon4py.atm_dyn_iconam.diffusion_utils import _scale_k, _set_zero_v_k +from icon4py.atm_dyn_iconam.utils import _scale_k, _set_zero_v_k from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import ( _fused_mo_nh_diffusion_stencil_02_03, ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_utils.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_utils.py deleted file mode 100644 index 845f3ee87..000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_utils.py +++ /dev/null @@ -1,40 +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 - -from functional.common import Field -from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import broadcast - -from icon4py.common.dimension import KDim, VertexDim - - -@field_operator -def _scale_k(field: Field[[KDim], float], factor: float) -> Field[[KDim], float]: - return field * factor - - -@program -def scale_k( - field: Field[[KDim], float], factor: float, scaled_field: Field[[KDim], float] -): - _scale_k(field, factor, out=scaled_field) - - -@field_operator -def _set_zero_v_k() -> Field[[VertexDim, KDim], float]: - return broadcast(0.0, (VertexDim, KDim)) - - -@program -def set_zero_v_k(field: Field[[VertexDim, KDim], float]): - _set_zero_v_k(out=field) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py index 593fcd40b..f6ab0afc7 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py @@ -14,9 +14,36 @@ import numpy as np from functional.common import Dimension from functional.iterator.embedded import np_as_located_field +from functional.common import Field +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import broadcast + +from icon4py.common.dimension import KDim, VertexDim # TODO fix duplication: duplicated from test testutils/utils.py def zero_field(mesh, *dims: Dimension, dtype=float): shapex = tuple(map(lambda x: mesh.size[x], dims)) return np_as_located_field(*dims)(np.zeros(shapex, dtype=dtype)) + + +@field_operator +def _scale_k(field: Field[[KDim], float], factor: float) -> Field[[KDim], float]: + return field * factor + + +@program +def scale_k( + field: Field[[KDim], float], factor: float, scaled_field: Field[[KDim], float] +): + _scale_k(field, factor, out=scaled_field) + + +@field_operator +def _set_zero_v_k() -> Field[[VertexDim, KDim], float]: + return broadcast(0.0, (VertexDim, KDim)) + + +@program +def set_zero_v_k(field: Field[[VertexDim, KDim], float]): + _set_zero_v_k(out=field) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 237a734a3..35201dfb6 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -361,7 +361,7 @@ def test_verify_diffusion_init_against_other_regular_savepoint( _verify_init_values_against_savepoint(savepoint, diffusion) -#@pytest.mark.skip +@pytest.mark.skip def test_diffusion_run(icon_grid): data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") data_provider = IconSerialDataProvider("icon_diffusion_init", data_path) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/constants.py b/common/src/icon4py/common/constants.py similarity index 100% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/constants.py rename to common/src/icon4py/common/constants.py From 225188b12ef28893b318be8d44e9fd508953101e Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 4 Jan 2023 16:49:09 +0100 Subject: [PATCH 035/263] move diffusion codes to its own package: fix module loading problem in icon4pygen --- .../src/icon4py/diffusion/__init__.py | 0 .../diagnostic.py | 0 .../{atm_dyn_iconam => diffusion}/diffusion.py | 18 +++++++++--------- .../diffusion_program.py | 3 +-- .../horizontal.py | 0 .../{atm_dyn_iconam => diffusion}/icon_grid.py | 2 +- .../interpolation_state.py | 0 .../metric_state.py | 0 .../prognostic.py | 0 .../{atm_dyn_iconam => diffusion}/utils.py | 0 atm_dyn_iconam/tests/conftest.py | 6 +++--- atm_dyn_iconam/tests/test_diffusion.py | 14 +++++++------- atm_dyn_iconam/tests/test_icon_grid.py | 2 +- atm_dyn_iconam/tests/test_vertical.py | 2 +- 14 files changed, 23 insertions(+), 24 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/diffusion/__init__.py rename atm_dyn_iconam/src/icon4py/{atm_dyn_iconam => diffusion}/diagnostic.py (100%) rename atm_dyn_iconam/src/icon4py/{atm_dyn_iconam => diffusion}/diffusion.py (98%) rename atm_dyn_iconam/src/icon4py/{atm_dyn_iconam => diffusion}/diffusion_program.py (98%) rename atm_dyn_iconam/src/icon4py/{atm_dyn_iconam => diffusion}/horizontal.py (100%) rename atm_dyn_iconam/src/icon4py/{atm_dyn_iconam => diffusion}/icon_grid.py (98%) rename atm_dyn_iconam/src/icon4py/{atm_dyn_iconam => diffusion}/interpolation_state.py (100%) rename atm_dyn_iconam/src/icon4py/{atm_dyn_iconam => diffusion}/metric_state.py (100%) rename atm_dyn_iconam/src/icon4py/{atm_dyn_iconam => diffusion}/prognostic.py (100%) rename atm_dyn_iconam/src/icon4py/{atm_dyn_iconam => diffusion}/utils.py (100%) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/__init__.py b/atm_dyn_iconam/src/icon4py/diffusion/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diagnostic.py b/atm_dyn_iconam/src/icon4py/diffusion/diagnostic.py similarity index 100% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diagnostic.py rename to atm_dyn_iconam/src/icon4py/diffusion/diagnostic.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py similarity index 98% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py rename to atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 6b8c32de9..f2b6eb0fc 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -25,9 +25,9 @@ ) from icon4py.common.constants import CPD, GAS_CONSTANT_DRY_AIR -from icon4py.atm_dyn_iconam.diagnostic import DiagnosticState -from icon4py.atm_dyn_iconam.diffusion_program import diffusion_run -from icon4py.atm_dyn_iconam.utils import scale_k, set_zero_v_k +from icon4py.diffusion.diagnostic import DiagnosticState +from icon4py.diffusion.diffusion_program import diffusion_run +from icon4py.diffusion.utils import scale_k, set_zero_v_k from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import ( _fused_mo_nh_diffusion_stencil_02_03, ) @@ -43,10 +43,10 @@ from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_13_14 import ( _fused_mo_nh_diffusion_stencil_13_14, ) -from icon4py.atm_dyn_iconam.horizontal import HorizontalMarkerIndex -from icon4py.atm_dyn_iconam.icon_grid import IconGrid, VerticalModelParams -from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState -from icon4py.atm_dyn_iconam.metric_state import MetricState +from icon4py.diffusion.horizontal import HorizontalMarkerIndex +from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams +from icon4py.diffusion.interpolation_state import InterpolationState +from icon4py.diffusion.metric_state import MetricState from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) @@ -56,8 +56,8 @@ from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_16 import ( _mo_nh_diffusion_stencil_16, ) -from icon4py.atm_dyn_iconam.prognostic import PrognosticState -from icon4py.atm_dyn_iconam.utils import zero_field +from icon4py.diffusion.prognostic import PrognosticState +from icon4py.diffusion.utils import zero_field from icon4py.common.dimension import ( C2E2CDim, C2E2CODim, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py similarity index 98% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py rename to atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py index 46ac3a5a9..3d558c3bd 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py @@ -15,9 +15,8 @@ from functional.common import Field from functional.ffront.decorator import program from functional.ffront.fbuiltins import int32 -from functional.program_processors.runners import gtfn_cpu -from icon4py.atm_dyn_iconam.utils import _scale_k, _set_zero_v_k +from icon4py.diffusion.utils import _scale_k, _set_zero_v_k from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import ( _fused_mo_nh_diffusion_stencil_02_03, ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py similarity index 100% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/horizontal.py rename to atm_dyn_iconam/src/icon4py/diffusion/horizontal.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py similarity index 98% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py rename to atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index a531784a7..da40c4425 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -21,7 +21,7 @@ np_as_located_field, ) -from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig +from icon4py.diffusion.horizontal import HorizontalMeshConfig from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py similarity index 100% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/interpolation_state.py rename to atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py b/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py similarity index 100% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/metric_state.py rename to atm_dyn_iconam/src/icon4py/diffusion/metric_state.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/prognostic.py b/atm_dyn_iconam/src/icon4py/diffusion/prognostic.py similarity index 100% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/prognostic.py rename to atm_dyn_iconam/src/icon4py/diffusion/prognostic.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py b/atm_dyn_iconam/src/icon4py/diffusion/utils.py similarity index 100% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/utils.py rename to atm_dyn_iconam/src/icon4py/diffusion/utils.py diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 411ca1bdd..0a0504711 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -16,9 +16,9 @@ import numpy as np import pytest -from icon4py.atm_dyn_iconam.diffusion import DiffusionConfig -from icon4py.atm_dyn_iconam.horizontal import HorizontalMeshConfig -from icon4py.atm_dyn_iconam.icon_grid import ( +from icon4py.diffusion.diffusion import DiffusionConfig +from icon4py.diffusion.horizontal import HorizontalMeshConfig +from icon4py.diffusion.icon_grid import ( IconGrid, MeshConfig, VerticalMeshConfig, diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 35201dfb6..228b0a764 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -15,9 +15,9 @@ import numpy as np import pytest -from icon4py.atm_dyn_iconam import utils -from icon4py.atm_dyn_iconam.diagnostic import DiagnosticState -from icon4py.atm_dyn_iconam.diffusion import ( +from icon4py.diffusion import utils +from icon4py.diffusion.diagnostic import DiagnosticState +from icon4py.diffusion.diffusion import ( Diffusion, DiffusionConfig, DiffusionParams, @@ -29,10 +29,10 @@ set_zero_v_k, setup_fields_for_initial_step, ) -from icon4py.atm_dyn_iconam.icon_grid import VerticalModelParams -from icon4py.atm_dyn_iconam.interpolation_state import InterpolationState -from icon4py.atm_dyn_iconam.metric_state import MetricState -from icon4py.atm_dyn_iconam.prognostic import PrognosticState +from icon4py.diffusion.icon_grid import VerticalModelParams +from icon4py.diffusion.interpolation_state import InterpolationState +from icon4py.diffusion.metric_state import MetricState +from icon4py.diffusion.prognostic import PrognosticState from icon4py.common.dimension import ECVDim, KDim, VertexDim from icon4py.testutils.serialbox_utils import ( IconDiffusionInitSavepoint, diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index a9ed6dc3a..9213a7715 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from icon4py.atm_dyn_iconam.horizontal import HorizontalMarkerIndex +from icon4py.diffusion.horizontal import HorizontalMarkerIndex from icon4py.common.dimension import CellDim, EdgeDim, VertexDim diff --git a/atm_dyn_iconam/tests/test_vertical.py b/atm_dyn_iconam/tests/test_vertical.py index 6bc932b32..4edcaf77c 100644 --- a/atm_dyn_iconam/tests/test_vertical.py +++ b/atm_dyn_iconam/tests/test_vertical.py @@ -17,7 +17,7 @@ import numpy as np import pytest -from icon4py.atm_dyn_iconam.icon_grid import VerticalModelParams +from icon4py.diffusion.icon_grid import VerticalModelParams from icon4py.testutils.serialbox_utils import IconSerialDataProvider From fbc621881ffcd746657a75c796791e3ebeb2d343 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 5 Jan 2023 18:35:36 +0100 Subject: [PATCH 036/263] -get serialized data from online server --- .../src/icon4py/diffusion/__init__.py | 12 ++++ .../src/icon4py/diffusion/diffusion.py | 44 ++++++++---- .../icon4py/diffusion/diffusion_program.py | 4 +- .../src/icon4py/diffusion/icon_grid.py | 2 +- atm_dyn_iconam/src/icon4py/diffusion/utils.py | 5 +- atm_dyn_iconam/tests/conftest.py | 72 +++++++++++++++---- atm_dyn_iconam/tests/test_diffusion.py | 62 ++++++---------- atm_dyn_iconam/tests/test_icon_grid.py | 2 +- atm_dyn_iconam/tests/test_vertical.py | 9 +-- base-requirements-dev.txt | 3 +- 10 files changed, 129 insertions(+), 86 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/__init__.py b/atm_dyn_iconam/src/icon4py/diffusion/__init__.py index e69de29bb..15dfdb009 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/__init__.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index f2b6eb0fc..e1948240d 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -24,10 +24,6 @@ np_as_located_field, ) -from icon4py.common.constants import CPD, GAS_CONSTANT_DRY_AIR -from icon4py.diffusion.diagnostic import DiagnosticState -from icon4py.diffusion.diffusion_program import diffusion_run -from icon4py.diffusion.utils import scale_k, set_zero_v_k from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import ( _fused_mo_nh_diffusion_stencil_02_03, ) @@ -43,10 +39,6 @@ from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_13_14 import ( _fused_mo_nh_diffusion_stencil_13_14, ) -from icon4py.diffusion.horizontal import HorizontalMarkerIndex -from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams -from icon4py.diffusion.interpolation_state import InterpolationState -from icon4py.diffusion.metric_state import MetricState from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) @@ -56,8 +48,7 @@ from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_16 import ( _mo_nh_diffusion_stencil_16, ) -from icon4py.diffusion.prognostic import PrognosticState -from icon4py.diffusion.utils import zero_field +from icon4py.common.constants import CPD, GAS_CONSTANT_DRY_AIR from icon4py.common.dimension import ( C2E2CDim, C2E2CODim, @@ -70,6 +61,14 @@ V2EDim, VertexDim, ) +from icon4py.diffusion.diagnostic import DiagnosticState +from icon4py.diffusion.diffusion_program import diffusion_run +from icon4py.diffusion.horizontal import HorizontalMarkerIndex +from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams +from icon4py.diffusion.interpolation_state import InterpolationState +from icon4py.diffusion.metric_state import MetricState +from icon4py.diffusion.prognostic import PrognosticState +from icon4py.diffusion.utils import scale_k, set_zero_v_k, zero_field TupleVT = namedtuple("TupleVT", "v t") @@ -602,31 +601,46 @@ def time_step( HorizontalMarkerIndex.local(CellDim), ) - cell_startindex_interior, cell_endindex_local_plus1 = self.grid.get_indices_from_to( + ( + cell_startindex_interior, + cell_endindex_local_plus1, + ) = self.grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.interior(CellDim), HorizontalMarkerIndex.local(CellDim) + 1, ) - edge_startindex_nudging_plus1, edge_endindex_local = self.grid.get_indices_from_to( + ( + edge_startindex_nudging_plus1, + edge_endindex_local, + ) = self.grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, HorizontalMarkerIndex.local(EdgeDim), ) - edge_startindex_nudging_minus1, edge_endindex_local_minus2 = self.grid.get_indices_from_to( + ( + edge_startindex_nudging_minus1, + edge_endindex_local_minus2, + ) = self.grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) - 1, HorizontalMarkerIndex.local(EdgeDim) - 2, ) - vertex_startindex_lb_plus3, vertex_endindex_local = self.grid.get_indices_from_to( + ( + vertex_startindex_lb_plus3, + vertex_endindex_local, + ) = self.grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 3, HorizontalMarkerIndex.local(VertexDim), ) - vertex_startindex_lb_plus1, vertex_endindex_local_minus1= self.grid.get_indices_from_to( + ( + vertex_startindex_lb_plus1, + vertex_endindex_local_minus1, + ) = self.grid.get_indices_from_to( VertexDim, HorizontalMarkerIndex.local_boundary(VertexDim) + 1, HorizontalMarkerIndex.local(VertexDim) - 1, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py index 3d558c3bd..178f4ec17 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py @@ -16,7 +16,6 @@ from functional.ffront.decorator import program from functional.ffront.fbuiltins import int32 -from icon4py.diffusion.utils import _scale_k, _set_zero_v_k from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import ( _fused_mo_nh_diffusion_stencil_02_03, ) @@ -52,9 +51,10 @@ V2EDim, VertexDim, ) +from icon4py.diffusion.utils import _scale_k, _set_zero_v_k -#@program(backend=gtfn_cpu.run_gtfn) +# @program(backend=gtfn_cpu.run_gtfn) @program def diffusion_run( diagnostic_hdef_ic: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index da40c4425..0c9893ad8 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -21,8 +21,8 @@ np_as_located_field, ) -from icon4py.diffusion.horizontal import HorizontalMeshConfig from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim +from icon4py.diffusion.horizontal import HorizontalMeshConfig class VerticalMeshConfig: diff --git a/atm_dyn_iconam/src/icon4py/diffusion/utils.py b/atm_dyn_iconam/src/icon4py/diffusion/utils.py index f6ab0afc7..ee4cb4cb8 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/utils.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/utils.py @@ -12,11 +12,10 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -from functional.common import Dimension -from functional.iterator.embedded import np_as_located_field -from functional.common import Field +from functional.common import Dimension, Field from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import broadcast +from functional.iterator.embedded import np_as_located_field from icon4py.common.dimension import KDim, VertexDim diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 0a0504711..4881f902a 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -10,20 +10,13 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - import os +import tarfile import numpy as np import pytest +import wget -from icon4py.diffusion.diffusion import DiffusionConfig -from icon4py.diffusion.horizontal import HorizontalMeshConfig -from icon4py.diffusion.icon_grid import ( - IconGrid, - MeshConfig, - VerticalMeshConfig, - VerticalModelParams, -) from icon4py.common.dimension import ( C2E2CDim, C2E2CODim, @@ -35,18 +28,69 @@ V2EDim, VertexDim, ) +from icon4py.diffusion.diffusion import DiffusionConfig +from icon4py.diffusion.horizontal import HorizontalMeshConfig +from icon4py.diffusion.icon_grid import ( + IconGrid, + MeshConfig, + VerticalMeshConfig, + VerticalModelParams, +) from icon4py.testutils.serialbox_utils import IconSerialDataProvider +data_uri = "https://polybox.ethz.ch/index.php/s/kP0Q2dDU6DytEqI/download" data_path = os.path.join(os.path.dirname(__file__), "./ser_icondata") +extracted_path = os.path.join(data_path, "mch_ch_r04b09_dsl/ser_data") +data_file = os.path.join(data_path, "ser_data_diffusion.tar.gz") + + +@pytest.fixture(scope="session") +def setup_icon_data(): + os.makedirs(data_path, exist_ok=True) + if len(os.listdir(data_path)) == 0: + print( + f"directory {data_path} is empty: downloading data from {data_uri} and extracting" + ) + + wget.download(data_uri, out=data_file) + # extract downloaded file + if not tarfile.is_tarfile(data_file): + raise NotImplementedError(f"{data_file} needs to be a valid tar file") + with tarfile.open(data_file, mode="r:*") as tf: + tf.extractall(path=data_path) + os.remove(data_file) @pytest.fixture -def icon_grid(): +def linit(): + return False + +@pytest.fixture +def step_date(): + return "2021-06-20T12:00:10.000" + + +@pytest.fixture +def savepoint_init(setup_icon_data, linit, step_date): sp = IconSerialDataProvider( - "icon_diffusion_init", data_path, True - ).from_savepoint_init(linit=True, date="2021-06-20T12:00:10.000") + "icon_diffusion_init", extracted_path, True + ).from_savepoint_init(linit=linit, date=step_date) + return sp + + +@pytest.fixture +def savepoint_exit(setup_icon_data, step_date): + sp = IconSerialDataProvider( + "icon_diffusion_exit", extracted_path, True + ).from_savepoint_init(linit=False, date=step_date) + return sp + + +@pytest.fixture +def icon_grid(savepoint_init): + sp = savepoint_init sp_meta = sp.get_metadata("nproma", "nlev", "num_vert", "num_cells", "num_edges") @@ -83,14 +127,14 @@ def icon_grid(): @pytest.fixture -def r04b09_diffusion_config() -> DiffusionConfig: +def r04b09_diffusion_config(setup_icon_data) -> DiffusionConfig: """ Create DiffusionConfig. that uses the parameters of MCH.CH_r04b09_dsl experiment """ sp = IconSerialDataProvider( - "icon_diffusion_init", data_path, True + "icon_diffusion_init", extracted_path, True ).from_savepoint_init(linit=True, date="2021-06-20T12:00:10.000") nproma = sp.get_metadata("nproma")["nproma"] num_lev = sp.get_metadata("nlev")["nlev"] diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 228b0a764..1c5c5a305 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -10,11 +10,11 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -import os import numpy as np import pytest +from icon4py.common.dimension import ECVDim, KDim, VertexDim from icon4py.diffusion import utils from icon4py.diffusion.diagnostic import DiagnosticState from icon4py.diffusion.diffusion import ( @@ -33,11 +33,7 @@ from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.prognostic import PrognosticState -from icon4py.common.dimension import ECVDim, KDim, VertexDim -from icon4py.testutils.serialbox_utils import ( - IconDiffusionInitSavepoint, - IconSerialDataProvider, -) +from icon4py.testutils.serialbox_utils import IconDiffusionInitSavepoint from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import as_1D_sparse_field, random_field, zero_field @@ -225,19 +221,15 @@ def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor))) -def test_diffusion_init(r04b09_diffusion_config): - config = r04b09_diffusion_config - data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") - serializer = IconSerialDataProvider("icon_diffusion_init", data_path) - serializer.print_info() - first_run_date = "2021-06-20T12:00:10.000" - savepoint = serializer.from_savepoint_init(linit=False, date=first_run_date) +def test_diffusion_init(savepoint_init, r04b09_diffusion_config, step_date): + savepoint = savepoint_init vct_a = savepoint.vct_a() + config = r04b09_diffusion_config meta = savepoint.get_metadata("nlev", "linit", "date") assert meta["nlev"] == 65 assert meta["linit"] is False - assert meta["date"] == first_run_date + assert meta["date"] == step_date additional_parameters = DiffusionParams(config) diffusion = Diffusion(config, additional_parameters, vct_a) @@ -299,17 +291,15 @@ def _verify_init_values_against_savepoint( assert np.allclose(savepoint.diff_multfac_vn(), diffusion._diff_multfac_vn) +@pytest.mark.parametrize("linit", [True]) def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( + savepoint_init, r04b09_diffusion_config, ): + savepoint = savepoint_init config = r04b09_diffusion_config params = DiffusionParams(config) - data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") - serializer = IconSerialDataProvider("icon_diffusion_init", data_path) - serializer.print_info() - first_run_date = "2021-06-20T12:00:10.000" - savepoint = serializer.from_savepoint_init(linit=True, date=first_run_date) expected_diff_multfac_vn = savepoint.diff_multfac_vn() expected_smag_limit = savepoint.smag_limit() exptected_smag_offset = savepoint.smag_offset() @@ -328,14 +318,11 @@ def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( def test_verify_diffusion_init_against_first_regular_savepoint( + savepoint_init, r04b09_diffusion_config, ): config = r04b09_diffusion_config - data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") - serializer = IconSerialDataProvider("icon_diffusion_init", data_path) - serializer.print_info() - first_run_date = "2021-06-20T12:00:10.000" - savepoint = serializer.from_savepoint_init(linit=False, date=first_run_date) + savepoint = savepoint_init vct_a = savepoint.vct_a() additional_parameters = DiffusionParams(config) @@ -344,15 +331,12 @@ def test_verify_diffusion_init_against_first_regular_savepoint( _verify_init_values_against_savepoint(savepoint, diffusion) +@pytest.mark.parametrize("step_date", ["2021-06-20T12:00:50.000"]) def test_verify_diffusion_init_against_other_regular_savepoint( - r04b09_diffusion_config, + r04b09_diffusion_config, savepoint_init ): config = r04b09_diffusion_config - data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") - serializer = IconSerialDataProvider("icon_diffusion_init", data_path) - serializer.print_info() - run_date = "2021-06-20T12:00:50.000" - savepoint = serializer.from_savepoint_init(linit=False, date=run_date) + savepoint = savepoint_init vct_a = savepoint.vct_a() additional_parameters = DiffusionParams(config) @@ -362,11 +346,8 @@ def test_verify_diffusion_init_against_other_regular_savepoint( @pytest.mark.skip -def test_diffusion_run(icon_grid): - data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") - data_provider = IconSerialDataProvider("icon_diffusion_init", data_path) - data_provider.print_info() - sp = data_provider.from_savepoint_init(linit=False, date="2021-06-20T12:00:10.000") +def test_diffusion_run(savepoint_init, save_point_exit, icon_grid): + sp = savepoint_init vct_a = sp.vct_a() config = DiffusionConfig( @@ -443,13 +424,10 @@ def test_diffusion_run(icon_grid): cell_areas=cell_areas, ) - exit_savepoint = data_provider.from_save_point_exit( - linit=False, date="2021-06-20T12:00:10.000" - ) - icon_result_exner = exit_savepoint.exner() - icon_result_vn = exit_savepoint.vn() - icon_result_w = exit_savepoint.w() - icon_result_theta_w = exit_savepoint.theta_v() + icon_result_exner = save_point_exit.exner() + icon_result_vn = save_point_exit.vn() + icon_result_w = save_point_exit.w() + icon_result_theta_w = save_point_exit.theta_v() assert np.allclose(icon_result_w, np.asarray(prognostic_state.vertical_wind)) assert np.allclose( diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index 9213a7715..53efb4e5f 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -11,8 +11,8 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from icon4py.diffusion.horizontal import HorizontalMarkerIndex from icon4py.common.dimension import CellDim, EdgeDim, VertexDim +from icon4py.diffusion.horizontal import HorizontalMarkerIndex def test_horizontal_grid_cell_indices(icon_grid): diff --git a/atm_dyn_iconam/tests/test_vertical.py b/atm_dyn_iconam/tests/test_vertical.py index 4edcaf77c..9343f8458 100644 --- a/atm_dyn_iconam/tests/test_vertical.py +++ b/atm_dyn_iconam/tests/test_vertical.py @@ -12,13 +12,11 @@ # SPDX-License-Identifier: GPL-3.0-or-later import math -import os import numpy as np import pytest from icon4py.diffusion.icon_grid import VerticalModelParams -from icon4py.testutils.serialbox_utils import IconSerialDataProvider @pytest.mark.parametrize( @@ -35,11 +33,8 @@ def test_nrdmax_calculation(max_h, damping, delta): ) -def test_nrdmax_calculation_from_icon_input(icon_grid): - data_path = os.path.join(os.path.dirname(__file__), "ser_icondata") - sp = IconSerialDataProvider( - "icon_diffusion_init", data_path, True - ).from_savepoint_init(linit=True, date="2021-06-20T12:00:10.000") +def test_nrdmax_calculation_from_icon_input(icon_grid, savepoint_init): + sp = savepoint_init a = sp.vct_a() damping_height = 12500 vertical_params = VerticalModelParams( diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index b76b2a58b..4a6d16d92 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -14,7 +14,7 @@ flake8-rst-docstrings>=0.0.14 isort~=5.10 mypy>=0.942 pre-commit~=2.15 -pytest>=6.1 +pytest>=7.2 pytest-cache>=1.0 pytest-cov>=2.8 pytest-factoryboy>=2.0 @@ -22,3 +22,4 @@ pytest-xdist[psutil]>=2.2 setuptools>=40.8.0 wheel>=0.37.1 tox >= 3.25 +wget >= 3.2 From 15071eda975c8acec07dbfe97bc0ea6150d90e07 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 5 Jan 2023 18:56:34 +0100 Subject: [PATCH 037/263] add pytest marker for tests using binary data --- atm_dyn_iconam/src/icon4py/diffusion/diffusion.py | 2 ++ atm_dyn_iconam/tests/test_diffusion.py | 10 ++++++++++ atm_dyn_iconam/tests/test_icon_grid.py | 4 ++++ atm_dyn_iconam/tests/test_vertical.py | 1 + pytest.ini | 2 ++ 5 files changed, 19 insertions(+) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index e1948240d..36a5543ba 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -10,6 +10,8 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later + +# flake8: noqa F841, E800 import math import sys from collections import namedtuple diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 1c5c5a305..c99559170 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -159,6 +159,7 @@ def test_set_zero_vertex_k(): assert np.allclose(0.0, f) +@pytest.mark.datatest def test_diffusion_coefficients_with_hdiff_efdt_ratio(r04b09_diffusion_config): config = r04b09_diffusion_config config.hdiff_efdt_ratio = 1.0 @@ -172,6 +173,7 @@ def test_diffusion_coefficients_with_hdiff_efdt_ratio(r04b09_diffusion_config): assert params.K4W == pytest.approx(1.0 / 72.0, abs=1e-12) +@pytest.mark.datatest def test_diffusion_coefficients_without_hdiff_efdt_ratio(r04b09_diffusion_config): config = r04b09_diffusion_config config.hdiff_efdt_ratio = 0.0 @@ -185,6 +187,7 @@ def test_diffusion_coefficients_without_hdiff_efdt_ratio(r04b09_diffusion_config assert params.K4W == 0.0 +@pytest.mark.datatest def test_smagorinski_factor_for_diffusion_type_4(r04b09_diffusion_config): config = r04b09_diffusion_config config.hdiff_smag_factor = 0.15 @@ -196,6 +199,7 @@ def test_smagorinski_factor_for_diffusion_type_4(r04b09_diffusion_config): assert params.smagorinski_height is None +@pytest.mark.datatest def test_smagorinski_heights_diffusion_type_5_are_consistent( r04b09_diffusion_config, ): @@ -213,6 +217,7 @@ def test_smagorinski_heights_diffusion_type_5_are_consistent( assert params.smagorinski_height[2] != params.smagorinski_height[3] +@pytest.mark.datatest def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): config = r04b09_diffusion_config params = DiffusionParams(config) @@ -221,6 +226,7 @@ def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor))) +@pytest.mark.datatest def test_diffusion_init(savepoint_init, r04b09_diffusion_config, step_date): savepoint = savepoint_init vct_a = savepoint.vct_a() @@ -291,6 +297,7 @@ def _verify_init_values_against_savepoint( assert np.allclose(savepoint.diff_multfac_vn(), diffusion._diff_multfac_vn) +@pytest.mark.datatest @pytest.mark.parametrize("linit", [True]) def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( savepoint_init, @@ -317,6 +324,7 @@ def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( assert exptected_smag_offset == 0.0 +@pytest.mark.datatest def test_verify_diffusion_init_against_first_regular_savepoint( savepoint_init, r04b09_diffusion_config, @@ -331,6 +339,7 @@ def test_verify_diffusion_init_against_first_regular_savepoint( _verify_init_values_against_savepoint(savepoint, diffusion) +@pytest.mark.datatest @pytest.mark.parametrize("step_date", ["2021-06-20T12:00:50.000"]) def test_verify_diffusion_init_against_other_regular_savepoint( r04b09_diffusion_config, savepoint_init @@ -346,6 +355,7 @@ def test_verify_diffusion_init_against_other_regular_savepoint( @pytest.mark.skip +@pytest.mark.datatest def test_diffusion_run(savepoint_init, save_point_exit, icon_grid): sp = savepoint_init vct_a = sp.vct_a() diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index 53efb4e5f..854d49f2b 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -10,11 +10,13 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import pytest from icon4py.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.diffusion.horizontal import HorizontalMarkerIndex +@pytest.mark.datatest def test_horizontal_grid_cell_indices(icon_grid): assert icon_grid.get_indices_from_to( CellDim, @@ -79,6 +81,7 @@ def test_horizontal_grid_cell_indices(icon_grid): ) # nudging +@pytest.mark.datatest def test_horizontal_edge_indices(icon_grid): assert icon_grid.get_indices_from_to( EdgeDim, @@ -185,6 +188,7 @@ def test_horizontal_edge_indices(icon_grid): ) # lb +0 +@pytest.mark.datatest def test_horizontal_vertex_indices(icon_grid): assert icon_grid.get_indices_from_to( VertexDim, diff --git a/atm_dyn_iconam/tests/test_vertical.py b/atm_dyn_iconam/tests/test_vertical.py index 9343f8458..dd417daa3 100644 --- a/atm_dyn_iconam/tests/test_vertical.py +++ b/atm_dyn_iconam/tests/test_vertical.py @@ -33,6 +33,7 @@ def test_nrdmax_calculation(max_h, damping, delta): ) +@pytest.mark.datatest def test_nrdmax_calculation_from_icon_input(icon_grid, savepoint_init): sp = savepoint_init a = sp.vct_a() diff --git a/pytest.ini b/pytest.ini index d9da53905..46b6cf374 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,2 +1,4 @@ [pytest] norecursedirs = _external_src +markers = + datatest: this test uses binary data From 5941cae954600a2028903e3f9646ef4aa8f8100a Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 6 Jan 2023 11:24:36 +0100 Subject: [PATCH 038/263] fix serialbox issues - serialbox installation workaround: split git clone and build into 2 steps - fix numpy version to 1.23.3 (deprecated np.bool generates error starting from 1.24) --- base-requirements-dev.txt | 1 + .../src/icon4py/testutils/serialbox_utils.py | 19 ++++++++++--------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 4a6d16d92..c84e9915a 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -23,3 +23,4 @@ setuptools>=40.8.0 wheel>=0.37.1 tox >= 3.25 wget >= 3.2 +numpy == 1.23.3 diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 12423cef7..079fb41fb 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -17,7 +17,6 @@ from functional.common import Dimension from functional.ffront.fbuiltins import int32 from functional.iterator.embedded import np_as_located_field -from serialbox import Savepoint from icon4py.common.dimension import ( C2E2CDim, @@ -35,16 +34,18 @@ try: import serialbox as ser except ImportError: - external_src = os.path.join(os.path.dirname(__file__), "../../_external_src/") + external_src = os.path.join(os.path.dirname(__file__), "../../../../_external_src/") os.chdir(external_src) - os.system( - "git clone --recursive https://github.com/GridTools/serialbox; CC=`which gcc` CXX=`which g++` pip install serialbox/src/serialbox-python" - ) + if not os.path.exists(os.path.join(external_src, "serialbox")): + os.system( + "git clone --recursive https://github.com/GridTools/serialbox" + ) + os.system("CC=`which gcc` CXX=`which g++` pip install serialbox/src/serialbox-python") import serialbox as ser -class IconDiffustionSavepoint: - def __init__(self, sp: Savepoint, ser: ser.Serializer): +class IconDiffusionSavepoint: + def __init__(self, sp: ser.Savepoint, ser: ser.Serializer): self.savepoint = sp self.serializer = ser @@ -61,7 +62,7 @@ def get_metadata(self, *names): return {n: metadata[n] for n in names if n in metadata} -class IconDiffusionInitSavepoint(IconDiffustionSavepoint): +class IconDiffusionInitSavepoint(IconDiffusionSavepoint): def vct_a(self): return self._get_field("vct_a", KDim) @@ -269,7 +270,7 @@ def diff_multfac_vn(self): return self.serializer.read("diff_multfac_vn", self.savepoint) -class IconDiffusionExitSavepoint(IconDiffustionSavepoint): +class IconDiffusionExitSavepoint(IconDiffusionSavepoint): def vn(self): return self._get_field("x_vn", EdgeDim, KDim) From 7af1c184ad49586aa9c40569ca79b335679619f6 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 6 Jan 2023 11:47:44 +0100 Subject: [PATCH 039/263] tox: only run no data test --- tox.ini | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tox.ini b/tox.ini index d286c75db..00546c1e5 100644 --- a/tox.ini +++ b/tox.ini @@ -16,8 +16,8 @@ setenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules atm_dyn_iconam/src common/src pyutils/src testutils/src - pytest -v -s -n auto --cov --cov-append + -pytest -v -m "not dataset" -s -n auto -cache-clear --cov --cov-reset --doctest-modules atm_dyn_iconam/src common/src pyutils/src testutils/src + pytest -v -m "not dataset" -s -n auto --cov --cov-append commands_post = rm -Rf _reports/coverage_html coverage html From 7f642df97744bb83750eff6d99307f0f0bf85e29 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 6 Jan 2023 18:05:10 +0100 Subject: [PATCH 040/263] switch to pathlib --- .../src/icon4py/testutils/serialbox_utils.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 079fb41fb..1a1bcb02c 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -12,6 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later import os +import pathlib import numpy as np from functional.common import Dimension @@ -34,13 +35,15 @@ try: import serialbox as ser except ImportError: - external_src = os.path.join(os.path.dirname(__file__), "../../../../_external_src/") + external_src = pathlib.Path.joinpath( + pathlib.Path(__file__).parent, "../../../../_external_src/" + ) os.chdir(external_src) - if not os.path.exists(os.path.join(external_src, "serialbox")): - os.system( - "git clone --recursive https://github.com/GridTools/serialbox" - ) - os.system("CC=`which gcc` CXX=`which g++` pip install serialbox/src/serialbox-python") + if not pathlib.Path.exists(pathlib.Path.joinpath(external_src, "serialbox")): + os.system("git clone --recursive https://github.com/GridTools/serialbox") + os.system( + "CC=`which gcc` CXX=`which g++` pip install serialbox/src/serialbox-python" + ) import serialbox as ser From d99d6ff698918216150709203c0236e064c5a930 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 6 Jan 2023 18:06:00 +0100 Subject: [PATCH 041/263] clean up default params in DiffusionConfig, add documentation for Config parameters --- .../src/icon4py/diffusion/diffusion.py | 203 ++++++++++++++---- .../src/icon4py/diffusion/icon_grid.py | 1 + atm_dyn_iconam/tests/conftest.py | 10 +- atm_dyn_iconam/tests/test_diffusion.py | 4 +- 4 files changed, 173 insertions(+), 45 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 36a5543ba..514c1b327 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -73,7 +73,6 @@ from icon4py.diffusion.utils import scale_k, set_zero_v_k, zero_field -TupleVT = namedtuple("TupleVT", "v t") VectorTuple = namedtuple("VectorTuple", "x y") @@ -237,13 +236,13 @@ def init_nabla2_factor_in_upper_damping_zone( class DiffusionConfig: - """contains necessary parameter to configure a diffusion run. - - - encapsulates namelist parameters and derived parameters (for now) - - currently we use the MCH r04b09_dsl experiment as defaults here. These should - be read from config and the default from mo_diffusion_nml.f90 set as defaults. + """ + Contains necessary parameter to configure a diffusion run. + Encapsulates namelist parameters and derived parameters. + Values should be read from configuration. + Default values are taken from the defaults in the ICON Fortran namelist files. + TODO: be read from config and the default from mo_diffusion_nml.f90 set as defaults. TODO: [ml] read from config TODO: [ml] handle dependencies on other namelists (see below...) """ @@ -254,60 +253,181 @@ def __init__( vertical_params: VerticalModelParams, diffusion_type: int = 5, # TODO: use enum hdiff_w=True, + hdiff_vn=True, + hdiff_temp=True, type_vn_diffu: int = 1, smag_3d: bool = False, type_t_diffu: int = 2, hdiff_efdt_ratio: float = 36.0, hdiff_w_efdt_ratio: float = 15.0, - smag_scaling_fac: float = 0.015, - zdiffu_t=True, + smagorinski_scaling_factor: float = 0.015, + n_substeps: int = 5, + zdiffu_t: bool = True, + hdiff_rcf: bool = True, + velocity_boundary_diffusion_denom: float = 200.0, + temperature_boundary_diffusion_denom: float = 135.0, + max_nudging_coeff: float = 0.02, + nudging_decay_rate: float = 2.0, ): # TODO [ml]: move external stuff out: grid related stuff, other than diffusion namelists (see below self.grid = grid self.vertical_params = vertical_params # parameters from namelist diffusion_nml - self.diffusion_type = ( - diffusion_type # hdiff_order : order of nabla operator for diffusion - ) - self.apply_to_vertical_wind = ( - hdiff_w # lhdiff_w, diffusion on the vertical wind field + self.diffusion_type: int = diffusion_type + """ + Order of nabla operator for diffusion. + + Called `hdiff_order` in mo_diffusion_nml.f90. + Possible values are: + - -1: no diffusion + - 2: 2nd order linear diffusion on all vertical levels + - 3: Smagorinsky diffusion without background diffusion + - 4: 4th order linear diffusion on all vertical levels + - 5: Smagorinsky diffusion with fourth-order background diffusion + + We only support type 5. + """ + + self.apply_to_vertical_wind: bool = hdiff_w + """Ff True, apply diffusion on the vertical wind field + + Called `lhdiff_w` in in mo_diffusion_nml.f90 + """ + + self.apply_to_horizontal_wind = hdiff_vn + """ + If True apply diffusion on the horizontal wind field, is ONLY used in mo_nh_stepping.f90 + + Called `lhdiff_vn` in in mo_diffusion_nml.f90 + """ + + self.apply_to_temperature: bool = hdiff_temp + """ + If True, apply horizontal diffusion to temperature field + + Called `lhdiff_temp` in in mo_diffusion_nml.f90 + """ + + self.compute_3d_smag_coeff: bool = smag_3d + """ + If True, compute 3D Smagorinsky diffusion coefficient. + + Called `lsmag_3d` in in mo_diffusion_nml.f90 + """ + + self.type_vn_diffu: int = type_vn_diffu + """ + Options for discretizing the Smagorinsky momentum diffusion. + + Called `itype_vn_diffu` in in mo_diffusion_nml.f90 + """ + + self.type_t_diffu = type_t_diffu + """ + Options for discretizing the Smagorinsky temperature diffusion. + + Called `itype_t_diffu` in in mo_diffusion_nml.f90 + """ + + self.hdiff_efdt_ratio: float = hdiff_efdt_ratio + """ + Ratio of e-folding time to (2*)time step. + + Called `hdiff_efdt_ratio` in in mo_diffusion_nml.f90. + """ + + self.hdiff_w_efdt_ratio: float = hdiff_w_efdt_ratio + """ + Ratio of e-folding time to time step for w diffusion (NH only). + + Called `hdiff_w_efdt_ratio` in in mo_diffusion_nml.f90. + """ + + self.smagorinski_scaling_factor: float = smagorinski_scaling_factor + """ + Scaling factor for Smagorinsky diffusion at height hdiff_smag_z and below. + + Called `hdiff_smag_fac` in in mo_diffusion_nml.f90. + """ + + self.apply_zdiffusion_t: bool = zdiffu_t + """ + If True, apply truly horizontal temperature diffusion over steep slopes. + + From parent namelist mo_nonhydrostatic_nml.f90, but is only used in diffusion, + and in mo_vertical_grid.f90>prepare_zdiffu. + Called 'l_zdiffu_t' in mo_nonhydrostatic_nml.f90. + """ + + # from other namelists + + # from parent namelist mo_nonhydrostatic_nml + self.ndyn_substeps: int = n_substeps + """ + Number of dynamics substeps per fast-physics step. + + Called 'ndyn_substeps' in mo_nonhydrostatic_nml.f90. + """ + + self.lhdiff_rcf: bool = hdiff_rcf + """ + If True, compute horizontal diffusion only at the large time step. + + Called 'lhdiff_rcf' in mo_nonhydrostatic_nml.f90. + """ + + # namelist mo_gridref_nml.f90 + self.temperature_boundary_diffusion_denominator: float = ( + temperature_boundary_diffusion_denom ) - self.apply_to_horizontal_wind = True # lhdiff_vn, diffusion on horizonal wind field, ONLY used in mo_nh_stepping.f90 - self.apply_to_temperature = ( - True # lhdiff_temp, apply horizontal diffusion to temperature + """ + Denominator for temperature boundary diffusion. + + Called 'denom_diffu_t' in mo_gridref_nml.f90. + """ + + self.velocity_boundary_diffusion_denominator: float = ( + velocity_boundary_diffusion_denom ) + """ + Denominator for velocity boundary diffusion. + + Called 'denom_diffu_v' in mo_gridref_nml.f90. + """ - self.compute_3d_smag_coeff = smag_3d # lsmag_3d, if `true`, compute 3D Smagorinsky diffusion coefficient. - self.itype_vn_diffu = type_vn_diffu # itype_vn_diffu, options for discretizing the Smagorinsky momentum diffusion - self.itype_t_diffu = type_t_diffu # itype_t_diffu, options for discretizing the Smagorinsky temperature diffusion - self.hdiff_efdt_ratio = hdiff_efdt_ratio # hdiff_efdt_ratio, ratio of e-folding time to (2*)time step - self.hdiff_w_efdt_ratio = hdiff_w_efdt_ratio # hdiff_w_efdt_ratio, ratio of e-folding time to time step for w diffusion (NH only) - self.hdiff_smag_factor = smag_scaling_fac # hdiff_smag_fac, scaling factor for Smagorinsky diffusion at height hdiff_smag_z and below + # namelist: mo_interpol_nml.f90 + self.nudge_max_coeff: float = max_nudging_coeff + """ + Parameter describing the lateral boundary nudging in limited area mode. - self.zdiffu_t = zdiffu_t # l_zdiffu_t, apply truly horizontal temperature diffusion, from parent namelist nonhydrostatic_nml, but is only used in diffusion, and in mo_vertical_grid.prepare_zdiffu. + Maximal value of the nudging coefficients used cell row bordering the boundary + interpolation zone, from there nudging coefficients decay exponentially with + `nudge_efold_width` in units of cell rows. - # from other namelists - # from parent namelist nonhydrostatic_nml - self.ndyn_substeps = 5 - self.lhdiff_rcf = True + Called `nudge_max_coeff` in mo_interpol_nml.f90 + """ - # namelist gridref_nml - # default is v=200.0, t=135.0 - self.lateral_boundary_denominator = TupleVT(v=150.0, t=135.0) + self.nudge_efold_width: float = nudging_decay_rate + """ + Exponential decay rate (in units of cell rows) of the lateral boundary nudging coefficients. + + Called `nudge_efold_width` in mo_interpol_nml.f90 + """ - # name list: interpol_nml - self.nudge_max_coeff = 0.075 self._validate() def _validate(self): + """ + Apply consistency checks and validation on configuration parameters. + """ if self.diffusion_type != 5: raise NotImplementedError( - "only diffusion type 5 : `Smagorinsky diffusion with fourth-order background diffusion` is implemented" + "Only diffusion type 5 = `Smagorinsky diffusion with fourth-order background diffusion` is implemented" ) if not self.grid.limited_area: - raise NotImplementedError("only limited area mode is implemented") + raise NotImplementedError("Only limited area mode is implemented") if self.diffusion_type < 0: self.apply_to_temperature = False @@ -317,7 +437,7 @@ def _validate(self): self.apply_to_temperature = True self.apply_to_horizontal_wind = True - if not self.zdiffu_t: + if not self.apply_zdiffusion_t: raise NotImplementedError( "zdiffu_t = False is not implemented (leaves out stencil_15)" ) @@ -371,7 +491,9 @@ def determine_smagorinski_factor(self, config: DiffusionConfig): # according to mo_nh_diffusion.f90 this isn't used anywhere the factor is only # used for diffusion_type (3,5) but the defaults are only defined for iequations=3 smagorinski_factor = ( - config.hdiff_smag_factor if config.hdiff_smag_factor else 0.15, + config.smagorinski_scaling_factor + if config.smagorinski_scaling_factor + else 0.15, ) smagorinski_height = None case _: @@ -387,7 +509,7 @@ def _diffusion_type_5_smagorinski_factor(config: DiffusionConfig): magic_sqrt = math.sqrt(1600.0 * (1600 + 50000.0)) magic_fac2_value = 2e-6 * (1600.0 + 25000.0 + magic_sqrt) magic_z2 = 1600.0 + 50000.0 + magic_sqrt - factor = (config.hdiff_smag_factor, magic_fac2_value, 0.0, 1.0) + factor = (config.smagorinski_scaling_factor, magic_fac2_value, 0.0, 1.0) heights = (32500.0, magic_z2, 50000.0, 90000.0) return factor, heights @@ -445,9 +567,10 @@ def __init__( ) self.bdy_diff: float = 0.015 / (config.nudge_max_coeff + sys.float_info.epsilon) self.fac_bdydiff_v: float = ( - math.sqrt(config.substep_as_float()) / config.lateral_boundary_denominator.v + math.sqrt(config.substep_as_float()) + / config.velocity_boundary_diffusion_denominator if config.lhdiff_rcf - else 1.0 / config.lateral_boundary_denominator.v + else 1.0 / config.velocity_boundary_diffusion_denominator ) self.thresh_tdiff: float = ( -5.0 diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index 0c9893ad8..4972438f5 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -113,6 +113,7 @@ def with_connectivities(self, connectivity: Dict[Dimension, np.ndarray]): self.size.update({d: t.shape[1] for d, t in connectivity.items()}) def limited_area(self): + # TODO defined in mo_grid_nml.f90 return self.config.limited_area def n_lev(self): diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 4881f902a..0c0ac37be 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -129,9 +129,10 @@ def icon_grid(savepoint_init): @pytest.fixture def r04b09_diffusion_config(setup_icon_data) -> DiffusionConfig: """ - Create DiffusionConfig. + Create DiffusionConfig matching MCH_CH_r04b09_dsl. - that uses the parameters of MCH.CH_r04b09_dsl experiment + Sets values to the ones used in the MCH_CH_r04b09_dsl experiment where they differ + from the default. """ sp = IconSerialDataProvider( "icon_diffusion_init", extracted_path, True @@ -156,10 +157,13 @@ def r04b09_diffusion_config(setup_icon_data) -> DiffusionConfig: vertical_params=verticalParams, diffusion_type=5, hdiff_w=True, + hdiff_vn=True, type_t_diffu=2, type_vn_diffu=1, hdiff_efdt_ratio=24.0, hdiff_w_efdt_ratio=15.0, - smag_scaling_fac=0.025, + smagorinski_scaling_factor=0.025, zdiffu_t=True, + velocity_boundary_diffusion_denom=150.0, + max_nudging_coeff=0.075, ) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index c99559170..aef0b5ff1 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -190,7 +190,7 @@ def test_diffusion_coefficients_without_hdiff_efdt_ratio(r04b09_diffusion_config @pytest.mark.datatest def test_smagorinski_factor_for_diffusion_type_4(r04b09_diffusion_config): config = r04b09_diffusion_config - config.hdiff_smag_factor = 0.15 + config.smagorinski_scaling_factor = 0.15 config.diffusion_type = 4 params = DiffusionParams(config) @@ -204,7 +204,7 @@ def test_smagorinski_heights_diffusion_type_5_are_consistent( r04b09_diffusion_config, ): config = r04b09_diffusion_config - config.hdiff_smag_factor = 0.15 + config.smagorinski_scaling_factor = 0.15 config.diffusion_type = 5 params = DiffusionParams(config) From 42ac762cff12b5dbb4f228f58ed0c9c059e5e94e Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 6 Jan 2023 18:06:41 +0100 Subject: [PATCH 042/263] global mode: add global (non limited area) version of stencil mo_diffusion_stencil_05 --- .../mo_nh_diffusion_stencil_05.py | 43 +++++++++++++++ .../tests/test_mo_nh_diffusion_stencil_05.py | 52 +++++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_05.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_05.py index 064fcf6d1..11f1ed973 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_05.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_05.py @@ -17,6 +17,49 @@ from icon4py.common.dimension import EdgeDim, KDim +@field_operator +def _mo_nh_diffusion_stencil_05_global_mode( + area_edge: Field[[EdgeDim], float], + kh_smag_e: Field[[EdgeDim, KDim], float], + z_nabla2_e: Field[[EdgeDim, KDim], float], + z_nabla_e2: Field[[EdgeDim, KDim], float], + diff_multfac_vn: Field[[KDim], float], + vn: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + z_d_vn_hdf = area_edge * ( + kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla_e2 * area_edge + ) + return vn + z_d_vn_hdf + + +@program +def mo_nh_diffusion_stencil_05_global_mode( + area_edge: Field[[EdgeDim], float], + kh_smag_e: Field[[EdgeDim, KDim], float], + z_nabla2_e: Field[[EdgeDim, KDim], float], + z_nabla4_e2: Field[[EdgeDim, KDim], float], + diff_multfac_vn: Field[[KDim], float], + vn: Field[[EdgeDim, KDim], float], + horizontal_start: int, + horizontal_end: int, + vertical_start: int, + vertical_end: int, +): + _mo_nh_diffusion_stencil_05_global_mode( + area_edge, + kh_smag_e, + z_nabla2_e, + z_nabla4_e2, + diff_multfac_vn, + vn, + out=vn, + domain={ + EdgeDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) + + @field_operator def _mo_nh_diffusion_stencil_05( area_edge: Field[[EdgeDim], float], diff --git a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_05.py b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_05.py index 9e6df6fa2..78c4a61ae 100644 --- a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_05.py +++ b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_05.py @@ -15,6 +15,7 @@ from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_05 import ( mo_nh_diffusion_stencil_05, + mo_nh_diffusion_stencil_05_global_mode, ) from icon4py.common.dimension import EdgeDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh @@ -76,3 +77,54 @@ def test_mo_nh_diffusion_stencil_05(): offset_provider={}, ) assert np.allclose(vn, vn_ref) + + +def mo_nh_diffusion_stencil_05_global_mode_numpy( + area_edge: np.array, + kh_smag_e: np.array, + z_nabla2_e: np.array, + z_nabla4_e2: np.array, + diff_multfac_vn: np.array, + vn: np.array, +): + area_edge = np.expand_dims(area_edge, axis=-1) + diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0) + vn = vn + area_edge * ( + kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge + ) + return vn + + +def test_mo_nh_diffusion_stencil_05_global_mode(): + mesh = SimpleMesh() + area_edge = random_field(mesh, EdgeDim) + kh_smag_e = random_field(mesh, EdgeDim, KDim) + z_nabla2_e = random_field(mesh, EdgeDim, KDim) + z_nabla4_e2 = random_field(mesh, EdgeDim, KDim) + diff_multfac_vn = random_field(mesh, KDim) + vn = random_field(mesh, EdgeDim, KDim) + + vn_ref_np = mo_nh_diffusion_stencil_05_global_mode_numpy( + np.asarray(area_edge), + np.asarray(kh_smag_e), + np.asarray(z_nabla2_e), + np.asarray(z_nabla4_e2), + np.asarray(diff_multfac_vn), + np.asarray(vn), + ) + + mo_nh_diffusion_stencil_05_global_mode( + area_edge, + kh_smag_e, + z_nabla2_e, + z_nabla4_e2, + diff_multfac_vn, + vn, + 0, + mesh.n_edges, + 0, + mesh.k_level, + offset_provider={}, + ) + + assert np.allclose(vn, vn_ref_np) From a937df885a2bd1c26ef9b3aecbdb41961aad9260 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 6 Jan 2023 18:26:58 +0100 Subject: [PATCH 043/263] small fixes, cleanups, remove obsolete TODO --- .../src/icon4py/diffusion/diffusion.py | 31 ++++++++++--------- atm_dyn_iconam/tests/test_diffusion.py | 6 ++-- 2 files changed, 20 insertions(+), 17 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 514c1b327..e7e967f7f 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -99,7 +99,7 @@ def _setup_initial_diff_multfac_vn( @field_operator -def setup_fields_for_initial_step(k4: float, hdiff_efdt_ratio: float): +def _setup_fields_for_initial_step(k4: float, hdiff_efdt_ratio: float): diff_multfac_vn = _setup_initial_diff_multfac_vn(k4, hdiff_efdt_ratio) smag_limit = _setup_smag_limit(diff_multfac_vn) return diff_multfac_vn, smag_limit @@ -135,7 +135,7 @@ def _en_smag_fac_for_zero_nshift( @field_operator -def _init_diffusion_local_fields( +def _init_diffusion_local_fields_for_regular_timestemp( k4: float, dyn_substeps: float, hdiff_smag_fac: float, @@ -169,7 +169,7 @@ def _init_diffusion_local_fields( @program -def init_diffusion_local_fields( +def init_diffusion_local_fields_for_regular_timestep( k4: float, dyn_substeps: float, hdiff_smag_fac: float, @@ -185,7 +185,7 @@ def init_diffusion_local_fields( smag_limit: Field[[KDim], float], enh_smag_fac: Field[[KDim], float], ): - _init_diffusion_local_fields( + _init_diffusion_local_fields_for_regular_timestemp( k4, dyn_substeps, hdiff_smag_fac, @@ -211,7 +211,8 @@ def init_nabla2_factor_in_upper_damping_zone( """ Calculate diff_multfac_n2w. - numpy version gt4py does not allow non-constant indexing into fields + numpy version, since gt4py does not allow non-constant indexing into fields + TODO: [ml] fix this once IndexedFields are implemented Args k_size: number of vertical levels @@ -241,9 +242,8 @@ class DiffusionConfig: Encapsulates namelist parameters and derived parameters. Values should be read from configuration. - Default values are taken from the defaults in the ICON Fortran namelist files. - TODO: be read from config and the default from mo_diffusion_nml.f90 set as defaults. - TODO: [ml] read from config + Default values are taken from the defaults in the corresponding ICON Fortran namelist files. + TODO: [ml] to be read from config TODO: [ml] handle dependencies on other namelists (see below...) """ @@ -251,7 +251,7 @@ def __init__( self, grid: IconGrid, vertical_params: VerticalModelParams, - diffusion_type: int = 5, # TODO: use enum + diffusion_type: int = 5, hdiff_w=True, hdiff_vn=True, hdiff_temp=True, @@ -269,7 +269,6 @@ def __init__( max_nudging_coeff: float = 0.02, nudging_decay_rate: float = 2.0, ): - # TODO [ml]: move external stuff out: grid related stuff, other than diffusion namelists (see below self.grid = grid self.vertical_params = vertical_params @@ -287,6 +286,7 @@ def __init__( - 5: Smagorinsky diffusion with fourth-order background diffusion We only support type 5. + TODO: [ml] use enum """ self.apply_to_vertical_wind: bool = hdiff_w @@ -505,7 +505,11 @@ def determine_smagorinski_factor(self, config: DiffusionConfig): @staticmethod def _diffusion_type_5_smagorinski_factor(config: DiffusionConfig): - # initial values from mo_diffusion_nml.f90 + """ + Initialize smagorinski factors used in diffusion type 5. + + The calculation and magic numbers are taken from mo_diffusion_nml.f90 + """ magic_sqrt = math.sqrt(1600.0 * (1600 + 50000.0)) magic_fac2_value = 2e-6 * (1600.0 + 25000.0 + magic_sqrt) magic_z2 = 1600.0 + 50000.0 + magic_sqrt @@ -583,8 +587,7 @@ def __init__( self._allocate_local_fields() - # TODO different for initial run!, through diff_multfac_vn - init_diffusion_local_fields( + init_diffusion_local_fields_for_regular_timestep( params.K4, config.substep_as_float(), *params.smagorinski_factor, @@ -656,7 +659,7 @@ def initial_step( diff_multfac_vn = zero_field(self.grid, KDim) smag_limit = zero_field(self.grid, KDim) - setup_fields_for_initial_step( + _setup_fields_for_initial_step( self.params.K4, self.config.hdiff_efdt_ratio, out=(diff_multfac_vn, smag_limit), diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index aef0b5ff1..741bfd104 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -23,11 +23,11 @@ DiffusionParams, VectorTuple, _en_smag_fac_for_zero_nshift, + _setup_fields_for_initial_step, _setup_runtime_diff_multfac_vn, _setup_smag_limit, scale_k, set_zero_v_k, - setup_fields_for_initial_step, ) from icon4py.diffusion.icon_grid import VerticalModelParams from icon4py.diffusion.interpolation_state import InterpolationState @@ -68,7 +68,7 @@ def test_diff_multfac_vn_and_smag_limit_for_initial_step(): initial_diff_multfac_vn_numpy, shape, k4, efdt_ratio ) - setup_fields_for_initial_step( + _setup_fields_for_initial_step( k4, efdt_ratio, out=(diff_multfac_vn_init, smag_limit_init), offset_provider={} ) @@ -313,7 +313,7 @@ def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( diff_multfac_vn = utils.zero_field(config.grid, KDim) smag_limit = utils.zero_field(config.grid, KDim) - setup_fields_for_initial_step( + _setup_fields_for_initial_step( params.K4, config.hdiff_efdt_ratio, out=(diff_multfac_vn, smag_limit), From 73d76d607722235c7d72dd3d5355a97eb1837c76 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 10 Jan 2023 09:27:40 +0100 Subject: [PATCH 044/263] remove gt4py standard functional branch from base-requirements-dev.txt --- base-requirements-dev.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index c84e9915a..47c869d56 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -1,5 +1,4 @@ # VCS -#-e git+https://github.com/GridTools/gt4py.git@functional#egg=gt4py-functional -e git+https://github.com/tehrengruber/gt4py.git@reduce_foast_lowering_complexity#egg=gt4py-functional # PyPI From 631577079fd43eb1073dea3903e73debfe78f8bc Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 10 Jan 2023 09:39:38 +0100 Subject: [PATCH 045/263] - remove gt4py standard functional branch from base-requirements-dev.txt - remove numpy version restriction after update of serialbox --- base-requirements-dev.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 47c869d56..cb97c337d 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -22,4 +22,3 @@ setuptools>=40.8.0 wheel>=0.37.1 tox >= 3.25 wget >= 3.2 -numpy == 1.23.3 From bcc230de68cd3978c8e720af9a6ab903ade7a5a2 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 10 Jan 2023 09:54:16 +0100 Subject: [PATCH 046/263] add docstrings for fixtures in conftest.py --- atm_dyn_iconam/tests/conftest.py | 37 +++++++++++++++++++++++++- atm_dyn_iconam/tests/test_diffusion.py | 1 - 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 0c0ac37be..05da1ab87 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -47,6 +47,11 @@ @pytest.fixture(scope="session") def setup_icon_data(): + """ + Get the binary ICON data from a remote server. + + Session scoped fixture which is a prerequisite of all the other fixtures in this file. + """ os.makedirs(data_path, exist_ok=True) if len(os.listdir(data_path)) == 0: print( @@ -64,16 +69,34 @@ def setup_icon_data(): @pytest.fixture def linit(): + """ + Set the 'linit' flag for the ICON diffusion data savepoint. + + Defaults to False + """ return False @pytest.fixture def step_date(): + """ + Set the step date for the loaded ICON time stamp. + + Defaults to 2021-06-20T12:00:10.000' + """ return "2021-06-20T12:00:10.000" @pytest.fixture def savepoint_init(setup_icon_data, linit, step_date): + """ + Load data from ICON savepoint at start of diffusion module. + + date of the timestamp to be selected can be set seperately by overriding the 'step_data' + fixture, passing 'step_data=' + + linit flag can be set by overriding the 'linit' fixture + """ sp = IconSerialDataProvider( "icon_diffusion_init", extracted_path, True ).from_savepoint_init(linit=linit, date=step_date) @@ -82,6 +105,12 @@ def savepoint_init(setup_icon_data, linit, step_date): @pytest.fixture def savepoint_exit(setup_icon_data, step_date): + """ + Load data from ICON savepoint at exist of diffusion module. + + date of the timestamp to be selected can be set seperately by overriding the 'step_data' + fixture, passing 'step_data=' + """ sp = IconSerialDataProvider( "icon_diffusion_exit", extracted_path, True ).from_savepoint_init(linit=False, date=step_date) @@ -90,6 +119,12 @@ def savepoint_exit(setup_icon_data, step_date): @pytest.fixture def icon_grid(savepoint_init): + """ + Load the icon grid from an ICON savepoint. + + Uses the default save_point from 'savepoint_init' fixture, however these data don't change for + different time steps. + """ sp = savepoint_init sp_meta = sp.get_metadata("nproma", "nlev", "num_vert", "num_cells", "num_edges") @@ -131,7 +166,7 @@ def r04b09_diffusion_config(setup_icon_data) -> DiffusionConfig: """ Create DiffusionConfig matching MCH_CH_r04b09_dsl. - Sets values to the ones used in the MCH_CH_r04b09_dsl experiment where they differ + Set values to the ones used in the MCH_CH_r04b09_dsl experiment where they differ from the default. """ sp = IconSerialDataProvider( diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 741bfd104..a0896cb7f 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -354,7 +354,6 @@ def test_verify_diffusion_init_against_other_regular_savepoint( _verify_init_values_against_savepoint(savepoint, diffusion) -@pytest.mark.skip @pytest.mark.datatest def test_diffusion_run(savepoint_init, save_point_exit, icon_grid): sp = savepoint_init From b5496b75712f73f4fef5cc2219706505295fd376 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 10 Jan 2023 11:03:09 +0100 Subject: [PATCH 047/263] fix loading of savepoint_exit --- atm_dyn_iconam/tests/conftest.py | 4 ++-- atm_dyn_iconam/tests/test_diffusion.py | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 05da1ab87..6644b4828 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -112,8 +112,8 @@ def savepoint_exit(setup_icon_data, step_date): fixture, passing 'step_data=' """ sp = IconSerialDataProvider( - "icon_diffusion_exit", extracted_path, True - ).from_savepoint_init(linit=False, date=step_date) + "icon_diffusion_init", extracted_path, True + ).from_save_point_exit(linit=False, date=step_date) return sp diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index a0896cb7f..a1de3603d 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -355,7 +355,7 @@ def test_verify_diffusion_init_against_other_regular_savepoint( @pytest.mark.datatest -def test_diffusion_run(savepoint_init, save_point_exit, icon_grid): +def test_diffusion_run(savepoint_init, savepoint_exit, icon_grid): sp = savepoint_init vct_a = sp.vct_a() @@ -433,10 +433,10 @@ def test_diffusion_run(savepoint_init, save_point_exit, icon_grid): cell_areas=cell_areas, ) - icon_result_exner = save_point_exit.exner() - icon_result_vn = save_point_exit.vn() - icon_result_w = save_point_exit.w() - icon_result_theta_w = save_point_exit.theta_v() + icon_result_exner = savepoint_exit.exner() + icon_result_vn = savepoint_exit.vn() + icon_result_w = savepoint_exit.w() + icon_result_theta_w = savepoint_exit.theta_v() assert np.allclose(icon_result_w, np.asarray(prognostic_state.vertical_wind)) assert np.allclose( From 34fa7702b44b23e39fbfccad35b39bd36d9ec889 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 10 Jan 2023 12:02:50 +0100 Subject: [PATCH 048/263] add a test skeleton that runs 5 setup, init_step and 5 timesteps --- atm_dyn_iconam/tests/conftest.py | 23 +++++++++++++++------ atm_dyn_iconam/tests/test_diffusion.py | 28 +++++++++++++++++++++++--- 2 files changed, 42 insertions(+), 9 deletions(-) diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 6644b4828..6bd7307ed 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -78,9 +78,9 @@ def linit(): @pytest.fixture -def step_date(): +def step_date_init(): """ - Set the step date for the loaded ICON time stamp. + Set the step date for the loaded ICON time stamp at start of module. Defaults to 2021-06-20T12:00:10.000' """ @@ -88,7 +88,18 @@ def step_date(): @pytest.fixture -def savepoint_init(setup_icon_data, linit, step_date): +def step_date_exit(): + """ + Set the step date for the loaded ICON time stamp at the end of module. + + Defaults to 2021-06-20T12:00:10.000' + """ + return "2021-06-20T12:00:10.000" + + + +@pytest.fixture +def savepoint_init(setup_icon_data, linit, step_date_init): """ Load data from ICON savepoint at start of diffusion module. @@ -99,12 +110,12 @@ def savepoint_init(setup_icon_data, linit, step_date): """ sp = IconSerialDataProvider( "icon_diffusion_init", extracted_path, True - ).from_savepoint_init(linit=linit, date=step_date) + ).from_savepoint_init(linit=linit, date=step_date_init) return sp @pytest.fixture -def savepoint_exit(setup_icon_data, step_date): +def savepoint_exit(setup_icon_data, step_date_exit): """ Load data from ICON savepoint at exist of diffusion module. @@ -113,7 +124,7 @@ def savepoint_exit(setup_icon_data, step_date): """ sp = IconSerialDataProvider( "icon_diffusion_init", extracted_path, True - ).from_save_point_exit(linit=False, date=step_date) + ).from_save_point_exit(linit=False, date=step_date_exit) return sp diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index a1de3603d..35f4d1912 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -227,7 +227,7 @@ def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): @pytest.mark.datatest -def test_diffusion_init(savepoint_init, r04b09_diffusion_config, step_date): +def test_diffusion_init(savepoint_init, r04b09_diffusion_config, step_date_init): savepoint = savepoint_init vct_a = savepoint.vct_a() config = r04b09_diffusion_config @@ -235,7 +235,7 @@ def test_diffusion_init(savepoint_init, r04b09_diffusion_config, step_date): assert meta["nlev"] == 65 assert meta["linit"] is False - assert meta["date"] == step_date + assert meta["date"] == step_date_init additional_parameters = DiffusionParams(config) diffusion = Diffusion(config, additional_parameters, vct_a) @@ -355,7 +355,7 @@ def test_verify_diffusion_init_against_other_regular_savepoint( @pytest.mark.datatest -def test_diffusion_run(savepoint_init, savepoint_exit, icon_grid): +def test_run_diffusion_single_step(savepoint_init, savepoint_exit, icon_grid): sp = savepoint_init vct_a = sp.vct_a() @@ -448,3 +448,25 @@ def test_diffusion_run(savepoint_init, savepoint_exit, icon_grid): assert np.allclose( np.asarray(icon_result_exner), np.asarray(prognostic_state.exner_pressure) ) + +@pytest.mark.skip +def test_diffusion(icon_grid, savepoint_init, savepoint_exit, linit = True, step_date_exit = '2021-06-20T12:01:00.000'): + sp = savepoint_init + config = DiffusionConfig( + icon_grid, + vertical_params=VerticalModelParams( + vct_a=(sp.vct_a()), rayleigh_damping_height=12500.0 + ), + ) + + additional_parameters = DiffusionParams(config) + diffusion = Diffusion(config, additional_parameters, sp.vct_a()) + + + diffusion.initial_step(...) + for i in range(4): + diffusion.time_step(...) + + sp_exit = savepoint_exit + # TODO assert exit values + From cecf1811bcd43cffda2bf2b52a056b94fb35ee4e Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 10 Jan 2023 13:46:25 +0100 Subject: [PATCH 049/263] move field_operators and programs to utils.py --- .../src/icon4py/diffusion/diffusion.py | 175 +---------------- atm_dyn_iconam/src/icon4py/diffusion/utils.py | 179 +++++++++++++++++- atm_dyn_iconam/tests/conftest.py | 1 - atm_dyn_iconam/tests/test_diffusion.py | 44 +++-- 4 files changed, 214 insertions(+), 185 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index e7e967f7f..256d80ee3 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -19,8 +19,7 @@ import numpy as np from functional.common import Dimension -from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, broadcast, int32, maximum, minimum +from functional.ffront.fbuiltins import Field, int32 from functional.iterator.embedded import ( StridedNeighborOffsetProvider, np_as_located_field, @@ -59,7 +58,6 @@ ECVDim, EdgeDim, KDim, - Koff, V2EDim, VertexDim, ) @@ -70,172 +68,19 @@ from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.prognostic import PrognosticState -from icon4py.diffusion.utils import scale_k, set_zero_v_k, zero_field +from icon4py.diffusion.utils import ( + init_diffusion_local_fields_for_regular_timestep, + init_nabla2_factor_in_upper_damping_zone, + scale_k, + set_zero_v_k, + setup_fields_for_initial_step, + zero_field, +) VectorTuple = namedtuple("VectorTuple", "x y") -@field_operator -def _setup_smag_limit(diff_multfac_vn: Field[[KDim], float]) -> Field[[KDim], float]: - return 0.125 - 4.0 * diff_multfac_vn - - -@field_operator -def _setup_runtime_diff_multfac_vn( - k4: float, dyn_substeps: float -) -> Field[[KDim], float]: - con = 1.0 / 128.0 - dyn = k4 * dyn_substeps / 3.0 - return broadcast(minimum(con, dyn), (KDim,)) - - -# @field_operator(backend=run_gtfn) -@field_operator -def _setup_initial_diff_multfac_vn( - k4: float, hdiff_efdt_ratio: float -) -> Field[[KDim], float]: - return broadcast(k4 / 3.0 * hdiff_efdt_ratio, (KDim,)) - - -@field_operator -def _setup_fields_for_initial_step(k4: float, hdiff_efdt_ratio: float): - diff_multfac_vn = _setup_initial_diff_multfac_vn(k4, hdiff_efdt_ratio) - smag_limit = _setup_smag_limit(diff_multfac_vn) - return diff_multfac_vn, smag_limit - - -@field_operator -def _en_smag_fac_for_zero_nshift( - vect_a: Field[[KDim], float], - hdiff_smag_fac: float, - hdiff_smag_fac2: float, - hdiff_smag_fac3: float, - hdiff_smag_fac4: float, - hdiff_smag_z: float, - hdiff_smag_z2: float, - hdiff_smag_z3: float, - hdiff_smag_z4: float, -) -> Field[[KDim], float]: - dz21 = hdiff_smag_z2 - hdiff_smag_z - alin = (hdiff_smag_fac2 - hdiff_smag_fac) / dz21 - df32 = hdiff_smag_fac3 - hdiff_smag_fac2 - df42 = hdiff_smag_fac4 - hdiff_smag_fac2 - dz32 = hdiff_smag_z3 - hdiff_smag_z2 - dz42 = hdiff_smag_z4 - hdiff_smag_z2 - - bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32)) - aqdr = df32 / dz32 - bqdr * dz32 - zf = 0.5 * (vect_a + vect_a(Koff[1])) - - dzlin = minimum(dz21, maximum(0.0, zf - hdiff_smag_z)) - dzqdr = minimum(dz42, maximum(0.0, zf - hdiff_smag_z2)) - enh_smag_fac = hdiff_smag_fac + (dzlin * alin) + dzqdr * (aqdr + dzqdr * bqdr) - return enh_smag_fac - - -@field_operator -def _init_diffusion_local_fields_for_regular_timestemp( - k4: float, - dyn_substeps: float, - hdiff_smag_fac: float, - hdiff_smag_fac2: float, - hdiff_smag_fac3: float, - hdiff_smag_fac4: float, - hdiff_smag_z: float, - hdiff_smag_z2: float, - hdiff_smag_z3: float, - hdiff_smag_z4: float, - vect_a: Field[[KDim], float], -) -> tuple[Field[[KDim], float], Field[[KDim], float], Field[[KDim], float]]: - diff_multfac_vn = _setup_runtime_diff_multfac_vn(k4, dyn_substeps) - smag_limit = _setup_smag_limit(diff_multfac_vn) - enh_smag_fac = _en_smag_fac_for_zero_nshift( - vect_a, - hdiff_smag_fac, - hdiff_smag_fac2, - hdiff_smag_fac3, - hdiff_smag_fac4, - hdiff_smag_z, - hdiff_smag_z2, - hdiff_smag_z3, - hdiff_smag_z4, - ) - return ( - diff_multfac_vn, - smag_limit, - enh_smag_fac, - ) - - -@program -def init_diffusion_local_fields_for_regular_timestep( - k4: float, - dyn_substeps: float, - hdiff_smag_fac: float, - hdiff_smag_fac2: float, - hdiff_smag_fac3: float, - hdiff_smag_fac4: float, - hdiff_smag_z: float, - hdiff_smag_z2: float, - hdiff_smag_z3: float, - hdiff_smag_z4: float, - vect_a: Field[[KDim], float], - diff_multfac_vn: Field[[KDim], float], - smag_limit: Field[[KDim], float], - enh_smag_fac: Field[[KDim], float], -): - _init_diffusion_local_fields_for_regular_timestemp( - k4, - dyn_substeps, - hdiff_smag_fac, - hdiff_smag_fac2, - hdiff_smag_fac3, - hdiff_smag_fac4, - hdiff_smag_z, - hdiff_smag_z2, - hdiff_smag_z3, - hdiff_smag_z4, - vect_a, - out=( - diff_multfac_vn, - smag_limit, - enh_smag_fac, - ), - ) - - -def init_nabla2_factor_in_upper_damping_zone( - k_size: int, nrdmax: int, nshift: int, physical_heights: np.ndarray -) -> Field[[KDim], float]: - """ - Calculate diff_multfac_n2w. - - numpy version, since gt4py does not allow non-constant indexing into fields - TODO: [ml] fix this once IndexedFields are implemented - - Args - k_size: number of vertical levels - nrdmax: index of the level where rayleigh dampint starts - nshift: - physcial_heights: vector of physical heights [m] of the height levels - """ - buffer = np.zeros(k_size) - buffer[1 : nrdmax + 1] = ( - 1.0 - / 12.0 - * ( - ( - physical_heights[1 + nshift : nrdmax + 1 + nshift] - - physical_heights[nshift + nrdmax + 1] - ) - / (physical_heights[1] - physical_heights[nshift + nrdmax + 1]) - ) - ** 4 - ) - return np_as_located_field(KDim)(buffer) - - class DiffusionConfig: """ Contains necessary parameter to configure a diffusion run. @@ -659,7 +504,7 @@ def initial_step( diff_multfac_vn = zero_field(self.grid, KDim) smag_limit = zero_field(self.grid, KDim) - _setup_fields_for_initial_step( + setup_fields_for_initial_step( self.params.K4, self.config.hdiff_efdt_ratio, out=(diff_multfac_vn, smag_limit), diff --git a/atm_dyn_iconam/src/icon4py/diffusion/utils.py b/atm_dyn_iconam/src/icon4py/diffusion/utils.py index ee4cb4cb8..b6040838b 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/utils.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/utils.py @@ -10,14 +10,15 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +from typing import Tuple import numpy as np from functional.common import Dimension, Field from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import broadcast +from functional.ffront.fbuiltins import broadcast, maximum, minimum from functional.iterator.embedded import np_as_located_field -from icon4py.common.dimension import KDim, VertexDim +from icon4py.common.dimension import KDim, Koff, VertexDim # TODO fix duplication: duplicated from test testutils/utils.py @@ -46,3 +47,177 @@ def _set_zero_v_k() -> Field[[VertexDim, KDim], float]: @program def set_zero_v_k(field: Field[[VertexDim, KDim], float]): _set_zero_v_k(out=field) + + +@field_operator +def _setup_smag_limit(diff_multfac_vn: Field[[KDim], float]) -> Field[[KDim], float]: + return 0.125 - 4.0 * diff_multfac_vn + + +@field_operator +def _setup_runtime_diff_multfac_vn( + k4: float, dyn_substeps: float +) -> Field[[KDim], float]: + con = 1.0 / 128.0 + dyn = k4 * dyn_substeps / 3.0 + return broadcast(minimum(con, dyn), (KDim,)) + + +# @field_operator(backend=run_gtfn) +@field_operator +def _setup_initial_diff_multfac_vn( + k4: float, hdiff_efdt_ratio: float +) -> Field[[KDim], float]: + return broadcast(k4 / 3.0 * hdiff_efdt_ratio, (KDim,)) + + +@field_operator +def _setup_fields_for_initial_step( + k4: float, hdiff_efdt_ratio: float +) -> Tuple[Field[[KDim], float], Field[[KDim], float]]: + diff_multfac_vn = _setup_initial_diff_multfac_vn(k4, hdiff_efdt_ratio) + smag_limit = _setup_smag_limit(diff_multfac_vn) + return diff_multfac_vn, smag_limit + + +@program +def setup_fields_for_initial_step( + k4: float, + hdiff_efdt_ratio: float, + diff_multfac_vn: Field[[KDim], float], + smag_limit: Field[[KDim], float], +): + _setup_fields_for_initial_step( + k4, hdiff_efdt_ratio, out=(diff_multfac_vn, smag_limit) + ) + + +@field_operator +def _en_smag_fac_for_zero_nshift( + vect_a: Field[[KDim], float], + hdiff_smag_fac: float, + hdiff_smag_fac2: float, + hdiff_smag_fac3: float, + hdiff_smag_fac4: float, + hdiff_smag_z: float, + hdiff_smag_z2: float, + hdiff_smag_z3: float, + hdiff_smag_z4: float, +) -> Field[[KDim], float]: + dz21 = hdiff_smag_z2 - hdiff_smag_z + alin = (hdiff_smag_fac2 - hdiff_smag_fac) / dz21 + df32 = hdiff_smag_fac3 - hdiff_smag_fac2 + df42 = hdiff_smag_fac4 - hdiff_smag_fac2 + dz32 = hdiff_smag_z3 - hdiff_smag_z2 + dz42 = hdiff_smag_z4 - hdiff_smag_z2 + + bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32)) + aqdr = df32 / dz32 - bqdr * dz32 + zf = 0.5 * (vect_a + vect_a(Koff[1])) + + dzlin = minimum(dz21, maximum(0.0, zf - hdiff_smag_z)) + dzqdr = minimum(dz42, maximum(0.0, zf - hdiff_smag_z2)) + enh_smag_fac = hdiff_smag_fac + (dzlin * alin) + dzqdr * (aqdr + dzqdr * bqdr) + return enh_smag_fac + + +@field_operator +def _init_diffusion_local_fields_for_regular_timestemp( + k4: float, + dyn_substeps: float, + hdiff_smag_fac: float, + hdiff_smag_fac2: float, + hdiff_smag_fac3: float, + hdiff_smag_fac4: float, + hdiff_smag_z: float, + hdiff_smag_z2: float, + hdiff_smag_z3: float, + hdiff_smag_z4: float, + vect_a: Field[[KDim], float], +) -> tuple[Field[[KDim], float], Field[[KDim], float], Field[[KDim], float]]: + diff_multfac_vn = _setup_runtime_diff_multfac_vn(k4, dyn_substeps) + smag_limit = _setup_smag_limit(diff_multfac_vn) + enh_smag_fac = _en_smag_fac_for_zero_nshift( + vect_a, + hdiff_smag_fac, + hdiff_smag_fac2, + hdiff_smag_fac3, + hdiff_smag_fac4, + hdiff_smag_z, + hdiff_smag_z2, + hdiff_smag_z3, + hdiff_smag_z4, + ) + return ( + diff_multfac_vn, + smag_limit, + enh_smag_fac, + ) + + +@program +def init_diffusion_local_fields_for_regular_timestep( + k4: float, + dyn_substeps: float, + hdiff_smag_fac: float, + hdiff_smag_fac2: float, + hdiff_smag_fac3: float, + hdiff_smag_fac4: float, + hdiff_smag_z: float, + hdiff_smag_z2: float, + hdiff_smag_z3: float, + hdiff_smag_z4: float, + vect_a: Field[[KDim], float], + diff_multfac_vn: Field[[KDim], float], + smag_limit: Field[[KDim], float], + enh_smag_fac: Field[[KDim], float], +): + _init_diffusion_local_fields_for_regular_timestemp( + k4, + dyn_substeps, + hdiff_smag_fac, + hdiff_smag_fac2, + hdiff_smag_fac3, + hdiff_smag_fac4, + hdiff_smag_z, + hdiff_smag_z2, + hdiff_smag_z3, + hdiff_smag_z4, + vect_a, + out=( + diff_multfac_vn, + smag_limit, + enh_smag_fac, + ), + ) + + +def init_nabla2_factor_in_upper_damping_zone( + k_size: int, nrdmax: int, nshift: int, physical_heights: np.ndarray +) -> Field[[KDim], float]: + """ + Calculate diff_multfac_n2w. + + numpy version, since gt4py does not allow non-constant indexing into fields + TODO: [ml] fix this once IndexedFields are implemented + + Args + k_size: number of vertical levels + nrdmax: index of the level where rayleigh dampint starts + nshift: + physcial_heights: vector of physical heights [m] of the height levels + """ + buffer = np.zeros(k_size) + buffer[1 : nrdmax + 1] = ( + 1.0 + / 12.0 + * ( + ( + physical_heights[1 + nshift : nrdmax + 1 + nshift] + - physical_heights[nshift + nrdmax + 1] + ) + / (physical_heights[1] - physical_heights[nshift + nrdmax + 1]) + ) + ** 4 + ) + return np_as_located_field(KDim)(buffer) diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 6bd7307ed..1c268af63 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -97,7 +97,6 @@ def step_date_exit(): return "2021-06-20T12:00:10.000" - @pytest.fixture def savepoint_init(setup_icon_data, linit, step_date_init): """ diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 35f4d1912..0f66168b8 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -15,24 +15,25 @@ import pytest from icon4py.common.dimension import ECVDim, KDim, VertexDim -from icon4py.diffusion import utils from icon4py.diffusion.diagnostic import DiagnosticState from icon4py.diffusion.diffusion import ( Diffusion, DiffusionConfig, DiffusionParams, VectorTuple, +) +from icon4py.diffusion.icon_grid import VerticalModelParams +from icon4py.diffusion.interpolation_state import InterpolationState +from icon4py.diffusion.metric_state import MetricState +from icon4py.diffusion.prognostic import PrognosticState +from icon4py.diffusion.utils import ( _en_smag_fac_for_zero_nshift, - _setup_fields_for_initial_step, _setup_runtime_diff_multfac_vn, _setup_smag_limit, scale_k, set_zero_v_k, + setup_fields_for_initial_step, ) -from icon4py.diffusion.icon_grid import VerticalModelParams -from icon4py.diffusion.interpolation_state import InterpolationState -from icon4py.diffusion.metric_state import MetricState -from icon4py.diffusion.prognostic import PrognosticState from icon4py.testutils.serialbox_utils import IconDiffusionInitSavepoint from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import as_1D_sparse_field, random_field, zero_field @@ -68,8 +69,8 @@ def test_diff_multfac_vn_and_smag_limit_for_initial_step(): initial_diff_multfac_vn_numpy, shape, k4, efdt_ratio ) - _setup_fields_for_initial_step( - k4, efdt_ratio, out=(diff_multfac_vn_init, smag_limit_init), offset_provider={} + setup_fields_for_initial_step( + k4, efdt_ratio, diff_multfac_vn_init, smag_limit_init, offset_provider={} ) assert np.allclose(expected_diff_multfac_vn_init, diff_multfac_vn_init) @@ -311,12 +312,13 @@ def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( expected_smag_limit = savepoint.smag_limit() exptected_smag_offset = savepoint.smag_offset() - diff_multfac_vn = utils.zero_field(config.grid, KDim) - smag_limit = utils.zero_field(config.grid, KDim) - _setup_fields_for_initial_step( + diff_multfac_vn = zero_field(config.grid, KDim) + smag_limit = zero_field(config.grid, KDim) + setup_fields_for_initial_step( params.K4, config.hdiff_efdt_ratio, - out=(diff_multfac_vn, smag_limit), + diff_multfac_vn, + smag_limit, offset_provider={}, ) assert np.allclose(expected_smag_limit, smag_limit) @@ -340,7 +342,7 @@ def test_verify_diffusion_init_against_first_regular_savepoint( @pytest.mark.datatest -@pytest.mark.parametrize("step_date", ["2021-06-20T12:00:50.000"]) +@pytest.mark.parametrize("step_date_init", ["2021-06-20T12:00:50.000"]) def test_verify_diffusion_init_against_other_regular_savepoint( r04b09_diffusion_config, savepoint_init ): @@ -354,6 +356,7 @@ def test_verify_diffusion_init_against_other_regular_savepoint( _verify_init_values_against_savepoint(savepoint, diffusion) +@pytest.mark.skip @pytest.mark.datatest def test_run_diffusion_single_step(savepoint_init, savepoint_exit, icon_grid): sp = savepoint_init @@ -449,8 +452,15 @@ def test_run_diffusion_single_step(savepoint_init, savepoint_exit, icon_grid): np.asarray(icon_result_exner), np.asarray(prognostic_state.exner_pressure) ) + @pytest.mark.skip -def test_diffusion(icon_grid, savepoint_init, savepoint_exit, linit = True, step_date_exit = '2021-06-20T12:01:00.000'): +def test_diffusion( + icon_grid, + savepoint_init, + savepoint_exit, + linit=True, + step_date_exit="2021-06-20T12:01:00.000", +): sp = savepoint_init config = DiffusionConfig( icon_grid, @@ -462,11 +472,11 @@ def test_diffusion(icon_grid, savepoint_init, savepoint_exit, linit = True, step additional_parameters = DiffusionParams(config) diffusion = Diffusion(config, additional_parameters, sp.vct_a()) - diffusion.initial_step(...) - for i in range(4): + for _ in range(4): diffusion.time_step(...) sp_exit = savepoint_exit - # TODO assert exit values + # TODO assert exit values + sp_exit.w() From 6c6ec36cee5186a70147d53a537245745ff56f53 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 11 Jan 2023 14:15:44 +0100 Subject: [PATCH 050/263] - add output domain args to fused stencil programs - add option to switch from single fused stencil call in diffusion to gt4py program - load grid indices as int (instead of int32) --- .../fused_mo_nh_diffusion_stencil_02_03.py | 11 +- .../fused_mo_nh_diffusion_stencil_04_05_06.py | 11 +- ...sed_mo_nh_diffusion_stencil_07_08_09_10.py | 11 +- .../fused_mo_nh_diffusion_stencil_11_12.py | 17 +- .../fused_mo_nh_diffusion_stencil_13_14.py | 17 +- .../mo_nh_diffusion_stencil_16.py | 15 +- .../src/icon4py/diffusion/diffusion.py | 489 +++++++++--------- .../icon4py/diffusion/diffusion_program.py | 8 +- .../src/icon4py/diffusion/icon_grid.py | 6 +- atm_dyn_iconam/tests/test_diffusion.py | 118 ++++- ...est_mo_intp_rbf_rbf_vec_interpol_vertex.py | 14 +- .../tests/test_mo_nh_diffusion_stencil_16.py | 4 + 12 files changed, 428 insertions(+), 293 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py index 6b32d5a67..8648f9491 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -13,6 +13,7 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field +from functional.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import ( _mo_nh_diffusion_stencil_02, @@ -39,7 +40,7 @@ def _fused_mo_nh_diffusion_stencil_02_03( return div_ic, hdef_ic -@program +@program(backend=gtfn_cpu.run_gtfn) def fused_mo_nh_diffusion_stencil_02_03( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], @@ -49,6 +50,10 @@ def fused_mo_nh_diffusion_stencil_02_03( wgtfac_c: Field[[CellDim, KDim], float], div_ic: Field[[CellDim, KDim], float], hdef_ic: Field[[CellDim, KDim], float], + horizontal_start: int, + horizontal_end: int, + vertical_start: int, + vertical_end: int, ): _fused_mo_nh_diffusion_stencil_02_03( kh_smag_ec, @@ -58,4 +63,8 @@ def fused_mo_nh_diffusion_stencil_02_03( diff_multfac_smag, wgtfac_c, out=(div_ic, hdef_ic), + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py index 58e334501..8bb9768b3 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py @@ -13,6 +13,7 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, int32, where +from functional.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_04 import ( _mo_nh_diffusion_stencil_04, @@ -74,7 +75,7 @@ def _fused_mo_nh_diffusion_stencil_04_05_06( return vn -@program +@program(backend=gtfn_cpu.run_gtfn) def fused_mo_nh_diffusion_stencil_04_05_06( u_vert: Field[[VertexDim, KDim], float], v_vert: Field[[VertexDim, KDim], float], @@ -92,6 +93,10 @@ def fused_mo_nh_diffusion_stencil_04_05_06( nudgezone_diff: float, fac_bdydiff_v: float, start_2nd_nudge_line_idx_e: int32, + horizontal_start: int, + horizontal_end: int, + vertical_start: int, + vertical_end: int, ): _fused_mo_nh_diffusion_stencil_04_05_06( u_vert, @@ -111,4 +116,8 @@ def fused_mo_nh_diffusion_stencil_04_05_06( fac_bdydiff_v, start_2nd_nudge_line_idx_e, out=vn, + domain={ + EdgeDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py index a2a611a42..62a02aca4 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -13,6 +13,7 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, broadcast, int32, where +from functional.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_07 import ( _mo_nh_diffusion_stencil_07, @@ -82,7 +83,7 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( return w, dwdx, dwdy -@program +@program(backend=gtfn_cpu.run_gtfn) def fused_mo_nh_diffusion_stencil_07_08_09_10( area: Field[[CellDim], float], geofac_n2s: Field[[CellDim, C2E2CODim], float], @@ -99,6 +100,10 @@ def fused_mo_nh_diffusion_stencil_07_08_09_10( nrdmax: int32, interior_idx: int32, halo_idx: int32, + horizontal_start: int, + horizontal_end: int, + vertical_start: int, + vertical_end: int, ): _fused_mo_nh_diffusion_stencil_07_08_09_10( area, @@ -117,4 +122,8 @@ def fused_mo_nh_diffusion_stencil_07_08_09_10( interior_idx, halo_idx, out=(w, dwdx, dwdy), + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py index 375f8031d..77d66e537 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py @@ -13,6 +13,7 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field +from functional.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_11 import ( _mo_nh_diffusion_stencil_11, @@ -35,13 +36,25 @@ def _fused_mo_nh_diffusion_stencil_11_12( return kh_smag_e -@program +@program(backend=gtfn_cpu.run_gtfn) def fused_mo_nh_diffusion_stencil_11_12( theta_v: Field[[CellDim, KDim], float], theta_ref_mc: Field[[CellDim, KDim], float], thresh_tdiff: float, kh_smag_e: Field[[EdgeDim, KDim], float], + horizontal_start: int, + horizontal_end: int, + vertical_start: int, + vertical_end: int, ): _fused_mo_nh_diffusion_stencil_11_12( - theta_v, theta_ref_mc, thresh_tdiff, kh_smag_e, out=kh_smag_e + theta_v, + theta_ref_mc, + thresh_tdiff, + kh_smag_e, + out=kh_smag_e, + domain={ + EdgeDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py index 6ebe54825..d3edc5f56 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -13,6 +13,7 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field +from functional.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_13 import ( _mo_nh_diffusion_stencil_13, @@ -35,14 +36,26 @@ def _fused_mo_nh_diffusion_stencil_13_14( return z_temp -@program +@program(backend=gtfn_cpu.run_gtfn) def fused_mo_nh_diffusion_stencil_13_14( kh_smag_e: Field[[EdgeDim, KDim], float], inv_dual_edge_length: Field[[EdgeDim], float], theta_v: Field[[CellDim, KDim], float], geofac_div: Field[[CellDim, C2EDim], float], z_temp: Field[[CellDim, KDim], float], + horizontal_start: int, + horizontal_end: int, + vertical_start: int, + vertical_end: int, ): _fused_mo_nh_diffusion_stencil_13_14( - kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp + kh_smag_e, + inv_dual_edge_length, + theta_v, + geofac_div, + out=z_temp, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_16.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_16.py index 03d8a463a..eea2f719c 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_16.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_16.py @@ -38,7 +38,20 @@ def mo_nh_diffusion_stencil_16( theta_v: Field[[CellDim, KDim], float], exner: Field[[CellDim, KDim], float], rd_o_cvd: float, + horizontal_start: int, + horizontal_end: int, + vertical_start: int, + vertical_end: int, ): _mo_nh_diffusion_stencil_16( - z_temp, area, theta_v, exner, rd_o_cvd, out=(theta_v, exner) + z_temp, + area, + theta_v, + exner, + rd_o_cvd, + out=(theta_v, exner), + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 256d80ee3..3859e85bc 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -# flake8: noqa F841, E800 + import math import sys from collections import namedtuple @@ -25,29 +25,30 @@ np_as_located_field, ) +import icon4py.diffusion.diffusion_program as diff_prog from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import ( - _fused_mo_nh_diffusion_stencil_02_03, + fused_mo_nh_diffusion_stencil_02_03, ) from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import ( - _fused_mo_nh_diffusion_stencil_04_05_06, + fused_mo_nh_diffusion_stencil_04_05_06, ) from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_07_08_09_10 import ( - _fused_mo_nh_diffusion_stencil_07_08_09_10, + fused_mo_nh_diffusion_stencil_07_08_09_10, ) from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_11_12 import ( - _fused_mo_nh_diffusion_stencil_11_12, + fused_mo_nh_diffusion_stencil_11_12, ) from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_13_14 import ( - _fused_mo_nh_diffusion_stencil_13_14, + fused_mo_nh_diffusion_stencil_13_14, ) from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import ( - _mo_nh_diffusion_stencil_01, + mo_nh_diffusion_stencil_01, ) from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_16 import ( - _mo_nh_diffusion_stencil_16, + mo_nh_diffusion_stencil_16, ) from icon4py.common.constants import CPD, GAS_CONSTANT_DRY_AIR from icon4py.common.dimension import ( @@ -62,7 +63,6 @@ VertexDim, ) from icon4py.diffusion.diagnostic import DiagnosticState -from icon4py.diffusion.diffusion_program import diffusion_run from icon4py.diffusion.horizontal import HorizontalMarkerIndex from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams from icon4py.diffusion.interpolation_state import InterpolationState @@ -241,7 +241,7 @@ def __init__( Called 'denom_diffu_v' in mo_gridref_nml.f90. """ - # namelist: mo_interpol_nml.f90 + # parameters from namelist: mo_interpol_nml.f90 self.nudge_max_coeff: float = max_nudging_coeff """ Parameter describing the lateral boundary nudging in limited area mode. @@ -263,12 +263,11 @@ def __init__( self._validate() def _validate(self): - """ - Apply consistency checks and validation on configuration parameters. - """ + """Apply consistency checks and validation on configuration parameters.""" if self.diffusion_type != 5: raise NotImplementedError( - "Only diffusion type 5 = `Smagorinsky diffusion with fourth-order background diffusion` is implemented" + "Only diffusion type 5 = `Smagorinsky diffusion with fourth-order background " + "diffusion` is implemented" ) if not self.grid.limited_area: @@ -295,7 +294,7 @@ class DiffusionParams: """Calculates derived quantities depending on the diffusion config.""" def __init__(self, config: DiffusionConfig): - self.boundary_diffusion_start_index_edges = int32( + self.boundary_diffusion_start_index_edges = ( 5 # mo_nh_diffusion.start_bdydiff_e - 1 = 5 -1 ) @@ -401,12 +400,14 @@ def __init__( config: DiffusionConfig, params: DiffusionParams, vct_a: Field[[KDim], float], + run_program=True, ): """ Initialize Diffusion granule. calculates all local fields that are used in diffusion within the time loop """ + self._run_program = run_program self.params: DiffusionParams = params self.config: DiffusionConfig = config self.grid = config.grid @@ -425,7 +426,7 @@ def __init__( -5.0 ) # threshold temperature deviation from neighboring grid points hat activates extra diffusion against runaway cooling - self._smag_offset: float = 0.25 * params.K4 * config.substep_as_float() + self.smag_offset: float = 0.25 * params.K4 * config.substep_as_float() self.diff_multfac_w: float = min( 1.0 / 48.0, params.K4W * config.substep_as_float() ) @@ -438,8 +439,8 @@ def __init__( *params.smagorinski_factor, *params.smagorinski_height, vct_a, - self._diff_multfac_vn, - self._smag_limit, + self.diff_multfac_vn, + self.smag_limit, self.enh_smag_fac, offset_provider={"Koff": KDim}, ) @@ -459,9 +460,9 @@ def _index_field(dim: Dimension, size=None): size = size if size else self.grid.size[dim] return np_as_located_field(dim)(np.arange(size, dtype=int32)) - self._diff_multfac_vn = _allocate(KDim) + self.diff_multfac_vn = _allocate(KDim) - self._smag_limit = _allocate(KDim) + self.smag_limit = _allocate(KDim) self.enh_smag_fac = _allocate(KDim) self.u_vert = _allocate(VertexDim, KDim) self.v_vert = _allocate(VertexDim, KDim) @@ -549,25 +550,6 @@ def time_step( runs a diffusion step for the parameter linit=False, within regular time loop. """ - # self._do_diffusion_step( - # diagnostic_state, - # prognostic_state, - # metric_state, - # interpolation_state, - # dtime, - # tangent_orientation, - # inverse_primal_edge_lengths, - # inverse_dual_edge_length, - # inverse_vertical_vertex_lengths, - # primal_normal_vert, - # dual_normal_vert, - # edge_areas, - # cell_areas, - # self._diff_multfac_vn, - # self._smag_limit, - # self._smag_offset, - # ) - cell_startindex_nudging, cell_endindex_local = self.grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.nudging(CellDim), @@ -619,91 +601,111 @@ def time_step( HorizontalMarkerIndex.local(VertexDim) - 1, ) - diffusion_run( - diagnostic_state.hdef_ic, - diagnostic_state.div_ic, - diagnostic_state.dwdx, - diagnostic_state.dwdy, - prognostic_state.vertical_wind, - prognostic_state.normal_wind, - prognostic_state.exner_pressure, - prognostic_state.theta_v, - metric_state.theta_ref_mc, - metric_state.wgtfac_c, - metric_state.mask_hdiff, - metric_state.zd_vertidx, - metric_state.zd_diffcoef, - metric_state.zd_intcoef, - interpolation_state.e_bln_c_s, - interpolation_state.rbf_coeff_1, - interpolation_state.rbf_coeff_2, - interpolation_state.geofac_div, - interpolation_state.geofac_grg_x, - interpolation_state.geofac_grg_y, - interpolation_state.nudgecoeff_e, - interpolation_state.geofac_n2s, - tangent_orientation, - inverse_primal_edge_lengths, - inverse_dual_edge_length, - inverse_vertical_vertex_lengths, - primal_normal_vert[0], - primal_normal_vert[1], - dual_normal_vert[0], - dual_normal_vert[1], - edge_areas, - cell_areas, - self._diff_multfac_vn, - dtime, - self.rd_o_cvd, - self.thresh_tdiff, - self._smag_limit, - self.u_vert, - self.v_vert, - self.enh_smag_fac, - self.kh_smag_e, - self.kh_smag_ec, - self.z_nabla2_e, - self.z_temp, - self.diff_multfac_smag, - self.diff_multfac_n2w, - self._smag_offset, - self.nudgezone_diff, - self.fac_bdydiff_v, - self.diff_multfac_w, - self.vertical_index, - self.horizontal_cell_index, - self.horizontal_edge_index, - cell_startindex_interior, - cell_startindex_nudging, - cell_endindex_local_plus1, - cell_endindex_local, - edge_startindex_nudging_plus1, - edge_startindex_nudging_minus1, - edge_endindex_local, - edge_endindex_local_minus2, - vertex_startindex_lb_plus3, - vertex_startindex_lb_plus1, - vertex_endindex_local, - vertex_endindex_local_minus1, - self.config.vertical_params.index_of_damping_height, - self.grid.n_lev(), - self.params.boundary_diffusion_start_index_edges, - offset_provider={ - "V2E": self.grid.get_v2e_connectivity(), - "V2EDim": V2EDim, - "E2C2V": self.grid.get_e2c2v_connectivity(), - "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, 4), - "C2E": self.grid.get_c2e_connectivity(), - "C2EDim": C2EDim, - "E2C": self.grid.get_e2c_connectivity(), - "C2E2C": self.grid.get_c2e2c_connectivity(), - "C2E2CDim": C2E2CDim, - "ECVDim": ECVDim, - "C2E2CODim": C2E2CODim, - "C2E2CO": self.grid.get_c2e2co_connectivity(), - "Koff": KDim, - }, - ) + if not self._run_program: + self._do_diffusion_step( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + metric_state=metric_state, + interpolation_state=interpolation_state, + dtime=dtime, + tangent_orientation=tangent_orientation, + inverse_primal_edge_lengths=inverse_primal_edge_lengths, + inverse_dual_edge_length=inverse_dual_edge_length, + inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, + primal_normal_vert=primal_normal_vert, + dual_normal_vert=dual_normal_vert, + edge_areas=edge_areas, + cell_areas=cell_areas, + diff_multfac_vn=self.diff_multfac_vn, + smag_limit=self.smag_limit, + smag_offset=self.smag_offset, + ) + else: + diff_prog.diffusion_run( + diagnostic_hdef_ic=diagnostic_state.hdef_ic, + diagnostic_div_ic=diagnostic_state.div_ic, + diagnostic_dwdx=diagnostic_state.dwdx, + diagnostic_dwdy=diagnostic_state.dwdy, + prognostic_vertical_wind=prognostic_state.vertical_wind, + prognostic_normal_wind=prognostic_state.normal_wind, + prognostic_exner_pressure=prognostic_state.exner_pressure, + prognostic_theta_v=prognostic_state.theta_v, + metric_theta_ref_mc=metric_state.theta_ref_mc, + metric_wgtfac_c=metric_state.wgtfac_c, + metric_mask_hdiff=metric_state.mask_hdiff, + metric_zd_vertidx=metric_state.zd_vertidx, + metric_zd_diffcoef=metric_state.zd_diffcoef, + metric_zd_intcoef=metric_state.zd_intcoef, + interpolation_e_bln_c_s=interpolation_state.e_bln_c_s, + interpolation_rbf_coeff_1=interpolation_state.rbf_coeff_1, + interpolation_rbf_coeff_2=interpolation_state.rbf_coeff_2, + interpolation_geofac_div=interpolation_state.geofac_div, + interpolation_geofac_grg_x=interpolation_state.geofac_grg_x, + interpolation_geofac_grg_y=interpolation_state.geofac_grg_y, + interpolation_nudgecoeff_e=interpolation_state.nudgecoeff_e, + interpolation_geofac_n2s=interpolation_state.geofac_n2s, + tangent_orientation=tangent_orientation, + inverse_primal_edge_lengths=inverse_primal_edge_lengths, + inverse_dual_edge_lengths=inverse_dual_edge_length, + inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, + primal_normal_vert_1=primal_normal_vert[0], + primal_normal_vert_2=primal_normal_vert[1], + dual_normal_vert_1=dual_normal_vert[0], + dual_normal_vert_2=dual_normal_vert[1], + edge_areas=edge_areas, + cell_areas=cell_areas, + diff_multfac_vn=self.diff_multfac_vn, + dtime=dtime, + rd_o_cvd=self.rd_o_cvd, + local_thresh_tdiff=self.thresh_tdiff, + local_smag_limit=self.smag_limit, + local_u_vert=self.u_vert, + local_v_vert=self.v_vert, + local_enh_smag_fac=self.enh_smag_fac, + local_kh_smag_e=self.kh_smag_e, + local_kh_smag_ec=self.kh_smag_ec, + local_z_nabla2_e=self.z_nabla2_e, + local_z_temp=self.z_temp, + local_diff_multfac_smag=self.diff_multfac_smag, + local_diff_multfac_n2w=self.diff_multfac_n2w, + local_smag_offset=self.smag_offset, + local_nudgezone_diff=self.nudgezone_diff, + local_fac_bdydiff_v=self.fac_bdydiff_v, + local_diff_multfac_w=self.diff_multfac_w, + local_vertical_index=self.vertical_index, + local_horizontal_cell_index=self.horizontal_cell_index, + local_horizontal_edge_index=self.horizontal_edge_index, + cell_startindex_interior=cell_startindex_interior, + cell_startindex_nudging=cell_startindex_nudging, + cell_endindex_local_plus1=cell_endindex_local_plus1, + cell_endindex_local=cell_endindex_local, + edge_startindex_nudging_plus1=edge_startindex_nudging_plus1, + edge_startindex_nudging_minus1=edge_startindex_nudging_minus1, + edge_endindex_local=edge_endindex_local, + edge_endindex_local_minus2=edge_endindex_local_minus2, + vertex_startindex_lb_plus3=vertex_startindex_lb_plus3, + vertex_startindex_lb_plus1=vertex_startindex_lb_plus1, + vertex_endindex_local=vertex_endindex_local, + vertex_endindex_local_minus1=vertex_endindex_local_minus1, + index_of_damping_height=self.config.vertical_params.index_of_damping_height, + nlev=self.grid.n_lev(), + boundary_diffusion_start_index_edges=self.params.boundary_diffusion_start_index_edges, + offset_provider={ + "V2E": self.grid.get_v2e_connectivity(), + "V2EDim": V2EDim, + "E2C2V": self.grid.get_e2c2v_connectivity(), + "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, 4), + "C2E": self.grid.get_c2e_connectivity(), + "C2EDim": C2EDim, + "E2C": self.grid.get_e2c_connectivity(), + "C2E2C": self.grid.get_c2e2c_connectivity(), + "C2E2CDim": C2E2CDim, + "ECVDim": ECVDim, + "C2E2CODim": C2E2CODim, + "C2E2CO": self.grid.get_c2e2co_connectivity(), + "Koff": KDim, + }, + ) def _do_diffusion_step( self, @@ -746,9 +748,6 @@ def _do_diffusion_step( smag_offset: """ - # ------- - # OUTLINE - # ------- klevels = self.grid.n_lev() k_start_end_minus2 = klevels - 2 @@ -758,13 +757,13 @@ def _do_diffusion_step( HorizontalMarkerIndex.local(CellDim) + 1, ) - cell_start_interior, cell_end_halo = self.grid.get_indices_from_to( + cell_start_interior, cell_end_local = self.grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.interior(CellDim), HorizontalMarkerIndex.local(CellDim), ) - cell_start_nudging, cell_end_local = self.grid.get_indices_from_to( + cell_start_nudging, _ = self.grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.nudging(CellDim), HorizontalMarkerIndex.local(CellDim), @@ -813,88 +812,78 @@ def _do_diffusion_step( set_zero_v_k(self.v_vert, offset_provider={}) # # 1. CALL rbf_vec_interpol_vertex - - v_start = vertex_start_local_boundary_plus1 - v_end = vertex_end_local_minus1 - k_start = 0 - k_end = klevels mo_intp_rbf_rbf_vec_interpol_vertex( - prognostic_state.normal_wind, - interpolation_state.rbf_coeff_1, - interpolation_state.rbf_coeff_2, - self.u_vert, - self.v_vert, - v_start, - v_end, - k_start, - k_end, + p_e_in=prognostic_state.normal_wind, + ptr_coeff_1=interpolation_state.rbf_coeff_1, + ptr_coeff_2=interpolation_state.rbf_coeff_2, + p_u_out=self.u_vert, + p_v_out=self.v_vert, + horizontal_start=vertex_start_local_boundary_plus3, + horizontal_end=vertex_end_local_minus1, + vertical_start=0, + vertical_end=klevels, offset_provider={"V2E": self.grid.get_v2e_connectivity(), "V2EDim": V2EDim}, ) # 2. HALO EXCHANGE -- CALL sync_patch_array_mult # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 - e_start = self.params.boundary_diffusion_start_index_edges - e_end = edge_end_local_minus2 - k_start = 0 - k_end = klevels - _mo_nh_diffusion_stencil_01( - self.diff_multfac_smag, - tangent_orientation, - inverse_primal_edge_lengths, - inverse_vertical_vertex_lengths, - self.u_vert, - self.v_vert, - primal_normal_vert[0], - primal_normal_vert[1], - dual_normal_vert[0], - dual_normal_vert[1], - prognostic_state.normal_wind, - smag_limit, - smag_offset, - out=(self.kh_smag_e, self.kh_smag_ec, self.z_nabla2_e), + mo_nh_diffusion_stencil_01( + diff_multfac_smag=self.diff_multfac_smag, + tangent_orientation=tangent_orientation, + inv_primal_edge_length=inverse_primal_edge_lengths, + inv_vert_vert_length=inverse_vertical_vertex_lengths, + u_vert=self.u_vert, + v_vert=self.v_vert, + primal_normal_vert_x=primal_normal_vert[0], + primal_normal_vert_y=primal_normal_vert[1], + dual_normal_vert_x=dual_normal_vert[0], + dual_normal_vert_y=dual_normal_vert[1], + vn=prognostic_state.normal_wind, + smag_limit=smag_limit, + kh_smag_e=self.kh_smag_e, + kh_smag_ec=self.kh_smag_ec, + z_nabla2_e=self.z_nabla2_e, + smag_offset=smag_offset, + horizontal_start=self.params.boundary_diffusion_start_index_edges, + horizontal_end=edge_end_local_minus2, + vertical_start=0, + vertical_end=klevels, offset_provider={ "E2C2V": self.grid.get_e2c2v_connectivity(), "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, 4), }, ) - c_start = (cell_start_nudging,) - c_end = (cell_end_local,) - k_start = 0 - k_end = klevels - _fused_mo_nh_diffusion_stencil_02_03( - self.kh_smag_ec, - prognostic_state.normal_wind, - interpolation_state.e_bln_c_s, - interpolation_state.geofac_div, - self.diff_multfac_smag, - metric_state.wgtfac_c, - out=( - diagnostic_state.div_ic, - diagnostic_state.hdef_ic, - ), + fused_mo_nh_diffusion_stencil_02_03( + kh_smag_ec=self.kh_smag_ec, + vn=prognostic_state.normal_wind, + e_bln_c_s=interpolation_state.e_bln_c_s, + geofac_div=interpolation_state.geofac_div, + diff_multfac_smag=self.diff_multfac_smag, + wgtfac_c=metric_state.wgtfac_c, + div_ic=diagnostic_state.div_ic, + hdef_ic=diagnostic_state.hdef_ic, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local, + vertical_start=0, + vertical_end=klevels, offset_provider={"C2E": self.grid.get_c2e_connectivity(), "C2EDim": C2EDim}, ) # # # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH # # # 5. CALL rbf_vec_interpol_vertex_wp - - k_start = 0 - k_end = klevels - v_start = vertex_start_local_boundary_plus3 - v_end = vertex_end_local mo_intp_rbf_rbf_vec_interpol_vertex( - self.z_nabla2_e, - interpolation_state.rbf_coeff_1, - interpolation_state.rbf_coeff_2, - self.u_vert, - self.v_vert, - v_start, - v_end, - k_start, - k_end, + p_e_in=self.z_nabla2_e, + ptr_coeff_1=interpolation_state.rbf_coeff_1, + ptr_coeff_2=interpolation_state.rbf_coeff_2, + p_u_out=self.u_vert, + p_v_out=self.v_vert, + horizontal_start=vertex_start_local_boundary_plus3, + horizontal_end=vertex_end_local, + vertical_start=0, + vertical_end=klevels, offset_provider={"V2E": self.grid.get_v2e_connectivity(), "V2EDim": V2EDim}, ) # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult @@ -902,12 +891,7 @@ def _do_diffusion_step( # # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 # # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 # - - e_start = edge_start_nudging_plus_one - e_end = edge_end_local - k_start = 0 - k_end = klevels - _fused_mo_nh_diffusion_stencil_04_05_06( + fused_mo_nh_diffusion_stencil_04_05_06( self.u_vert, self.v_vert, primal_normal_vert[0], @@ -924,7 +908,11 @@ def _do_diffusion_step( self.nudgezone_diff, self.fac_bdydiff_v, edge_start_nudging_minus1, - out=prognostic_state.normal_wind, + prognostic_state.normal_wind, + horizontal_start=edge_start_nudging_plus_one, + horizontal_end=edge_end_local, + vertical_start=0, + vertical_end=klevels, offset_provider={ "E2C2V": self.grid.get_e2c2v_connectivity(), "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, 4), @@ -934,31 +922,26 @@ def _do_diffusion_step( # # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, # # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 - c_start = cell_start_nudging - c_end = cell_end_local_plus1 - k_start = 0 - k_end = klevels - _fused_mo_nh_diffusion_stencil_07_08_09_10( - cell_areas, - interpolation_state.geofac_n2s, - interpolation_state.geofac_grg_x, - interpolation_state.geofac_grg_y, - prognostic_state.vertical_wind, - prognostic_state.vertical_wind, - diagnostic_state.dwdx, - diagnostic_state.dwdy, - self.diff_multfac_w, - self.diff_multfac_n2w, - self.vertical_index, - self.horizontal_cell_index, - self.config.vertical_params.index_of_damping_height, - cell_start_interior, - cell_end_halo, - out=( - prognostic_state.vertical_wind, - diagnostic_state.dwdx, - diagnostic_state.dwdy, - ), + fused_mo_nh_diffusion_stencil_07_08_09_10( + area=cell_areas, + geofac_n2s=interpolation_state.geofac_n2s, + geofac_grg_x=interpolation_state.geofac_grg_x, + geofac_grg_y=interpolation_state.geofac_grg_y, + w_old=prognostic_state.vertical_wind, + w=prognostic_state.vertical_wind, + dwdx=diagnostic_state.dwdx, + dwdy=diagnostic_state.dwdy, + diff_multfac_w=self.diff_multfac_w, + diff_multfac_n2w=self.diff_multfac_n2w, + vert_idx=self.vertical_index, + horz_idx=self.horizontal_cell_index, + nrdmax=self.config.vertical_params.index_of_damping_height, + interior_idx=cell_start_interior, # h end index for stencil_09 and stencil_10 + halo_idx=cell_end_local, # h end index for stencil_09 and stencil_10, + horizontal_start=cell_start_nudging, # h start index for stencil_07 and stencil_08 + horizontal_end=cell_end_local_plus1, # h end index for stencil_07 and stencil_08 + vertical_start=0, + vertical_end=klevels, offset_provider={ "C2E2CO": self.grid.get_c2e2co_connectivity(), "C2E2CODim": C2E2CODim, @@ -970,35 +953,31 @@ def _do_diffusion_step( # # # TODO check: kh_smag_e is an out field, should not be calculated in init? # - - e_start = edge_start_nudging_plus_one - e_end = edge_end_local - c_start = cell_start_nudging_minus1 - c_end = cell_end_local_plus1 - k_start = k_start_end_minus2 - k_end = klevels - - _fused_mo_nh_diffusion_stencil_11_12( - prognostic_state.theta_v, - metric_state.theta_ref_mc, - self.thresh_tdiff, - self.kh_smag_e, - out=self.kh_smag_e, + fused_mo_nh_diffusion_stencil_11_12( + theta_v=prognostic_state.theta_v, + theta_ref_mc=metric_state.theta_ref_mc, + thresh_tdiff=self.thresh_tdiff, + kh_smag_e=self.kh_smag_e, + horizontal_start=edge_start_nudging_plus_one, + horizontal_end=edge_end_local, + vertical_start=k_start_end_minus2, + vertical_end=klevels, offset_provider={ "E2C": self.grid.get_e2c_connectivity(), "C2E2C": self.grid.get_c2e2c_connectivity(), "C2E2CDim": C2E2CDim, }, ) - - c_start = (cell_start_nudging,) - c_end = cell_end_local - _fused_mo_nh_diffusion_stencil_13_14( - self.kh_smag_e, - inverse_dual_edge_length, - prognostic_state.theta_v, - interpolation_state.geofac_div, - out=self.z_temp, + fused_mo_nh_diffusion_stencil_13_14( + kh_smag_e=self.kh_smag_e, + inv_dual_edge_length=inverse_dual_edge_length, + theta_v=prognostic_state.theta_v, + geofac_div=interpolation_state.geofac_div, + z_temp=self.z_temp, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local, + vertical_start=0, + vertical_end=klevels, offset_provider={ "C2E": self.grid.get_c2e_connectivity(), "E2C": self.grid.get_e2c_connectivity(), @@ -1024,14 +1003,16 @@ def _do_diffusion_step( # offset_provider={}, # ) - c_start = cell_start_nudging - c_end = cell_end_local - _mo_nh_diffusion_stencil_16( - self.z_temp, - cell_areas, - prognostic_state.theta_v, - prognostic_state.exner_pressure, - self.rd_o_cvd, + mo_nh_diffusion_stencil_16( + z_temp=self.z_temp, + area=cell_areas, + theta_v=prognostic_state.theta_v, + exner=prognostic_state.exner_pressure, + rd_o_cvd=self.rd_o_cvd, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local, + vertical_start=0, + vertical_end=klevels, offset_provider={}, ) # 10. HALO EXCHANGE sync_patch_array diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py index 178f4ec17..f17aefb8a 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py @@ -15,6 +15,7 @@ from functional.common import Field from functional.ffront.decorator import program from functional.ffront.fbuiltins import int32 +from functional.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import ( _fused_mo_nh_diffusion_stencil_02_03, @@ -54,8 +55,7 @@ from icon4py.diffusion.utils import _scale_k, _set_zero_v_k -# @program(backend=gtfn_cpu.run_gtfn) -@program +@program(backend=gtfn_cpu.run_gtfn) def diffusion_run( diagnostic_hdef_ic: Field[[CellDim, KDim], float], diagnostic_div_ic: Field[[CellDim, KDim], float], @@ -81,7 +81,7 @@ def diffusion_run( interpolation_geofac_n2s: Field[[CellDim, C2E2CODim], float], tangent_orientation: Field[[EdgeDim], float], inverse_primal_edge_lengths: Field[[EdgeDim], float], - inverse_dual_edge_length: Field[[EdgeDim], float], + inverse_dual_edge_lengths: Field[[EdgeDim], float], inverse_vertical_vertex_lengths: Field[[EdgeDim], float], primal_normal_vert_1: Field[[ECVDim], float], primal_normal_vert_2: Field[[ECVDim], float], @@ -275,7 +275,7 @@ def diffusion_run( ) _fused_mo_nh_diffusion_stencil_13_14( local_kh_smag_e, - inverse_dual_edge_length, + inverse_dual_edge_lengths, prognostic_theta_v, interpolation_geofac_div, out=local_z_temp, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index 4972438f5..979fd70bf 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -102,8 +102,8 @@ def with_config(self, config: MeshConfig): def with_start_end_indices( self, dim: Dimension, start_indices: np.ndarray, end_indices: np.ndarray ): - self.start_indices[dim] = start_indices - self.end_indices[dim] = end_indices + self.start_indices[dim] = start_indices.astype(int) + self.end_indices[dim] = end_indices.astype(int) @builder def with_connectivities(self, connectivity: Dict[Dimension, np.ndarray]): @@ -130,7 +130,7 @@ def num_edges(self): def get_indices_from_to( self, dim: Dimension, start_marker: int, end_marker: int - ) -> Tuple[int32, int32]: + ) -> Tuple[int, int]: """ Use to specifzy domains of a field for field_operator. diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 0f66168b8..b2ea47ce4 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -259,15 +259,15 @@ def test_diffusion_init(savepoint_init, r04b09_diffusion_config, step_date_init) ) assert ( - diffusion._smag_offset + diffusion.smag_offset == 0.25 * additional_parameters.K4 * config.substep_as_float() ) - assert np.allclose(expected_smag_limit, diffusion._smag_limit) + assert np.allclose(expected_smag_limit, diffusion.smag_limit) expected_diff_multfac_vn = diff_multfac_vn_numpy( shape_k, additional_parameters.K4, config.substep_as_float() ) - assert np.allclose(expected_diff_multfac_vn, diffusion._diff_multfac_vn) + assert np.allclose(expected_diff_multfac_vn, diffusion.diff_multfac_vn) expected_enh_smag_fac = enhanced_smagorinski_factor_numpy( additional_parameters.smagorinski_factor, additional_parameters.smagorinski_height, @@ -284,7 +284,7 @@ def _verify_init_values_against_savepoint( assert savepoint.nudgezone_diff() == diffusion.nudgezone_diff assert savepoint.bdy_diff() == diffusion.bdy_diff assert savepoint.fac_bdydiff_v() == diffusion.fac_bdydiff_v - assert savepoint.smag_offset() == diffusion._smag_offset + assert savepoint.smag_offset() == diffusion.smag_offset assert savepoint.diff_multfac_w() == diffusion.diff_multfac_w # this is done in diffusion.run(...) because it depends on the dtime @@ -293,9 +293,9 @@ def _verify_init_values_against_savepoint( ) assert np.allclose(savepoint.diff_multfac_smag(), diffusion.diff_multfac_smag) - assert np.allclose(savepoint.smag_limit(), diffusion._smag_limit) + assert np.allclose(savepoint.smag_limit(), diffusion.smag_limit) np.allclose(savepoint.diff_multfac_n2w(), np.asarray(diffusion.diff_multfac_n2w)) - assert np.allclose(savepoint.diff_multfac_vn(), diffusion._diff_multfac_vn) + assert np.allclose(savepoint.diff_multfac_vn(), diffusion.diff_multfac_vn) @pytest.mark.datatest @@ -356,8 +356,8 @@ def test_verify_diffusion_init_against_other_regular_savepoint( _verify_init_values_against_savepoint(savepoint, diffusion) -@pytest.mark.skip @pytest.mark.datatest +@pytest.mark.skip def test_run_diffusion_single_step(savepoint_init, savepoint_exit, icon_grid): sp = savepoint_init vct_a = sp.vct_a() @@ -371,7 +371,7 @@ def test_run_diffusion_single_step(savepoint_init, savepoint_exit, icon_grid): additional_parameters = DiffusionParams(config) - diffusion = Diffusion(config, additional_parameters, vct_a) + diffusion = Diffusion(config, additional_parameters, vct_a, run_program=False) diagnostic_state = DiagnosticState( hdef_ic=sp.hdef_ic(), div_ic=sp.div_ic(), dwdx=sp.dwdx(), dwdy=sp.dwdy() @@ -454,7 +454,8 @@ def test_run_diffusion_single_step(savepoint_init, savepoint_exit, icon_grid): @pytest.mark.skip -def test_diffusion( +@pytest.mark.datatest +def test_diffusion_five_steps( icon_grid, savepoint_init, savepoint_exit, @@ -462,21 +463,104 @@ def test_diffusion( step_date_exit="2021-06-20T12:01:00.000", ): sp = savepoint_init + + diagnostic_state = DiagnosticState( + hdef_ic=sp.hdef_ic(), div_ic=sp.div_ic(), dwdx=sp.dwdx(), dwdy=sp.dwdy() + ) + prognostic_state = PrognosticState( + vertical_wind=sp.w(), + normal_wind=sp.vn(), + exner_pressure=sp.exner(), + theta_v=sp.theta_v(), + ) + grg = sp.geofac_grg() + + interpolation_state = InterpolationState( + e_bln_c_s=sp.e_bln_c_s(), + rbf_coeff_1=sp.rbf_vec_coeff_v1(), + rbf_coeff_2=sp.rbf_vec_coeff_v2(), + geofac_div=sp.geofac_div(), + geofac_n2s=sp.geofac_n2s(), + geofac_grg_x=grg[0], + geofac_grg_y=grg[1], + nudgecoeff_e=sp.nudgecoeff_e(), + ) + + metric_state = MetricState( + mask_hdiff=sp.mask_diff(), + theta_ref_mc=sp.theta_ref_mc(), + wgtfac_c=sp.wgtfac_c(), + zd_intcoef=sp.zd_intcoef(), + zd_vertidx=sp.zd_vertidx(), + zd_diffcoef=sp.zd_diffcoef(), + ) + dtime = sp.get_metadata("dtime").get("dtime") + orientation = sp.tangent_orientation() + inverse_primal_edge_lengths = sp.inverse_primal_edge_lengths() + inverse_vertical_vertex_lengths = sp.inv_vert_vert_length() + inverse_dual_edge_length = sp.inv_dual_edge_length() + primal_normal_vert: VectorTuple = ( + as_1D_sparse_field(sp.primal_normal_vert_x(), ECVDim), + as_1D_sparse_field(sp.primal_normal_vert_y(), ECVDim), + ) + dual_normal_vert: VectorTuple = ( + as_1D_sparse_field(sp.dual_normal_vert_x(), ECVDim), + as_1D_sparse_field(sp.dual_normal_vert_y(), ECVDim), + ) + edge_areas = sp.edge_areas() + cell_areas = sp.cell_areas() + config = DiffusionConfig( icon_grid, vertical_params=VerticalModelParams( vct_a=(sp.vct_a()), rayleigh_damping_height=12500.0 ), ) - additional_parameters = DiffusionParams(config) diffusion = Diffusion(config, additional_parameters, sp.vct_a()) - - diffusion.initial_step(...) + diffusion.initial_step( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + metric_state=metric_state, + interpolation_state=interpolation_state, + dtime=dtime, + tangent_orientation=orientation, + inverse_primal_edge_lengths=inverse_primal_edge_lengths, + inverse_dual_edge_length=inverse_dual_edge_length, + inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, + primal_normal_vert=primal_normal_vert, + dual_normal_vert=dual_normal_vert, + edge_areas=edge_areas, + cell_areas=cell_areas, + ) for _ in range(4): - diffusion.time_step(...) + diffusion.time_step( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + metric_state=metric_state, + interpolation_state=interpolation_state, + dtime=dtime, + tangent_orientation=orientation, + inverse_primal_edge_lengths=inverse_primal_edge_lengths, + inverse_dual_edge_length=inverse_dual_edge_length, + inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, + primal_normal_vert=primal_normal_vert, + dual_normal_vert=dual_normal_vert, + edge_areas=edge_areas, + cell_areas=cell_areas, + ) - sp_exit = savepoint_exit - - # TODO assert exit values - sp_exit.w() + icon_result_exner = savepoint_exit.exner() + icon_result_vn = savepoint_exit.vn() + icon_result_w = savepoint_exit.w() + icon_result_theta_w = savepoint_exit.theta_v() + assert np.allclose(icon_result_w, np.asarray(prognostic_state.vertical_wind)) + assert np.allclose( + np.asarray(icon_result_vn), np.asarray(prognostic_state.normal_wind) + ) + assert np.allclose( + np.asarray(icon_result_theta_w), np.asarray(prognostic_state.theta_v) + ) + assert np.allclose( + np.asarray(icon_result_exner), np.asarray(prognostic_state.exner_pressure) + ) diff --git a/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index 1466ec08d..2eb90a21a 100644 --- a/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -25,12 +25,12 @@ def mo_intp_rbf_rbf_vec_interpol_vertex_numpy( v2e: np.array, p_e_in: np.array, ptr_coeff_1: np.array, ptr_coeff_2: np.array ) -> tuple[np.array]: ptr_coeff_1 = np.expand_dims(ptr_coeff_1, axis=-1) - p_v_out = np.sum(p_e_in[v2e] * ptr_coeff_1, axis=1) + p_u_out = np.sum(p_e_in[v2e] * ptr_coeff_1, axis=1) ptr_coeff_2 = np.expand_dims(ptr_coeff_2, axis=-1) - p_u_out = np.sum(p_e_in[v2e] * ptr_coeff_2, axis=1) + p_v_out = np.sum(p_e_in[v2e] * ptr_coeff_2, axis=1) - return p_v_out, p_u_out + return p_u_out, p_v_out def test_mo_intp_rbf_rbf_vec_interpol_vertex(): @@ -42,21 +42,21 @@ def test_mo_intp_rbf_rbf_vec_interpol_vertex(): p_v_out = zero_field(mesh, VertexDim, KDim) p_u_out = zero_field(mesh, VertexDim, KDim) - p_v_out_ref, p_u_out_ref = mo_intp_rbf_rbf_vec_interpol_vertex_numpy( - mesh.v2e, np.asarray(p_e_in), np.asarray(ptr_coeff_1), np.asarray(ptr_coeff_2) - ) mo_intp_rbf_rbf_vec_interpol_vertex( p_e_in, ptr_coeff_1, ptr_coeff_2, - p_v_out, p_u_out, + p_v_out, 0, mesh.n_vertices, 0, mesh.k_level, offset_provider={"V2E": mesh.get_v2e_offset_provider(), "V2EDim": V2EDim}, ) + p_u_out_ref, p_v_out_ref = mo_intp_rbf_rbf_vec_interpol_vertex_numpy( + mesh.v2e, np.asarray(p_e_in), np.asarray(ptr_coeff_1), np.asarray(ptr_coeff_2) + ) assert np.allclose(p_v_out, p_v_out_ref) assert np.allclose(p_u_out, p_u_out_ref) diff --git a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_16.py b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_16.py index 48af1e45b..1020d07d7 100644 --- a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_16.py +++ b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_16.py @@ -57,6 +57,10 @@ def test_mo_nh_diffusion_stencil_16(): theta_v, exner, rd_o_cvd, + horizontal_start=0, + horizontal_end=mesh.n_cells, + vertical_start=0, + vertical_end=mesh.k_level, offset_provider={}, ) assert np.allclose(theta_v, theta_v_ref) From a0007e77f053f62a2449067c9ab8fb576ec229b1 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 12 Jan 2023 10:40:50 +0100 Subject: [PATCH 051/263] fix pre-commit issues --- .../src/icon4py/diffusion/diffusion.py | 34 +++++++++---------- .../src/icon4py/diffusion/horizontal.py | 1 - 2 files changed, 17 insertions(+), 18 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 3859e85bc..3cff7a2de 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -985,23 +985,23 @@ def _do_diffusion_step( }, ) - # mo_nh_diffusion_stencil_15_numpy( - # mask_hdiff=metric_state.mask_hdiff, - # zd_vertidx=metric_state.zd_vertidx, - # vcoef=metric_state.zd_diffcoef, - # zd_diffcoef=metric_state.zd_diffcoef, - # geofac_n2s=interpolation_state.geofac_n2s, - # theta_v=prognostic_state.theta_v, - # z_temp=self.z_temp, - # domain={ - # CellDim: self.grid.get_indices_from_to( - # CellDim, - # HorizontalMarkerIndex.nudging(CellDim), - # HorizontalMarkerIndex.halo(CellDim), - # ), - # }, - # offset_provider={}, - # ) + mo_nh_diffusion_stencil_15_numpy( + mask_hdiff=metric_state.mask_hdiff, + zd_vertidx=metric_state.zd_vertidx, + vcoef=metric_state.zd_diffcoef, + zd_diffcoef=metric_state.zd_diffcoef, + geofac_n2s=interpolation_state.geofac_n2s, + theta_v=prognostic_state.theta_v, + z_temp=self.z_temp, + domain={ + CellDim: self.grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.nudging(CellDim), + HorizontalMarkerIndex.halo(CellDim), + ), + }, + offset_provider={}, + ) mo_nh_diffusion_stencil_16( z_temp=self.z_temp, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py index 846933f32..36df836bb 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py @@ -67,7 +67,6 @@ def local_boundary(cls, dim: Dimension) -> int: @classmethod def local(cls, dim: Dimension) -> int: - """ """ match (dim): case (dimension.CellDim): return cls._HALO_CELLS From b49aa41cc73bf533d5f1f0e5b3533448bf553fda Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 17 Jan 2023 08:55:48 +0100 Subject: [PATCH 052/263] remove serialbox workaround --- base-requirements-dev.txt | 4 +++- .../src/icon4py/testutils/serialbox_utils.py | 19 +------------------ 2 files changed, 4 insertions(+), 19 deletions(-) diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index cb97c337d..570c27208 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -1,5 +1,7 @@ # VCS --e git+https://github.com/tehrengruber/gt4py.git@reduce_foast_lowering_complexity#egg=gt4py-functional +-e git+https://github.com/tehrengruber/gt4py.git@2a9dcb0367a1ec1ef6d78951358d700424867b07#egg=gt4py-functional +git+https://github.com/halungge/serialbox@fix_python_build#egg=serialbox&subdirectory=src/serialbox-python + # PyPI flake8>=3.8 diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 1a1bcb02c..53b5945ad 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -11,10 +11,8 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import os -import pathlib - import numpy as np +import serialbox as ser from functional.common import Dimension from functional.ffront.fbuiltins import int32 from functional.iterator.embedded import np_as_located_field @@ -32,21 +30,6 @@ ) -try: - import serialbox as ser -except ImportError: - external_src = pathlib.Path.joinpath( - pathlib.Path(__file__).parent, "../../../../_external_src/" - ) - os.chdir(external_src) - if not pathlib.Path.exists(pathlib.Path.joinpath(external_src, "serialbox")): - os.system("git clone --recursive https://github.com/GridTools/serialbox") - os.system( - "CC=`which gcc` CXX=`which g++` pip install serialbox/src/serialbox-python" - ) - import serialbox as ser - - class IconDiffusionSavepoint: def __init__(self, sp: ser.Savepoint, ser: ser.Serializer): self.savepoint = sp From 14335e13a93f822582220d6a7e4d07e87f796175 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 17 Jan 2023 09:46:52 +0100 Subject: [PATCH 053/263] add boost to workflow runner --- .github/workflows/tox.yml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.github/workflows/tox.yml b/.github/workflows/tox.yml index c5bf4f7a9..e459c76ac 100644 --- a/.github/workflows/tox.yml +++ b/.github/workflows/tox.yml @@ -18,6 +18,11 @@ jobs: steps: - uses: actions/checkout@v2 + - name: install boost + run: | + sudo apt-get update + sudo apt-get install libboost-all-dev + - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v2 with: From 79272cc3a1141748f339db6b2e3d911a21688a9a Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 17 Jan 2023 10:57:28 +0100 Subject: [PATCH 054/263] fix tox.ini: wrong naming of datatest marker --- tox.ini | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tox.ini b/tox.ini index 1d6e2dae7..76b886571 100644 --- a/tox.ini +++ b/tox.ini @@ -16,8 +16,8 @@ setenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -m "not dataset" -s -n auto -cache-clear --cov --cov-reset --doctest-modules atm_dyn_iconam/src common/src pyutils/src testutils/src - pytest -v -m "not dataset" -s -n auto --cov --cov-append + -pytest -v -m "not datatest" -s -n auto -cache-clear --cov --cov-reset --doctest-modules atm_dyn_iconam/src common/src pyutils/src testutils/src + pytest -v -m "not datatest" -s -n auto --cov --cov-append commands_post = rm -Rf _reports/coverage_html coverage html From ee6dfe1a477f3323c6dc0cac82ab750e72f5b5ee Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 19 Jan 2023 13:38:52 +0100 Subject: [PATCH 055/263] switch to pathlib when loading serialbox data --- atm_dyn_iconam/tests/conftest.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 1c268af63..ce77d7d33 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -10,8 +10,8 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -import os import tarfile +from pathlib import Path import numpy as np import pytest @@ -40,9 +40,9 @@ data_uri = "https://polybox.ethz.ch/index.php/s/kP0Q2dDU6DytEqI/download" -data_path = os.path.join(os.path.dirname(__file__), "./ser_icondata") -extracted_path = os.path.join(data_path, "mch_ch_r04b09_dsl/ser_data") -data_file = os.path.join(data_path, "ser_data_diffusion.tar.gz") +data_path = Path(__file__).parent.joinpath("ser_icondata") +extracted_path = data_path.joinpath("mch_ch_r04b09_dsl/ser_data") +data_file = data_path.joinpath("ser_data_diffusion.tar.gz").name @pytest.fixture(scope="session") @@ -52,19 +52,18 @@ def setup_icon_data(): Session scoped fixture which is a prerequisite of all the other fixtures in this file. """ - os.makedirs(data_path, exist_ok=True) - if len(os.listdir(data_path)) == 0: + data_path.mkdir(parents=True, exist_ok=True) + if not any(data_path.iterdir()): print( f"directory {data_path} is empty: downloading data from {data_uri} and extracting" ) - wget.download(data_uri, out=data_file) # extract downloaded file if not tarfile.is_tarfile(data_file): raise NotImplementedError(f"{data_file} needs to be a valid tar file") with tarfile.open(data_file, mode="r:*") as tf: tf.extractall(path=data_path) - os.remove(data_file) + Path(data_file).unlink(missing_ok=True) @pytest.fixture From 1e00c05fdaf6ad6153475f29f6109ec663893251 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Thu, 24 Nov 2022 09:04:58 +0100 Subject: [PATCH 056/263] First version of fused diffusion stencils --- .../fused_mo_nh_diffusion_stencil_02_03.py | 47 ++++++++++ .../fused_mo_nh_diffusion_stencil_04_05_06.py | 82 ++++++++++++++++++ ...sed_mo_nh_diffusion_stencil_07_08_09_10.py | 86 +++++++++++++++++++ .../fused_mo_nh_diffusion_stencil_11_12.py | 40 +++++++++ .../fused_mo_nh_diffusion_stencil_13_14.py | 67 +++++++++++++++ .../mo_nh_diffusion_stencil_10.py | 8 +- 6 files changed, 326 insertions(+), 4 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py new file mode 100644 index 000000000..8c111b5d4 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -0,0 +1,47 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum + +from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim, Koff + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import _mo_nh_diffusion_stencil_02 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import _mo_nh_diffusion_stencil_03 + +@field_operator +def _fused_mo_nh_diffusion_stencil_02_03( + kh_smag_ec: Field[[EdgeDim, KDim], float], + vn: Field[[EdgeDim, KDim], float], + e_bln_c_s: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + diff_multfac_smag: Field[[KDim], float], + wgtfac_c: Field[[CellDim, KDim], float], +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + kh_c, div = _mo_nh_diffusion_stencil_02(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag) + div_ic, hdef_ic = _mo_nh_diffusion_stencil_03(div, kh_c, wgtfac_c) + return div_ic, hdef_ic + + +@program +def fused_mo_nh_diffusion_stencil_02_03( + kh_smag_ec: Field[[EdgeDim, KDim], float], + vn: Field[[EdgeDim, KDim], float], + e_bln_c_s: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + diff_multfac_smag: Field[[KDim], float], + wgtfac_c: Field[[CellDim, KDim], float], + div_ic: Field[[CellDim, KDim], float], + hdef_ic: Field[[CellDim, KDim], float], +): + _fused_mo_nh_diffusion_stencil_02_03(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag, wgtfac_c, out=(div_ic, hdef_ic)) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py new file mode 100644 index 000000000..1438e07c1 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py @@ -0,0 +1,82 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum, maximum, where, int32 + +from icon4py.common.dimension import ( + E2C2V, + E2ECV, + ECVDim, + EdgeDim, + KDim, + VertexDim, +) + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_04 import _mo_nh_diffusion_stencil_04 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_05 import _mo_nh_diffusion_stencil_05 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_06 import _mo_nh_diffusion_stencil_06 + +@field_operator +def _fused_mo_nh_diffusion_stencil_04_05_06( + u_vert: Field[[VertexDim, KDim], float], + v_vert: Field[[VertexDim, KDim], float], + primal_normal_vert_v1: Field[[ECVDim], float], + primal_normal_vert_v2: Field[[ECVDim], float], + z_nabla2_e: Field[[EdgeDim, KDim], float], + inv_vert_vert_length: Field[[EdgeDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + area_edge: Field[[EdgeDim], float], + kh_smag_e: Field[[EdgeDim, KDim], float], + diff_multfac_vn: Field[[KDim], float], + nudgecoeff_e: Field[[EdgeDim], float], + vn: Field[[EdgeDim, KDim], float], + horz_idx: Field[[EdgeDim], int32], + nudgezone_diff: float, + fac_bdydiff_v: float, + start_2nd_nudge_line_idx_e: int32, +) -> Field[[EdgeDim, KDim], float]: + + z_nabla4_e2 = _mo_nh_diffusion_stencil_04(u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2, z_nabla2_e, inv_vert_vert_length, inv_primal_edge_length) + + vn = where( + horz_idx >= start_2nd_nudge_line_idx_e, + _mo_nh_diffusion_stencil_05(area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, nudgecoeff_e, vn, nudgezone_diff), + _mo_nh_diffusion_stencil_06(z_nabla2_e, area_edge, vn, fac_bdydiff_v) + ) + + return vn + + +@program +def fused_mo_nh_diffusion_stencil_04_05_06( + u_vert: Field[[VertexDim, KDim], float], + v_vert: Field[[VertexDim, KDim], float], + primal_normal_vert_v1: Field[[ECVDim], float], + primal_normal_vert_v2: Field[[ECVDim], float], + z_nabla2_e: Field[[EdgeDim, KDim], float], + inv_vert_vert_length: Field[[EdgeDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + area_edge: Field[[EdgeDim], float], + kh_smag_e: Field[[EdgeDim, KDim], float], + diff_multfac_vn: Field[[KDim], float], + nudgecoeff_e: Field[[EdgeDim], float], + vn: Field[[EdgeDim, KDim], float], + horz_idx: Field[[EdgeDim], int32], + nudgezone_diff: float, + fac_bdydiff_v: float, + start_2nd_nudge_line_idx_e: int32, +): + _fused_mo_nh_diffusion_stencil_04_05_06(u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2, z_nabla2_e, inv_vert_vert_length, + inv_primal_edge_length, area_edge, kh_smag_e, diff_multfac_vn, nudgecoeff_e, vn, horz_idx, nudgezone_diff, fac_bdydiff_v, + start_2nd_nudge_line_idx_e, out=vn) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py new file mode 100644 index 000000000..a1e305ae5 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -0,0 +1,86 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum, where, int32, broadcast + +from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_07 import _mo_nh_diffusion_stencil_07 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_08 import _mo_nh_diffusion_stencil_08 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_09 import _mo_nh_diffusion_stencil_09 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_10 import _mo_nh_diffusion_stencil_10 + +@field_operator +def _fused_mo_nh_diffusion_stencil_07_08_09_10( + area: Field[[CellDim], float], + geofac_n2s: Field[[CellDim, C2E2CODim], float], + geofac_grg_x: Field[[CellDim, C2E2CODim], float], + geofac_grg_y: Field[[CellDim, C2E2CODim], float], + w_old: Field[[CellDim, KDim], float], + w: Field[[CellDim, KDim], float], + dwdx: Field[[CellDim, KDim], float], + dwdy: Field[[CellDim, KDim], float], + diff_multfac_w: float, + diff_multfac_n2w: Field[[KDim], float], + vert_idx: Field[[KDim], int32], + horz_idx: Field[[CellDim], int32], + nrdmax: int32, + interior_idx: int32, + halo_idx: int32, +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + + vert_idx = broadcast(vert_idx, (CellDim, KDim)) + + dwdx, dwdy = where( + vert_idx > int32(0), + _mo_nh_diffusion_stencil_08(w_old, geofac_grg_x, geofac_grg_y), + (dwdx, dwdy) + ) + + z_nabla2_c = _mo_nh_diffusion_stencil_07(w_old, geofac_n2s) + + w = where( + (horz_idx >= interior_idx) & (horz_idx < halo_idx), + _mo_nh_diffusion_stencil_09(area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w), + w_old, + ) + + w = where( + (vert_idx > int32(0)) & (vert_idx < nrdmax) & (horz_idx >= interior_idx) & (horz_idx < halo_idx), + _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, area, z_nabla2_c), + w, + ) + + return w, dwdx, dwdy + + +@program +def fused_mo_nh_diffusion_stencil_07_08_09_10( + area: Field[[CellDim], float], + geofac_n2s: Field[[CellDim, C2E2CODim], float], + geofac_grg_x: Field[[CellDim, C2E2CODim], float], + geofac_grg_y: Field[[CellDim, C2E2CODim], float], + w_old: Field[[CellDim, KDim], float], + w: Field[[CellDim, KDim], float], + dwdx: Field[[CellDim, KDim], float], + dwdy: Field[[CellDim, KDim], float], + diff_multfac_w: float, + diff_multfac_n2w: Field[[KDim], float], + vert_idx: Field[[KDim], int32], + horz_idx: Field[[CellDim], int32], + nrdmax: int32, + interior_idx: int32, + halo_idx: int32, +): + _fused_mo_nh_diffusion_stencil_07_08_09_10(area, geofac_n2s, geofac_grg_x, geofac_grg_y, w_old, w, dwdx, dwdy, diff_multfac_w, diff_multfac_n2w, vert_idx, horz_idx, nrdmax, interior_idx, halo_idx, out = (w, dwdx, dwdy)) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py new file mode 100644 index 000000000..02fd00530 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py @@ -0,0 +1,40 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, where, neighbor_sum, max_over, maximum + +from icon4py.common.dimension import E2C, E2CDim, C2E2C, C2E2CDim, CellDim, EdgeDim, KDim + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_11 import _mo_nh_diffusion_stencil_11 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_12 import _mo_nh_diffusion_stencil_12 + +@field_operator +def _fused_mo_nh_diffusion_stencil_11_12( + theta_v: Field[[CellDim, KDim], float], + theta_ref_mc: Field[[CellDim, KDim], float], + thresh_tdiff: float, + kh_smag_e: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + enh_diffu_3d = _mo_nh_diffusion_stencil_11(theta_v, theta_ref_mc, thresh_tdiff) + kh_smag_e = _mo_nh_diffusion_stencil_12(kh_smag_e, enh_diffu_3d) + return kh_smag_e + +@program +def fused_mo_nh_diffusion_stencil_11_12( + theta_v: Field[[CellDim, KDim], float], + theta_ref_mc: Field[[CellDim, KDim], float], + thresh_tdiff: float, + kh_smag_e: Field[[EdgeDim, KDim], float], +): + _fused_mo_nh_diffusion_stencil_11_12(theta_v, theta_ref_mc, thresh_tdiff, kh_smag_e, out=kh_smag_e) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py new file mode 100644 index 000000000..40faa10b5 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -0,0 +1,67 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum + +from icon4py.common.dimension import ( + E2C, + E2CDim, + C2E2C, + C2E2CDim, + C2CE, + C2E, + C2EDim, + CEDim, + CellDim, + EdgeDim, + KDim, +) + +@field_operator +def _mo_nh_diffusion_stencil_13( + kh_smag_e: Field[[EdgeDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + theta_v: Field[[CellDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + z_nabla2_e = kh_smag_e * inv_dual_edge_length * (theta_v(E2C[1]) - theta_v(E2C[0])) + return z_nabla2_e + +@field_operator +def _mo_nh_diffusion_stencil_14( + z_nabla2_e: Field[[EdgeDim, KDim], float], + geofac_div: Field[[CEDim], float], +) -> Field[[CellDim, KDim], float]: + z_temp = neighbor_sum(z_nabla2_e(C2E) * geofac_div(C2CE), axis=C2EDim) + return z_temp + +@field_operator +def _fused_mo_nh_diffusion_stencil_13_14( + kh_smag_e: Field[[EdgeDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + theta_v: Field[[CellDim, KDim], float], + geofac_div: Field[[CEDim], float], +) -> Field[[CellDim, KDim], float]: + z_nabla2_e = _mo_nh_diffusion_stencil_13(kh_smag_e, inv_dual_edge_length, theta_v) + z_temp = _mo_nh_diffusion_stencil_14(z_nabla2_e, geofac_div) + return z_temp + +@program +def fused_mo_nh_diffusion_stencil_13_14( + kh_smag_e: Field[[EdgeDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + theta_v: Field[[CellDim, KDim], float], + geofac_div: Field[[CEDim], float], + z_temp: Field[[CellDim, KDim], float], +): + _fused_mo_nh_diffusion_stencil_13_14(kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py index e439f742e..fed277e85 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py @@ -21,10 +21,10 @@ def _mo_nh_diffusion_stencil_10( w: Field[[CellDim, KDim], float], diff_multfac_n2w: Field[[KDim], float], - cell_area: Field[[CellDim], float], + area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - w = w + diff_multfac_n2w * (cell_area * z_nabla2_c) + w = w + diff_multfac_n2w * (area * z_nabla2_c) return w @@ -32,7 +32,7 @@ def _mo_nh_diffusion_stencil_10( def mo_nh_diffusion_stencil_10( w: Field[[CellDim, KDim], float], diff_multfac_n2w: Field[[KDim], float], - cell_area: Field[[CellDim], float], + area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], ): - _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w) + _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, area, z_nabla2_c, out=w) From 2abd8b93738c6193255b9b42b3c242b72474f7b4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Tue, 29 Nov 2022 12:30:45 +0100 Subject: [PATCH 057/263] Added prototype of fused_mo_nh_diffusion_stencil_01_02_03_rbf Currently this stencil does not work since it would require elements of a tuple of output variables of a field_operator to have different types. --- ...ed_mo_nh_diffusion_stencil_01_02_03_rbf.py | 92 +++++++++++++++++++ 1 file changed, 92 insertions(+) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py new file mode 100644 index 000000000..ae2b30777 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py @@ -0,0 +1,92 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field + +from icon4py.common.dimension import ECVDim, C2EDim, V2EDim, CellDim, EdgeDim, VertexDim, KDim + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import _mo_nh_diffusion_stencil_01 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import _mo_nh_diffusion_stencil_02 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import _mo_nh_diffusion_stencil_03 +from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import _mo_intp_rbf_rbf_vec_interpol_vertex + +@field_operator +def _fused_mo_nh_diffusion_stencil_01_02_03_rbf( + diff_multfac_smag: Field[[KDim], float], + tangent_orientation: Field[[EdgeDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + inv_vert_vert_length: Field[[EdgeDim], float], + u_vert_old: Field[[VertexDim, KDim], float], + v_vert_old: Field[[VertexDim, KDim], float], + primal_normal_vert_x: Field[[ECVDim], float], + primal_normal_vert_y: Field[[ECVDim], float], + dual_normal_vert_x: Field[[ECVDim], float], + dual_normal_vert_y: Field[[ECVDim], float], + vn: Field[[EdgeDim, KDim], float], + smag_limit: Field[[KDim], float], + smag_offset: float, + e_bln_c_s: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + wgtfac_c: Field[[CellDim, KDim], float], + ptr_coeff_1: Field[[VertexDim, V2EDim], float], + ptr_coeff_2: Field[[VertexDim, V2EDim], float], +) -> tuple[ + Field[[EdgeDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[VertexDim, KDim], float], + Field[[VertexDim, KDim], float], + ]: + + kh_smag_e, kh_smag_ec, z_nabla2_e = _mo_nh_diffusion_stencil_01( diff_multfac_smag, + tangent_orientation, inv_primal_edge_length, inv_vert_vert_length, u_vert_old, v_vert_old, + primal_normal_vert_x, primal_normal_vert_y, dual_normal_vert_x, dual_normal_vert_y, + vn, smag_limit, smag_offset) + + kh_c, div = _mo_nh_diffusion_stencil_02(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag) + + div_ic, hdef_ic = _mo_nh_diffusion_stencil_03(div, kh_c, wgtfac_c) + + u_vert, v_vert = _mo_intp_rbf_rbf_vec_interpol_vertex(z_nabla2_e, ptr_coeff_1, ptr_coeff_2) + + return kh_smag_e, div_ic, hdef_ic, u_vert, v_vert + + +@program +def fused_mo_nh_diffusion_stencil_01_02_03_rbf( + diff_multfac_smag: Field[[KDim], float], + tangent_orientation: Field[[EdgeDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + inv_vert_vert_length: Field[[EdgeDim], float], + u_vert_old: Field[[VertexDim, KDim], float], + v_vert_old: Field[[VertexDim, KDim], float], + primal_normal_vert_x: Field[[ECVDim], float], + primal_normal_vert_y: Field[[ECVDim], float], + dual_normal_vert_x: Field[[ECVDim], float], + dual_normal_vert_y: Field[[ECVDim], float], + vn: Field[[EdgeDim, KDim], float], + smag_limit: Field[[KDim], float], + smag_offset: float, + e_bln_c_s: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + wgtfac_c: Field[[CellDim, KDim], float], + ptr_coeff_1: Field[[VertexDim, V2EDim], float], + ptr_coeff_2: Field[[VertexDim, V2EDim], float], + kh_smag_e: Field[[EdgeDim, KDim], float], + div_ic: Field[[CellDim, KDim], float], + hdef_ic: Field[[CellDim, KDim], float], + u_vert: Field[[VertexDim, KDim], float], + v_vert: Field[[VertexDim, KDim], float], +): + _fused_mo_nh_diffusion_stencil_01_02_03_rbf(diff_multfac_smag, tangent_orientation, inv_primal_edge_length, inv_vert_vert_length, u_vert_old, v_vert_old, primal_normal_vert_x, primal_normal_vert_y, dual_normal_vert_x, dual_normal_vert_y, vn, smag_limit, smag_offset, e_bln_c_s, geofac_div, wgtfac_c, ptr_coeff_1, ptr_coeff_2, out=(kh_smag_e, div_ic, hdef_ic, u_vert, v_vert)) From ee4e21e7ce44e43212cecb61cc27ed89b4c065e2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Wed, 7 Dec 2022 15:19:32 +0100 Subject: [PATCH 058/263] Removed duplicate stencil definition from fused stencil 13 and 14 --- ...used_mo_nh_diffusion_stencil_07_08_09_10.py | 2 +- .../fused_mo_nh_diffusion_stencil_13_14.py | 18 ++---------------- 2 files changed, 3 insertions(+), 17 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py index a1e305ae5..a97a75d21 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -40,7 +40,7 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( halo_idx: int32, ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - vert_idx = broadcast(vert_idx, (CellDim, KDim)) + vert_idx = broadcast(vert_idx, (CellDim, KDim)) dwdx, dwdy = where( vert_idx > int32(0), diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py index 40faa10b5..8d2092008 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -28,22 +28,8 @@ KDim, ) -@field_operator -def _mo_nh_diffusion_stencil_13( - kh_smag_e: Field[[EdgeDim, KDim], float], - inv_dual_edge_length: Field[[EdgeDim], float], - theta_v: Field[[CellDim, KDim], float], -) -> Field[[EdgeDim, KDim], float]: - z_nabla2_e = kh_smag_e * inv_dual_edge_length * (theta_v(E2C[1]) - theta_v(E2C[0])) - return z_nabla2_e - -@field_operator -def _mo_nh_diffusion_stencil_14( - z_nabla2_e: Field[[EdgeDim, KDim], float], - geofac_div: Field[[CEDim], float], -) -> Field[[CellDim, KDim], float]: - z_temp = neighbor_sum(z_nabla2_e(C2E) * geofac_div(C2CE), axis=C2EDim) - return z_temp +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_13 import _mo_nh_diffusion_stencil_13 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_14 import _mo_nh_diffusion_stencil_14 @field_operator def _fused_mo_nh_diffusion_stencil_13_14( From 381be77f2e55c5f1138929b8623a7b0d2e98d060 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Fri, 20 Jan 2023 14:07:29 +0100 Subject: [PATCH 059/263] Cleanup --- .../fused_mo_nh_diffusion_stencil_02_03.py | 2 +- .../fused_mo_nh_diffusion_stencil_04_05_06.py | 11 ++--------- .../fused_mo_nh_diffusion_stencil_07_08_09_10.py | 4 ++-- .../fused_mo_nh_diffusion_stencil_11_12.py | 4 ++-- .../fused_mo_nh_diffusion_stencil_13_14.py | 16 ++-------------- .../atm_dyn_iconam/mo_nh_diffusion_stencil_10.py | 8 ++++---- 6 files changed, 13 insertions(+), 32 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py index 8c111b5d4..f034e5a62 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -14,7 +14,7 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim, Koff +from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import _mo_nh_diffusion_stencil_02 from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import _mo_nh_diffusion_stencil_03 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py index 1438e07c1..b74c52b4e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py @@ -12,16 +12,9 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, neighbor_sum, maximum, where, int32 +from functional.ffront.fbuiltins import Field, where, int32 -from icon4py.common.dimension import ( - E2C2V, - E2ECV, - ECVDim, - EdgeDim, - KDim, - VertexDim, -) +from icon4py.common.dimension import ECVDim, EdgeDim, KDim, VertexDim from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_04 import _mo_nh_diffusion_stencil_04 from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_05 import _mo_nh_diffusion_stencil_05 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py index a97a75d21..0bbcce058 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -12,9 +12,9 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, neighbor_sum, where, int32, broadcast +from functional.ffront.fbuiltins import Field, where, int32, broadcast -from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim +from icon4py.common.dimension import C2E2CODim, CellDim, KDim from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_07 import _mo_nh_diffusion_stencil_07 from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_08 import _mo_nh_diffusion_stencil_08 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py index 02fd00530..cda940f52 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py @@ -12,9 +12,9 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, where, neighbor_sum, max_over, maximum +from functional.ffront.fbuiltins import Field -from icon4py.common.dimension import E2C, E2CDim, C2E2C, C2E2CDim, CellDim, EdgeDim, KDim +from icon4py.common.dimension import CellDim, EdgeDim, KDim from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_11 import _mo_nh_diffusion_stencil_11 from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_12 import _mo_nh_diffusion_stencil_12 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py index 8d2092008..ae317277c 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -12,21 +12,9 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, neighbor_sum +from functional.ffront.fbuiltins import Field -from icon4py.common.dimension import ( - E2C, - E2CDim, - C2E2C, - C2E2CDim, - C2CE, - C2E, - C2EDim, - CEDim, - CellDim, - EdgeDim, - KDim, -) +from icon4py.common.dimension import CEDim, CellDim, EdgeDim, KDim from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_13 import _mo_nh_diffusion_stencil_13 from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_14 import _mo_nh_diffusion_stencil_14 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py index fed277e85..e439f742e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py @@ -21,10 +21,10 @@ def _mo_nh_diffusion_stencil_10( w: Field[[CellDim, KDim], float], diff_multfac_n2w: Field[[KDim], float], - area: Field[[CellDim], float], + cell_area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - w = w + diff_multfac_n2w * (area * z_nabla2_c) + w = w + diff_multfac_n2w * (cell_area * z_nabla2_c) return w @@ -32,7 +32,7 @@ def _mo_nh_diffusion_stencil_10( def mo_nh_diffusion_stencil_10( w: Field[[CellDim, KDim], float], diff_multfac_n2w: Field[[KDim], float], - area: Field[[CellDim], float], + cell_area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], ): - _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, area, z_nabla2_c, out=w) + _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w) From f5ad668dd24bdb55aa338fc6dfe50822bdb482b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Mon, 23 Jan 2023 01:26:50 +0100 Subject: [PATCH 060/263] Cleanup --- ...ed_mo_nh_diffusion_stencil_01_02_03_rbf.py | 87 +++++++++++++++---- .../fused_mo_nh_diffusion_stencil_02_03.py | 22 ++++- .../fused_mo_nh_diffusion_stencil_04_05_06.py | 59 ++++++++++--- ...sed_mo_nh_diffusion_stencil_07_08_09_10.py | 54 +++++++++--- .../fused_mo_nh_diffusion_stencil_11_12.py | 13 ++- .../fused_mo_nh_diffusion_stencil_13_14.py | 13 ++- 6 files changed, 200 insertions(+), 48 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py index ae2b30777..fb3ecaf12 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py @@ -14,12 +14,28 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field -from icon4py.common.dimension import ECVDim, C2EDim, V2EDim, CellDim, EdgeDim, VertexDim, KDim +from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( + _mo_intp_rbf_rbf_vec_interpol_vertex, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import ( + _mo_nh_diffusion_stencil_01, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import ( + _mo_nh_diffusion_stencil_02, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import ( + _mo_nh_diffusion_stencil_03, +) +from icon4py.common.dimension import ( + C2EDim, + CellDim, + ECVDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import _mo_nh_diffusion_stencil_01 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import _mo_nh_diffusion_stencil_02 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import _mo_nh_diffusion_stencil_03 -from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import _mo_intp_rbf_rbf_vec_interpol_vertex @field_operator def _fused_mo_nh_diffusion_stencil_01_02_03_rbf( @@ -42,23 +58,38 @@ def _fused_mo_nh_diffusion_stencil_01_02_03_rbf( ptr_coeff_1: Field[[VertexDim, V2EDim], float], ptr_coeff_2: Field[[VertexDim, V2EDim], float], ) -> tuple[ - Field[[EdgeDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[VertexDim, KDim], float], - Field[[VertexDim, KDim], float], - ]: + Field[[EdgeDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[VertexDim, KDim], float], + Field[[VertexDim, KDim], float], +]: - kh_smag_e, kh_smag_ec, z_nabla2_e = _mo_nh_diffusion_stencil_01( diff_multfac_smag, - tangent_orientation, inv_primal_edge_length, inv_vert_vert_length, u_vert_old, v_vert_old, - primal_normal_vert_x, primal_normal_vert_y, dual_normal_vert_x, dual_normal_vert_y, - vn, smag_limit, smag_offset) + kh_smag_e, kh_smag_ec, z_nabla2_e = _mo_nh_diffusion_stencil_01( + diff_multfac_smag, + tangent_orientation, + inv_primal_edge_length, + inv_vert_vert_length, + u_vert_old, + v_vert_old, + primal_normal_vert_x, + primal_normal_vert_y, + dual_normal_vert_x, + dual_normal_vert_y, + vn, + smag_limit, + smag_offset, + ) - kh_c, div = _mo_nh_diffusion_stencil_02(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag) + kh_c, div = _mo_nh_diffusion_stencil_02( + kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag + ) div_ic, hdef_ic = _mo_nh_diffusion_stencil_03(div, kh_c, wgtfac_c) - u_vert, v_vert = _mo_intp_rbf_rbf_vec_interpol_vertex(z_nabla2_e, ptr_coeff_1, ptr_coeff_2) + u_vert, v_vert = _mo_intp_rbf_rbf_vec_interpol_vertex( + z_nabla2_e, ptr_coeff_1, ptr_coeff_2 + ) return kh_smag_e, div_ic, hdef_ic, u_vert, v_vert @@ -89,4 +120,24 @@ def fused_mo_nh_diffusion_stencil_01_02_03_rbf( u_vert: Field[[VertexDim, KDim], float], v_vert: Field[[VertexDim, KDim], float], ): - _fused_mo_nh_diffusion_stencil_01_02_03_rbf(diff_multfac_smag, tangent_orientation, inv_primal_edge_length, inv_vert_vert_length, u_vert_old, v_vert_old, primal_normal_vert_x, primal_normal_vert_y, dual_normal_vert_x, dual_normal_vert_y, vn, smag_limit, smag_offset, e_bln_c_s, geofac_div, wgtfac_c, ptr_coeff_1, ptr_coeff_2, out=(kh_smag_e, div_ic, hdef_ic, u_vert, v_vert)) + _fused_mo_nh_diffusion_stencil_01_02_03_rbf( + diff_multfac_smag, + tangent_orientation, + inv_primal_edge_length, + inv_vert_vert_length, + u_vert_old, + v_vert_old, + primal_normal_vert_x, + primal_normal_vert_y, + dual_normal_vert_x, + dual_normal_vert_y, + vn, + smag_limit, + smag_offset, + e_bln_c_s, + geofac_div, + wgtfac_c, + ptr_coeff_1, + ptr_coeff_2, + out=(kh_smag_e, div_ic, hdef_ic, u_vert, v_vert), + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py index f034e5a62..c3a392b22 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -14,10 +14,14 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, neighbor_sum +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import ( + _mo_nh_diffusion_stencil_02, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import ( + _mo_nh_diffusion_stencil_03, +) from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import _mo_nh_diffusion_stencil_02 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import _mo_nh_diffusion_stencil_03 @field_operator def _fused_mo_nh_diffusion_stencil_02_03( @@ -28,7 +32,9 @@ def _fused_mo_nh_diffusion_stencil_02_03( diff_multfac_smag: Field[[KDim], float], wgtfac_c: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - kh_c, div = _mo_nh_diffusion_stencil_02(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag) + kh_c, div = _mo_nh_diffusion_stencil_02( + kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag + ) div_ic, hdef_ic = _mo_nh_diffusion_stencil_03(div, kh_c, wgtfac_c) return div_ic, hdef_ic @@ -44,4 +50,12 @@ def fused_mo_nh_diffusion_stencil_02_03( div_ic: Field[[CellDim, KDim], float], hdef_ic: Field[[CellDim, KDim], float], ): - _fused_mo_nh_diffusion_stencil_02_03(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag, wgtfac_c, out=(div_ic, hdef_ic)) + _fused_mo_nh_diffusion_stencil_02_03( + kh_smag_ec, + vn, + e_bln_c_s, + geofac_div, + diff_multfac_smag, + wgtfac_c, + out=(div_ic, hdef_ic), + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py index b74c52b4e..58e334501 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py @@ -12,13 +12,19 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, where, int32 +from functional.ffront.fbuiltins import Field, int32, where +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_04 import ( + _mo_nh_diffusion_stencil_04, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_05 import ( + _mo_nh_diffusion_stencil_05, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_06 import ( + _mo_nh_diffusion_stencil_06, +) from icon4py.common.dimension import ECVDim, EdgeDim, KDim, VertexDim -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_04 import _mo_nh_diffusion_stencil_04 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_05 import _mo_nh_diffusion_stencil_05 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_06 import _mo_nh_diffusion_stencil_06 @field_operator def _fused_mo_nh_diffusion_stencil_04_05_06( @@ -40,12 +46,29 @@ def _fused_mo_nh_diffusion_stencil_04_05_06( start_2nd_nudge_line_idx_e: int32, ) -> Field[[EdgeDim, KDim], float]: - z_nabla4_e2 = _mo_nh_diffusion_stencil_04(u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2, z_nabla2_e, inv_vert_vert_length, inv_primal_edge_length) + z_nabla4_e2 = _mo_nh_diffusion_stencil_04( + u_vert, + v_vert, + primal_normal_vert_v1, + primal_normal_vert_v2, + z_nabla2_e, + inv_vert_vert_length, + inv_primal_edge_length, + ) vn = where( horz_idx >= start_2nd_nudge_line_idx_e, - _mo_nh_diffusion_stencil_05(area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, nudgecoeff_e, vn, nudgezone_diff), - _mo_nh_diffusion_stencil_06(z_nabla2_e, area_edge, vn, fac_bdydiff_v) + _mo_nh_diffusion_stencil_05( + area_edge, + kh_smag_e, + z_nabla2_e, + z_nabla4_e2, + diff_multfac_vn, + nudgecoeff_e, + vn, + nudgezone_diff, + ), + _mo_nh_diffusion_stencil_06(z_nabla2_e, area_edge, vn, fac_bdydiff_v), ) return vn @@ -70,6 +93,22 @@ def fused_mo_nh_diffusion_stencil_04_05_06( fac_bdydiff_v: float, start_2nd_nudge_line_idx_e: int32, ): - _fused_mo_nh_diffusion_stencil_04_05_06(u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2, z_nabla2_e, inv_vert_vert_length, - inv_primal_edge_length, area_edge, kh_smag_e, diff_multfac_vn, nudgecoeff_e, vn, horz_idx, nudgezone_diff, fac_bdydiff_v, - start_2nd_nudge_line_idx_e, out=vn) + _fused_mo_nh_diffusion_stencil_04_05_06( + u_vert, + v_vert, + primal_normal_vert_v1, + primal_normal_vert_v2, + z_nabla2_e, + inv_vert_vert_length, + inv_primal_edge_length, + area_edge, + kh_smag_e, + diff_multfac_vn, + nudgecoeff_e, + vn, + horz_idx, + nudgezone_diff, + fac_bdydiff_v, + start_2nd_nudge_line_idx_e, + out=vn, + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py index 0bbcce058..a2a611a42 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -12,14 +12,22 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, where, int32, broadcast +from functional.ffront.fbuiltins import Field, broadcast, int32, where +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_07 import ( + _mo_nh_diffusion_stencil_07, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_08 import ( + _mo_nh_diffusion_stencil_08, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_09 import ( + _mo_nh_diffusion_stencil_09, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_10 import ( + _mo_nh_diffusion_stencil_10, +) from icon4py.common.dimension import C2E2CODim, CellDim, KDim -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_07 import _mo_nh_diffusion_stencil_07 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_08 import _mo_nh_diffusion_stencil_08 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_09 import _mo_nh_diffusion_stencil_09 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_10 import _mo_nh_diffusion_stencil_10 @field_operator def _fused_mo_nh_diffusion_stencil_07_08_09_10( @@ -38,26 +46,35 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( nrdmax: int32, interior_idx: int32, halo_idx: int32, -) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: +) -> tuple[ + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], +]: vert_idx = broadcast(vert_idx, (CellDim, KDim)) dwdx, dwdy = where( vert_idx > int32(0), _mo_nh_diffusion_stencil_08(w_old, geofac_grg_x, geofac_grg_y), - (dwdx, dwdy) + (dwdx, dwdy), ) z_nabla2_c = _mo_nh_diffusion_stencil_07(w_old, geofac_n2s) w = where( (horz_idx >= interior_idx) & (horz_idx < halo_idx), - _mo_nh_diffusion_stencil_09(area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w), + _mo_nh_diffusion_stencil_09( + area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w + ), w_old, ) w = where( - (vert_idx > int32(0)) & (vert_idx < nrdmax) & (horz_idx >= interior_idx) & (horz_idx < halo_idx), + (vert_idx > int32(0)) + & (vert_idx < nrdmax) + & (horz_idx >= interior_idx) + & (horz_idx < halo_idx), _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, area, z_nabla2_c), w, ) @@ -83,4 +100,21 @@ def fused_mo_nh_diffusion_stencil_07_08_09_10( interior_idx: int32, halo_idx: int32, ): - _fused_mo_nh_diffusion_stencil_07_08_09_10(area, geofac_n2s, geofac_grg_x, geofac_grg_y, w_old, w, dwdx, dwdy, diff_multfac_w, diff_multfac_n2w, vert_idx, horz_idx, nrdmax, interior_idx, halo_idx, out = (w, dwdx, dwdy)) + _fused_mo_nh_diffusion_stencil_07_08_09_10( + area, + geofac_n2s, + geofac_grg_x, + geofac_grg_y, + w_old, + w, + dwdx, + dwdy, + diff_multfac_w, + diff_multfac_n2w, + vert_idx, + horz_idx, + nrdmax, + interior_idx, + halo_idx, + out=(w, dwdx, dwdy), + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py index cda940f52..375f8031d 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py @@ -14,10 +14,14 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_11 import ( + _mo_nh_diffusion_stencil_11, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_12 import ( + _mo_nh_diffusion_stencil_12, +) from icon4py.common.dimension import CellDim, EdgeDim, KDim -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_11 import _mo_nh_diffusion_stencil_11 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_12 import _mo_nh_diffusion_stencil_12 @field_operator def _fused_mo_nh_diffusion_stencil_11_12( @@ -30,6 +34,7 @@ def _fused_mo_nh_diffusion_stencil_11_12( kh_smag_e = _mo_nh_diffusion_stencil_12(kh_smag_e, enh_diffu_3d) return kh_smag_e + @program def fused_mo_nh_diffusion_stencil_11_12( theta_v: Field[[CellDim, KDim], float], @@ -37,4 +42,6 @@ def fused_mo_nh_diffusion_stencil_11_12( thresh_tdiff: float, kh_smag_e: Field[[EdgeDim, KDim], float], ): - _fused_mo_nh_diffusion_stencil_11_12(theta_v, theta_ref_mc, thresh_tdiff, kh_smag_e, out=kh_smag_e) + _fused_mo_nh_diffusion_stencil_11_12( + theta_v, theta_ref_mc, thresh_tdiff, kh_smag_e, out=kh_smag_e + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py index ae317277c..0941bbd77 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -14,10 +14,14 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_13 import ( + _mo_nh_diffusion_stencil_13, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_14 import ( + _mo_nh_diffusion_stencil_14, +) from icon4py.common.dimension import CEDim, CellDim, EdgeDim, KDim -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_13 import _mo_nh_diffusion_stencil_13 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_14 import _mo_nh_diffusion_stencil_14 @field_operator def _fused_mo_nh_diffusion_stencil_13_14( @@ -30,6 +34,7 @@ def _fused_mo_nh_diffusion_stencil_13_14( z_temp = _mo_nh_diffusion_stencil_14(z_nabla2_e, geofac_div) return z_temp + @program def fused_mo_nh_diffusion_stencil_13_14( kh_smag_e: Field[[EdgeDim, KDim], float], @@ -38,4 +43,6 @@ def fused_mo_nh_diffusion_stencil_13_14( geofac_div: Field[[CEDim], float], z_temp: Field[[CellDim, KDim], float], ): - _fused_mo_nh_diffusion_stencil_13_14(kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp) + _fused_mo_nh_diffusion_stencil_13_14( + kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp + ) From ae10ca27736694b17dd982d49b5ef601e32e6bb9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Mon, 23 Jan 2023 10:22:55 +0100 Subject: [PATCH 061/263] Change includes, remove 01_02_03_rbf from tox tests --- ..._02_03_rbf.py => fused_mo_nh_diffusion_stencil_01_02_03_rbf} | 0 .../atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) rename atm_dyn_iconam/src/icon4py/atm_dyn_iconam/{fused_mo_nh_diffusion_stencil_01_02_03_rbf.py => fused_mo_nh_diffusion_stencil_01_02_03_rbf} (100%) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf similarity index 100% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py rename to atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py index c3a392b22..6b32d5a67 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, neighbor_sum +from functional.ffront.fbuiltins import Field from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import ( _mo_nh_diffusion_stencil_02, From aa291b35d9bf2263345e1a6542bcf3e3af111827 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 25 Jan 2023 09:38:44 +0100 Subject: [PATCH 062/263] rename files containt prognostic and diagnostic state. Diffusion: split construction and initialization into two steps --- .../{diagnostic.py => diagnostic_state.py} | 0 .../src/icon4py/diffusion/diffusion.py | 128 ++++++++++-------- .../icon4py/diffusion/diffusion_wrapper.py | 68 ++++++++++ .../{prognostic.py => prognostic_state.py} | 2 +- atm_dyn_iconam/tests/test_diffusion.py | 88 ++++++++---- 5 files changed, 205 insertions(+), 81 deletions(-) rename atm_dyn_iconam/src/icon4py/diffusion/{diagnostic.py => diagnostic_state.py} (100%) create mode 100644 atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py rename atm_dyn_iconam/src/icon4py/diffusion/{prognostic.py => prognostic_state.py} (90%) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diagnostic.py b/atm_dyn_iconam/src/icon4py/diffusion/diagnostic_state.py similarity index 100% rename from atm_dyn_iconam/src/icon4py/diffusion/diagnostic.py rename to atm_dyn_iconam/src/icon4py/diffusion/diagnostic_state.py diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 3cff7a2de..2e9d95cd3 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -15,7 +15,7 @@ import math import sys from collections import namedtuple -from typing import Final, Tuple +from typing import Final, Tuple, Optional import numpy as np from functional.common import Dimension @@ -62,12 +62,12 @@ V2EDim, VertexDim, ) -from icon4py.diffusion.diagnostic import DiagnosticState +from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.horizontal import HorizontalMarkerIndex from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState -from icon4py.diffusion.prognostic import PrognosticState +from icon4py.diffusion.prognostic_state import PrognosticState from icon4py.diffusion.utils import ( init_diffusion_local_fields_for_regular_timestep, init_nabla2_factor_in_upper_damping_zone, @@ -395,23 +395,48 @@ def mo_nh_diffusion_stencil_15_numpy( class Diffusion: """Class that configures diffusion and does one diffusion step.""" - def __init__( + def __init__(self, run_program = True): + + self._initialized = False + self.rd_o_cvd: float = GAS_CONSTANT_DRY_AIR / (CPD - GAS_CONSTANT_DRY_AIR) + self.thresh_tdiff: float = ( + -5.0 + ) # threshold temperature deviation from neighboring grid points hat activates extra diffusion against runaway cooling + self._run_program = run_program + self.grid: Optional[IconGrid] = None + self.config:Optional[DiffusionConfig] = None + self.params:Optional[DiffusionParams] = None + self.interpolation_state = None + self.metric_state = None + + self.diff_multfac_w: Optional[float] = None + self.smag_offset: Optional[float] = None + self.fac_bdydiff_v: Optional[float] = None + self.bdy_diff: Optional[float] = None + self.nudgezone_diff: Optional[float] = None + self.diff_multfac_n2w:Optional[float] = None + self._allocate_local_fields() + + def init( self, config: DiffusionConfig, params: DiffusionParams, vct_a: Field[[KDim], float], - run_program=True, + metric_state: MetricState, + interpolation_state: InterpolationState ): """ - Initialize Diffusion granule. + Initialize Diffusion granule with configuration. calculates all local fields that are used in diffusion within the time loop """ - self._run_program = run_program + self.params: DiffusionParams = params self.config: DiffusionConfig = config + self.metric_state: MetricState = metric_state + self.interpolation_state: InterpolationState = interpolation_state self.grid = config.grid - self.rd_o_cvd: float = GAS_CONSTANT_DRY_AIR / (CPD - GAS_CONSTANT_DRY_AIR) + self.nudgezone_diff: float = 0.04 / ( config.nudge_max_coeff + sys.float_info.epsilon ) @@ -422,17 +447,13 @@ def __init__( if config.lhdiff_rcf else 1.0 / config.velocity_boundary_diffusion_denominator ) - self.thresh_tdiff: float = ( - -5.0 - ) # threshold temperature deviation from neighboring grid points hat activates extra diffusion against runaway cooling + self.smag_offset: float = 0.25 * params.K4 * config.substep_as_float() self.diff_multfac_w: float = min( 1.0 / 48.0, params.K4W * config.substep_as_float() ) - self._allocate_local_fields() - init_diffusion_local_fields_for_regular_timestep( params.K4, config.substep_as_float(), @@ -451,6 +472,11 @@ def __init__( physical_heights=np.asarray(vct_a), nrdmax=self.config.vertical_params.index_of_damping_height, ) + self._initialized = True + + @property + def initialized(self): + return self._initialized def _allocate_local_fields(self): def _allocate(*dims: Dimension): @@ -479,8 +505,6 @@ def initial_step( self, diagnostic_state: DiagnosticState, prognostic_state: PrognosticState, - metric_state: MetricState, - interpolation_state: InterpolationState, dtime: float, tangent_orientation: Field[[EdgeDim], float], inverse_primal_edge_lengths: Field[[EdgeDim], float], @@ -513,8 +537,6 @@ def initial_step( self._do_diffusion_step( diagnostic_state, prognostic_state, - metric_state, - interpolation_state, dtime, tangent_orientation, inverse_primal_edge_lengths, @@ -533,8 +555,6 @@ def time_step( self, diagnostic_state: DiagnosticState, prognostic_state: PrognosticState, - metric_state: MetricState, - interpolation_state: InterpolationState, dtime: float, tangent_orientation: Field[[EdgeDim], float], inverse_primal_edge_lengths: Field[[EdgeDim], float], @@ -605,8 +625,6 @@ def time_step( self._do_diffusion_step( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, - metric_state=metric_state, - interpolation_state=interpolation_state, dtime=dtime, tangent_orientation=tangent_orientation, inverse_primal_edge_lengths=inverse_primal_edge_lengths, @@ -630,20 +648,20 @@ def time_step( prognostic_normal_wind=prognostic_state.normal_wind, prognostic_exner_pressure=prognostic_state.exner_pressure, prognostic_theta_v=prognostic_state.theta_v, - metric_theta_ref_mc=metric_state.theta_ref_mc, - metric_wgtfac_c=metric_state.wgtfac_c, - metric_mask_hdiff=metric_state.mask_hdiff, - metric_zd_vertidx=metric_state.zd_vertidx, - metric_zd_diffcoef=metric_state.zd_diffcoef, - metric_zd_intcoef=metric_state.zd_intcoef, - interpolation_e_bln_c_s=interpolation_state.e_bln_c_s, - interpolation_rbf_coeff_1=interpolation_state.rbf_coeff_1, - interpolation_rbf_coeff_2=interpolation_state.rbf_coeff_2, - interpolation_geofac_div=interpolation_state.geofac_div, - interpolation_geofac_grg_x=interpolation_state.geofac_grg_x, - interpolation_geofac_grg_y=interpolation_state.geofac_grg_y, - interpolation_nudgecoeff_e=interpolation_state.nudgecoeff_e, - interpolation_geofac_n2s=interpolation_state.geofac_n2s, + metric_theta_ref_mc=self.metric_state.theta_ref_mc, + metric_wgtfac_c=self.metric_state.wgtfac_c, + metric_mask_hdiff=self.metric_state.mask_hdiff, + metric_zd_vertidx=self.metric_state.zd_vertidx, + metric_zd_diffcoef=self.metric_state.zd_diffcoef, + metric_zd_intcoef=self.metric_state.zd_intcoef, + interpolation_e_bln_c_s=self.interpolation_state.e_bln_c_s, + interpolation_rbf_coeff_1=self.interpolation_state.rbf_coeff_1, + interpolation_rbf_coeff_2=self.interpolation_state.rbf_coeff_2, + interpolation_geofac_div=self.interpolation_state.geofac_div, + interpolation_geofac_grg_x=self.interpolation_state.geofac_grg_x, + interpolation_geofac_grg_y=self.interpolation_state.geofac_grg_y, + interpolation_nudgecoeff_e=self.interpolation_state.nudgecoeff_e, + interpolation_geofac_n2s=self.interpolation_state.geofac_n2s, tangent_orientation=tangent_orientation, inverse_primal_edge_lengths=inverse_primal_edge_lengths, inverse_dual_edge_lengths=inverse_dual_edge_length, @@ -711,8 +729,6 @@ def _do_diffusion_step( self, diagnostic_state: DiagnosticState, prognostic_state: PrognosticState, - metric_state: MetricState, - interpolation_state: InterpolationState, dtime: float, tangent_orientation: Field[[EdgeDim], float], inverse_primal_edge_lengths: Field[[EdgeDim], float], @@ -732,8 +748,6 @@ def _do_diffusion_step( Args: diagnostic_state: output argument, data class that contains diagnostic variables prognostic_state: output argument, data class that contains prognostic variables - metric_state: - interpolation_state: dtime: the time step, tangent_orientation: inverse_primal_edge_lengths: @@ -814,8 +828,8 @@ def _do_diffusion_step( # # 1. CALL rbf_vec_interpol_vertex mo_intp_rbf_rbf_vec_interpol_vertex( p_e_in=prognostic_state.normal_wind, - ptr_coeff_1=interpolation_state.rbf_coeff_1, - ptr_coeff_2=interpolation_state.rbf_coeff_2, + ptr_coeff_1=self.interpolation_state.rbf_coeff_1, + ptr_coeff_2=self.interpolation_state.rbf_coeff_2, p_u_out=self.u_vert, p_v_out=self.v_vert, horizontal_start=vertex_start_local_boundary_plus3, @@ -858,10 +872,10 @@ def _do_diffusion_step( fused_mo_nh_diffusion_stencil_02_03( kh_smag_ec=self.kh_smag_ec, vn=prognostic_state.normal_wind, - e_bln_c_s=interpolation_state.e_bln_c_s, - geofac_div=interpolation_state.geofac_div, + e_bln_c_s=self.interpolation_state.e_bln_c_s, + geofac_div=self.interpolation_state.geofac_div, diff_multfac_smag=self.diff_multfac_smag, - wgtfac_c=metric_state.wgtfac_c, + wgtfac_c=self.metric_state.wgtfac_c, div_ic=diagnostic_state.div_ic, hdef_ic=diagnostic_state.hdef_ic, horizontal_start=cell_start_nudging, @@ -876,8 +890,8 @@ def _do_diffusion_step( # # 5. CALL rbf_vec_interpol_vertex_wp mo_intp_rbf_rbf_vec_interpol_vertex( p_e_in=self.z_nabla2_e, - ptr_coeff_1=interpolation_state.rbf_coeff_1, - ptr_coeff_2=interpolation_state.rbf_coeff_2, + ptr_coeff_1=self.interpolation_state.rbf_coeff_1, + ptr_coeff_2=self.interpolation_state.rbf_coeff_2, p_u_out=self.u_vert, p_v_out=self.v_vert, horizontal_start=vertex_start_local_boundary_plus3, @@ -902,7 +916,7 @@ def _do_diffusion_step( edge_areas, self.kh_smag_e, diff_multfac_vn, - interpolation_state.nudgecoeff_e, + self.interpolation_state.nudgecoeff_e, prognostic_state.normal_wind, self.horizontal_edge_index, self.nudgezone_diff, @@ -924,9 +938,9 @@ def _do_diffusion_step( fused_mo_nh_diffusion_stencil_07_08_09_10( area=cell_areas, - geofac_n2s=interpolation_state.geofac_n2s, - geofac_grg_x=interpolation_state.geofac_grg_x, - geofac_grg_y=interpolation_state.geofac_grg_y, + geofac_n2s=self.interpolation_state.geofac_n2s, + geofac_grg_x=self.interpolation_state.geofac_grg_x, + geofac_grg_y=self.interpolation_state.geofac_grg_y, w_old=prognostic_state.vertical_wind, w=prognostic_state.vertical_wind, dwdx=diagnostic_state.dwdx, @@ -955,7 +969,7 @@ def _do_diffusion_step( # fused_mo_nh_diffusion_stencil_11_12( theta_v=prognostic_state.theta_v, - theta_ref_mc=metric_state.theta_ref_mc, + theta_ref_mc=self.metric_state.theta_ref_mc, thresh_tdiff=self.thresh_tdiff, kh_smag_e=self.kh_smag_e, horizontal_start=edge_start_nudging_plus_one, @@ -972,7 +986,7 @@ def _do_diffusion_step( kh_smag_e=self.kh_smag_e, inv_dual_edge_length=inverse_dual_edge_length, theta_v=prognostic_state.theta_v, - geofac_div=interpolation_state.geofac_div, + geofac_div=self.interpolation_state.geofac_div, z_temp=self.z_temp, horizontal_start=cell_start_nudging, horizontal_end=cell_end_local, @@ -986,11 +1000,11 @@ def _do_diffusion_step( ) mo_nh_diffusion_stencil_15_numpy( - mask_hdiff=metric_state.mask_hdiff, - zd_vertidx=metric_state.zd_vertidx, - vcoef=metric_state.zd_diffcoef, - zd_diffcoef=metric_state.zd_diffcoef, - geofac_n2s=interpolation_state.geofac_n2s, + mask_hdiff=self.metric_state.mask_hdiff, + zd_vertidx=self.metric_state.zd_vertidx, + vcoef=self.metric_state.zd_diffcoef, + zd_diffcoef=self.metric_state.zd_diffcoef, + geofac_n2s=self.interpolation_state.geofac_n2s, theta_v=prognostic_state.theta_v, z_temp=self.z_temp, domain={ diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py new file mode 100644 index 000000000..ae5a97045 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py @@ -0,0 +1,68 @@ +from typing import Optional + +import numpy as np + +from icon4py.diffusion.diagnostic_state import DiagnosticState +from icon4py.diffusion.diffusion import Diffusion, DiffusionConfig, DiffusionParams +from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams +from icon4py.diffusion.interpolation_state import InterpolationState +from icon4py.diffusion.metric_state import MetricState +from icon4py.diffusion.prognostic_state import PrognosticState + +diffusion: Diffusion(run_program=True) + +def diffusion_init(vct_a: np.ndarray, nrdmax: float): + """ + Instantiate and Initialize the diffusion object. + + should only accept simple fields as arguments for compatibility with the standalone + Fortran ICON Diffusion component (aka Diffusion granule) + + """ + grid = IconGrid() + + vertical_params = VerticalModelParams(vct_a=vct_a, rayleigh_damping_height=nrdmax) + config: DiffusionConfig = DiffusionConfig(grid, vertical_params) + derived_diffusion_params = DiffusionParams(config) + metric_state = MetricState() + interpolation_state = InterpolationState() + + diffusion.init(config=config, params=derived_diffusion_params, vct_a=vertical_params.vct_a,) +#ELSE IF ((diffu_type == 3 .OR. diffu_type == 5) .AND. discr_vn == 1 .AND. .NOT. diffusion_config(jg)%lsmag_3d) +# mch branch +# ELSE IF ((diffu_type == 3 .OR. diffu_type == 5) .AND. discr_vn == 1) THEN t_nh_diag%vt AND t_nh_diag%theta_v_ic are used +# matching but not mch branch +# and I couldn't find discr_vn = mo_diffusion_nml:itype_vn_diffu != 1 in any experiment +def diffusion_run(dtime:float, + linit:bool, + p_int: InterpolationState, + vn: np.ndarray, w:np.ndarray, + theta_v:np.ndarray, + exner:np.ndarray, + #vt: np.ndarray(EdgeDim, KDim) + #theta_v_ic: np.ndarray (CellDim, KHalfDim) + div_ic: np.ndarray, + hdef_ic:np.ndarray, + dwdx: np.ndarray, + dwdy:np.ndarray, + ): + diagnostic_state = DiagnosticState(hdef_ic, div_ic, dwdx, dwdy) + prognostic_state = PrognosticState(vertical_wind=w, normal_wind=vn, exner_pressure=exner, theta_v=theta_v) + if linit: + diffusion.initial_step(diagnostic_state, + prognostic_state, + dtime, + ) + else: + diffusion.time_step(diagnostic_state, + prognostic_state, + dtime) + + + + + + +class DuplicateInitializationException(Exception): + """ Raised if the component is already initilalized""" + pass diff --git a/atm_dyn_iconam/src/icon4py/diffusion/prognostic.py b/atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py similarity index 90% rename from atm_dyn_iconam/src/icon4py/diffusion/prognostic.py rename to atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py index 5495e5004..6109a0e6b 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/prognostic.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py @@ -25,7 +25,7 @@ class PrognosticState: corresponds to ICON t_nh_prog """ - vertical_wind: Field[[CellDim, KDim], float] # w(nproma, nlevp1, nblks_c) [m/s] + vertical_wind: Field[[CellDim, KDim], float] #vertical_wind field w(nproma, nlevp1, nblks_c) [m/s] normal_wind: Field[[EdgeDim, KDim], float] # vn(nproma, nlev, nblks_e) [m/s] exner_pressure: Field[[CellDim, KDim], float] # exner(nrpoma, nlev, nblks_c) theta_v: Field[[CellDim, KDim], float] # (nproma, nlev, nlbks_c) [K] diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index b2ea47ce4..4b28fae29 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -15,7 +15,7 @@ import pytest from icon4py.common.dimension import ECVDim, KDim, VertexDim -from icon4py.diffusion.diagnostic import DiagnosticState +from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.diffusion import ( Diffusion, DiffusionConfig, @@ -25,7 +25,7 @@ from icon4py.diffusion.icon_grid import VerticalModelParams from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState -from icon4py.diffusion.prognostic import PrognosticState +from icon4py.diffusion.prognostic_state import PrognosticState from icon4py.diffusion.utils import ( _en_smag_fac_for_zero_nshift, _setup_runtime_diff_multfac_vn, @@ -332,11 +332,32 @@ def test_verify_diffusion_init_against_first_regular_savepoint( r04b09_diffusion_config, ): config = r04b09_diffusion_config + additional_parameters = DiffusionParams(config) + savepoint = savepoint_init vct_a = savepoint.vct_a() - additional_parameters = DiffusionParams(config) - diffusion = Diffusion(config, additional_parameters, vct_a) + grg = savepoint.geofac_grg() + interpolation_state = InterpolationState( + e_bln_c_s=savepoint.e_bln_c_s(), + rbf_coeff_1=savepoint.rbf_vec_coeff_v1(), + rbf_coeff_2=savepoint.rbf_vec_coeff_v2(), + geofac_div=savepoint.geofac_div(), + geofac_n2s=savepoint.geofac_n2s(), + geofac_grg_x=grg[0], + geofac_grg_y=grg[1], + nudgecoeff_e=savepoint.nudgecoeff_e(), + ) + metric_state = MetricState( + mask_hdiff=savepoint.mask_diff(), + theta_ref_mc=savepoint.theta_ref_mc(), + wgtfac_c=savepoint.wgtfac_c(), + zd_intcoef=savepoint.zd_intcoef(), + zd_vertidx=savepoint.zd_vertidx(), + zd_diffcoef=savepoint.zd_diffcoef(), + ) + diffusion = Diffusion() + diffusion.init(config, additional_parameters, vct_a, metric_state, interpolation_state) _verify_init_values_against_savepoint(savepoint, diffusion) @@ -347,11 +368,32 @@ def test_verify_diffusion_init_against_other_regular_savepoint( r04b09_diffusion_config, savepoint_init ): config = r04b09_diffusion_config + additional_parameters = DiffusionParams(config) + savepoint = savepoint_init vct_a = savepoint.vct_a() + grg = savepoint.geofac_grg() + interpolation_state = InterpolationState( + e_bln_c_s=savepoint.e_bln_c_s(), + rbf_coeff_1=savepoint.rbf_vec_coeff_v1(), + rbf_coeff_2=savepoint.rbf_vec_coeff_v2(), + geofac_div=savepoint.geofac_div(), + geofac_n2s=savepoint.geofac_n2s(), + geofac_grg_x=grg[0], + geofac_grg_y=grg[1], + nudgecoeff_e=savepoint.nudgecoeff_e(), + ) + metric_state = MetricState( + mask_hdiff=savepoint.mask_diff(), + theta_ref_mc=savepoint.theta_ref_mc(), + wgtfac_c=savepoint.wgtfac_c(), + zd_intcoef=savepoint.zd_intcoef(), + zd_vertidx=savepoint.zd_vertidx(), + zd_diffcoef=savepoint.zd_diffcoef(), + ) - additional_parameters = DiffusionParams(config) - diffusion = Diffusion(config, additional_parameters, vct_a) + diffusion = Diffusion() + diffusion.init(config, additional_parameters, vct_a, metric_state, interpolation_state) _verify_init_values_against_savepoint(savepoint, diffusion) @@ -371,19 +413,7 @@ def test_run_diffusion_single_step(savepoint_init, savepoint_exit, icon_grid): additional_parameters = DiffusionParams(config) - diffusion = Diffusion(config, additional_parameters, vct_a, run_program=False) - - diagnostic_state = DiagnosticState( - hdef_ic=sp.hdef_ic(), div_ic=sp.div_ic(), dwdx=sp.dwdx(), dwdy=sp.dwdy() - ) - prognostic_state = PrognosticState( - vertical_wind=sp.w(), - normal_wind=sp.vn(), - exner_pressure=sp.exner(), - theta_v=sp.theta_v(), - ) grg = sp.geofac_grg() - interpolation_state = InterpolationState( e_bln_c_s=sp.e_bln_c_s(), rbf_coeff_1=sp.rbf_vec_coeff_v1(), @@ -403,6 +433,20 @@ def test_run_diffusion_single_step(savepoint_init, savepoint_exit, icon_grid): zd_vertidx=sp.zd_vertidx(), zd_diffcoef=sp.zd_diffcoef(), ) + + diffusion = Diffusion(run_program=False) + diffusion.init(config, additional_parameters, vct_a, metric_state, interpolation_state) + + diagnostic_state = DiagnosticState( + hdef_ic=sp.hdef_ic(), div_ic=sp.div_ic(), dwdx=sp.dwdx(), dwdy=sp.dwdy() + ) + prognostic_state = PrognosticState( + vertical_wind=sp.w(), + normal_wind=sp.vn(), + exner_pressure=sp.exner(), + theta_v=sp.theta_v(), + ) + dtime = sp.get_metadata("dtime").get("dtime") orientation = sp.tangent_orientation() @@ -517,12 +561,12 @@ def test_diffusion_five_steps( ), ) additional_parameters = DiffusionParams(config) - diffusion = Diffusion(config, additional_parameters, sp.vct_a()) + + diffusion = Diffusion() + diffusion.init(config, additional_parameters, sp.vct_a(), metric_state, interpolation_state) diffusion.initial_step( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, - metric_state=metric_state, - interpolation_state=interpolation_state, dtime=dtime, tangent_orientation=orientation, inverse_primal_edge_lengths=inverse_primal_edge_lengths, @@ -537,8 +581,6 @@ def test_diffusion_five_steps( diffusion.time_step( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, - metric_state=metric_state, - interpolation_state=interpolation_state, dtime=dtime, tangent_orientation=orientation, inverse_primal_edge_lengths=inverse_primal_edge_lengths, From 97ee6830bb1287d04b1a132e160a868d36a54429 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 25 Jan 2023 15:08:59 +0100 Subject: [PATCH 063/263] - disentangle grid and vertical params from config - add dataclass for Edge and Cell based properties --- .flake8 | 1 + .../src/icon4py/diffusion/diffusion.py | 77 +++++----- .../icon4py/diffusion/diffusion_wrapper.py | 83 +++++++---- .../src/icon4py/diffusion/horizontal.py | 121 +++++++++++++--- .../src/icon4py/diffusion/icon_grid.py | 29 ++-- .../icon4py/diffusion/interpolation_state.py | 2 +- .../src/icon4py/diffusion/prognostic_state.py | 4 +- atm_dyn_iconam/tests/conftest.py | 38 ++--- atm_dyn_iconam/tests/test_diffusion.py | 131 ++++++++++++------ atm_dyn_iconam/tests/test_vertical.py | 4 +- .../src/icon4py/testutils/serialbox_utils.py | 3 - 11 files changed, 314 insertions(+), 179 deletions(-) diff --git a/.flake8 b/.flake8 index 1caa5971d..4fff82ba1 100644 --- a/.flake8 +++ b/.flake8 @@ -37,3 +37,4 @@ rst-roles = per-file-ignores = pyutils/src/icon4py/pyutils/icochainsize.py:E800 + atm_dyn_iconam/src/icon4py/diffusion/horizontal.py:W605 diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 2e9d95cd3..e8985e323 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -15,7 +15,7 @@ import math import sys from collections import namedtuple -from typing import Final, Tuple, Optional +from typing import Final, Optional, Tuple import numpy as np from functional.common import Dimension @@ -94,8 +94,6 @@ class DiffusionConfig: def __init__( self, - grid: IconGrid, - vertical_params: VerticalModelParams, diffusion_type: int = 5, hdiff_w=True, hdiff_vn=True, @@ -114,8 +112,6 @@ def __init__( max_nudging_coeff: float = 0.02, nudging_decay_rate: float = 2.0, ): - self.grid = grid - self.vertical_params = vertical_params # parameters from namelist diffusion_nml self.diffusion_type: int = diffusion_type @@ -135,7 +131,8 @@ def __init__( """ self.apply_to_vertical_wind: bool = hdiff_w - """Ff True, apply diffusion on the vertical wind field + """ + If True, apply diffusion on the vertical wind field Called `lhdiff_w` in in mo_diffusion_nml.f90 """ @@ -270,9 +267,6 @@ def _validate(self): "diffusion` is implemented" ) - if not self.grid.limited_area: - raise NotImplementedError("Only limited area mode is implemented") - if self.diffusion_type < 0: self.apply_to_temperature = False self.apply_to_horizontal_wind = False @@ -395,7 +389,7 @@ def mo_nh_diffusion_stencil_15_numpy( class Diffusion: """Class that configures diffusion and does one diffusion step.""" - def __init__(self, run_program = True): + def __init__(self, run_program=True): self._initialized = False self.rd_o_cvd: float = GAS_CONSTANT_DRY_AIR / (CPD - GAS_CONSTANT_DRY_AIR) @@ -404,38 +398,40 @@ def __init__(self, run_program = True): ) # threshold temperature deviation from neighboring grid points hat activates extra diffusion against runaway cooling self._run_program = run_program self.grid: Optional[IconGrid] = None - self.config:Optional[DiffusionConfig] = None - self.params:Optional[DiffusionParams] = None + self.config: Optional[DiffusionConfig] = None + self.params: Optional[DiffusionParams] = None + self.vertical_params: Optional[VerticalModelParams] = None self.interpolation_state = None self.metric_state = None - self.diff_multfac_w: Optional[float] = None self.smag_offset: Optional[float] = None self.fac_bdydiff_v: Optional[float] = None self.bdy_diff: Optional[float] = None self.nudgezone_diff: Optional[float] = None - self.diff_multfac_n2w:Optional[float] = None - self._allocate_local_fields() + self.diff_multfac_n2w: Optional[float] = None def init( self, + grid: IconGrid, config: DiffusionConfig, params: DiffusionParams, - vct_a: Field[[KDim], float], + vertical_params: VerticalModelParams, metric_state: MetricState, - interpolation_state: InterpolationState + interpolation_state: InterpolationState, ): """ Initialize Diffusion granule with configuration. calculates all local fields that are used in diffusion within the time loop """ - - self.params: DiffusionParams = params self.config: DiffusionConfig = config + self.params: DiffusionParams = params + self.grid = grid + self.vertical_params = vertical_params self.metric_state: MetricState = metric_state self.interpolation_state: InterpolationState = interpolation_state - self.grid = config.grid + + self._allocate_local_fields() self.nudgezone_diff: float = 0.04 / ( config.nudge_max_coeff + sys.float_info.epsilon @@ -448,7 +444,6 @@ def init( else 1.0 / config.velocity_boundary_diffusion_denominator ) - self.smag_offset: float = 0.25 * params.K4 * config.substep_as_float() self.diff_multfac_w: float = min( 1.0 / 48.0, params.K4W * config.substep_as_float() @@ -459,7 +454,7 @@ def init( config.substep_as_float(), *params.smagorinski_factor, *params.smagorinski_height, - vct_a, + self.vertical_params.physical_heights, self.diff_multfac_vn, self.smag_limit, self.enh_smag_fac, @@ -467,10 +462,10 @@ def init( ) self.diff_multfac_n2w = init_nabla2_factor_in_upper_damping_zone( - k_size=config.grid.n_lev(), + k_size=self.grid.n_lev(), nshift=0, - physical_heights=np.asarray(vct_a), - nrdmax=self.config.vertical_params.index_of_damping_height, + physical_heights=np.asarray(self.vertical_params.physical_heights), + nrdmax=self.vertical_params.index_of_damping_layer, ) self._initialized = True @@ -629,7 +624,7 @@ def time_step( tangent_orientation=tangent_orientation, inverse_primal_edge_lengths=inverse_primal_edge_lengths, inverse_dual_edge_length=inverse_dual_edge_length, - inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, + inverse_vertex_vertex_lengths=inverse_vertical_vertex_lengths, primal_normal_vert=primal_normal_vert, dual_normal_vert=dual_normal_vert, edge_areas=edge_areas, @@ -705,7 +700,7 @@ def time_step( vertex_startindex_lb_plus1=vertex_startindex_lb_plus1, vertex_endindex_local=vertex_endindex_local, vertex_endindex_local_minus1=vertex_endindex_local_minus1, - index_of_damping_height=self.config.vertical_params.index_of_damping_height, + index_of_damping_height=self.config.vertical_params._index_of_damping_height, nlev=self.grid.n_lev(), boundary_diffusion_start_index_edges=self.params.boundary_diffusion_start_index_edges, offset_provider={ @@ -730,14 +725,18 @@ def _do_diffusion_step( diagnostic_state: DiagnosticState, prognostic_state: PrognosticState, dtime: float, - tangent_orientation: Field[[EdgeDim], float], - inverse_primal_edge_lengths: Field[[EdgeDim], float], - inverse_dual_edge_length: Field[[EdgeDim], float], - inverse_vertical_vertex_lengths: Field[[EdgeDim], float], - primal_normal_vert: Tuple[Field[[ECVDim], float], Field[[ECVDim], float]], - dual_normal_vert: Tuple[Field[[ECVDim], float], Field[[ECVDim], float]], - edge_areas: Field[[EdgeDim], float], - cell_areas: Field[[CellDim], float], + tangent_orientation: Field[[EdgeDim], float], # from p_patch%edges + inverse_primal_edge_lengths: Field[[EdgeDim], float], # from p_patch%edges + inverse_dual_edge_length: Field[[EdgeDim], float], # from p_patch%edges + inverse_vertex_vertex_lengths: Field[[EdgeDim], float], # from p_patch%edges + primal_normal_vert: Tuple[ + Field[[ECVDim], float], Field[[ECVDim], float] + ], # from p_patch%edges + dual_normal_vert: Tuple[ + Field[[ECVDim], float], Field[[ECVDim], float] + ], # from p_patch%edges + edge_areas: Field[[EdgeDim], float], # from p_patch%edges + cell_areas: Field[[CellDim], float], # from p_atch%cells diff_multfac_vn: Field[[KDim], float], smag_limit: Field[[KDim], float], smag_offset: float, @@ -752,7 +751,7 @@ def _do_diffusion_step( tangent_orientation: inverse_primal_edge_lengths: inverse_dual_edge_length: - inverse_vertical_vertex_lengths: + inverse_vertex_vertex_lengths: primal_normal_vert: dual_normal_vert: edge_areas: @@ -846,7 +845,7 @@ def _do_diffusion_step( diff_multfac_smag=self.diff_multfac_smag, tangent_orientation=tangent_orientation, inv_primal_edge_length=inverse_primal_edge_lengths, - inv_vert_vert_length=inverse_vertical_vertex_lengths, + inv_vert_vert_length=inverse_vertex_vertex_lengths, u_vert=self.u_vert, v_vert=self.v_vert, primal_normal_vert_x=primal_normal_vert[0], @@ -911,7 +910,7 @@ def _do_diffusion_step( primal_normal_vert[0], primal_normal_vert[1], self.z_nabla2_e, - inverse_vertical_vertex_lengths, + inverse_vertex_vertex_lengths, inverse_primal_edge_lengths, edge_areas, self.kh_smag_e, @@ -949,7 +948,7 @@ def _do_diffusion_step( diff_multfac_n2w=self.diff_multfac_n2w, vert_idx=self.vertical_index, horz_idx=self.horizontal_cell_index, - nrdmax=self.config.vertical_params.index_of_damping_height, + nrdmax=self.config.vertical_params._index_of_damping_height, interior_idx=cell_start_interior, # h end index for stencil_09 and stencil_10 halo_idx=cell_end_local, # h end index for stencil_09 and stencil_10, horizontal_start=cell_start_nudging, # h start index for stencil_07 and stencil_08 diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py index ae5a97045..216478027 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py @@ -1,16 +1,34 @@ -from typing import Optional +# 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 +# flake8: noqa import numpy as np from icon4py.diffusion.diagnostic_state import DiagnosticState -from icon4py.diffusion.diffusion import Diffusion, DiffusionConfig, DiffusionParams +from icon4py.diffusion.diffusion import ( + Diffusion, + DiffusionConfig, + DiffusionParams, +) from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.prognostic_state import PrognosticState + diffusion: Diffusion(run_program=True) + def diffusion_init(vct_a: np.ndarray, nrdmax: float): """ Instantiate and Initialize the diffusion object. @@ -27,42 +45,47 @@ def diffusion_init(vct_a: np.ndarray, nrdmax: float): metric_state = MetricState() interpolation_state = InterpolationState() - diffusion.init(config=config, params=derived_diffusion_params, vct_a=vertical_params.vct_a,) -#ELSE IF ((diffu_type == 3 .OR. diffu_type == 5) .AND. discr_vn == 1 .AND. .NOT. diffusion_config(jg)%lsmag_3d) + diffusion.init( + config=config, + params=derived_diffusion_params, + vct_a=vertical_params._vct_a, + ) + + +# ELSE IF ((diffu_type == 3 .OR. diffu_type == 5) .AND. discr_vn == 1 .AND. .NOT. diffusion_config(jg)%lsmag_3d) # mch branch # ELSE IF ((diffu_type == 3 .OR. diffu_type == 5) .AND. discr_vn == 1) THEN t_nh_diag%vt AND t_nh_diag%theta_v_ic are used # matching but not mch branch # and I couldn't find discr_vn = mo_diffusion_nml:itype_vn_diffu != 1 in any experiment -def diffusion_run(dtime:float, - linit:bool, - p_int: InterpolationState, - vn: np.ndarray, w:np.ndarray, - theta_v:np.ndarray, - exner:np.ndarray, - #vt: np.ndarray(EdgeDim, KDim) - #theta_v_ic: np.ndarray (CellDim, KHalfDim) - div_ic: np.ndarray, - hdef_ic:np.ndarray, - dwdx: np.ndarray, - dwdy:np.ndarray, - ): +def diffusion_run( + dtime: float, + linit: bool, + vn: np.ndarray, + w: np.ndarray, + theta_v: np.ndarray, + exner: np.ndarray, + # vt: np.ndarray(EdgeDim, KDim) # not used in mch_ch_r04b09 see above + # theta_v_ic: np.ndarray (CellDim, KHalfDim) # not used in mch_ch_r04b09 see above + div_ic: np.ndarray, + hdef_ic: np.ndarray, + dwdx: np.ndarray, + dwdy: np.ndarray, +): diagnostic_state = DiagnosticState(hdef_ic, div_ic, dwdx, dwdy) - prognostic_state = PrognosticState(vertical_wind=w, normal_wind=vn, exner_pressure=exner, theta_v=theta_v) + prognostic_state = PrognosticState( + vertical_wind=w, normal_wind=vn, exner_pressure=exner, theta_v=theta_v + ) if linit: - diffusion.initial_step(diagnostic_state, - prognostic_state, - dtime, - ) + diffusion.initial_step( + diagnostic_state, + prognostic_state, + dtime, + ) else: - diffusion.time_step(diagnostic_state, - prognostic_state, - dtime) - - - - + diffusion.time_step(diagnostic_state, prognostic_state, dtime) class DuplicateInitializationException(Exception): - """ Raised if the component is already initilalized""" + """Raised if the component is already initilalized.""" + pass diff --git a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py index 36df836bb..b5b95d5b1 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py @@ -10,11 +10,13 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -from typing import Final +from dataclasses import dataclass +from typing import Final, Tuple -from functional.common import Dimension +from functional.common import Dimension, Field from icon4py.common import dimension +from icon4py.common.dimension import CellDim, ECVDim, EdgeDim class HorizontalMarkerIndex: @@ -106,20 +108,107 @@ def end(cls, dim: Dimension) -> int: return cls._END_VERTICES -class HorizontalMeshConfig: - def __init__(self, num_vertices: int, num_edges: int, num_cells: int): - self._num_vertices = num_vertices - self._num_edges = num_edges - self._num_cells = num_cells +@dataclass(frozen=True) +class HorizontalMeshSize: + num_vertices: int + num_edges: int + num_cells: int - @property - def num_vertices(self): - return self._num_vertices - @property - def num_edges(self): - return self._num_edges +@dataclass(frozen=True) +class EdgeParams: + tangent_orientation: Field[[EdgeDim], float] + """ + Orientation of vector product of the edge and the adjacent cell centers + v3 + / \ + / \ + / c1 \ + / | \ + v1----|--->v2 + \ | / + \ v / + \ c2 / + \ / + v4 + +1 or -1 depending on whether the vector product of + (v2-v1) x (c2-c1) points outside (+) or inside (-) the sphere - @property - def num_cells(self): - return self._num_cells + defined in ICON in mo_model_domain.f90:t_grid_edges%tangent_orientation + """ + + primal_edge_lengths: Field[[EdgeDim], float] + """ + Length of the triangle edge. + + defined int ICON in mo_model_domain.f90:t_grid_edges%primal_edge_length + """ + + inverse_primal_edge_lengths: Field[[EdgeDim], float] + """ + Inverse of the triangle edge length: 1.0/primal_edge_length. + + defined int ICON in mo_model_domain.f90:t_grid_edges%inv_primal_edge_length + """ + + dual_edge_lengths: Field[[EdgeDim], float] + """ + Length of the hexagon/pentagon edge. + + defined int ICON in mo_model_domain.f90:t_grid_edges%dual_edge_length + """ + + inverse_dual_edge_lengths: Field[[EdgeDim], float] + """ + Inverse of hexagon/pentagon edge length: 1.0/dual_edge_length. + + defined int ICON in mo_model_domain.f90:t_grid_edges%inv_dual_edge_length + """ + + inverse_vertex_vertex_lengths: Field[[EdgeDim], float] + """ + Inverse distance between outer vertices of adjacent cells. + + v1-------- + | /| + | / | + | e | + | / | + |/ | + --------v2 + + inverse_vertex_vertex_length(e) = 1.0/|v2-v1| + + defined int ICON in mo_model_domain.f90:t_grid_edges%inv_vert_vert_length + """ + + primal_normal_vert: Tuple[Field[[ECVDim], float], Field[[ECVDim], float]] + """ + Normal of the triangle edge, projected onto the location of the vertices + + defined int ICON in mo_model_domain.f90:t_grid_edges%primal_normal_vert + """ + + dual_normal_vert: Tuple[Field[[ECVDim], float], Field[[ECVDim], float]] + """ + Tangent to the triangle edge, projected onto the location of vertices. + + defined int ICON in mo_model_domain.f90:t_grid_edges%dual_normal_vert + """ + + edge_areas: Field[[EdgeDim], float] + """ + Area of the quadrilateral (two triangles) adjacent to the edge. + + defined int ICON in mo_model_domain.f90:t_grid_edges%area_edge + """ + + +@dataclass(frozen=True) +class CellParams: + area: Field[[CellDim], float] + """ + Area of a cell. + + defined int ICON in mo_model_domain.f90:t_grid_cells%area + """ diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index 979fd70bf..cf6c89434 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -16,13 +16,10 @@ import numpy as np from functional.common import Dimension, DimensionKind, Field from functional.ffront.fbuiltins import int32 -from functional.iterator.embedded import ( - NeighborTableOffsetProvider, - np_as_located_field, -) +from functional.iterator.embedded import NeighborTableOffsetProvider from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim -from icon4py.diffusion.horizontal import HorizontalMeshConfig +from icon4py.diffusion.horizontal import HorizontalMeshSize class VerticalMeshConfig: @@ -37,7 +34,7 @@ def num_lev(self) -> int: class MeshConfig: def __init__( self, - horizontal_config: HorizontalMeshConfig, + horizontal_config: HorizontalMeshSize, vertical_config: VerticalMeshConfig, limited_area=True, ): @@ -178,22 +175,24 @@ def get_v2e_connectivity(self): class VerticalModelParams: - def __init__(self, vct_a: np.ndarray, rayleigh_damping_height: float): + def __init__(self, vct_a: Field[[KDim], float], rayleigh_damping_height: float): """ - Contains physical parameters defined on the grid. + Contains vertical physical parameters defined on the grid. Args: vct_a: field containing the physical heights of the k level rayleigh_damping_height: height of rayleigh damping in [m] mo_nonhydro_nml """ self.rayleigh_damping_height = rayleigh_damping_height - self.vct_a = vct_a - self.index_of_damping_height = int32( - np.argmax(np.where(np.asarray(self.vct_a) >= self.rayleigh_damping_height)) + self._vct_a = vct_a + self._index_of_damping_height = int32( + np.argmax(np.where(np.asarray(self._vct_a) >= self.rayleigh_damping_height)) ) - def get_index_of_damping_layer(self): - return self.index_of_damping_height + @property + def index_of_damping_layer(self): + return self._index_of_damping_height - def get_physical_heights(self) -> Field[[KDim], float]: - return np_as_located_field(KDim)(self.vct_a) + @property + def physical_heights(self) -> Field[[KDim], float]: + return self._vct_a diff --git a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py index 2d9a4035c..a4cbdd388 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py @@ -30,7 +30,7 @@ class InterpolationState: """ represents the ICON interpolation state. - TODO: keep? does this state make sense at all? + TODO [ml]: keep? does this state make sense at all? """ e_bln_c_s: Field[ diff --git a/atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py b/atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py index 6109a0e6b..d82aa083b 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py @@ -25,7 +25,9 @@ class PrognosticState: corresponds to ICON t_nh_prog """ - vertical_wind: Field[[CellDim, KDim], float] #vertical_wind field w(nproma, nlevp1, nblks_c) [m/s] + vertical_wind: Field[ + [CellDim, KDim], float + ] # vertical_wind field w(nproma, nlevp1, nblks_c) [m/s] normal_wind: Field[[EdgeDim, KDim], float] # vn(nproma, nlev, nblks_e) [m/s] exner_pressure: Field[[CellDim, KDim], float] # exner(nrpoma, nlev, nblks_c) theta_v: Field[[CellDim, KDim], float] # (nproma, nlev, nlbks_c) [K] diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index ce77d7d33..7e48620aa 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -29,13 +29,8 @@ VertexDim, ) from icon4py.diffusion.diffusion import DiffusionConfig -from icon4py.diffusion.horizontal import HorizontalMeshConfig -from icon4py.diffusion.icon_grid import ( - IconGrid, - MeshConfig, - VerticalMeshConfig, - VerticalModelParams, -) +from icon4py.diffusion.horizontal import HorizontalMeshSize +from icon4py.diffusion.icon_grid import IconGrid, MeshConfig, VerticalMeshConfig from icon4py.testutils.serialbox_utils import IconSerialDataProvider @@ -107,7 +102,7 @@ def savepoint_init(setup_icon_data, linit, step_date_init): linit flag can be set by overriding the 'linit' fixture """ sp = IconSerialDataProvider( - "icon_diffusion_init", extracted_path, True + "icon_diffusion_init", str(extracted_path), True ).from_savepoint_init(linit=linit, date=step_date_init) return sp @@ -146,7 +141,7 @@ def icon_grid(savepoint_init): edge_ends = sp.edge_end_index() config = MeshConfig( - HorizontalMeshConfig( + HorizontalMeshSize( num_vertices=sp_meta["nproma"], # or rather "num_vert" num_cells=sp_meta["nproma"], # or rather "num_cells" num_edges=sp_meta["nproma"], # or rather "num_edges" @@ -178,27 +173,7 @@ def r04b09_diffusion_config(setup_icon_data) -> DiffusionConfig: Set values to the ones used in the MCH_CH_r04b09_dsl experiment where they differ from the default. """ - sp = IconSerialDataProvider( - "icon_diffusion_init", extracted_path, True - ).from_savepoint_init(linit=True, date="2021-06-20T12:00:10.000") - nproma = sp.get_metadata("nproma")["nproma"] - num_lev = sp.get_metadata("nlev")["nlev"] - horizontal_config = HorizontalMeshConfig( - num_vertices=nproma, num_cells=nproma, num_edges=nproma - ) - vertical_config = VerticalMeshConfig(num_lev=num_lev) - - grid = IconGrid().with_config( - MeshConfig(horizontal_config=horizontal_config, vertical_config=vertical_config) - ) - - verticalParams = VerticalModelParams( - rayleigh_damping_height=12500, vct_a=sp.vct_a() - ) - return DiffusionConfig( - grid=grid, - vertical_params=verticalParams, diffusion_type=5, hdiff_w=True, hdiff_vn=True, @@ -211,3 +186,8 @@ def r04b09_diffusion_config(setup_icon_data) -> DiffusionConfig: velocity_boundary_diffusion_denom=150.0, max_nudging_coeff=0.075, ) + + +@pytest.fixture +def damping_height(): + return 12500 diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 4b28fae29..d816a6c45 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -16,12 +16,7 @@ from icon4py.common.dimension import ECVDim, KDim, VertexDim from icon4py.diffusion.diagnostic_state import DiagnosticState -from icon4py.diffusion.diffusion import ( - Diffusion, - DiffusionConfig, - DiffusionParams, - VectorTuple, -) +from icon4py.diffusion.diffusion import Diffusion, DiffusionParams, VectorTuple from icon4py.diffusion.icon_grid import VerticalModelParams from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState @@ -200,7 +195,6 @@ def test_smagorinski_factor_for_diffusion_type_4(r04b09_diffusion_config): assert params.smagorinski_height is None -@pytest.mark.datatest def test_smagorinski_heights_diffusion_type_5_are_consistent( r04b09_diffusion_config, ): @@ -218,7 +212,6 @@ def test_smagorinski_heights_diffusion_type_5_are_consistent( assert params.smagorinski_height[2] != params.smagorinski_height[3] -@pytest.mark.datatest def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): config = r04b09_diffusion_config params = DiffusionParams(config) @@ -228,18 +221,49 @@ def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): @pytest.mark.datatest -def test_diffusion_init(savepoint_init, r04b09_diffusion_config, step_date_init): +def test_diffusion_init( + savepoint_init, icon_grid, r04b09_diffusion_config, step_date_init, damping_height +): savepoint = savepoint_init - vct_a = savepoint.vct_a() config = r04b09_diffusion_config + additional_parameters = DiffusionParams(config) + vertical_params = VerticalModelParams(savepoint.vct_a(), damping_height) + meta = savepoint.get_metadata("nlev", "linit", "date") assert meta["nlev"] == 65 assert meta["linit"] is False assert meta["date"] == step_date_init - additional_parameters = DiffusionParams(config) - diffusion = Diffusion(config, additional_parameters, vct_a) + grg = savepoint.geofac_grg() + interpolation_state = InterpolationState( + e_bln_c_s=savepoint.e_bln_c_s(), + rbf_coeff_1=savepoint.rbf_vec_coeff_v1(), + rbf_coeff_2=savepoint.rbf_vec_coeff_v2(), + geofac_div=savepoint.geofac_div(), + geofac_n2s=savepoint.geofac_n2s(), + geofac_grg_x=grg[0], + geofac_grg_y=grg[1], + nudgecoeff_e=savepoint.nudgecoeff_e(), + ) + metric_state = MetricState( + mask_hdiff=savepoint.mask_diff(), + theta_ref_mc=savepoint.theta_ref_mc(), + wgtfac_c=savepoint.wgtfac_c(), + zd_intcoef=savepoint.zd_intcoef(), + zd_vertidx=savepoint.zd_vertidx(), + zd_diffcoef=savepoint.zd_diffcoef(), + ) + + diffusion = Diffusion() + diffusion.init( + grid=icon_grid, + vertical_params=vertical_params, + config=config, + params=additional_parameters, + metric_state=metric_state, + interpolation_state=interpolation_state, + ) assert diffusion.diff_multfac_w == min( 1.0 / 48.0, additional_parameters.K4W * config.substep_as_float() @@ -250,7 +274,7 @@ def test_diffusion_init(savepoint_init, r04b09_diffusion_config, step_date_init) assert np.allclose(0.0, np.asarray(diffusion.kh_smag_ec)) assert np.allclose(0.0, np.asarray(diffusion.kh_smag_e)) - shape_k = (config.grid.n_lev(),) + shape_k = (icon_grid.n_lev(),) expected_smag_limit = smag_limit_numpy( diff_multfac_vn_numpy, shape_k, @@ -271,7 +295,7 @@ def test_diffusion_init(savepoint_init, r04b09_diffusion_config, step_date_init) expected_enh_smag_fac = enhanced_smagorinski_factor_numpy( additional_parameters.smagorinski_factor, additional_parameters.smagorinski_height, - vct_a, + savepoint.vct_a(), ) assert np.allclose(expected_enh_smag_fac, np.asarray(diffusion.enh_smag_fac)) @@ -301,8 +325,7 @@ def _verify_init_values_against_savepoint( @pytest.mark.datatest @pytest.mark.parametrize("linit", [True]) def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( - savepoint_init, - r04b09_diffusion_config, + savepoint_init, r04b09_diffusion_config, icon_grid ): savepoint = savepoint_init config = r04b09_diffusion_config @@ -312,8 +335,8 @@ def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( expected_smag_limit = savepoint.smag_limit() exptected_smag_offset = savepoint.smag_offset() - diff_multfac_vn = zero_field(config.grid, KDim) - smag_limit = zero_field(config.grid, KDim) + diff_multfac_vn = zero_field(icon_grid, KDim) + smag_limit = zero_field(icon_grid, KDim) setup_fields_for_initial_step( params.K4, config.hdiff_efdt_ratio, @@ -328,8 +351,7 @@ def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( @pytest.mark.datatest def test_verify_diffusion_init_against_first_regular_savepoint( - savepoint_init, - r04b09_diffusion_config, + savepoint_init, r04b09_diffusion_config, icon_grid, damping_height ): config = r04b09_diffusion_config additional_parameters = DiffusionParams(config) @@ -357,7 +379,14 @@ def test_verify_diffusion_init_against_first_regular_savepoint( zd_diffcoef=savepoint.zd_diffcoef(), ) diffusion = Diffusion() - diffusion.init(config, additional_parameters, vct_a, metric_state, interpolation_state) + diffusion.init( + config=config, + grid=icon_grid, + params=additional_parameters, + vertical_params=VerticalModelParams(vct_a, damping_height), + metric_state=metric_state, + interpolation_state=interpolation_state, + ) _verify_init_values_against_savepoint(savepoint, diffusion) @@ -365,13 +394,13 @@ def test_verify_diffusion_init_against_first_regular_savepoint( @pytest.mark.datatest @pytest.mark.parametrize("step_date_init", ["2021-06-20T12:00:50.000"]) def test_verify_diffusion_init_against_other_regular_savepoint( - r04b09_diffusion_config, savepoint_init + r04b09_diffusion_config, icon_grid, savepoint_init, damping_height ): config = r04b09_diffusion_config additional_parameters = DiffusionParams(config) savepoint = savepoint_init - vct_a = savepoint.vct_a() + vertical_params = VerticalModelParams(savepoint.vct_a(), damping_height) grg = savepoint.geofac_grg() interpolation_state = InterpolationState( e_bln_c_s=savepoint.e_bln_c_s(), @@ -393,25 +422,30 @@ def test_verify_diffusion_init_against_other_regular_savepoint( ) diffusion = Diffusion() - diffusion.init(config, additional_parameters, vct_a, metric_state, interpolation_state) + diffusion.init( + icon_grid, + config, + additional_parameters, + vertical_params, + metric_state, + interpolation_state, + ) _verify_init_values_against_savepoint(savepoint, diffusion) @pytest.mark.datatest @pytest.mark.skip -def test_run_diffusion_single_step(savepoint_init, savepoint_exit, icon_grid): +def test_run_diffusion_single_step( + savepoint_init, savepoint_exit, icon_grid, r04b09_diffusion_config, damping_height +): sp = savepoint_init - vct_a = sp.vct_a() - - config = DiffusionConfig( - icon_grid, - vertical_params=VerticalModelParams( - vct_a=vct_a, rayleigh_damping_height=12500.0 - ), + vct_a = sp._vct_a() + vertical_params = VerticalModelParams( + vct_a=vct_a, rayleigh_damping_height=damping_height ) - additional_parameters = DiffusionParams(config) + additional_parameters = DiffusionParams(r04b09_diffusion_config) grg = sp.geofac_grg() interpolation_state = InterpolationState( @@ -435,7 +469,14 @@ def test_run_diffusion_single_step(savepoint_init, savepoint_exit, icon_grid): ) diffusion = Diffusion(run_program=False) - diffusion.init(config, additional_parameters, vct_a, metric_state, interpolation_state) + diffusion.init( + grid=icon_grid, + config=r04b09_diffusion_config, + params=additional_parameters, + vertical_params=vertical_params, + metric_state=metric_state, + interpolation_state=interpolation_state, + ) diagnostic_state = DiagnosticState( hdef_ic=sp.hdef_ic(), div_ic=sp.div_ic(), dwdx=sp.dwdx(), dwdy=sp.dwdy() @@ -467,8 +508,6 @@ def test_run_diffusion_single_step(savepoint_init, savepoint_exit, icon_grid): diffusion.time_step( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, - metric_state=metric_state, - interpolation_state=interpolation_state, dtime=dtime, tangent_orientation=orientation, inverse_primal_edge_lengths=inverse_primal_edge_lengths, @@ -500,6 +539,7 @@ def test_run_diffusion_single_step(savepoint_init, savepoint_exit, icon_grid): @pytest.mark.skip @pytest.mark.datatest def test_diffusion_five_steps( + r04b09_diffusion_config, icon_grid, savepoint_init, savepoint_exit, @@ -554,16 +594,21 @@ def test_diffusion_five_steps( edge_areas = sp.edge_areas() cell_areas = sp.cell_areas() - config = DiffusionConfig( - icon_grid, - vertical_params=VerticalModelParams( - vct_a=(sp.vct_a()), rayleigh_damping_height=12500.0 - ), + vertical_params = VerticalModelParams( + vct_a=(sp._vct_a()), rayleigh_damping_height=12500.0 ) - additional_parameters = DiffusionParams(config) + + additional_parameters = DiffusionParams(r04b09_diffusion_config) diffusion = Diffusion() - diffusion.init(config, additional_parameters, sp.vct_a(), metric_state, interpolation_state) + diffusion.init( + grid=icon_grid, + config=r04b09_diffusion_config, + params=additional_parameters, + vertical_params=vertical_params, + metric_state=metric_state, + interpolation_state=interpolation_state, + ) diffusion.initial_step( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, diff --git a/atm_dyn_iconam/tests/test_vertical.py b/atm_dyn_iconam/tests/test_vertical.py index dd417daa3..7da77063b 100644 --- a/atm_dyn_iconam/tests/test_vertical.py +++ b/atm_dyn_iconam/tests/test_vertical.py @@ -28,7 +28,7 @@ def test_nrdmax_calculation(max_h, damping, delta): vct_a = vct_a[::-1] vertical_params = VerticalModelParams(rayleigh_damping_height=damping, vct_a=vct_a) assert ( - vertical_params.get_index_of_damping_layer() + vertical_params.index_of_damping_layer == vct_a.shape[0] - math.ceil(damping / delta) - 1 ) @@ -41,7 +41,7 @@ def test_nrdmax_calculation_from_icon_input(icon_grid, savepoint_init): vertical_params = VerticalModelParams( rayleigh_damping_height=damping_height, vct_a=a ) - assert 9 == vertical_params.get_index_of_damping_layer() + assert 9 == vertical_params.index_of_damping_layer a_array = np.asarray(a) assert a_array[9] > damping_height assert a_array[10] < damping_height diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 53b5945ad..72ded87b9 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -234,9 +234,6 @@ def smag_limit(self): def diff_multfac_n2w(self): return np.squeeze(self.serializer.read("diff_multfac_n2w", self.savepoint)) - def nrdmax(self): - return self.serializer.read("nrdmax", self.savepoint)[0] - def nudgezone_diff(self) -> int: return self.serializer.read("nudgezone_diff", self.savepoint)[0] From 40de762a61ffe19e10a818e0dfd502ff06248a32 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 26 Jan 2023 10:09:32 +0100 Subject: [PATCH 064/263] add params to diffusion init --- .../icon4py/diffusion/diffusion_wrapper.py | 66 +++++++++++++++++-- .../src/icon4py/diffusion/horizontal.py | 4 +- common/src/icon4py/common/constants.py | 1 + 3 files changed, 63 insertions(+), 8 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py index 216478027..885cafc49 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py @@ -14,12 +14,16 @@ # flake8: noqa import numpy as np +from functional.common import Field +from icon4py.common.dimension import KDim, CellDim, C2E2CDim, C2EDim, V2EDim, VertexDim, C2E2CODim, \ + EdgeDim, ECVDim from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.diffusion import ( Diffusion, DiffusionConfig, DiffusionParams, ) +from icon4py.diffusion.horizontal import EdgeParams, CellParams from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState @@ -29,7 +33,34 @@ diffusion: Diffusion(run_program=True) -def diffusion_init(vct_a: np.ndarray, nrdmax: float): +def diffusion_init(vct_a:Field[[KDim], float], + nrdmax: float, # in config + theta_ref_mc:Field[[CellDim, KDim], float], + wgtfac_c:Field[[CellDim, KDim], float], + mask_hdiff: Field[[CellDim, KDim], int], + zd_vertidx: Field[[CellDim, C2E2CDim, KDim], int], + zd_diffcoef: Field[[CellDim, KDim], float], + zd_intcoef: Field[[CellDim, C2E2CDim, KDim], float], + e_bln_c_s: Field[[CellDim, C2EDim], float], + rbf_coeff_1: Field[[VertexDim, V2EDim], float], + rbf_coeff_2: Field[[VertexDim, V2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + geofac_n2s: Field[[CellDim, C2E2CODim], float], + geofac_grg_x: Field[[CellDim, C2E2CODim], float], + geofac_grg_y: Field[[CellDim, C2E2CODim], float], + nudgecoeff_e: Field[[EdgeDim], float], + tangent_orientation: Field[[EdgeDim], float], + primal_edge_lengths: Field[[EdgeDim], float], + inverse_primal_edge_lengths: Field[[EdgeDim], float], + dual_edge_lengths: Field[[EdgeDim], float], + inverse_dual_edge_lengths: Field[[EdgeDim], float], + inverse_vertex_vertex_lengths: Field[[EdgeDim], float], + primal_normal_vert: tuple[Field[[ECVDim], float], Field[[ECVDim], float]], + dual_normal_vert: tuple[Field[[ECVDim], float], Field[[ECVDim], float]], + edge_areas: Field[[EdgeDim], float], + cell_areas: Field[[CellDim], float] + + ): """ Instantiate and Initialize the diffusion object. @@ -37,18 +68,41 @@ def diffusion_init(vct_a: np.ndarray, nrdmax: float): Fortran ICON Diffusion component (aka Diffusion granule) """ - grid = IconGrid() - + grid = IconGrid() # TODO where to get this from + edges_params = EdgeParams(tangent_orientation=tangent_orientation, + primal_edge_lengths=primal_edge_lengths, + inverse_primal_edge_lengths=inverse_primal_edge_lengths, + dual_normal_vert=dual_normal_vert, + dual_edge_lengths=dual_edge_lengths, + inverse_dual_edge_lengths=inverse_dual_edge_lengths, + inverse_vertex_vertex_lengths=inverse_vertex_vertex_lengths, + primal_normal_vert=primal_normal_vert, + edge_areas=edge_areas, + ) + cell_params = CellParams(cell_areas) vertical_params = VerticalModelParams(vct_a=vct_a, rayleigh_damping_height=nrdmax) config: DiffusionConfig = DiffusionConfig(grid, vertical_params) derived_diffusion_params = DiffusionParams(config) - metric_state = MetricState() - interpolation_state = InterpolationState() + metric_state = MetricState(theta_ref_mc, wgtfac_c, mask_hdiff, zd_vertidx, zd_diffcoef, zd_intcoef) + interpolation_state = InterpolationState(e_bln_c_s, + rbf_coeff_1, + rbf_coeff_2, + geofac_div , + geofac_n2s, + geofac_grg_x, + geofac_grg_y, + nudgecoeff_e) diffusion.init( + grid= grid, + cell_params = cell_params, + edges_params = edges_params, config=config, params=derived_diffusion_params, - vct_a=vertical_params._vct_a, + vertical_params=vertical_params, + metric_state=metric_state, + interpolation_state=interpolation_state, + ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py index b5b95d5b1..8857626c2 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py @@ -182,14 +182,14 @@ class EdgeParams: defined int ICON in mo_model_domain.f90:t_grid_edges%inv_vert_vert_length """ - primal_normal_vert: Tuple[Field[[ECVDim], float], Field[[ECVDim], float]] + primal_normal_vert: tuple[Field[[ECVDim], float], Field[[ECVDim], float]] """ Normal of the triangle edge, projected onto the location of the vertices defined int ICON in mo_model_domain.f90:t_grid_edges%primal_normal_vert """ - dual_normal_vert: Tuple[Field[[ECVDim], float], Field[[ECVDim], float]] + dual_normal_vert: tuple[Field[[ECVDim], float], Field[[ECVDim], float]] """ Tangent to the triangle edge, projected onto the location of vertices. diff --git a/common/src/icon4py/common/constants.py b/common/src/icon4py/common/constants.py index 15cc0e379..0918c2a10 100644 --- a/common/src/icon4py/common/constants.py +++ b/common/src/icon4py/common/constants.py @@ -22,3 +22,4 @@ GAS_CONSTANT_WATER_VAPOR: Annotated[ float, "gas constant for water vapor [J/K/kg], rv in Icon" ] = 461.51 +GRAVITATIONAL_ACCELERATION: Annotated[float, "av. gravitational acceleration [m/s^2]"] = 9.8066 From f30f1b086517267118af9fdf425d422a04ae94be Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 1 Feb 2023 08:29:55 +0100 Subject: [PATCH 065/263] use renamed diffusion stencils in fused stencils --- ...fused_mo_nh_diffusion_stencil_01_02_03_rbf | 18 +++++------ .../fused_mo_nh_diffusion_stencil_02_03.py | 12 ++++---- .../fused_mo_nh_diffusion_stencil_04_05_06.py | 20 ++++++------- ...sed_mo_nh_diffusion_stencil_07_08_09_10.py | 30 +++++++++---------- .../fused_mo_nh_diffusion_stencil_11_12.py | 16 ++++++---- .../fused_mo_nh_diffusion_stencil_13_14.py | 12 ++++---- 6 files changed, 54 insertions(+), 54 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf index fb3ecaf12..e33e884f8 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf @@ -17,14 +17,14 @@ from functional.ffront.fbuiltins import Field from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( _mo_intp_rbf_rbf_vec_interpol_vertex, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import ( - _mo_nh_diffusion_stencil_01, +from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( + _calculate_nabla2_and_smag_coefficients_for_vn, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import ( - _mo_nh_diffusion_stencil_02, +from icon4py.atm_dyn_iconam.temporary_fields_for_turbulance_diagnostics import ( + _temporary_fields_for_turbulance_diagnostics, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import ( - _mo_nh_diffusion_stencil_03, +from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulance import ( + _calculate_diagnostics_for_turbulance, ) from icon4py.common.dimension import ( C2EDim, @@ -65,7 +65,7 @@ def _fused_mo_nh_diffusion_stencil_01_02_03_rbf( Field[[VertexDim, KDim], float], ]: - kh_smag_e, kh_smag_ec, z_nabla2_e = _mo_nh_diffusion_stencil_01( + kh_smag_e, kh_smag_ec, z_nabla2_e = _calculate_nabla2_and_smag_coefficients_for_vn( diff_multfac_smag, tangent_orientation, inv_primal_edge_length, @@ -81,11 +81,11 @@ def _fused_mo_nh_diffusion_stencil_01_02_03_rbf( smag_offset, ) - kh_c, div = _mo_nh_diffusion_stencil_02( + kh_c, div = _temporary_fields_for_turbulance_diagnostics( kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag ) - div_ic, hdef_ic = _mo_nh_diffusion_stencil_03(div, kh_c, wgtfac_c) + div_ic, hdef_ic = _calculate_diagnostics_for_turbulance(div, kh_c, wgtfac_c) u_vert, v_vert = _mo_intp_rbf_rbf_vec_interpol_vertex( z_nabla2_e, ptr_coeff_1, ptr_coeff_2 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py index 6b32d5a67..2d74f63e9 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -14,11 +14,11 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import ( - _mo_nh_diffusion_stencil_02, +from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulance import ( + _calculate_diagnostics_for_turbulance, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import ( - _mo_nh_diffusion_stencil_03, +from icon4py.atm_dyn_iconam.temporary_fields_for_turbulance_diagnostics import ( + _temporary_fields_for_turbulance_diagnostics, ) from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim @@ -32,10 +32,10 @@ def _fused_mo_nh_diffusion_stencil_02_03( diff_multfac_smag: Field[[KDim], float], wgtfac_c: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - kh_c, div = _mo_nh_diffusion_stencil_02( + kh_c, div = _temporary_fields_for_turbulance_diagnostics( kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag ) - div_ic, hdef_ic = _mo_nh_diffusion_stencil_03(div, kh_c, wgtfac_c) + div_ic, hdef_ic = _calculate_diagnostics_for_turbulance(div, kh_c, wgtfac_c) return div_ic, hdef_ic diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py index 58e334501..a2a00a427 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py @@ -14,15 +14,13 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, int32, where -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_04 import ( - _mo_nh_diffusion_stencil_04, +from icon4py.atm_dyn_iconam.apply_nabla2_and_nabla4_to_vn import ( + _apply_nabla2_and_nabla4_to_vn, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_05 import ( - _mo_nh_diffusion_stencil_05, -) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_06 import ( - _mo_nh_diffusion_stencil_06, +from icon4py.atm_dyn_iconam.apply_nabla2_to_vn_in_lateral_boundary import ( + _apply_nabla2_to_vn_in_lateral_boundary, ) +from icon4py.atm_dyn_iconam.calculate_nabla4 import _calculate_nabla4 from icon4py.common.dimension import ECVDim, EdgeDim, KDim, VertexDim @@ -46,7 +44,7 @@ def _fused_mo_nh_diffusion_stencil_04_05_06( start_2nd_nudge_line_idx_e: int32, ) -> Field[[EdgeDim, KDim], float]: - z_nabla4_e2 = _mo_nh_diffusion_stencil_04( + z_nabla4_e2 = _calculate_nabla4( u_vert, v_vert, primal_normal_vert_v1, @@ -58,7 +56,7 @@ def _fused_mo_nh_diffusion_stencil_04_05_06( vn = where( horz_idx >= start_2nd_nudge_line_idx_e, - _mo_nh_diffusion_stencil_05( + _apply_nabla2_and_nabla4_to_vn( area_edge, kh_smag_e, z_nabla2_e, @@ -68,7 +66,9 @@ def _fused_mo_nh_diffusion_stencil_04_05_06( vn, nudgezone_diff, ), - _mo_nh_diffusion_stencil_06(z_nabla2_e, area_edge, vn, fac_bdydiff_v), + _apply_nabla2_to_vn_in_lateral_boundary( + z_nabla2_e, area_edge, vn, fac_bdydiff_v + ), ) return vn diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py index a2a611a42..0849dfe56 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -14,18 +14,14 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, broadcast, int32, where -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_07 import ( - _mo_nh_diffusion_stencil_07, +from icon4py.atm_dyn_iconam.apply_nabla2_to_w import _apply_nabla2_to_w +from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import ( + _apply_nabla2_to_w_in_upper_damping_layer, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_08 import ( - _mo_nh_diffusion_stencil_08, -) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_09 import ( - _mo_nh_diffusion_stencil_09, -) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_10 import ( - _mo_nh_diffusion_stencil_10, +from icon4py.atm_dyn_iconam.calculate_horizontal_gradients_for_turbulance import ( + _calculate_horizontal_gradients_for_turbulance, ) +from icon4py.atm_dyn_iconam.calculate_nabla2_for_w import _calculate_nabla2_for_w from icon4py.common.dimension import C2E2CODim, CellDim, KDim @@ -56,17 +52,17 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( dwdx, dwdy = where( vert_idx > int32(0), - _mo_nh_diffusion_stencil_08(w_old, geofac_grg_x, geofac_grg_y), + _calculate_horizontal_gradients_for_turbulance( + w_old, geofac_grg_x, geofac_grg_y + ), (dwdx, dwdy), ) - z_nabla2_c = _mo_nh_diffusion_stencil_07(w_old, geofac_n2s) + z_nabla2_c = _calculate_nabla2_for_w(w_old, geofac_n2s) w = where( (horz_idx >= interior_idx) & (horz_idx < halo_idx), - _mo_nh_diffusion_stencil_09( - area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w - ), + _apply_nabla2_to_w(area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w), w_old, ) @@ -75,7 +71,9 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( & (vert_idx < nrdmax) & (horz_idx >= interior_idx) & (horz_idx < halo_idx), - _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, area, z_nabla2_c), + _apply_nabla2_to_w_in_upper_damping_layer( + w, diff_multfac_n2w, area, z_nabla2_c + ), w, ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py index 375f8031d..8374f29a9 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py @@ -14,11 +14,11 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_11 import ( - _mo_nh_diffusion_stencil_11, +from icon4py.atm_dyn_iconam.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( + _enhance_diffusion_coefficient_for_grid_point_cold_pools, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_12 import ( - _mo_nh_diffusion_stencil_12, +from icon4py.atm_dyn_iconam.temporary_field_for_grid_point_cold_pools_enhancement import ( + _temporary_field_for_grid_point_cold_pools_enhancement, ) from icon4py.common.dimension import CellDim, EdgeDim, KDim @@ -30,8 +30,12 @@ def _fused_mo_nh_diffusion_stencil_11_12( thresh_tdiff: float, kh_smag_e: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - enh_diffu_3d = _mo_nh_diffusion_stencil_11(theta_v, theta_ref_mc, thresh_tdiff) - kh_smag_e = _mo_nh_diffusion_stencil_12(kh_smag_e, enh_diffu_3d) + enh_diffu_3d = _temporary_field_for_grid_point_cold_pools_enhancement( + theta_v, theta_ref_mc, thresh_tdiff + ) + kh_smag_e = _enhance_diffusion_coefficient_for_grid_point_cold_pools( + kh_smag_e, enh_diffu_3d + ) return kh_smag_e diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py index 0941bbd77..1b286e0f0 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -14,11 +14,9 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_13 import ( - _mo_nh_diffusion_stencil_13, -) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_14 import ( - _mo_nh_diffusion_stencil_14, +from icon4py.atm_dyn_iconam.calculate_nabla2_for_z import _calculate_nabla2_for_z +from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( + _calculate_nabla2_of_theta, ) from icon4py.common.dimension import CEDim, CellDim, EdgeDim, KDim @@ -30,8 +28,8 @@ def _fused_mo_nh_diffusion_stencil_13_14( theta_v: Field[[CellDim, KDim], float], geofac_div: Field[[CEDim], float], ) -> Field[[CellDim, KDim], float]: - z_nabla2_e = _mo_nh_diffusion_stencil_13(kh_smag_e, inv_dual_edge_length, theta_v) - z_temp = _mo_nh_diffusion_stencil_14(z_nabla2_e, geofac_div) + z_nabla2_e = _calculate_nabla2_for_z(kh_smag_e, inv_dual_edge_length, theta_v) + z_temp = _calculate_nabla2_of_theta(z_nabla2_e, geofac_div) return z_temp From 1bf3381c26c92c1a9082097ab251a93bfe7f1813 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 1 Feb 2023 15:07:41 +0100 Subject: [PATCH 066/263] cleanup: revert to ICON naming for prognistic winds --- .../src/icon4py/diffusion/diffusion.py | 18 ++++++++--------- .../icon4py/diffusion/diffusion_program.py | 20 +++++++++---------- .../icon4py/diffusion/diffusion_wrapper.py | 2 +- .../src/icon4py/diffusion/prognostic_state.py | 6 +++--- atm_dyn_iconam/tests/test_diffusion.py | 20 ++++++++----------- base-requirements-dev.txt | 2 +- 6 files changed, 32 insertions(+), 36 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index ddea6756c..7e9f90131 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -637,8 +637,8 @@ def time_step( diagnostic_div_ic=diagnostic_state.div_ic, diagnostic_dwdx=diagnostic_state.dwdx, diagnostic_dwdy=diagnostic_state.dwdy, - prognostic_vertical_wind=prognostic_state.vertical_wind, - prognostic_normal_wind=prognostic_state.normal_wind, + prognostic_w=prognostic_state.w, + prognostic_vn=prognostic_state.vn, prognostic_exner_pressure=prognostic_state.exner_pressure, prognostic_theta_v=prognostic_state.theta_v, metric_theta_ref_mc=self.metric_state.theta_ref_mc, @@ -824,7 +824,7 @@ def _do_diffusion_step( # # 1. CALL rbf_vec_interpol_vertex mo_intp_rbf_rbf_vec_interpol_vertex( - p_e_in=prognostic_state.normal_wind, + p_e_in=prognostic_state.vn, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, ptr_coeff_2=self.interpolation_state.rbf_coeff_2, p_u_out=self.u_vert, @@ -850,7 +850,7 @@ def _do_diffusion_step( primal_normal_vert_y=primal_normal_vert[1], dual_normal_vert_x=dual_normal_vert[0], dual_normal_vert_y=dual_normal_vert[1], - vn=prognostic_state.normal_wind, + vn=prognostic_state.vn, smag_limit=smag_limit, kh_smag_e=self.kh_smag_e, kh_smag_ec=self.kh_smag_ec, @@ -868,7 +868,7 @@ def _do_diffusion_step( fused_mo_nh_diffusion_stencil_02_03( kh_smag_ec=self.kh_smag_ec, - vn=prognostic_state.normal_wind, + vn=prognostic_state.vn, e_bln_c_s=self.interpolation_state.e_bln_c_s, geofac_div=self.interpolation_state.geofac_div, diff_multfac_smag=self.diff_multfac_smag, @@ -914,12 +914,12 @@ def _do_diffusion_step( self.kh_smag_e, diff_multfac_vn, self.interpolation_state.nudgecoeff_e, - prognostic_state.normal_wind, + prognostic_state.vn, self.horizontal_edge_index, self.nudgezone_diff, self.fac_bdydiff_v, edge_start_nudging_minus1, - prognostic_state.normal_wind, + prognostic_state.vn, horizontal_start=edge_start_nudging_plus_one, horizontal_end=edge_end_local, vertical_start=0, @@ -938,8 +938,8 @@ def _do_diffusion_step( geofac_n2s=self.interpolation_state.geofac_n2s, geofac_grg_x=self.interpolation_state.geofac_grg_x, geofac_grg_y=self.interpolation_state.geofac_grg_y, - w_old=prognostic_state.vertical_wind, - w=prognostic_state.vertical_wind, + w_old=prognostic_state.w, + w=prognostic_state.w, dwdx=diagnostic_state.dwdx, dwdy=diagnostic_state.dwdy, diff_multfac_w=self.diff_multfac_w, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py index c5785073d..426a3d93b 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py @@ -59,8 +59,8 @@ def diffusion_run( diagnostic_div_ic: Field[[CellDim, KDim], float], diagnostic_dwdx: Field[[CellDim, KDim], float], diagnostic_dwdy: Field[[CellDim, KDim], float], - prognostic_vertical_wind: Field[[CellDim, KDim], float], - prognostic_normal_wind: Field[[EdgeDim, KDim], float], + prognostic_w: Field[[CellDim, KDim], float], + prognostic_vn: Field[[EdgeDim, KDim], float], prognostic_exner_pressure: Field[[CellDim, KDim], float], prognostic_theta_v: Field[[CellDim, KDim], float], metric_theta_ref_mc: Field[[CellDim, KDim], float], @@ -132,7 +132,7 @@ def diffusion_run( # # 1. CALL rbf_vec_interpol_vertex _mo_intp_rbf_rbf_vec_interpol_vertex( - prognostic_normal_wind, + prognostic_vn, interpolation_rbf_coeff_1, interpolation_rbf_coeff_2, out=(local_u_vert, local_v_vert), @@ -156,7 +156,7 @@ def diffusion_run( primal_normal_vert_2, dual_normal_vert_1, dual_normal_vert_2, - prognostic_normal_wind, + prognostic_vn, local_smag_limit, local_smag_offset, out=(local_kh_smag_e, local_kh_smag_ec, local_z_nabla2_e), @@ -168,7 +168,7 @@ def diffusion_run( _fused_mo_nh_diffusion_stencil_02_03( local_kh_smag_ec, - prognostic_normal_wind, + prognostic_vn, interpolation_e_bln_c_s, interpolation_geofac_div, local_diff_multfac_smag, @@ -213,12 +213,12 @@ def diffusion_run( local_kh_smag_e, diff_multfac_vn, interpolation_nudgecoeff_e, - prognostic_normal_wind, + prognostic_vn, local_horizontal_edge_index, local_nudgezone_diff, local_fac_bdydiff_v, edge_startindex_nudging_minus1, - out=prognostic_normal_wind, + out=prognostic_vn, domain={ EdgeDim: (edge_startindex_nudging_plus1, edge_endindex_local), KDim: (0, nlev), @@ -232,8 +232,8 @@ def diffusion_run( interpolation_geofac_n2s, interpolation_geofac_grg_x, interpolation_geofac_grg_y, - prognostic_vertical_wind, - prognostic_vertical_wind, + prognostic_w, + prognostic_w, diagnostic_dwdx, diagnostic_dwdy, local_diff_multfac_w, @@ -244,7 +244,7 @@ def diffusion_run( cell_startindex_interior, cell_endindex_local, out=( - prognostic_vertical_wind, + prognostic_w, diagnostic_dwdx, diagnostic_dwdy, ), diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py index 303a1744a..5187e5e22 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py @@ -140,7 +140,7 @@ def diffusion_run( ): diagnostic_state = DiagnosticState(hdef_ic, div_ic, dwdx, dwdy) prognostic_state = PrognosticState( - vertical_wind=w, normal_wind=vn, exner_pressure=exner, theta_v=theta_v + w=w, vn=vn, exner_pressure=exner, theta_v=theta_v ) if linit: diffusion.initial_step( diff --git a/atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py b/atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py index d82aa083b..88eb9a987 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py @@ -25,9 +25,9 @@ class PrognosticState: corresponds to ICON t_nh_prog """ - vertical_wind: Field[ + w: Field[ [CellDim, KDim], float - ] # vertical_wind field w(nproma, nlevp1, nblks_c) [m/s] - normal_wind: Field[[EdgeDim, KDim], float] # vn(nproma, nlev, nblks_e) [m/s] + ] # vertical_wind field, w(nproma, nlevp1, nblks_c) [m/s] + vn: Field[[EdgeDim, KDim], float] # vn(nproma, nlev, nblks_e) [m/s] exner_pressure: Field[[CellDim, KDim], float] # exner(nrpoma, nlev, nblks_c) theta_v: Field[[CellDim, KDim], float] # (nproma, nlev, nlbks_c) [K] diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index d816a6c45..2b1ebeb2e 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -482,8 +482,8 @@ def test_run_diffusion_single_step( hdef_ic=sp.hdef_ic(), div_ic=sp.div_ic(), dwdx=sp.dwdx(), dwdy=sp.dwdy() ) prognostic_state = PrognosticState( - vertical_wind=sp.w(), - normal_wind=sp.vn(), + w=sp.w(), + vn=sp.vn(), exner_pressure=sp.exner(), theta_v=sp.theta_v(), ) @@ -524,10 +524,8 @@ def test_run_diffusion_single_step( icon_result_w = savepoint_exit.w() icon_result_theta_w = savepoint_exit.theta_v() - assert np.allclose(icon_result_w, np.asarray(prognostic_state.vertical_wind)) - assert np.allclose( - np.asarray(icon_result_vn), np.asarray(prognostic_state.normal_wind) - ) + assert np.allclose(icon_result_w, np.asarray(prognostic_state.w)) + assert np.allclose(np.asarray(icon_result_vn), np.asarray(prognostic_state.vn)) assert np.allclose( np.asarray(icon_result_theta_w), np.asarray(prognostic_state.theta_v) ) @@ -552,8 +550,8 @@ def test_diffusion_five_steps( hdef_ic=sp.hdef_ic(), div_ic=sp.div_ic(), dwdx=sp.dwdx(), dwdy=sp.dwdy() ) prognostic_state = PrognosticState( - vertical_wind=sp.w(), - normal_wind=sp.vn(), + w=sp.w(), + vn=sp.vn(), exner_pressure=sp.exner(), theta_v=sp.theta_v(), ) @@ -641,10 +639,8 @@ def test_diffusion_five_steps( icon_result_vn = savepoint_exit.vn() icon_result_w = savepoint_exit.w() icon_result_theta_w = savepoint_exit.theta_v() - assert np.allclose(icon_result_w, np.asarray(prognostic_state.vertical_wind)) - assert np.allclose( - np.asarray(icon_result_vn), np.asarray(prognostic_state.normal_wind) - ) + assert np.allclose(icon_result_w, np.asarray(prognostic_state.w)) + assert np.allclose(np.asarray(icon_result_vn), np.asarray(prognostic_state.vn)) assert np.allclose( np.asarray(icon_result_theta_w), np.asarray(prognostic_state.theta_v) ) diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 570c27208..29ee005a5 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -1,6 +1,6 @@ # VCS -e git+https://github.com/tehrengruber/gt4py.git@2a9dcb0367a1ec1ef6d78951358d700424867b07#egg=gt4py-functional -git+https://github.com/halungge/serialbox@fix_python_build#egg=serialbox&subdirectory=src/serialbox-python +git+https://github.com/GridTools/serialbox#egg=serialbox&subdirectory=src/serialbox-python # PyPI From 7e2e1079a7e7dbd16b7462e0116c037c75a01e66 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 1 Feb 2023 15:41:25 +0100 Subject: [PATCH 067/263] add cffi decorator wrapper: --- .../src/icon4py/diffusion/wrapper/__init__.py | 0 .../src/icon4py/diffusion/wrapper/decorators.py | 7 +++++++ .../diffusion/{ => wrapper}/diffusion_wrapper.py | 13 ++++--------- 3 files changed, 11 insertions(+), 9 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/diffusion/wrapper/__init__.py create mode 100644 atm_dyn_iconam/src/icon4py/diffusion/wrapper/decorators.py rename atm_dyn_iconam/src/icon4py/diffusion/{ => wrapper}/diffusion_wrapper.py (88%) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/__init__.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/decorators.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/decorators.py new file mode 100644 index 000000000..810008a5d --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/decorators.py @@ -0,0 +1,7 @@ +def cffi_plugin(plugin_name): + from plugin_name import ffi + def cffi_plugin_decorator(func): + def plugin_wrapper(*args, **kwargs): + return ffi.def_extern(func(*args, **kwargs)) + + return plugin_wrapper diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py rename to atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py index 5187e5e22..a75d09725 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_wrapper.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py @@ -14,6 +14,7 @@ # flake8: noqa import numpy as np from functional.common import Field +from functional.iterator.embedded import np_as_located_field from icon4py.common.dimension import ( C2E2CDim, @@ -119,11 +120,6 @@ def diffusion_init( ) -# ELSE IF ((diffu_type == 3 .OR. diffu_type == 5) .AND. discr_vn == 1 .AND. .NOT. diffusion_config(jg)%lsmag_3d) -# mch branch -# ELSE IF ((diffu_type == 3 .OR. diffu_type == 5) .AND. discr_vn == 1) THEN t_nh_diag%vt AND t_nh_diag%theta_v_ic are used -# matching but not mch branch -# and I couldn't find discr_vn = mo_diffusion_nml:itype_vn_diffu != 1 in any experiment def diffusion_run( dtime: float, linit: bool, @@ -131,16 +127,15 @@ def diffusion_run( w: np.ndarray, theta_v: np.ndarray, exner: np.ndarray, - # vt: np.ndarray(EdgeDim, KDim) # not used in mch_ch_r04b09 see above - # theta_v_ic: np.ndarray (CellDim, KHalfDim) # not used in mch_ch_r04b09 see above div_ic: np.ndarray, hdef_ic: np.ndarray, dwdx: np.ndarray, dwdy: np.ndarray, ): - diagnostic_state = DiagnosticState(hdef_ic, div_ic, dwdx, dwdy) + as_cell_k = np_as_located_field(CellDim, KDim) + diagnostic_state = DiagnosticState(as_cell_k(hdef_ic), as_cell_k(div_ic), as_cell_k(dwdx), as_cell_k(dwdy)) prognostic_state = PrognosticState( - w=w, vn=vn, exner_pressure=exner, theta_v=theta_v + w=as_cell_k(w), vn=np_as_located_field(EdgeDim, KDim)(vn), exner_pressure=as_cell_k(exner), theta_v=as_cell_k(theta_v) ) if linit: diffusion.initial_step( From 0b68cae5b8f3f9a4385df5879ca56be6e8372fd4 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 3 Feb 2023 10:50:54 +0100 Subject: [PATCH 068/263] add c header and fortran generation for wrapper functions --- .../src/icon4py/diffusion/wrapper/__init__.py | 12 ++ .../src/icon4py/diffusion/wrapper/binding.py | 119 ++++++++++++++++ .../icon4py/diffusion/wrapper/decorators.py | 15 ++ .../diffusion/wrapper/diffusion_wrapper.py | 35 +++-- .../src/icon4py/diffusion/wrapper/parsing.py | 51 +++++++ atm_dyn_iconam/tests/test_parsing.py | 23 +++ atm_dyn_iconam/tests/test_wrapper.py | 132 ++++++++++++++++++ 7 files changed, 372 insertions(+), 15 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py create mode 100644 atm_dyn_iconam/src/icon4py/diffusion/wrapper/parsing.py create mode 100644 atm_dyn_iconam/tests/test_parsing.py create mode 100644 atm_dyn_iconam/tests/test_wrapper.py diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/__init__.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/__init__.py index e69de29bb..15dfdb009 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/__init__.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py new file mode 100644 index 000000000..3bfa55070 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py @@ -0,0 +1,119 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from typing import Sequence + +import eve +from eve.codegen import JinjaTemplate as as_jinja +from eve.codegen import TemplatedGenerator +from functional.type_system.type_specifications import ScalarKind + +from icon4py.bindings.codegen.type_conversion import ( + BUILTIN_TO_CPP_TYPE, + BUILTIN_TO_ISO_C_TYPE, +) + + +class DimensionType(eve.Node): + name: str + length: int + + +class FuncParameter(eve.Node): + name: str + d_type: ScalarKind + dimensions: Sequence[DimensionType] + + +class Func(eve.Node): + name: str + args: Sequence[FuncParameter] + + +class CffiPlugin(eve.Node): + name: str + functions: Sequence[Func] + + +def to_c_type(scalar_type: ScalarKind): + return BUILTIN_TO_CPP_TYPE[scalar_type] + + +def to_f_type(scalar_type: ScalarKind): + return BUILTIN_TO_ISO_C_TYPE[scalar_type] + + +def field_extension(param: FuncParameter, language: str) -> str: + size = len(param.dimensions) + if size == 0: + return "" + if "C" == language: + return "*" + else: + dims = ",".join(map(lambda x: ":", range(size))) + return f"({dims})" + + +class CHeaderGenerator(TemplatedGenerator): + CffiPlugin = as_jinja("""{% for func in functions: %}{{func}}\n{% endfor %}""") + + Func = as_jinja("""extern void {{name}}({{", ".join(args)}});""") + + def visit_FuncParameter(self, param: FuncParameter): + return self.generic_visit( + param, + rendered_type=to_c_type(param.d_type), + dim=field_extension(param, "C"), + ) + + FuncParameter = as_jinja("""{{rendered_type}}{{dim}} {{name}}""") + + +class FortranInterfaceGenerator(TemplatedGenerator): + CffiPlugin = as_jinja( + """ + module {{name}} + use, intrinsic:: iso_c_binding + implicit none + + public + interface + {% for func in functions: %}\ + {{func}}\ + {% endfor %}\ + end interface + end module + """ + ) + + def visit_Func(self, func: Func): + arg_names = ", ".join(map(lambda x: x.name, func.args)) + return self.generic_visit(func, param_names=arg_names) + + Func = as_jinja( + """subroutine {{name}}({{param_names}}) bind(c, name='{{name}}') + use iso_c_binding + {% for arg in args: %}\ + {{arg}}\ + {% endfor %}\ + end subroutine {{name}} + """ + ) + + def visit_FuncParameter(self, param: FuncParameter, param_names=""): + return self.generic_visit(param, rendered_type=to_f_type(param.d_type)) + + FuncParameter = as_jinja( + """{{rendered_type}}, intent(inout):: {{name}} + """ + ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/decorators.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/decorators.py index 810008a5d..03028e378 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/decorators.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/decorators.py @@ -1,5 +1,20 @@ +# 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 + + def cffi_plugin(plugin_name): from plugin_name import ffi + def cffi_plugin_decorator(func): def plugin_wrapper(*args, **kwargs): return ffi.def_extern(func(*args, **kwargs)) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py index a75d09725..c4a30b75a 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py @@ -38,6 +38,7 @@ from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.prognostic_state import PrognosticState +from icon4py.diffusion.wrapper.decorators import cffi_plugin diffusion: Diffusion(run_program=True) @@ -66,8 +67,10 @@ def diffusion_init( dual_edge_lengths: Field[[EdgeDim], float], inverse_dual_edge_lengths: Field[[EdgeDim], float], inverse_vertex_vertex_lengths: Field[[EdgeDim], float], - primal_normal_vert: tuple[Field[[ECVDim], float], Field[[ECVDim], float]], - dual_normal_vert: tuple[Field[[ECVDim], float], Field[[ECVDim], float]], + primal_normal_vert_1: Field[[ECVDim], float], + primal_normal_vert_2: Field[[ECVDim], float], + dual_normal_vert_1: Field[[ECVDim], float], + dual_normal_vert_2: Field[[ECVDim], float], edge_areas: Field[[EdgeDim], float], cell_areas: Field[[CellDim], float], ): @@ -83,11 +86,11 @@ def diffusion_init( tangent_orientation=tangent_orientation, primal_edge_lengths=primal_edge_lengths, inverse_primal_edge_lengths=inverse_primal_edge_lengths, - dual_normal_vert=dual_normal_vert, + dual_normal_vert=(dual_normal_vert_1, dual_normal_vert_2), dual_edge_lengths=dual_edge_lengths, inverse_dual_edge_lengths=inverse_dual_edge_lengths, inverse_vertex_vertex_lengths=inverse_vertex_vertex_lengths, - primal_normal_vert=primal_normal_vert, + primal_normal_vert=(primal_normal_vert_1, primal_normal_vert_2), edge_areas=edge_areas, ) cell_params = CellParams(cell_areas) @@ -123,19 +126,21 @@ def diffusion_init( def diffusion_run( dtime: float, linit: bool, - vn: np.ndarray, - w: np.ndarray, - theta_v: np.ndarray, - exner: np.ndarray, - div_ic: np.ndarray, - hdef_ic: np.ndarray, - dwdx: np.ndarray, - dwdy: np.ndarray, + vn: Field[[EdgeDim, KDim], float], + w: Field[[CellDim, KDim], float], + theta_v: Field[[CellDim, KDim], float], + exner: Field[[CellDim, KDim], float], + div_ic: Field[[CellDim, KDim], float], + hdef_ic: Field[[CellDim, KDim], float], + dwdx: Field[[CellDim, KDim], float], + dwdy: Field[[CellDim, KDim], float], ): - as_cell_k = np_as_located_field(CellDim, KDim) - diagnostic_state = DiagnosticState(as_cell_k(hdef_ic), as_cell_k(div_ic), as_cell_k(dwdx), as_cell_k(dwdy)) + diagnostic_state = DiagnosticState(hdef_ic, div_ic, dwdx, dwdy) prognostic_state = PrognosticState( - w=as_cell_k(w), vn=np_as_located_field(EdgeDim, KDim)(vn), exner_pressure=as_cell_k(exner), theta_v=as_cell_k(theta_v) + w=w, + vn=vn, + exner_pressure=exner, + theta_v=theta_v, ) if linit: diffusion.initial_step( diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/parsing.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/parsing.py new file mode 100644 index 000000000..93e484ab8 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/parsing.py @@ -0,0 +1,51 @@ +# 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 importlib +from inspect import signature + +from functional.type_system.type_specifications import ScalarType +from functional.type_system.type_translation import from_type_hint + +from icon4py.diffusion.wrapper.binding import ( + CffiPlugin, + DimensionType, + Func, + FuncParameter, +) + + +def parse_functions_from_module(module_name: str, func_names: list[str]) -> CffiPlugin: + module = importlib.import_module(module_name) + funcs = [_parse_function(module, fn) for fn in func_names] + return CffiPlugin(name=module_name, functions=funcs) + + +def _parse_function(module, s): + func = getattr(module, s) + params = [ + _parse_params(signature(func).parameters, p) + for p in (signature(func).parameters) + ] + return Func(name=s, args=params) + + +def _parse_params(params, s): + type_spec = from_type_hint(params[s].annotation) + if isinstance(type_spec, ScalarType): + dtype = type_spec.kind + dims = [] + else: + dtype = type_spec.dtype.kind + dims = [DimensionType(name=d.value, length=10) for d in type_spec.dims] + return FuncParameter(name=s, d_type=dtype, dimensions=dims) diff --git a/atm_dyn_iconam/tests/test_parsing.py b/atm_dyn_iconam/tests/test_parsing.py new file mode 100644 index 000000000..26b9de5e8 --- /dev/null +++ b/atm_dyn_iconam/tests/test_parsing.py @@ -0,0 +1,23 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from icon4py.diffusion.wrapper.parsing import parse_functions_from_module + + +def test_parse_functions(): + path = "icon4py.diffusion.wrapper.diffusion_wrapper" + plugin = parse_functions_from_module(path, ["diffusion_init", "diffusion_run"]) + assert plugin.name == path + assert len(plugin.functions) == 2 + assert "diffusion_init" in map(lambda f: f.name, plugin.functions) + assert "diffusion_run" in map(lambda f: f.name, plugin.functions) diff --git a/atm_dyn_iconam/tests/test_wrapper.py b/atm_dyn_iconam/tests/test_wrapper.py new file mode 100644 index 000000000..08bff644c --- /dev/null +++ b/atm_dyn_iconam/tests/test_wrapper.py @@ -0,0 +1,132 @@ +# 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 string + +import pytest +from functional.type_system.type_specifications import ScalarKind + +from icon4py.diffusion.wrapper.binding import ( + CffiPlugin, + CHeaderGenerator, + DimensionType, + FortranInterfaceGenerator, + Func, + FuncParameter, + field_extension, +) + + +fieldParam2d = FuncParameter( + name="name", + d_type=ScalarKind.FLOAT32, + dimensions=[DimensionType(name="K", length=13), DimensionType(name="J", length=13)], +) +fieldParam1d = FuncParameter( + name="name", + d_type=ScalarKind.FLOAT32, + dimensions=[DimensionType(name="K", length=13)], +) + +simpleType = FuncParameter(name="name", d_type=ScalarKind.FLOAT32, dimensions=[]) + + +@pytest.mark.parametrize(("lang", "expected"), (("C", "*"), ("F", "(:,:)"))) +def test_field_extension_2d(lang, expected): + assert field_extension(fieldParam2d, lang) == expected + + +@pytest.mark.parametrize(("lang", "expected"), (("C", "*"), ("F", "(:)"))) +def test_field_extension_1d(lang, expected): + assert field_extension(fieldParam1d, lang) == expected + + +@pytest.mark.parametrize("lang", ("C", "F")) +def test_is_field_simple_type(lang): + assert field_extension(simpleType, lang) == "" + + +foo = Func( + name="foo", + args=[ + FuncParameter(name="one", d_type=ScalarKind.INT32, dimensions=[]), + FuncParameter(name="two", d_type=ScalarKind.FLOAT64, dimensions=[]), + ], +) + +bar = Func( + name="bar", + args=[ + FuncParameter( + name="one", + d_type=ScalarKind.FLOAT32, + dimensions=[ + DimensionType(name="KDim", length=10), + DimensionType(name="VDim", length=50000), + ], + ), + FuncParameter(name="two", d_type=ScalarKind.INT32, dimensions=[]), + ], +) + + +def test_cheader_generation_for_single_function(): + functions = [foo] + plugin = CffiPlugin(name="libtest", functions=functions) + + header = CHeaderGenerator.apply(plugin) + assert header == "extern void foo(int one, double two);\n" + + +def test_cheader_for_pointer_args(): + functions = [bar] + plugin = CffiPlugin(name="libtest", functions=functions) + + header = CHeaderGenerator.apply(plugin) + assert header == "extern void bar(float* one, int two);\n" + + +def compare_ignore_whitespace(s1: str, s2: str): + no_whitespace = {ord(c): None for c in string.whitespace} + return s1.translate(no_whitespace) == s2.translate(no_whitespace) + + +def test_cheader_with_several_functions(): + functions = [bar, foo] + plugin = CffiPlugin(name="libtest", functions=functions) + header = CHeaderGenerator.apply(plugin) + assert ( + header + == """extern void bar(float* one, int two);\nextern void foo(int one, double two);\n""" + ) + + +def test_fortran_interface(): + functions = [foo] + plugin = CffiPlugin(name="libtest", functions=functions) + interface = FortranInterfaceGenerator.apply(plugin) + expteced = """ + module libtest + use, intrinsic:: iso_c_binding + implicit none + + public + interface + subroutine foo(one, two) bind(c, name='foo') + use iso_c_binding + integer(c_int), intent(inout):: one + real(c_double), intent(inout):: two + end subroutine foo + end interface + end module + """ + assert compare_ignore_whitespace(interface, expteced) From 5d62a62d77b9903314e7ef03feaee5fbf734f5c1 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 3 Feb 2023 15:16:31 +0100 Subject: [PATCH 069/263] add cffi library generation (version 1) --- .../src/icon4py/diffusion/wrapper/binding.py | 13 +++++- .../icon4py/diffusion/wrapper/decorators.py | 44 +++++++++++++++++++ .../diffusion/wrapper/diffusion_wrapper.py | 3 -- .../src/icon4py/diffusion/wrapper/parsing.py | 3 +- .../icon4py/diffusion/wrapper/py2f90gen.py | 24 ++++++++++ atm_dyn_iconam/tests/test_parsing.py | 2 +- atm_dyn_iconam/tests/test_wrapper.py | 4 +- 7 files changed, 85 insertions(+), 8 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/diffusion/wrapper/py2f90gen.py diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py index 3bfa55070..ddf14a836 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py @@ -14,6 +14,7 @@ from typing import Sequence import eve +from eve import codegen from eve.codegen import JinjaTemplate as as_jinja from eve.codegen import TemplatedGenerator from functional.type_system.type_specifications import ScalarKind @@ -22,6 +23,7 @@ BUILTIN_TO_CPP_TYPE, BUILTIN_TO_ISO_C_TYPE, ) +from icon4py.bindings.utils import write_string class DimensionType(eve.Node): @@ -79,7 +81,7 @@ def visit_FuncParameter(self, param: FuncParameter): FuncParameter = as_jinja("""{{rendered_type}}{{dim}} {{name}}""") -class FortranInterfaceGenerator(TemplatedGenerator): +class F90InterfaceGenerator(TemplatedGenerator): CffiPlugin = as_jinja( """ module {{name}} @@ -117,3 +119,12 @@ def visit_FuncParameter(self, param: FuncParameter, param_names=""): """{{rendered_type}}, intent(inout):: {{name}} """ ) + + +def generate_c_header(plugin: CffiPlugin)->str: + generated_code = CHeaderGenerator.apply(plugin) + return codegen.format_source("cpp", generated_code, style="LLVM") + +def generate_f90_interface(plugin: CffiPlugin)->str: + generated_code = F90InterfaceGenerator.apply(plugin) + return generated_code diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/decorators.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/decorators.py index 03028e378..8c4c530a6 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/decorators.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/decorators.py @@ -10,6 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import cffi def cffi_plugin(plugin_name): @@ -20,3 +21,46 @@ def plugin_wrapper(*args, **kwargs): return ffi.def_extern(func(*args, **kwargs)) return plugin_wrapper + + +def compile_cffi_plugin( + plugin_name: str, c_header: str, cffi_functions_file: str, build_path="." +): + """ + Create C shared library. + + Create a linkable C library for the functions in {cffi_functions_file} that are decorated + with '@ffi.def_extern' and correspond to a C signature in the header string + + Args: + plugin_name: name of the plugin, a linkable C library with the name 'lib{plugin_name}.so' will be + created in the build_path folder' + c_header: C type header signature for the python functions. + cffi_functions_file: input file that contains python functions correspondig to the signature in the '{c_header}' + string, these functions must be decorated with @ffi.def_extern() and the file must contain the import + 'from {plugin_name} import ffi' + build_path: *optional* path to build directory + + Returns: + """ + c_header_file = plugin_name + ".h" + with open("/".join([build_path, c_header_file]), "w") as f: + f.write(c_header) + + builder = cffi.FFI() + + builder.embedding_api(c_header) + builder.set_source(plugin_name, f'#include "{c_header_file}"') + + with open(cffi_functions_file) as f: + module = f.read() + + import_str = f"from {plugin_name} import ffi\n" + extern_decorator = "@ffi.def_extern()\n" + module = f"{import_str}{module}" + module.replace("def diffusion_init", extern_decorator + "def diffusion_init") + + + builder.embedding_init_code(module) + builder.compile(tmpdir=build_path, target=f"lib{plugin_name}.*", verbose=True) + diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py index c4a30b75a..2a37aae8f 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py @@ -40,10 +40,8 @@ from icon4py.diffusion.prognostic_state import PrognosticState from icon4py.diffusion.wrapper.decorators import cffi_plugin - diffusion: Diffusion(run_program=True) - def diffusion_init( vct_a: Field[[KDim], float], nrdmax: float, # in config @@ -122,7 +120,6 @@ def diffusion_init( interpolation_state=interpolation_state, ) - def diffusion_run( dtime: float, linit: bool, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/parsing.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/parsing.py index 93e484ab8..2b51d1ec3 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/parsing.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/parsing.py @@ -28,7 +28,8 @@ def parse_functions_from_module(module_name: str, func_names: list[str]) -> CffiPlugin: module = importlib.import_module(module_name) funcs = [_parse_function(module, fn) for fn in func_names] - return CffiPlugin(name=module_name, functions=funcs) + plugin_name=module_name.split(".")[-1] + return CffiPlugin(name=plugin_name, functions=funcs) def _parse_function(module, s): diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/py2f90gen.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/py2f90gen.py new file mode 100644 index 000000000..9d9781a82 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/py2f90gen.py @@ -0,0 +1,24 @@ +from pathlib import Path + +from icon4py.bindings.utils import write_string +from icon4py.diffusion.wrapper.binding import generate_c_header, \ + generate_f90_interface +from icon4py.diffusion.wrapper.decorators import compile_cffi_plugin +from icon4py.diffusion.wrapper.parsing import parse_functions_from_module + +def main(): + module_name = "icon4py.diffusion.wrapper.diffusion_wrapper" + python_src_file = "diffusion_wrapper.py" + build_path = Path("./build") + build_path.mkdir(exist_ok=True, parents=True) + plugin = parse_functions_from_module(module_name, ["diffusion_init", "diffusion_run"]) + c_header = generate_c_header(plugin) + f90_interface= generate_f90_interface(plugin) + write_string(f90_interface, build_path, f"{plugin.name}.f90") + + compile_cffi_plugin(plugin.name, c_header, python_src_file , str(build_path)) + + +if __name__ == '__main__': + main() + diff --git a/atm_dyn_iconam/tests/test_parsing.py b/atm_dyn_iconam/tests/test_parsing.py index 26b9de5e8..c706a2747 100644 --- a/atm_dyn_iconam/tests/test_parsing.py +++ b/atm_dyn_iconam/tests/test_parsing.py @@ -17,7 +17,7 @@ def test_parse_functions(): path = "icon4py.diffusion.wrapper.diffusion_wrapper" plugin = parse_functions_from_module(path, ["diffusion_init", "diffusion_run"]) - assert plugin.name == path + assert plugin.name == "diffusion_wrapper" assert len(plugin.functions) == 2 assert "diffusion_init" in map(lambda f: f.name, plugin.functions) assert "diffusion_run" in map(lambda f: f.name, plugin.functions) diff --git a/atm_dyn_iconam/tests/test_wrapper.py b/atm_dyn_iconam/tests/test_wrapper.py index 08bff644c..203b139ec 100644 --- a/atm_dyn_iconam/tests/test_wrapper.py +++ b/atm_dyn_iconam/tests/test_wrapper.py @@ -19,7 +19,7 @@ CffiPlugin, CHeaderGenerator, DimensionType, - FortranInterfaceGenerator, + F90InterfaceGenerator, Func, FuncParameter, field_extension, @@ -113,7 +113,7 @@ def test_cheader_with_several_functions(): def test_fortran_interface(): functions = [foo] plugin = CffiPlugin(name="libtest", functions=functions) - interface = FortranInterfaceGenerator.apply(plugin) + interface = F90InterfaceGenerator.apply(plugin) expteced = """ module libtest use, intrinsic:: iso_c_binding From b1c814e95a61f45e02b31827e631bfaeb11d4423 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 3 Feb 2023 15:27:02 +0100 Subject: [PATCH 070/263] add missing dimension to fortran interface --- atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py index ddf14a836..432e25969 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py @@ -23,7 +23,6 @@ BUILTIN_TO_CPP_TYPE, BUILTIN_TO_ISO_C_TYPE, ) -from icon4py.bindings.utils import write_string class DimensionType(eve.Node): @@ -113,10 +112,13 @@ def visit_Func(self, func: Func): ) def visit_FuncParameter(self, param: FuncParameter, param_names=""): - return self.generic_visit(param, rendered_type=to_f_type(param.d_type)) + return self.generic_visit(param, + rendered_type=to_f_type(param.d_type), + dim=field_extension(param, "F") + ) FuncParameter = as_jinja( - """{{rendered_type}}, intent(inout):: {{name}} + """{{rendered_type}}, intent(inout):: {{name}}{{dim}} """ ) From 9ab3cf53909b6304865b1d809be5667af8af3cce Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 3 Feb 2023 16:04:48 +0100 Subject: [PATCH 071/263] add target keyword for scalar params in generated fortran interface --- .../src/icon4py/diffusion/wrapper/binding.py | 23 ++++++++++++++----- atm_dyn_iconam/tests/test_wrapper.py | 19 +++++++++------ 2 files changed, 29 insertions(+), 13 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py index 432e25969..e044f76eb 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py @@ -46,15 +46,24 @@ class CffiPlugin(eve.Node): functions: Sequence[Func] -def to_c_type(scalar_type: ScalarKind): +def to_c_type(scalar_type: ScalarKind)->str: return BUILTIN_TO_CPP_TYPE[scalar_type] -def to_f_type(scalar_type: ScalarKind): +def to_f_type(scalar_type: ScalarKind)->str: return BUILTIN_TO_ISO_C_TYPE[scalar_type] +def as_f90_value(param: FuncParameter)->str: + """ + If param is a scalar type (dimension=0) then return the F90 'value' keyword. + + Used for F90 generation only. + """ + return "value, " if len(param.dimensions) == 0 else "" + + -def field_extension(param: FuncParameter, language: str) -> str: +def as_field(param: FuncParameter, language: str) -> str: size = len(param.dimensions) if size == 0: return "" @@ -74,7 +83,7 @@ def visit_FuncParameter(self, param: FuncParameter): return self.generic_visit( param, rendered_type=to_c_type(param.d_type), - dim=field_extension(param, "C"), + dim=as_field(param, "C"), ) FuncParameter = as_jinja("""{{rendered_type}}{{dim}} {{name}}""") @@ -112,13 +121,15 @@ def visit_Func(self, func: Func): ) def visit_FuncParameter(self, param: FuncParameter, param_names=""): + #kw-arg param_names needs to be present because it is present up the tree return self.generic_visit(param, + value = as_f90_value(param), rendered_type=to_f_type(param.d_type), - dim=field_extension(param, "F") + dim=as_field(param, "F") ) FuncParameter = as_jinja( - """{{rendered_type}}, intent(inout):: {{name}}{{dim}} + """{{rendered_type}}, {{value}} intent(inout):: {{name}}{{dim}} """ ) diff --git a/atm_dyn_iconam/tests/test_wrapper.py b/atm_dyn_iconam/tests/test_wrapper.py index 203b139ec..1643815d2 100644 --- a/atm_dyn_iconam/tests/test_wrapper.py +++ b/atm_dyn_iconam/tests/test_wrapper.py @@ -22,37 +22,42 @@ F90InterfaceGenerator, Func, FuncParameter, - field_extension, + as_field, as_f90_value, ) -fieldParam2d = FuncParameter( +field_2d = FuncParameter( name="name", d_type=ScalarKind.FLOAT32, dimensions=[DimensionType(name="K", length=13), DimensionType(name="J", length=13)], ) -fieldParam1d = FuncParameter( +field_1d = FuncParameter( name="name", d_type=ScalarKind.FLOAT32, dimensions=[DimensionType(name="K", length=13)], ) -simpleType = FuncParameter(name="name", d_type=ScalarKind.FLOAT32, dimensions=[]) +simple_type = FuncParameter(name="name", d_type=ScalarKind.FLOAT32, dimensions=[]) + +@pytest.mark.parametrize(("param", "expected"), ((simple_type, "value, "), (field_2d, ""), (field_1d, ""))) +def test_as_target(param, expected): + assert expected == as_f90_value(param) + @pytest.mark.parametrize(("lang", "expected"), (("C", "*"), ("F", "(:,:)"))) def test_field_extension_2d(lang, expected): - assert field_extension(fieldParam2d, lang) == expected + assert as_field(field_2d, lang) == expected @pytest.mark.parametrize(("lang", "expected"), (("C", "*"), ("F", "(:)"))) def test_field_extension_1d(lang, expected): - assert field_extension(fieldParam1d, lang) == expected + assert as_field(field_1d, lang) == expected @pytest.mark.parametrize("lang", ("C", "F")) def test_is_field_simple_type(lang): - assert field_extension(simpleType, lang) == "" + assert as_field(simple_type, lang) == "" foo = Func( From 196275bc3f7f01e35358a6aa716ea3f22c1467ce Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Mon, 6 Feb 2023 17:46:05 +0100 Subject: [PATCH 072/263] register functions for generation add python script for code generation and cffi compilation --- .../icon4py/diffusion/wrapper/cffi_utils.py | 89 +++++++++++++++++++ .../{binding.py => code_generation.py} | 31 ++++--- .../icon4py/diffusion/wrapper/decorators.py | 66 -------------- .../diffusion/wrapper/diffusion_wrapper.py | 7 +- .../src/icon4py/diffusion/wrapper/parsing.py | 14 +-- .../icon4py/diffusion/wrapper/py2f90gen.py | 64 +++++++++---- ...est_wrapper.py => test_code_generation.py} | 21 +++-- atm_dyn_iconam/tests/test_parsing.py | 19 +++- base-requirements.txt | 2 + 9 files changed, 200 insertions(+), 113 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/diffusion/wrapper/cffi_utils.py rename atm_dyn_iconam/src/icon4py/diffusion/wrapper/{binding.py => code_generation.py} (81%) delete mode 100644 atm_dyn_iconam/src/icon4py/diffusion/wrapper/decorators.py rename atm_dyn_iconam/tests/{test_wrapper.py => test_code_generation.py} (89%) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/cffi_utils.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/cffi_utils.py new file mode 100644 index 000000000..e66a14253 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/cffi_utils.py @@ -0,0 +1,89 @@ +# 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 functools + +import cffi + + +FFI_DEF_EXTERN_DECORATOR = "@ffi.def_extern()" + +CFFI_GEN_DECORATOR = "@CffiMethod.register" + + +class CffiMethod: + _registry = {} + + @classmethod + def register(cls, func): + try: + cls._registry[func.__module__].append(func.__name__) + except KeyError: + cls._registry[func.__module__] = [func.__name__] + + @functools.wraps(func) + def wrapper(*args, **kwargs): + return func(*args, **kwargs) + + return wrapper + + @classmethod + def get(cls, name: str): + return cls._registry[name] + + +def with_cffi_gen(func): + @functools.wraps(func) + def _cffi_gen(*args, **kwargs): + return func(*args, **kwargs) + + return _cffi_gen + + +def generate_and_compile_cffi_plugin( + plugin_name: str, c_header: str, module_name: str, build_path="." +): + """ + Create C shared library. + + Create a linkable C library and F90 interface for the functions in the python module + {module_name} that are decorated with '@CffiMethod.register'. + + Args: + plugin_name: name of the plugin, a linkable C library with the name + 'lib{plugin_name}.so' will be created in the {build_path} folder' + c_header: C type header signature for the python functions. + module_name: python module name that contains python functions corresponding + to the signature in the '{c_header}' string, these functions must be decorated + with @CffiMethod.register and the file must contain the import + build_path: *optional* path to build directory + + """ + python_src_file = f"{module_name.split('.')[-1]}.py" + c_header_file = plugin_name + ".h" + with open("/".join([build_path, c_header_file]), "w") as f: + f.write(c_header) + + builder = cffi.FFI() + + builder.embedding_api(c_header) + builder.set_source(plugin_name, f'#include "{c_header_file}"') + + with open(python_src_file) as f: + module = f.read() + + module = f"from {plugin_name} import ffi\n{module}".replace( + CFFI_GEN_DECORATOR, FFI_DEF_EXTERN_DECORATOR + ) + + builder.embedding_init_code(module) + builder.compile(tmpdir=build_path, target=f"lib{plugin_name}.*", verbose=True) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/code_generation.py similarity index 81% rename from atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py rename to atm_dyn_iconam/src/icon4py/diffusion/wrapper/code_generation.py index e044f76eb..22907e9e3 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/binding.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/code_generation.py @@ -23,6 +23,7 @@ BUILTIN_TO_CPP_TYPE, BUILTIN_TO_ISO_C_TYPE, ) +from icon4py.bindings.utils import write_string class DimensionType(eve.Node): @@ -46,14 +47,15 @@ class CffiPlugin(eve.Node): functions: Sequence[Func] -def to_c_type(scalar_type: ScalarKind)->str: +def to_c_type(scalar_type: ScalarKind) -> str: return BUILTIN_TO_CPP_TYPE[scalar_type] -def to_f_type(scalar_type: ScalarKind)->str: +def to_f_type(scalar_type: ScalarKind) -> str: return BUILTIN_TO_ISO_C_TYPE[scalar_type] -def as_f90_value(param: FuncParameter)->str: + +def as_f90_value(param: FuncParameter) -> str: """ If param is a scalar type (dimension=0) then return the F90 'value' keyword. @@ -62,7 +64,6 @@ def as_f90_value(param: FuncParameter)->str: return "value, " if len(param.dimensions) == 0 else "" - def as_field(param: FuncParameter, language: str) -> str: size = len(param.dimensions) if size == 0: @@ -121,23 +122,25 @@ def visit_Func(self, func: Func): ) def visit_FuncParameter(self, param: FuncParameter, param_names=""): - #kw-arg param_names needs to be present because it is present up the tree - return self.generic_visit(param, - value = as_f90_value(param), - rendered_type=to_f_type(param.d_type), - dim=as_field(param, "F") - ) + # kw-arg param_names needs to be present because it is present up the tree + return self.generic_visit( + param, + value=as_f90_value(param), + rendered_type=to_f_type(param.d_type), + dim=as_field(param, "F"), + ) FuncParameter = as_jinja( - """{{rendered_type}}, {{value}} intent(inout):: {{name}}{{dim}} + """{{rendered_type}}, {{value}} intent(inout):: {{name}}{{dim}} """ ) -def generate_c_header(plugin: CffiPlugin)->str: +def generate_c_header(plugin: CffiPlugin) -> str: generated_code = CHeaderGenerator.apply(plugin) return codegen.format_source("cpp", generated_code, style="LLVM") -def generate_f90_interface(plugin: CffiPlugin)->str: + +def generate_and_write_f90_interface(build_path: str, plugin: CffiPlugin): generated_code = F90InterfaceGenerator.apply(plugin) - return generated_code + write_string(generated_code, build_path, f"{plugin.name}.f90") diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/decorators.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/decorators.py deleted file mode 100644 index 8c4c530a6..000000000 --- a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/decorators.py +++ /dev/null @@ -1,66 +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 cffi - - -def cffi_plugin(plugin_name): - from plugin_name import ffi - - def cffi_plugin_decorator(func): - def plugin_wrapper(*args, **kwargs): - return ffi.def_extern(func(*args, **kwargs)) - - return plugin_wrapper - - -def compile_cffi_plugin( - plugin_name: str, c_header: str, cffi_functions_file: str, build_path="." -): - """ - Create C shared library. - - Create a linkable C library for the functions in {cffi_functions_file} that are decorated - with '@ffi.def_extern' and correspond to a C signature in the header string - - Args: - plugin_name: name of the plugin, a linkable C library with the name 'lib{plugin_name}.so' will be - created in the build_path folder' - c_header: C type header signature for the python functions. - cffi_functions_file: input file that contains python functions correspondig to the signature in the '{c_header}' - string, these functions must be decorated with @ffi.def_extern() and the file must contain the import - 'from {plugin_name} import ffi' - build_path: *optional* path to build directory - - Returns: - """ - c_header_file = plugin_name + ".h" - with open("/".join([build_path, c_header_file]), "w") as f: - f.write(c_header) - - builder = cffi.FFI() - - builder.embedding_api(c_header) - builder.set_source(plugin_name, f'#include "{c_header_file}"') - - with open(cffi_functions_file) as f: - module = f.read() - - import_str = f"from {plugin_name} import ffi\n" - extern_decorator = "@ffi.def_extern()\n" - module = f"{import_str}{module}" - module.replace("def diffusion_init", extern_decorator + "def diffusion_init") - - - builder.embedding_init_code(module) - builder.compile(tmpdir=build_path, target=f"lib{plugin_name}.*", verbose=True) - diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py index 2a37aae8f..9c0f2fa7d 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py @@ -38,10 +38,13 @@ from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.prognostic_state import PrognosticState -from icon4py.diffusion.wrapper.decorators import cffi_plugin +from icon4py.diffusion.wrapper.cffi_utils import CffiMethod, with_cffi_gen + diffusion: Diffusion(run_program=True) + +@CffiMethod.register def diffusion_init( vct_a: Field[[KDim], float], nrdmax: float, # in config @@ -120,6 +123,8 @@ def diffusion_init( interpolation_state=interpolation_state, ) + +@CffiMethod.register def diffusion_run( dtime: float, linit: bool, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/parsing.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/parsing.py index 2b51d1ec3..f3bae36fb 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/parsing.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/parsing.py @@ -12,12 +12,13 @@ # SPDX-License-Identifier: GPL-3.0-or-later import importlib -from inspect import signature +from inspect import signature, unwrap from functional.type_system.type_specifications import ScalarType from functional.type_system.type_translation import from_type_hint -from icon4py.diffusion.wrapper.binding import ( +from icon4py.diffusion.wrapper.cffi_utils import CffiMethod +from icon4py.diffusion.wrapper.code_generation import ( CffiPlugin, DimensionType, Func, @@ -25,17 +26,18 @@ ) -def parse_functions_from_module(module_name: str, func_names: list[str]) -> CffiPlugin: +def parse_functions_from_module(module_name: str) -> CffiPlugin: module = importlib.import_module(module_name) + func_names = CffiMethod.get(module_name) funcs = [_parse_function(module, fn) for fn in func_names] - plugin_name=module_name.split(".")[-1] + plugin_name = module_name.split(".")[-1] return CffiPlugin(name=plugin_name, functions=funcs) def _parse_function(module, s): - func = getattr(module, s) + func = unwrap(getattr(module, s)) params = [ - _parse_params(signature(func).parameters, p) + _parse_params(signature(func, follow_wrapped=False).parameters, p) for p in (signature(func).parameters) ] return Func(name=s, args=params) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/py2f90gen.py b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/py2f90gen.py index 9d9781a82..4d73fdf7a 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/py2f90gen.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/wrapper/py2f90gen.py @@ -1,24 +1,56 @@ -from pathlib import Path +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later -from icon4py.bindings.utils import write_string -from icon4py.diffusion.wrapper.binding import generate_c_header, \ - generate_f90_interface -from icon4py.diffusion.wrapper.decorators import compile_cffi_plugin +import pathlib + +import click + +from icon4py.diffusion.wrapper.cffi_utils import generate_and_compile_cffi_plugin +from icon4py.diffusion.wrapper.code_generation import ( + generate_and_write_f90_interface, + generate_c_header, +) from icon4py.diffusion.wrapper.parsing import parse_functions_from_module -def main(): - module_name = "icon4py.diffusion.wrapper.diffusion_wrapper" - python_src_file = "diffusion_wrapper.py" - build_path = Path("./build") + +@click.command( + "py2f90gen", +) +@click.argument("module", type=str) +@click.argument( + "build_path", + type=click.Path(dir_okay=True, resolve_path=True, path_type=pathlib.Path), + default=".", +) +def main(module: str, build_path: pathlib.Path) -> None: + """ + Generate C and F90 wrappers and C library for embedding the python MODULE in C and Fortran. + + Args: + - module: name of the python module containing the methods to be embedded. Those + methods have to be decoratoed with CffiMethod.register + + - build_path: directory where the generated code and compiled libraries are to be found. + """ + module_name = module build_path.mkdir(exist_ok=True, parents=True) - plugin = parse_functions_from_module(module_name, ["diffusion_init", "diffusion_run"]) + plugin = parse_functions_from_module(module_name) c_header = generate_c_header(plugin) - f90_interface= generate_f90_interface(plugin) - write_string(f90_interface, build_path, f"{plugin.name}.f90") - - compile_cffi_plugin(plugin.name, c_header, python_src_file , str(build_path)) + generate_and_compile_cffi_plugin( + plugin.name, c_header, module_name, str(build_path) + ) + generate_and_write_f90_interface(plugin) -if __name__ == '__main__': +if __name__ == "__main__": main() - diff --git a/atm_dyn_iconam/tests/test_wrapper.py b/atm_dyn_iconam/tests/test_code_generation.py similarity index 89% rename from atm_dyn_iconam/tests/test_wrapper.py rename to atm_dyn_iconam/tests/test_code_generation.py index 1643815d2..978b22afb 100644 --- a/atm_dyn_iconam/tests/test_wrapper.py +++ b/atm_dyn_iconam/tests/test_code_generation.py @@ -15,14 +15,15 @@ import pytest from functional.type_system.type_specifications import ScalarKind -from icon4py.diffusion.wrapper.binding import ( +from icon4py.diffusion.wrapper.code_generation import ( CffiPlugin, CHeaderGenerator, DimensionType, F90InterfaceGenerator, Func, FuncParameter, - as_field, as_f90_value, + as_f90_value, + as_field, ) @@ -40,11 +41,13 @@ simple_type = FuncParameter(name="name", d_type=ScalarKind.FLOAT32, dimensions=[]) - -@pytest.mark.parametrize(("param", "expected"), ((simple_type, "value, "), (field_2d, ""), (field_1d, ""))) +@pytest.mark.parametrize( + ("param", "expected"), ((simple_type, "value, "), (field_2d, ""), (field_1d, "")) +) def test_as_target(param, expected): assert expected == as_f90_value(param) + @pytest.mark.parametrize(("lang", "expected"), (("C", "*"), ("F", "(:,:)"))) def test_field_extension_2d(lang, expected): assert as_field(field_2d, lang) == expected @@ -105,7 +108,7 @@ def compare_ignore_whitespace(s1: str, s2: str): return s1.translate(no_whitespace) == s2.translate(no_whitespace) -def test_cheader_with_several_functions(): +def test_c_header_with_several_functions(): functions = [bar, foo] plugin = CffiPlugin(name="libtest", functions=functions) header = CHeaderGenerator.apply(plugin) @@ -119,7 +122,7 @@ def test_fortran_interface(): functions = [foo] plugin = CffiPlugin(name="libtest", functions=functions) interface = F90InterfaceGenerator.apply(plugin) - expteced = """ + expected = """ module libtest use, intrinsic:: iso_c_binding implicit none @@ -128,10 +131,10 @@ def test_fortran_interface(): interface subroutine foo(one, two) bind(c, name='foo') use iso_c_binding - integer(c_int), intent(inout):: one - real(c_double), intent(inout):: two + integer(c_int), value, intent(inout):: one + real(c_double), value, intent(inout):: two end subroutine foo end interface end module """ - assert compare_ignore_whitespace(interface, expteced) + assert compare_ignore_whitespace(interface, expected) diff --git a/atm_dyn_iconam/tests/test_parsing.py b/atm_dyn_iconam/tests/test_parsing.py index c706a2747..fc3915300 100644 --- a/atm_dyn_iconam/tests/test_parsing.py +++ b/atm_dyn_iconam/tests/test_parsing.py @@ -11,13 +11,30 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from icon4py.diffusion.wrapper.cffi_utils import CffiMethod from icon4py.diffusion.wrapper.parsing import parse_functions_from_module def test_parse_functions(): path = "icon4py.diffusion.wrapper.diffusion_wrapper" - plugin = parse_functions_from_module(path, ["diffusion_init", "diffusion_run"]) + plugin = parse_functions_from_module(path) + assert plugin.name == "diffusion_wrapper" assert len(plugin.functions) == 2 assert "diffusion_init" in map(lambda f: f.name, plugin.functions) assert "diffusion_run" in map(lambda f: f.name, plugin.functions) + + +@CffiMethod.register +def do_foo(foo: str): + return foo + + +@CffiMethod.register +def do_bar(): + return "bar" + + +def test_register_with_cffi(): + assert "do_foo" in CffiMethod.get(__name__) + assert "do_bar" in CffiMethod.get(__name__) diff --git a/base-requirements.txt b/base-requirements.txt index 046fd3ce2..9f82dad3d 100644 --- a/base-requirements.txt +++ b/base-requirements.txt @@ -1,2 +1,4 @@ # VCS gt4py-functional @ git+https://github.com/GridTools/gt4py.git@572b04caaf4932a23ac0a6a266faa9ecb228807b + +cffi>=1.5 From af499c22a94bda26328af00f8ee8d4d29be4f159 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 8 Feb 2023 12:48:42 +0100 Subject: [PATCH 073/263] move f90 interface wrappers to its own package. --- ...ed_mo_nh_diffusion_stencil_01_02_03_rbf.py | 143 ------------------ .../src/icon4py/diffusion/icon_grid.py | 2 +- base-requirements.txt | 1 - py2f/README.md | 20 +++ py2f/requirements-dev.txt | 4 + py2f/requirements.txt | 4 + py2f/setup.cfg | 58 +++++++ py2f/setup.py | 18 +++ .../src/icon4py/py2f}/__init__.py | 0 .../src/icon4py/py2f}/cffi_utils.py | 0 .../src/icon4py/py2f/codegen.py | 0 .../src/icon4py/py2f}/parsing.py | 9 +- .../src/icon4py/py2f/py2fgen.py | 6 +- py2f/src/icon4py/py2f/wrappers/__init__.py | 12 ++ .../py2f/wrappers}/diffusion_wrapper.py | 4 +- .../tests/test_code_generation.py | 2 +- .../tests/test_parsing.py | 7 +- requirements-dev.txt | 1 + requirements.txt | 1 + 19 files changed, 129 insertions(+), 163 deletions(-) delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py create mode 100644 py2f/README.md create mode 100644 py2f/requirements-dev.txt create mode 100644 py2f/requirements.txt create mode 100644 py2f/setup.cfg create mode 100644 py2f/setup.py rename {atm_dyn_iconam/src/icon4py/diffusion/wrapper => py2f/src/icon4py/py2f}/__init__.py (100%) rename {atm_dyn_iconam/src/icon4py/diffusion/wrapper => py2f/src/icon4py/py2f}/cffi_utils.py (100%) rename atm_dyn_iconam/src/icon4py/diffusion/wrapper/code_generation.py => py2f/src/icon4py/py2f/codegen.py (100%) rename {atm_dyn_iconam/src/icon4py/diffusion/wrapper => py2f/src/icon4py/py2f}/parsing.py (90%) rename atm_dyn_iconam/src/icon4py/diffusion/wrapper/py2f90gen.py => py2f/src/icon4py/py2f/py2fgen.py (88%) create mode 100644 py2f/src/icon4py/py2f/wrappers/__init__.py rename {atm_dyn_iconam/src/icon4py/diffusion/wrapper => py2f/src/icon4py/py2f/wrappers}/diffusion_wrapper.py (97%) rename {atm_dyn_iconam => py2f}/tests/test_code_generation.py (98%) rename {atm_dyn_iconam => py2f}/tests/test_parsing.py (85%) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py deleted file mode 100644 index fb3ecaf12..000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py +++ /dev/null @@ -1,143 +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 - -from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field - -from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( - _mo_intp_rbf_rbf_vec_interpol_vertex, -) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import ( - _mo_nh_diffusion_stencil_01, -) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import ( - _mo_nh_diffusion_stencil_02, -) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import ( - _mo_nh_diffusion_stencil_03, -) -from icon4py.common.dimension import ( - C2EDim, - CellDim, - ECVDim, - EdgeDim, - KDim, - V2EDim, - VertexDim, -) - - -@field_operator -def _fused_mo_nh_diffusion_stencil_01_02_03_rbf( - diff_multfac_smag: Field[[KDim], float], - tangent_orientation: Field[[EdgeDim], float], - inv_primal_edge_length: Field[[EdgeDim], float], - inv_vert_vert_length: Field[[EdgeDim], float], - u_vert_old: Field[[VertexDim, KDim], float], - v_vert_old: Field[[VertexDim, KDim], float], - primal_normal_vert_x: Field[[ECVDim], float], - primal_normal_vert_y: Field[[ECVDim], float], - dual_normal_vert_x: Field[[ECVDim], float], - dual_normal_vert_y: Field[[ECVDim], float], - vn: Field[[EdgeDim, KDim], float], - smag_limit: Field[[KDim], float], - smag_offset: float, - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], - wgtfac_c: Field[[CellDim, KDim], float], - ptr_coeff_1: Field[[VertexDim, V2EDim], float], - ptr_coeff_2: Field[[VertexDim, V2EDim], float], -) -> tuple[ - Field[[EdgeDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[VertexDim, KDim], float], - Field[[VertexDim, KDim], float], -]: - - kh_smag_e, kh_smag_ec, z_nabla2_e = _mo_nh_diffusion_stencil_01( - diff_multfac_smag, - tangent_orientation, - inv_primal_edge_length, - inv_vert_vert_length, - u_vert_old, - v_vert_old, - primal_normal_vert_x, - primal_normal_vert_y, - dual_normal_vert_x, - dual_normal_vert_y, - vn, - smag_limit, - smag_offset, - ) - - kh_c, div = _mo_nh_diffusion_stencil_02( - kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag - ) - - div_ic, hdef_ic = _mo_nh_diffusion_stencil_03(div, kh_c, wgtfac_c) - - u_vert, v_vert = _mo_intp_rbf_rbf_vec_interpol_vertex( - z_nabla2_e, ptr_coeff_1, ptr_coeff_2 - ) - - return kh_smag_e, div_ic, hdef_ic, u_vert, v_vert - - -@program -def fused_mo_nh_diffusion_stencil_01_02_03_rbf( - diff_multfac_smag: Field[[KDim], float], - tangent_orientation: Field[[EdgeDim], float], - inv_primal_edge_length: Field[[EdgeDim], float], - inv_vert_vert_length: Field[[EdgeDim], float], - u_vert_old: Field[[VertexDim, KDim], float], - v_vert_old: Field[[VertexDim, KDim], float], - primal_normal_vert_x: Field[[ECVDim], float], - primal_normal_vert_y: Field[[ECVDim], float], - dual_normal_vert_x: Field[[ECVDim], float], - dual_normal_vert_y: Field[[ECVDim], float], - vn: Field[[EdgeDim, KDim], float], - smag_limit: Field[[KDim], float], - smag_offset: float, - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], - wgtfac_c: Field[[CellDim, KDim], float], - ptr_coeff_1: Field[[VertexDim, V2EDim], float], - ptr_coeff_2: Field[[VertexDim, V2EDim], float], - kh_smag_e: Field[[EdgeDim, KDim], float], - div_ic: Field[[CellDim, KDim], float], - hdef_ic: Field[[CellDim, KDim], float], - u_vert: Field[[VertexDim, KDim], float], - v_vert: Field[[VertexDim, KDim], float], -): - _fused_mo_nh_diffusion_stencil_01_02_03_rbf( - diff_multfac_smag, - tangent_orientation, - inv_primal_edge_length, - inv_vert_vert_length, - u_vert_old, - v_vert_old, - primal_normal_vert_x, - primal_normal_vert_y, - dual_normal_vert_x, - dual_normal_vert_y, - vn, - smag_limit, - smag_offset, - e_bln_c_s, - geofac_div, - wgtfac_c, - ptr_coeff_1, - ptr_coeff_2, - out=(kh_smag_e, div_ic, hdef_ic, u_vert, v_vert), - ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index cf6c89434..dc6314708 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -129,7 +129,7 @@ def get_indices_from_to( self, dim: Dimension, start_marker: int, end_marker: int ) -> Tuple[int, int]: """ - Use to specifzy domains of a field for field_operator. + Use to specify domains of a field for field_operator. For a given dimension, returns the start and end index if a horizontal region in a field given by the markers. diff --git a/base-requirements.txt b/base-requirements.txt index 9f82dad3d..5f05f7206 100644 --- a/base-requirements.txt +++ b/base-requirements.txt @@ -1,4 +1,3 @@ # VCS gt4py-functional @ git+https://github.com/GridTools/gt4py.git@572b04caaf4932a23ac0a6a266faa9ecb228807b -cffi>=1.5 diff --git a/py2f/README.md b/py2f/README.md new file mode 100644 index 000000000..af6792103 --- /dev/null +++ b/py2f/README.md @@ -0,0 +1,20 @@ +# icon4py py2f90 + +## Description + +Python utilities for generating a C library and Fortran interface to call Python icon4py modules. The library [embeds python via CFFI ](https://cffi.readthedocs.io/en/latest/embedding.html) + +### py2f90gen + +Generates a C header file and a Fortran interface and compiles python functions into a C library embedding python. + +The functions need to be decorated with `CffiMethod.register` and have a signature with scalar arguments of gt4py fields: + +``` +@CffiMethod.register +def foo(i:int, param:float, field1: Field[[VertexDim, KDim], float], field2: Field[CellDim, KDim], float]) +``` + +## Installation instructions + +Check `README.md` file in the root of the repository. diff --git a/py2f/requirements-dev.txt b/py2f/requirements-dev.txt new file mode 100644 index 000000000..217c64d71 --- /dev/null +++ b/py2f/requirements-dev.txt @@ -0,0 +1,4 @@ +-r ../base-requirements-dev.txt +-e ../common +-e ../atm_dyn_iconam +-e . diff --git a/py2f/requirements.txt b/py2f/requirements.txt new file mode 100644 index 000000000..b6eac2aa8 --- /dev/null +++ b/py2f/requirements.txt @@ -0,0 +1,4 @@ +-r ../base-requirements.txt +../common +../atm_dyn_iconam +. diff --git a/py2f/setup.cfg b/py2f/setup.cfg new file mode 100644 index 000000000..f7d5ac4cd --- /dev/null +++ b/py2f/setup.cfg @@ -0,0 +1,58 @@ +# This file is mainly used to configure package creation with setuptools. +# Documentation: +# http://setuptools.readthedocs.io/en/latest/setuptools.html#configuring-setup-using-setup-cfg-files +# +[metadata] +name = icon4py_py2f +description = Icon inspired code in Python and GT4Py +long_description = file: README.md +long_description_content_type = text/markdown +url = https://github.com/C2SM/icon4py +author = ETH Zurich +author_email = gridtools@cscs.ch +license = gpl3 +license_files = LICENSE +platforms = Linux, Mac +classifiers = + Development Status :: 3 - Alpha + Intended Audience :: Science/Research + License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+) + Operating System :: POSIX + Programming Language :: Python + Programming Language :: Python :: 3 + Programming Language :: Python :: 3 :: Only + Programming Language :: Python :: 3.10 + Programming Language :: Python :: Implementation :: CPython + Topic :: Scientific/Engineering :: Atmospheric Science + Topic :: Scientific/Engineering :: Mathematics + Topic :: Scientific/Engineering :: Physics +project_urls = + Source Code = https://github.com/C2SM/icon4py + +[options] +packages = find_namespace: +install_requires = + icon4py-common + icon4py-atm_dyn_iconam + cffi>=1.5 + tabulate>=0.8.9 + fprettify>=0.3.7 +python_requires = >=3.10 +package_dir = + = src +zip_safe = False + +[options.package_data] +# References: +# https://setuptools.pypa.io/en/latest/userguide/datafiles.html +# https://github.com/abravalheri/experiment-setuptools-package-data +* = *.md, *.rst, *.toml, *.txt, py.typed + +[options.packages.find] +where = src +exclude = + tests + +[options.entry_points] +console_scripts = + py2fgen = icon4py.py2f.wrapper.py2fgen:main diff --git a/py2f/setup.py b/py2f/setup.py new file mode 100644 index 000000000..9c9f7b81c --- /dev/null +++ b/py2f/setup.py @@ -0,0 +1,18 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from setuptools import setup + + +if __name__ == "__main__": + setup() diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/__init__.py b/py2f/src/icon4py/py2f/__init__.py similarity index 100% rename from atm_dyn_iconam/src/icon4py/diffusion/wrapper/__init__.py rename to py2f/src/icon4py/py2f/__init__.py diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/cffi_utils.py b/py2f/src/icon4py/py2f/cffi_utils.py similarity index 100% rename from atm_dyn_iconam/src/icon4py/diffusion/wrapper/cffi_utils.py rename to py2f/src/icon4py/py2f/cffi_utils.py diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/code_generation.py b/py2f/src/icon4py/py2f/codegen.py similarity index 100% rename from atm_dyn_iconam/src/icon4py/diffusion/wrapper/code_generation.py rename to py2f/src/icon4py/py2f/codegen.py diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/parsing.py b/py2f/src/icon4py/py2f/parsing.py similarity index 90% rename from atm_dyn_iconam/src/icon4py/diffusion/wrapper/parsing.py rename to py2f/src/icon4py/py2f/parsing.py index f3bae36fb..763b554d7 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/parsing.py +++ b/py2f/src/icon4py/py2f/parsing.py @@ -17,13 +17,8 @@ from functional.type_system.type_specifications import ScalarType from functional.type_system.type_translation import from_type_hint -from icon4py.diffusion.wrapper.cffi_utils import CffiMethod -from icon4py.diffusion.wrapper.code_generation import ( - CffiPlugin, - DimensionType, - Func, - FuncParameter, -) +from icon4py.py2f.cffi_utils import CffiMethod +from icon4py.py2f.codegen import CffiPlugin, DimensionType, Func, FuncParameter def parse_functions_from_module(module_name: str) -> CffiPlugin: diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/py2f90gen.py b/py2f/src/icon4py/py2f/py2fgen.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/diffusion/wrapper/py2f90gen.py rename to py2f/src/icon4py/py2f/py2fgen.py index 4d73fdf7a..f638f6e28 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/py2f90gen.py +++ b/py2f/src/icon4py/py2f/py2fgen.py @@ -15,12 +15,12 @@ import click -from icon4py.diffusion.wrapper.cffi_utils import generate_and_compile_cffi_plugin -from icon4py.diffusion.wrapper.code_generation import ( +from icon4py.f2py.cffi_utils import generate_and_compile_cffi_plugin +from icon4py.f2py.codegen import ( generate_and_write_f90_interface, generate_c_header, ) -from icon4py.diffusion.wrapper.parsing import parse_functions_from_module +from icon4py.f2py.parsing import parse_functions_from_module @click.command( diff --git a/py2f/src/icon4py/py2f/wrappers/__init__.py b/py2f/src/icon4py/py2f/wrappers/__init__.py new file mode 100644 index 000000000..15dfdb009 --- /dev/null +++ b/py2f/src/icon4py/py2f/wrappers/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py similarity index 97% rename from atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py rename to py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py index 9c0f2fa7d..fb6f99e40 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/wrapper/diffusion_wrapper.py +++ b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py @@ -12,9 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later # flake8: noqa -import numpy as np from functional.common import Field -from functional.iterator.embedded import np_as_located_field from icon4py.common.dimension import ( C2E2CDim, @@ -38,7 +36,7 @@ from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.prognostic_state import PrognosticState -from icon4py.diffusion.wrapper.cffi_utils import CffiMethod, with_cffi_gen +from icon4py.py2f.cffi_utils import CffiMethod diffusion: Diffusion(run_program=True) diff --git a/atm_dyn_iconam/tests/test_code_generation.py b/py2f/tests/test_code_generation.py similarity index 98% rename from atm_dyn_iconam/tests/test_code_generation.py rename to py2f/tests/test_code_generation.py index 978b22afb..3270c07ec 100644 --- a/atm_dyn_iconam/tests/test_code_generation.py +++ b/py2f/tests/test_code_generation.py @@ -15,7 +15,7 @@ import pytest from functional.type_system.type_specifications import ScalarKind -from icon4py.diffusion.wrapper.code_generation import ( +from icon4py.py2f.codegen import ( CffiPlugin, CHeaderGenerator, DimensionType, diff --git a/atm_dyn_iconam/tests/test_parsing.py b/py2f/tests/test_parsing.py similarity index 85% rename from atm_dyn_iconam/tests/test_parsing.py rename to py2f/tests/test_parsing.py index fc3915300..07ccbaf6a 100644 --- a/atm_dyn_iconam/tests/test_parsing.py +++ b/py2f/tests/test_parsing.py @@ -10,13 +10,12 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - -from icon4py.diffusion.wrapper.cffi_utils import CffiMethod -from icon4py.diffusion.wrapper.parsing import parse_functions_from_module +from icon4py.py2f.cffi_utils import CffiMethod +from icon4py.py2f.parsing import parse_functions_from_module def test_parse_functions(): - path = "icon4py.diffusion.wrapper.diffusion_wrapper" + path = "icon4py.py2f.wrappers.diffusion_wrapper" plugin = parse_functions_from_module(path) assert plugin.name == "diffusion_wrapper" diff --git a/requirements-dev.txt b/requirements-dev.txt index bacaeba53..a625417b7 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -5,3 +5,4 @@ -e ./pyutils -e ./testutils -e ./atm_dyn_iconam +-e ./py2f diff --git a/requirements.txt b/requirements.txt index 4e67c2a53..74cf9c4bf 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,3 +5,4 @@ ./pyutils ./testutils ./atm_dyn_iconam +./py2f From 796843e24bbe4dca97837ee9cf4dbe4004d6dd00 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 16 Feb 2023 14:09:29 +0100 Subject: [PATCH 074/263] - fix import of gt4py merge of main and functional - switch to head of https://github.com/tehrengruber/gt4py.git@reduce_foast_lowering_complexity - skip tests: iterator ir stencils (offset provider containing local dims ) --- .../fused_mo_nh_diffusion_stencil_01_02_03_rbf | 4 ++-- .../fused_mo_nh_diffusion_stencil_02_03.py | 4 ++-- .../fused_mo_nh_diffusion_stencil_04_05_06.py | 6 +++--- ...fused_mo_nh_diffusion_stencil_07_08_09_10.py | 6 +++--- .../fused_mo_nh_diffusion_stencil_11_12.py | 6 +++--- .../fused_mo_nh_diffusion_stencil_13_14.py | 6 +++--- .../src/icon4py/diffusion/diagnostic_state.py | 2 +- .../src/icon4py/diffusion/diffusion.py | 6 +++--- .../src/icon4py/diffusion/diffusion_program.py | 8 ++++---- .../src/icon4py/diffusion/horizontal.py | 2 +- .../src/icon4py/diffusion/icon_grid.py | 6 +++--- .../icon4py/diffusion/interpolation_state.py | 2 +- .../src/icon4py/diffusion/metric_state.py | 2 +- .../src/icon4py/diffusion/prognostic_state.py | 2 +- atm_dyn_iconam/src/icon4py/diffusion/utils.py | 8 ++++---- atm_dyn_iconam/tests/conftest.py | 2 +- atm_dyn_iconam/tests/test_diffusion.py | 8 ++++---- .../tests/test_mo_nh_diffusion_stencil_15.py | 2 ++ .../tests/test_mo_solve_nonhydro_stencil_20.py | 2 ++ .../tests/test_mo_solve_nonhydro_stencil_21.py | 3 +++ base-requirements-dev.txt | 5 +++-- py2f/src/icon4py/py2f/codegen.py | 17 ++++++++--------- py2f/src/icon4py/py2f/parsing.py | 4 ++-- .../icon4py/py2f/wrappers/diffusion_wrapper.py | 2 +- py2f/tests/test_code_generation.py | 2 +- pyutils/tests/test_exceptions.py | 7 +++---- pyutils/tests/test_field_rendering.py | 3 ++- .../src/icon4py/testutils/serialbox_utils.py | 6 +++--- 28 files changed, 70 insertions(+), 63 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf index e33e884f8..88ecc9645 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf @@ -11,8 +11,8 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( _mo_intp_rbf_rbf_vec_interpol_vertex, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py index 23a407a31..ab7dcc0c4 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -11,8 +11,8 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulance import ( _calculate_diagnostics_for_turbulance, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py index e8943d417..15b75110c 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py @@ -11,9 +11,9 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, int32, where -from functional.program_processors.runners import gtfn_cpu +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, int32, where +from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.apply_nabla2_and_nabla4_to_vn import ( _apply_nabla2_and_nabla4_to_vn, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py index f3fddf390..0c647c1ce 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -11,9 +11,9 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, broadcast, int32, where -from functional.program_processors.runners import gtfn_cpu +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where +from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.apply_nabla2_to_w import _apply_nabla2_to_w from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import ( diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py index 6fdc44063..673dd6be5 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py @@ -11,9 +11,9 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field -from functional.program_processors.runners import gtfn_cpu +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( _enhance_diffusion_coefficient_for_grid_point_cold_pools, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py index 4e438df69..4c9d39693 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -11,9 +11,9 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field -from functional.program_processors.runners import gtfn_cpu +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.calculate_nabla2_for_z import _calculate_nabla2_for_z from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diagnostic_state.py b/atm_dyn_iconam/src/icon4py/diffusion/diagnostic_state.py index c3edbd21a..b74384af8 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diagnostic_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diagnostic_state.py @@ -13,7 +13,7 @@ from dataclasses import dataclass -from functional.common import Field +from gt4py.next.common import Field from icon4py.common.dimension import CellDim, KDim diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 7e9f90131..06692f857 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -18,9 +18,9 @@ from typing import Final, Optional, Tuple import numpy as np -from functional.common import Dimension -from functional.ffront.fbuiltins import Field, int32 -from functional.iterator.embedded import ( +from gt4py.next.common import Dimension +from gt4py.next.ffront.fbuiltins import Field, int32 +from gt4py.next.iterator.embedded import ( StridedNeighborOffsetProvider, np_as_located_field, ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py index 426a3d93b..75fb57932 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py @@ -12,10 +12,10 @@ # SPDX-License-Identifier: GPL-3.0-or-later -from functional.common import Field -from functional.ffront.decorator import program -from functional.ffront.fbuiltins import int32 -from functional.program_processors.runners import gtfn_cpu +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import program +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( _calculate_nabla2_and_smag_coefficients_for_vn, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py index a4228a941..531926929 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py @@ -13,7 +13,7 @@ from dataclasses import dataclass from typing import Final -from functional.common import Dimension, Field +from gt4py.next.common import Dimension, Field from icon4py.common import dimension from icon4py.common.dimension import CellDim, ECVDim, EdgeDim diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index dc6314708..9e49d22f2 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -14,9 +14,9 @@ from typing import Dict, Tuple import numpy as np -from functional.common import Dimension, DimensionKind, Field -from functional.ffront.fbuiltins import int32 -from functional.iterator.embedded import NeighborTableOffsetProvider +from gt4py.next.common import Dimension, DimensionKind, Field +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator.embedded import NeighborTableOffsetProvider from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim from icon4py.diffusion.horizontal import HorizontalMeshSize diff --git a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py index a4cbdd388..12620b7cf 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py @@ -13,7 +13,7 @@ from dataclasses import dataclass -from functional.common import Field +from gt4py.next.common import Field from icon4py.common.dimension import ( C2E2CODim, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py b/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py index faf172be5..0f28b2ca1 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py @@ -13,7 +13,7 @@ from dataclasses import dataclass -from functional.common import Field +from gt4py.next.common import Field from icon4py.common.dimension import C2E2CDim, CellDim, KDim diff --git a/atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py b/atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py index 88eb9a987..f659eaca9 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py @@ -13,7 +13,7 @@ from dataclasses import dataclass -from functional.common import Field +from gt4py.next.common import Field from icon4py.common.dimension import CellDim, EdgeDim, KDim diff --git a/atm_dyn_iconam/src/icon4py/diffusion/utils.py b/atm_dyn_iconam/src/icon4py/diffusion/utils.py index b6040838b..b3854d981 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/utils.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/utils.py @@ -13,10 +13,10 @@ from typing import Tuple import numpy as np -from functional.common import Dimension, Field -from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import broadcast, maximum, minimum -from functional.iterator.embedded import np_as_located_field +from gt4py.next.common import Dimension, Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import broadcast, maximum, minimum +from gt4py.next.iterator.embedded import np_as_located_field from icon4py.common.dimension import KDim, Koff, VertexDim diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 7e48620aa..44fe9afaf 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -116,7 +116,7 @@ def savepoint_exit(setup_icon_data, step_date_exit): fixture, passing 'step_data=' """ sp = IconSerialDataProvider( - "icon_diffusion_init", extracted_path, True + "icon_diffusion_init", str(extracted_path), True ).from_save_point_exit(linit=False, date=step_date_exit) return sp diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 2b1ebeb2e..8fe2641f8 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -115,7 +115,7 @@ def test_diff_multfac_vn_smag_limit_for_loop_run_with_k4_substeps(): def test_init_enh_smag_fac(): mesh = SimpleMesh() enh_smag_fac = zero_field(mesh, KDim) - a_vec = random_field(mesh, KDim, low=1.0, high=10.0) + a_vec = random_field(mesh, KDim, low=1.0, high=10.0, extend={KDim: 1}) fac = (0.67, 0.5, 1.3, 0.8) z = (0.1, 0.2, 0.3, 0.4) @@ -124,7 +124,7 @@ def test_init_enh_smag_fac(): _en_smag_fac_for_zero_nshift( a_vec, *fac, *z, out=enh_smag_fac, offset_provider={"Koff": KDim} ) - assert np.allclose(enhanced_smag_fac_np, np.asarray(enh_smag_fac[:-1])) + assert np.allclose(enhanced_smag_fac_np, np.asarray(enh_smag_fac)) def diff_multfac_vn_numpy(shape, k4, substeps): @@ -440,7 +440,7 @@ def test_run_diffusion_single_step( savepoint_init, savepoint_exit, icon_grid, r04b09_diffusion_config, damping_height ): sp = savepoint_init - vct_a = sp._vct_a() + vct_a = sp.vct_a() vertical_params = VerticalModelParams( vct_a=vct_a, rayleigh_damping_height=damping_height ) @@ -468,7 +468,7 @@ def test_run_diffusion_single_step( zd_diffcoef=sp.zd_diffcoef(), ) - diffusion = Diffusion(run_program=False) + diffusion = Diffusion(run_program=True) diffusion.init( grid=icon_grid, config=r04b09_diffusion_config, diff --git a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py index 829fb9391..4ae2fb0b1 100644 --- a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py @@ -12,6 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_15 import ( mo_nh_diffusion_stencil_15, @@ -60,6 +61,7 @@ def mo_nh_diffusion_stencil_15_numpy( ) +@pytest.mark.skip("new lowering: dims in offset provider") def test_mo_nh_diffusion_stencil_15(): mesh = SimpleMesh() diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py index 040a88ba2..4ee09f0e7 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py @@ -12,6 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_20 import ( mo_solve_nonhydro_stencil_20, @@ -64,6 +65,7 @@ def at_neighbor(i): return z_gradh_exner +@pytest.mark.skip("new lowering: dims in offset provider") def test_mo_solve_nonhydro_stencil_20(): mesh = SimpleMesh() diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py index d98f55326..9e76b51ec 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py @@ -12,6 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from gt4py.next.iterator.embedded import constant_field from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_21 import ( @@ -85,6 +86,7 @@ def _apply_index_field(shape, to_index, neighbor_table, offset_field): return z_hydro_corr +@pytest.mark.skip("new lowering: dims in offset provider") def test_mo_solve_nonhydro_stencil_21(): mesh = SimpleMesh() @@ -138,6 +140,7 @@ def test_mo_solve_nonhydro_stencil_21(): kend, offset_provider={ "E2C": mesh.get_e2c_offset_provider(), + "E2CDim": E2CDim, "Koff": KDim, }, ) diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 3d121fee7..29b446eac 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -1,5 +1,6 @@ # VCS --e git+https://github.com/tehrengruber/gt4py.git@2a9dcb0367a1ec1ef6d78951358d700424867b07#egg=gt4py-functional +#-e git+https://github.com/tehrengruber/gt4py.git@2a9dcb0367a1ec1ef6d78951358d700424867b07#egg=gt4py-functional +-e git+https://github.com/tehrengruber/gt4py.git@reduce_foast_lowering_complexity#egg=gt4py git+https://github.com/GridTools/serialbox#egg=serialbox&subdirectory=src/serialbox-python @@ -15,7 +16,7 @@ flake8-rst-docstrings>=0.0.14 isort~=5.10 mypy>=0.942 pre-commit~=2.15 -pytest>=7.2 +pytest>=6.1 pytest-cache>=1.0 pytest-cov>=2.8 pytest-factoryboy>=2.0 diff --git a/py2f/src/icon4py/py2f/codegen.py b/py2f/src/icon4py/py2f/codegen.py index 22907e9e3..f85592c3f 100644 --- a/py2f/src/icon4py/py2f/codegen.py +++ b/py2f/src/icon4py/py2f/codegen.py @@ -13,11 +13,10 @@ from typing import Sequence -import eve -from eve import codegen -from eve.codegen import JinjaTemplate as as_jinja -from eve.codegen import TemplatedGenerator -from functional.type_system.type_specifications import ScalarKind +from gt4py.eve import Node, codegen +from gt4py.eve.codegen import JinjaTemplate as as_jinja +from gt4py.eve.codegen import TemplatedGenerator +from gt4py.next.type_system.type_specifications import ScalarKind from icon4py.bindings.codegen.type_conversion import ( BUILTIN_TO_CPP_TYPE, @@ -26,23 +25,23 @@ from icon4py.bindings.utils import write_string -class DimensionType(eve.Node): +class DimensionType(Node): name: str length: int -class FuncParameter(eve.Node): +class FuncParameter(Node): name: str d_type: ScalarKind dimensions: Sequence[DimensionType] -class Func(eve.Node): +class Func(Node): name: str args: Sequence[FuncParameter] -class CffiPlugin(eve.Node): +class CffiPlugin(Node): name: str functions: Sequence[Func] diff --git a/py2f/src/icon4py/py2f/parsing.py b/py2f/src/icon4py/py2f/parsing.py index 763b554d7..0bf398d74 100644 --- a/py2f/src/icon4py/py2f/parsing.py +++ b/py2f/src/icon4py/py2f/parsing.py @@ -14,8 +14,8 @@ import importlib from inspect import signature, unwrap -from functional.type_system.type_specifications import ScalarType -from functional.type_system.type_translation import from_type_hint +from gt4py.next.type_system.type_specifications import ScalarType +from gt4py.next.type_system.type_translation import from_type_hint from icon4py.py2f.cffi_utils import CffiMethod from icon4py.py2f.codegen import CffiPlugin, DimensionType, Func, FuncParameter diff --git a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py index fb6f99e40..09e1dfce8 100644 --- a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py +++ b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later # flake8: noqa -from functional.common import Field +from gt4py.next.common import Field from icon4py.common.dimension import ( C2E2CDim, diff --git a/py2f/tests/test_code_generation.py b/py2f/tests/test_code_generation.py index 3270c07ec..bfe72b953 100644 --- a/py2f/tests/test_code_generation.py +++ b/py2f/tests/test_code_generation.py @@ -13,7 +13,7 @@ import string import pytest -from functional.type_system.type_specifications import ScalarKind +from gt4py.next.type_system.type_specifications import ScalarKind from icon4py.py2f.codegen import ( CffiPlugin, diff --git a/pyutils/tests/test_exceptions.py b/pyutils/tests/test_exceptions.py index a47c7a59c..0ac8b746c 100644 --- a/pyutils/tests/test_exceptions.py +++ b/pyutils/tests/test_exceptions.py @@ -12,8 +12,8 @@ # SPDX-License-Identifier: GPL-3.0-or-later import pytest -from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Dimension, Field, neighbor_sum +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Dimension, Field from icon4py.bindings.codegen.render.location import LocationRenderer from icon4py.bindings.entities import Offset, chain_from_str @@ -22,7 +22,7 @@ BindingsTypeConsistencyException, ) from icon4py.bindings.workflow import PyBindGen -from icon4py.common.dimension import E2CDim, EdgeDim, KDim +from icon4py.common.dimension import EdgeDim, KDim from icon4py.pyutils.metadata import get_stencil_info @@ -94,7 +94,6 @@ def bad_program( [field.renderer.render_ctype("py") for field in bindgen.fields] -@pytest.mark.skip("raises exception due to dims in offset provider") def test_scalar_sid_rendering_exception(): @field_operator def bad_stencil(a: Field[[EdgeDim], float], b: float) -> Field[[EdgeDim], float]: diff --git a/pyutils/tests/test_field_rendering.py b/pyutils/tests/test_field_rendering.py index e94d25104..a826597ce 100644 --- a/pyutils/tests/test_field_rendering.py +++ b/pyutils/tests/test_field_rendering.py @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - +import pytest from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum @@ -75,6 +75,7 @@ def identity_prog( ) +@pytest.mark.skip("new lowering: dims in offset provider") def test_vertical_sparse_field_sid_rendering(): @field_operator def reduction( diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 72ded87b9..d1d2bb9e9 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -13,9 +13,9 @@ import numpy as np import serialbox as ser -from functional.common import Dimension -from functional.ffront.fbuiltins import int32 -from functional.iterator.embedded import np_as_located_field +from gt4py.next.common import Dimension +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator.embedded import np_as_located_field from icon4py.common.dimension import ( C2E2CDim, From 7a62010c264fdd8b3c0423f0bf44ca0da1ea234c Mon Sep 17 00:00:00 2001 From: Magdalena Date: Thu, 23 Feb 2023 10:03:02 +0100 Subject: [PATCH 075/263] =?UTF-8?q?adapt=20serialbox=5Futils.py=20and=20fi?= =?UTF-8?q?xtures=20to=20new=20structure=20of=20serialized=20=E2=80=A6=20(?= =?UTF-8?q?#166)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * adapt serialbox_utils.py and fixtures to new structure of serialized data * update download url * unify names * change calculation of nudge_max_coeff in diffusion.init() due to changes in ICON. (mo_interpol_nml.f90: line 299) --- .../src/icon4py/diffusion/diffusion.py | 16 ++- atm_dyn_iconam/tests/conftest.py | 32 +++-- atm_dyn_iconam/tests/test_diffusion.py | 125 ++++++++++-------- atm_dyn_iconam/tests/test_vertical.py | 5 +- common/src/icon4py/common/constants.py | 3 + .../src/icon4py/testutils/serialbox_utils.py | 69 +++++----- 6 files changed, 146 insertions(+), 104 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 06692f857..fa19500cf 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -48,7 +48,11 @@ mo_intp_rbf_rbf_vec_interpol_vertex, ) from icon4py.atm_dyn_iconam.update_theta_and_exner import update_theta_and_exner -from icon4py.common.constants import CPD, GAS_CONSTANT_DRY_AIR +from icon4py.common.constants import ( + CPD, + DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO, + GAS_CONSTANT_DRY_AIR, +) from icon4py.common.dimension import ( C2E2CDim, C2E2CODim, @@ -308,6 +312,10 @@ def __init__(self, config: DiffusionConfig): self.smagorinski_factor, self.smagorinski_height, ) = self.determine_smagorinski_factor(config) + # see mo_interpol_nml.f90: + self.scaled_nudge_max_coeff = ( + config.nudge_max_coeff * DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO + ) def determine_smagorinski_factor(self, config: DiffusionConfig): """Enhanced Smagorinsky diffusion factor. @@ -432,9 +440,11 @@ def init( self._allocate_local_fields() self.nudgezone_diff: float = 0.04 / ( - config.nudge_max_coeff + sys.float_info.epsilon + params.scaled_nudge_max_coeff + sys.float_info.epsilon + ) + self.bdy_diff: float = 0.015 / ( + params.scaled_nudge_max_coeff + sys.float_info.epsilon ) - self.bdy_diff: float = 0.015 / (config.nudge_max_coeff + sys.float_info.epsilon) self.fac_bdydiff_v: float = ( math.sqrt(config.substep_as_float()) / config.velocity_boundary_diffusion_denominator diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 44fe9afaf..fb880d0c1 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -34,10 +34,10 @@ from icon4py.testutils.serialbox_utils import IconSerialDataProvider -data_uri = "https://polybox.ethz.ch/index.php/s/kP0Q2dDU6DytEqI/download" +data_uri = "https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download" data_path = Path(__file__).parent.joinpath("ser_icondata") extracted_path = data_path.joinpath("mch_ch_r04b09_dsl/ser_data") -data_file = data_path.joinpath("ser_data_diffusion.tar.gz").name +data_file = data_path.joinpath("mch_ch_r04b09_dsl_v2.tar.gz").name @pytest.fixture(scope="session") @@ -61,6 +61,11 @@ def setup_icon_data(): Path(data_file).unlink(missing_ok=True) +@pytest.fixture +def data_provider(setup_icon_data) -> IconSerialDataProvider: + return IconSerialDataProvider("icon_pydycore", str(extracted_path), True) + + @pytest.fixture def linit(): """ @@ -92,7 +97,7 @@ def step_date_exit(): @pytest.fixture -def savepoint_init(setup_icon_data, linit, step_date_init): +def diffusion_savepoint_init(data_provider, linit, step_date_init): """ Load data from ICON savepoint at start of diffusion module. @@ -101,36 +106,30 @@ def savepoint_init(setup_icon_data, linit, step_date_init): linit flag can be set by overriding the 'linit' fixture """ - sp = IconSerialDataProvider( - "icon_diffusion_init", str(extracted_path), True - ).from_savepoint_init(linit=linit, date=step_date_init) - return sp + return data_provider.from_savepoint_diffusion_init(linit=linit, date=step_date_init) @pytest.fixture -def savepoint_exit(setup_icon_data, step_date_exit): +def diffusion_savepoint_exit(data_provider, step_date_exit): """ Load data from ICON savepoint at exist of diffusion module. date of the timestamp to be selected can be set seperately by overriding the 'step_data' fixture, passing 'step_data=' """ - sp = IconSerialDataProvider( - "icon_diffusion_init", str(extracted_path), True - ).from_save_point_exit(linit=False, date=step_date_exit) + sp = data_provider.from_savepoint_diffusion_exit(linit=False, date=step_date_exit) return sp @pytest.fixture -def icon_grid(savepoint_init): +def icon_grid(data_provider): """ Load the icon grid from an ICON savepoint. Uses the default save_point from 'savepoint_init' fixture, however these data don't change for different time steps. """ - sp = savepoint_init - + sp = data_provider.from_savepoint_grid() sp_meta = sp.get_metadata("nproma", "nlev", "num_vert", "num_cells", "num_edges") cell_starts = sp.cells_start_index() @@ -165,6 +164,11 @@ def icon_grid(savepoint_init): return grid +@pytest.fixture +def grid_savepoint(data_provider): + return data_provider.from_savepoint_grid() + + @pytest.fixture def r04b09_diffusion_config(setup_icon_data) -> DiffusionConfig: """ diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 8fe2641f8..a4b0c9f17 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -222,37 +222,41 @@ def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): @pytest.mark.datatest def test_diffusion_init( - savepoint_init, icon_grid, r04b09_diffusion_config, step_date_init, damping_height + diffusion_savepoint_init, + grid_savepoint, + icon_grid, + r04b09_diffusion_config, + step_date_init, + damping_height, ): - savepoint = savepoint_init config = r04b09_diffusion_config additional_parameters = DiffusionParams(config) - vertical_params = VerticalModelParams(savepoint.vct_a(), damping_height) + vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) - meta = savepoint.get_metadata("nlev", "linit", "date") + meta = diffusion_savepoint_init.get_metadata("nlev", "linit", "date") assert meta["nlev"] == 65 assert meta["linit"] is False assert meta["date"] == step_date_init - grg = savepoint.geofac_grg() + grg = diffusion_savepoint_init.geofac_grg() interpolation_state = InterpolationState( - e_bln_c_s=savepoint.e_bln_c_s(), - rbf_coeff_1=savepoint.rbf_vec_coeff_v1(), - rbf_coeff_2=savepoint.rbf_vec_coeff_v2(), - geofac_div=savepoint.geofac_div(), - geofac_n2s=savepoint.geofac_n2s(), + e_bln_c_s=diffusion_savepoint_init.e_bln_c_s(), + rbf_coeff_1=diffusion_savepoint_init.rbf_vec_coeff_v1(), + rbf_coeff_2=diffusion_savepoint_init.rbf_vec_coeff_v2(), + geofac_div=diffusion_savepoint_init.geofac_div(), + geofac_n2s=diffusion_savepoint_init.geofac_n2s(), geofac_grg_x=grg[0], geofac_grg_y=grg[1], - nudgecoeff_e=savepoint.nudgecoeff_e(), + nudgecoeff_e=diffusion_savepoint_init.nudgecoeff_e(), ) metric_state = MetricState( - mask_hdiff=savepoint.mask_diff(), - theta_ref_mc=savepoint.theta_ref_mc(), - wgtfac_c=savepoint.wgtfac_c(), - zd_intcoef=savepoint.zd_intcoef(), - zd_vertidx=savepoint.zd_vertidx(), - zd_diffcoef=savepoint.zd_diffcoef(), + mask_hdiff=diffusion_savepoint_init.mask_diff(), + theta_ref_mc=diffusion_savepoint_init.theta_ref_mc(), + wgtfac_c=diffusion_savepoint_init.wgtfac_c(), + zd_intcoef=diffusion_savepoint_init.zd_intcoef(), + zd_vertidx=diffusion_savepoint_init.zd_vertidx(), + zd_diffcoef=diffusion_savepoint_init.zd_diffcoef(), ) diffusion = Diffusion() @@ -295,7 +299,7 @@ def test_diffusion_init( expected_enh_smag_fac = enhanced_smagorinski_factor_numpy( additional_parameters.smagorinski_factor, additional_parameters.smagorinski_height, - savepoint.vct_a(), + grid_savepoint.vct_a(), ) assert np.allclose(expected_enh_smag_fac, np.asarray(diffusion.enh_smag_fac)) @@ -318,16 +322,18 @@ def _verify_init_values_against_savepoint( assert np.allclose(savepoint.diff_multfac_smag(), diffusion.diff_multfac_smag) assert np.allclose(savepoint.smag_limit(), diffusion.smag_limit) - np.allclose(savepoint.diff_multfac_n2w(), np.asarray(diffusion.diff_multfac_n2w)) + assert np.allclose( + savepoint.diff_multfac_n2w(), np.asarray(diffusion.diff_multfac_n2w) + ) assert np.allclose(savepoint.diff_multfac_vn(), diffusion.diff_multfac_vn) @pytest.mark.datatest @pytest.mark.parametrize("linit", [True]) def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( - savepoint_init, r04b09_diffusion_config, icon_grid + diffusion_savepoint_init, r04b09_diffusion_config, icon_grid ): - savepoint = savepoint_init + savepoint = diffusion_savepoint_init config = r04b09_diffusion_config params = DiffusionParams(config) @@ -351,13 +357,17 @@ def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( @pytest.mark.datatest def test_verify_diffusion_init_against_first_regular_savepoint( - savepoint_init, r04b09_diffusion_config, icon_grid, damping_height + diffusion_savepoint_init, + grid_savepoint, + r04b09_diffusion_config, + icon_grid, + damping_height, ): config = r04b09_diffusion_config additional_parameters = DiffusionParams(config) - savepoint = savepoint_init - vct_a = savepoint.vct_a() + savepoint = diffusion_savepoint_init + vct_a = grid_savepoint.vct_a() grg = savepoint.geofac_grg() interpolation_state = InterpolationState( @@ -394,31 +404,34 @@ def test_verify_diffusion_init_against_first_regular_savepoint( @pytest.mark.datatest @pytest.mark.parametrize("step_date_init", ["2021-06-20T12:00:50.000"]) def test_verify_diffusion_init_against_other_regular_savepoint( - r04b09_diffusion_config, icon_grid, savepoint_init, damping_height + r04b09_diffusion_config, + grid_savepoint, + icon_grid, + diffusion_savepoint_init, + damping_height, ): config = r04b09_diffusion_config additional_parameters = DiffusionParams(config) - savepoint = savepoint_init - vertical_params = VerticalModelParams(savepoint.vct_a(), damping_height) - grg = savepoint.geofac_grg() + vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) + grg = diffusion_savepoint_init.geofac_grg() interpolation_state = InterpolationState( - e_bln_c_s=savepoint.e_bln_c_s(), - rbf_coeff_1=savepoint.rbf_vec_coeff_v1(), - rbf_coeff_2=savepoint.rbf_vec_coeff_v2(), - geofac_div=savepoint.geofac_div(), - geofac_n2s=savepoint.geofac_n2s(), + e_bln_c_s=diffusion_savepoint_init.e_bln_c_s(), + rbf_coeff_1=diffusion_savepoint_init.rbf_vec_coeff_v1(), + rbf_coeff_2=diffusion_savepoint_init.rbf_vec_coeff_v2(), + geofac_div=diffusion_savepoint_init.geofac_div(), + geofac_n2s=diffusion_savepoint_init.geofac_n2s(), geofac_grg_x=grg[0], geofac_grg_y=grg[1], - nudgecoeff_e=savepoint.nudgecoeff_e(), + nudgecoeff_e=diffusion_savepoint_init.nudgecoeff_e(), ) metric_state = MetricState( - mask_hdiff=savepoint.mask_diff(), - theta_ref_mc=savepoint.theta_ref_mc(), - wgtfac_c=savepoint.wgtfac_c(), - zd_intcoef=savepoint.zd_intcoef(), - zd_vertidx=savepoint.zd_vertidx(), - zd_diffcoef=savepoint.zd_diffcoef(), + mask_hdiff=diffusion_savepoint_init.mask_diff(), + theta_ref_mc=diffusion_savepoint_init.theta_ref_mc(), + wgtfac_c=diffusion_savepoint_init.wgtfac_c(), + zd_intcoef=diffusion_savepoint_init.zd_intcoef(), + zd_vertidx=diffusion_savepoint_init.zd_vertidx(), + zd_diffcoef=diffusion_savepoint_init.zd_diffcoef(), ) diffusion = Diffusion() @@ -431,15 +444,19 @@ def test_verify_diffusion_init_against_other_regular_savepoint( interpolation_state, ) - _verify_init_values_against_savepoint(savepoint, diffusion) + _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) @pytest.mark.datatest @pytest.mark.skip def test_run_diffusion_single_step( - savepoint_init, savepoint_exit, icon_grid, r04b09_diffusion_config, damping_height + diffusion_savepoint_init, + diffusion_savepoint_exit, + icon_grid, + r04b09_diffusion_config, + damping_height, ): - sp = savepoint_init + sp = diffusion_savepoint_init vct_a = sp.vct_a() vertical_params = VerticalModelParams( vct_a=vct_a, rayleigh_damping_height=damping_height @@ -519,10 +536,10 @@ def test_run_diffusion_single_step( cell_areas=cell_areas, ) - icon_result_exner = savepoint_exit.exner() - icon_result_vn = savepoint_exit.vn() - icon_result_w = savepoint_exit.w() - icon_result_theta_w = savepoint_exit.theta_v() + icon_result_exner = diffusion_savepoint_exit.exner() + icon_result_vn = diffusion_savepoint_exit.vn() + icon_result_w = diffusion_savepoint_exit.w() + icon_result_theta_w = diffusion_savepoint_exit.theta_v() assert np.allclose(icon_result_w, np.asarray(prognostic_state.w)) assert np.allclose(np.asarray(icon_result_vn), np.asarray(prognostic_state.vn)) @@ -539,12 +556,12 @@ def test_run_diffusion_single_step( def test_diffusion_five_steps( r04b09_diffusion_config, icon_grid, - savepoint_init, - savepoint_exit, + diffusion_savepoint_init, + diffusion_savepoint_exit, linit=True, step_date_exit="2021-06-20T12:01:00.000", ): - sp = savepoint_init + sp = diffusion_savepoint_init diagnostic_state = DiagnosticState( hdef_ic=sp.hdef_ic(), div_ic=sp.div_ic(), dwdx=sp.dwdx(), dwdy=sp.dwdy() @@ -635,10 +652,10 @@ def test_diffusion_five_steps( cell_areas=cell_areas, ) - icon_result_exner = savepoint_exit.exner() - icon_result_vn = savepoint_exit.vn() - icon_result_w = savepoint_exit.w() - icon_result_theta_w = savepoint_exit.theta_v() + icon_result_exner = diffusion_savepoint_exit.exner() + icon_result_vn = diffusion_savepoint_exit.vn() + icon_result_w = diffusion_savepoint_exit.w() + icon_result_theta_w = diffusion_savepoint_exit.theta_v() assert np.allclose(icon_result_w, np.asarray(prognostic_state.w)) assert np.allclose(np.asarray(icon_result_vn), np.asarray(prognostic_state.vn)) assert np.allclose( diff --git a/atm_dyn_iconam/tests/test_vertical.py b/atm_dyn_iconam/tests/test_vertical.py index 7da77063b..ccf0b7dbf 100644 --- a/atm_dyn_iconam/tests/test_vertical.py +++ b/atm_dyn_iconam/tests/test_vertical.py @@ -34,9 +34,8 @@ def test_nrdmax_calculation(max_h, damping, delta): @pytest.mark.datatest -def test_nrdmax_calculation_from_icon_input(icon_grid, savepoint_init): - sp = savepoint_init - a = sp.vct_a() +def test_nrdmax_calculation_from_icon_input(icon_grid, grid_savepoint): + a = grid_savepoint.vct_a() damping_height = 12500 vertical_params = VerticalModelParams( rayleigh_damping_height=damping_height, vct_a=a diff --git a/common/src/icon4py/common/constants.py b/common/src/icon4py/common/constants.py index e4d0cfb2c..8aa751bde 100644 --- a/common/src/icon4py/common/constants.py +++ b/common/src/icon4py/common/constants.py @@ -25,3 +25,6 @@ GRAVITATIONAL_ACCELERATION: Annotated[ float, "av. gravitational acceleration [m/s^2]" ] = 9.8066 +DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO: Annotated[ + float, "default physics to dynamics time step ratio" +] = 5.0 diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index d1d2bb9e9..9f19e769a 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -30,7 +30,7 @@ ) -class IconDiffusionSavepoint: +class IconSavepoint: def __init__(self, sp: ser.Savepoint, ser: ser.Serializer): self.savepoint = sp self.serializer = ser @@ -47,8 +47,32 @@ def get_metadata(self, *names): metadata = self.savepoint.metainfo.to_dict() return {n: metadata[n] for n in names if n in metadata} + def _read_int32_shift1(self, name: str): + """ + Read a index field and shift it by -1. + + use for start indeces: the shift accounts for the zero based python + values are converted to int32 + """ + return (self.serializer.read(name, self.savepoint) - 1).astype(int32) + + def _read_int32(self, name: str): + """ + Read a int field by name. + + use this for end indices: because FORTRAN slices are inclusive [from:to] _and_ one based + this accounts for being exclusive python exclusive bounds: [from:to) + field values are convert to int32 + """ + return self.serializer.read(name, self.savepoint).astype(int32) -class IconDiffusionInitSavepoint(IconDiffusionSavepoint): + def read_int(self, name: str): + buffer = self.serializer.read(name, self.savepoint).astype(int) + print(f"{name} {buffer.shape}") + return buffer + + +class IconGridSavePoint(IconSavepoint): def vct_a(self): return self._get_field("vct_a", KDim) @@ -85,33 +109,9 @@ def inv_dual_edge_length(self): def cells_start_index(self): return self._read_int32_shift1("c_start_index") - def _read_int32_shift1(self, name: str): - """ - Read a index field and shift it by -1. - - use for start indeces: the shift accounts for the zero based python - values are converted to int32 - """ - return (self.serializer.read(name, self.savepoint) - 1).astype(int32) - def cells_end_index(self): return self._read_int32("c_end_index") - def read_int(self, name: str): - buffer = self.serializer.read(name, self.savepoint).astype(int) - print(f"{name} {buffer.shape}") - return buffer - - def _read_int32(self, name: str): - """ - Read a int field by name. - - use this for end indices: because FORTRAN slices are inclusive [from:to] _and_ one based - this accounts for being exclusive python exclusive bounds: [from:to) - field values are convert to int32 - """ - return self.serializer.read(name, self.savepoint).astype(int32) - def vertex_start_index(self): return self._read_int32_shift1("v_start_index") @@ -156,6 +156,8 @@ def v2e(self): # subtract 1 to account for python being 0 based return self.serializer.read("v2e", self.savepoint) - 1 + +class IconDiffusionInitSavepoint(IconSavepoint): def hdef_ic(self): return self._get_field("hdef_ic", CellDim, KDim) @@ -217,7 +219,8 @@ def nudgecoeff_e(self): return self._get_field("nudgecoeff_e", EdgeDim) def zd_vertidx(self): - return self._get_field("zd_vertidx", CellDim, C2E2CDim, KDim, dtype=int) + # TODO fix this + return self._get_field("zd_vertidx", CellDim, C2E2CDim, dtype=int) def rbf_vec_coeff_v1(self): return self._get_field("rbf_vec_coeff_v1", VertexDim, V2EDim) @@ -253,7 +256,7 @@ def diff_multfac_vn(self): return self.serializer.read("diff_multfac_vn", self.savepoint) -class IconDiffusionExitSavepoint(IconDiffusionSavepoint): +class IconDiffusionExitSavepoint(IconSavepoint): def vn(self): return self._get_field("x_vn", EdgeDim, KDim) @@ -288,7 +291,13 @@ def print_info(self): print(f"SAVEPOINTS: {self.serializer.savepoint_list()}") print(f"FIELDNAMES: {self.serializer.fieldnames()}") - def from_savepoint_init(self, linit: bool, date: str) -> IconDiffusionInitSavepoint: + def from_savepoint_grid(self) -> IconGridSavePoint: + savepoint = self.serializer.savepoint["icon-grid"].id[1].as_savepoint() + return IconGridSavePoint(savepoint, self.serializer) + + def from_savepoint_diffusion_init( + self, linit: bool, date: str + ) -> IconDiffusionInitSavepoint: savepoint = ( self.serializer.savepoint["call-diffusion-init"] .linit[linit] @@ -297,7 +306,7 @@ def from_savepoint_init(self, linit: bool, date: str) -> IconDiffusionInitSavepo ) return IconDiffusionInitSavepoint(savepoint, self.serializer) - def from_save_point_exit( + def from_savepoint_diffusion_exit( self, linit: bool, date: str ) -> IconDiffusionExitSavepoint: savepoint = ( From 35abecca5190173ac867d1b4c5eb1b6152df06a6 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 23 Feb 2023 14:25:13 +0100 Subject: [PATCH 076/263] run get damping height index from correct location --- atm_dyn_iconam/src/icon4py/diffusion/diffusion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index fa19500cf..de86cdef8 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -708,7 +708,7 @@ def time_step( vertex_startindex_lb_plus1=vertex_startindex_lb_plus1, vertex_endindex_local=vertex_endindex_local, vertex_endindex_local_minus1=vertex_endindex_local_minus1, - index_of_damping_height=self.config.vertical_params._index_of_damping_height, + index_of_damping_height=self.vertical_params.index_of_damping_layer, nlev=self.grid.n_lev(), boundary_diffusion_start_index_edges=self.params.boundary_diffusion_start_index_edges, offset_provider={ From 00408e06d734575a32af8223d328200ef3857274 Mon Sep 17 00:00:00 2001 From: Magdalena Date: Thu, 2 Mar 2023 16:14:26 +0100 Subject: [PATCH 077/263] Fix diffusion tests (#169) Fixes for diffusion: - switch to gtfn backend - remove diffusion_stencil_15 numpy version - bugfixes: direct usage refin_ctrl marker for domain boundary, int32/int64 casts --- ...ate_nabla2_and_smag_coefficients_for_vn.py | 3 +- .../fused_mo_nh_diffusion_stencil_02_03.py | 3 +- .../fused_mo_nh_diffusion_stencil_11_12.py | 2 +- .../mo_intp_rbf_rbf_vec_interpol_vertex.py | 3 +- .../atm_dyn_iconam/update_theta_and_exner.py | 3 +- .../src/icon4py/diffusion/diffusion.py | 335 +++++++++--------- .../icon4py/diffusion/diffusion_program.py | 51 ++- .../src/icon4py/diffusion/icon_grid.py | 27 +- .../icon4py/diffusion/interpolation_state.py | 17 +- .../src/icon4py/diffusion/metric_state.py | 2 +- atm_dyn_iconam/src/icon4py/diffusion/utils.py | 13 +- atm_dyn_iconam/tests/conftest.py | 3 +- ...ate_nabla2_and_smag_coefficients_for_vn.py | 4 +- atm_dyn_iconam/tests/test_diffusion.py | 219 ++++++------ atm_dyn_iconam/tests/test_icon_grid.py | 7 +- .../src/icon4py/testutils/serialbox_utils.py | 36 +- .../src/icon4py/testutils/simple_mesh.py | 7 + 17 files changed, 414 insertions(+), 321 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py index 74becbc26..9da38e186 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py @@ -13,6 +13,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, maximum, minimum, sqrt +from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.common.dimension import ( E2C2V, @@ -135,7 +136,7 @@ def _calculate_nabla2_and_smag_coefficients_for_vn( return kh_smag_e, kh_smag_ec, z_nabla2_e -@program +@program(backend=gtfn_cpu.run_gtfn) def calculate_nabla2_and_smag_coefficients_for_vn( diff_multfac_smag: Field[[KDim], float], tangent_orientation: Field[[EdgeDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py index ab7dcc0c4..35cdf2127 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -13,6 +13,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulance import ( _calculate_diagnostics_for_turbulance, @@ -39,7 +40,7 @@ def _fused_mo_nh_diffusion_stencil_02_03( return div_ic, hdef_ic -@program +@program(backend=gtfn_cpu.run_gtfn) def fused_mo_nh_diffusion_stencil_02_03( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py index 673dd6be5..0ce4c8cdf 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py @@ -40,7 +40,7 @@ def _fused_mo_nh_diffusion_stencil_11_12( return kh_smag_e -@program(backend=gtfn_cpu) +@program(backend=gtfn_cpu.run_gtfn) def fused_mo_nh_diffusion_stencil_11_12( theta_v: Field[[CellDim, KDim], float], theta_ref_mc: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py index 2d34018a6..e062c3f53 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -13,6 +13,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum +from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.common.dimension import V2E, EdgeDim, KDim, V2EDim, VertexDim @@ -28,7 +29,7 @@ def _mo_intp_rbf_rbf_vec_interpol_vertex( return p_u_out, p_v_out -@program +@program(backend=gtfn_cpu.run_gtfn) def mo_intp_rbf_rbf_vec_interpol_vertex( p_e_in: Field[[EdgeDim, KDim], float], ptr_coeff_1: Field[[VertexDim, V2EDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py index c597ed3f2..7c73b3ca3 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py @@ -13,6 +13,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.common.dimension import CellDim, KDim @@ -31,7 +32,7 @@ def _update_theta_and_exner( return theta_v, exner -@program +@program(backend=gtfn_cpu.run_gtfn) def update_theta_and_exner( z_temp: Field[[CellDim, KDim], float], area: Field[[CellDim], float], diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index de86cdef8..2cba30b0e 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +# flake8: noqa import math import sys @@ -20,10 +21,7 @@ import numpy as np from gt4py.next.common import Dimension from gt4py.next.ffront.fbuiltins import Field, int32 -from gt4py.next.iterator.embedded import ( - StridedNeighborOffsetProvider, - np_as_located_field, -) +from gt4py.next.iterator.embedded import np_as_located_field import icon4py.diffusion.diffusion_program as diff_prog from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( @@ -47,6 +45,9 @@ from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_15 import ( + mo_nh_diffusion_stencil_15, +) from icon4py.atm_dyn_iconam.update_theta_and_exner import update_theta_and_exner from icon4py.common.constants import ( CPD, @@ -58,6 +59,8 @@ C2E2CODim, C2EDim, CellDim, + E2C2VDim, + E2CDim, ECVDim, EdgeDim, KDim, @@ -290,9 +293,6 @@ class DiffusionParams: """Calculates derived quantities depending on the diffusion config.""" def __init__(self, config: DiffusionConfig): - self.boundary_diffusion_start_index_edges = ( - 5 # mo_nh_diffusion.start_bdydiff_e - 1 = 5 -1 - ) self.K2: Final[float] = ( 1.0 / (config.hdiff_efdt_ratio * 8.0) @@ -362,36 +362,6 @@ def _diffusion_type_5_smagorinski_factor(config: DiffusionConfig): return factor, heights -def mo_nh_diffusion_stencil_15_numpy( - c2e2c, - mask_hdiff: Field[[CellDim, KDim], int], - zd_vertidx: Field[[C2E2CDim, KDim], int], - zd_diffcoef: Field[[CellDim, KDim], float], - geofac_n2s: Field[[CellDim, C2E2CODim], float], - vcoef: Field[[C2E2CDim, KDim], float], - theta_v: Field[[CellDim, KDim], float], - z_temp: Field[[CellDim, KDim], float], -): - geofac_n2s = np.asarray(geofac_n2s) - geofac_n2s_nbh = geofac_n2s[:, 1:] - geofac_n2s_c = geofac_n2s[:, 0] - mask_hdiff = np.asarray(mask_hdiff) - zd_vertidx = np.asarray(zd_vertidx) - zd_diffcoef = np.asarray(zd_diffcoef) - vcoef = np.asarray(vcoef) - second = (1 - vcoef) * theta_v[c2e2c][zd_vertidx + 1] - - z_temp = np.asarray(z_temp) - vertidx_ = vcoef * theta_v[c2e2c][zd_vertidx] - first = geofac_n2s_nbh * vertidx_ - summed = np.sum(first + second, axis=0) - - z_temp = np.where( - mask_hdiff, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + summed), z_temp - ) - return z_temp - - class Diffusion: """Class that configures diffusion and does one diffusion step.""" @@ -410,11 +380,11 @@ def __init__(self, run_program=True): self.interpolation_state = None self.metric_state = None self.diff_multfac_w: Optional[float] = None + self.diff_multfac_n2w: Field[[KDim], float] = None self.smag_offset: Optional[float] = None self.fac_bdydiff_v: Optional[float] = None self.bdy_diff: Optional[float] = None self.nudgezone_diff: Optional[float] = None - self.diff_multfac_n2w: Optional[float] = None def init( self, @@ -500,6 +470,7 @@ def _index_field(dim: Dimension, size=None): self.z_nabla2_e = _allocate(EdgeDim, KDim) self.z_temp = _allocate(CellDim, KDim) self.diff_multfac_smag = _allocate(KDim) + # TODO this is KHalfDim self.vertical_index = _index_field(KDim, self.grid.n_lev() + 1) self.horizontal_cell_index = _index_field(CellDim) self.horizontal_edge_index = _index_field(EdgeDim) @@ -512,7 +483,7 @@ def initial_step( tangent_orientation: Field[[EdgeDim], float], inverse_primal_edge_lengths: Field[[EdgeDim], float], inverse_dual_edge_length: Field[[EdgeDim], float], - inverse_vertical_vertex_lengths: Field[[EdgeDim], float], + inverse_vert_vert_lengths: Field[[EdgeDim], float], primal_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], dual_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], edge_areas: Field[[EdgeDim], float], @@ -535,7 +506,9 @@ def initial_step( setup_fields_for_initial_step( self.params.K4, self.config.hdiff_efdt_ratio, - out=(diff_multfac_vn, smag_limit), + diff_multfac_vn, + smag_limit, + offset_provider={}, ) self._do_diffusion_step( diagnostic_state, @@ -544,7 +517,7 @@ def initial_step( tangent_orientation, inverse_primal_edge_lengths, inverse_dual_edge_length, - inverse_vertical_vertex_lengths, + inverse_vert_vert_lengths, primal_normal_vert, dual_normal_vert, edge_areas, @@ -562,7 +535,7 @@ def time_step( tangent_orientation: Field[[EdgeDim], float], inverse_primal_edge_lengths: Field[[EdgeDim], float], inverse_dual_edge_length: Field[[EdgeDim], float], - inverse_vertical_vertex_lengths: Field[[EdgeDim], float], + inverse_vert_vert_lengths: Field[[EdgeDim], float], primal_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], dual_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], edge_areas: Field[[EdgeDim], float], @@ -573,56 +546,6 @@ def time_step( runs a diffusion step for the parameter linit=False, within regular time loop. """ - cell_startindex_nudging, cell_endindex_local = self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.local(CellDim), - ) - - ( - cell_startindex_interior, - cell_endindex_local_plus1, - ) = self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.interior(CellDim), - HorizontalMarkerIndex.local(CellDim) + 1, - ) - - ( - edge_startindex_nudging_plus1, - edge_endindex_local, - ) = self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim) + 1, - HorizontalMarkerIndex.local(EdgeDim), - ) - - ( - edge_startindex_nudging_minus1, - edge_endindex_local_minus2, - ) = self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim) - 1, - HorizontalMarkerIndex.local(EdgeDim) - 2, - ) - - ( - vertex_startindex_lb_plus3, - vertex_endindex_local, - ) = self.grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 3, - HorizontalMarkerIndex.local(VertexDim), - ) - - ( - vertex_startindex_lb_plus1, - vertex_endindex_local_minus1, - ) = self.grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 1, - HorizontalMarkerIndex.local(VertexDim) - 1, - ) if not self._run_program: self._do_diffusion_step( @@ -632,7 +555,7 @@ def time_step( tangent_orientation=tangent_orientation, inverse_primal_edge_lengths=inverse_primal_edge_lengths, inverse_dual_edge_length=inverse_dual_edge_length, - inverse_vertex_vertex_lengths=inverse_vertical_vertex_lengths, + inverse_vertex_vertex_lengths=inverse_vert_vert_lengths, primal_normal_vert=primal_normal_vert, dual_normal_vert=dual_normal_vert, edge_areas=edge_areas, @@ -642,6 +565,66 @@ def time_step( smag_offset=self.smag_offset, ) else: + print("run program") + ( + cell_startindex_nudging, + cell_endindex_local, + ) = self.grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.nudging(CellDim), + HorizontalMarkerIndex.local(CellDim), + ) + + ( + cell_startindex_interior, + cell_endindex_local_plus1, + ) = self.grid.get_indices_from_to( + CellDim, + HorizontalMarkerIndex.interior(CellDim), + HorizontalMarkerIndex.local(CellDim) - 1, + ) + + ( + edge_startindex_nudging_plus1, + edge_endindex_local, + ) = self.grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.nudging(EdgeDim) + 1, + HorizontalMarkerIndex.local(EdgeDim), + ) + + ( + edge_startindex_nudging_minus1, + edge_endindex_local_minus2, + ) = self.grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.nudging(EdgeDim) - 1, + HorizontalMarkerIndex.local(EdgeDim) - 2, + ) + + ( + vertex_startindex_lb_plus3, + vertex_endindex_local, + ) = self.grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.local_boundary(VertexDim) + 3, + HorizontalMarkerIndex.local(VertexDim), + ) + + ( + vertex_startindex_lb_plus1, + vertex_endindex_local_minus1, + ) = self.grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.local_boundary(VertexDim) + 1, + HorizontalMarkerIndex.local(VertexDim) - 1, + ) + edge_start_lb_plus4, _ = self.grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, + ) + diff_prog.diffusion_run( diagnostic_hdef_ic=diagnostic_state.hdef_ic, diagnostic_div_ic=diagnostic_state.div_ic, @@ -665,10 +648,12 @@ def time_step( interpolation_geofac_grg_y=self.interpolation_state.geofac_grg_y, interpolation_nudgecoeff_e=self.interpolation_state.nudgecoeff_e, interpolation_geofac_n2s=self.interpolation_state.geofac_n2s, + interpolation_geofac_n2s_c=self.interpolation_state.geofac_n2s_c, + interpolation_geofac_n2s_nbh=self.interpolation_state.geofac_n2s_nbh, tangent_orientation=tangent_orientation, inverse_primal_edge_lengths=inverse_primal_edge_lengths, inverse_dual_edge_lengths=inverse_dual_edge_length, - inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, + inverse_vert_vert_lengths=inverse_vert_vert_lengths, primal_normal_vert_1=primal_normal_vert[0], primal_normal_vert_2=primal_normal_vert[1], dual_normal_vert_1=dual_normal_vert[0], @@ -696,12 +681,13 @@ def time_step( local_vertical_index=self.vertical_index, local_horizontal_cell_index=self.horizontal_cell_index, local_horizontal_edge_index=self.horizontal_edge_index, - cell_startindex_interior=cell_startindex_interior, + cell_startindex_interior=int32(cell_startindex_interior), cell_startindex_nudging=cell_startindex_nudging, cell_endindex_local_plus1=cell_endindex_local_plus1, cell_endindex_local=cell_endindex_local, + cell_halo_idx=int32(cell_endindex_local), edge_startindex_nudging_plus1=edge_startindex_nudging_plus1, - edge_startindex_nudging_minus1=edge_startindex_nudging_minus1, + edge_startindex_nudging_minus1=int32(edge_startindex_nudging_minus1), edge_endindex_local=edge_endindex_local, edge_endindex_local_minus2=edge_endindex_local_minus2, vertex_startindex_lb_plus3=vertex_startindex_lb_plus3, @@ -710,18 +696,19 @@ def time_step( vertex_endindex_local_minus1=vertex_endindex_local_minus1, index_of_damping_height=self.vertical_params.index_of_damping_layer, nlev=self.grid.n_lev(), - boundary_diffusion_start_index_edges=self.params.boundary_diffusion_start_index_edges, + boundary_diffusion_start_index_edges=edge_start_lb_plus4, offset_provider={ "V2E": self.grid.get_v2e_connectivity(), "V2EDim": V2EDim, "E2C2V": self.grid.get_e2c2v_connectivity(), - "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, 4), + "E2C2VDim": E2C2VDim, + "E2CDim": E2CDim, + "E2ECV": self.grid.get_e2ecv_connectivity(), "C2E": self.grid.get_c2e_connectivity(), "C2EDim": C2EDim, "E2C": self.grid.get_e2c_connectivity(), "C2E2C": self.grid.get_c2e2c_connectivity(), "C2E2CDim": C2E2CDim, - "ECVDim": ECVDim, "C2E2CODim": C2E2CODim, "C2E2CO": self.grid.get_c2e2co_connectivity(), "Koff": KDim, @@ -733,18 +720,14 @@ def _do_diffusion_step( diagnostic_state: DiagnosticState, prognostic_state: PrognosticState, dtime: float, - tangent_orientation: Field[[EdgeDim], float], # from p_patch%edges - inverse_primal_edge_lengths: Field[[EdgeDim], float], # from p_patch%edges - inverse_dual_edge_length: Field[[EdgeDim], float], # from p_patch%edges - inverse_vertex_vertex_lengths: Field[[EdgeDim], float], # from p_patch%edges - primal_normal_vert: Tuple[ - Field[[ECVDim], float], Field[[ECVDim], float] - ], # from p_patch%edges - dual_normal_vert: Tuple[ - Field[[ECVDim], float], Field[[ECVDim], float] - ], # from p_patch%edges - edge_areas: Field[[EdgeDim], float], # from p_patch%edges - cell_areas: Field[[CellDim], float], # from p_atch%cells + tangent_orientation: Field[[EdgeDim], float], + inverse_primal_edge_lengths: Field[[EdgeDim], float], + inverse_dual_edge_length: Field[[EdgeDim], float], + inverse_vertex_vertex_lengths: Field[[EdgeDim], float], + primal_normal_vert: Tuple[Field[[ECVDim], float], Field[[ECVDim], float]], + dual_normal_vert: Tuple[Field[[ECVDim], float], Field[[ECVDim], float]], + edge_areas: Field[[EdgeDim], float], + cell_areas: Field[[CellDim], float], diff_multfac_vn: Field[[KDim], float], smag_limit: Field[[KDim], float], smag_offset: float, @@ -775,7 +758,7 @@ def _do_diffusion_step( cell_start_nudging_minus1, cell_end_local_plus1 = self.grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.nudging(CellDim) - 1, - HorizontalMarkerIndex.local(CellDim) + 1, + HorizontalMarkerIndex.local(CellDim) - 1, ) cell_start_interior, cell_end_local = self.grid.get_indices_from_to( @@ -796,6 +779,12 @@ def _do_diffusion_step( HorizontalMarkerIndex.local(EdgeDim), ) + edge_start_lb_plus4, _ = self.grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, + ) + ( edge_start_nudging_minus1, edge_end_local_minus2, @@ -831,7 +820,7 @@ def _do_diffusion_step( # TODO: is this needed?, if not remove set_zero_v_k(self.u_vert, offset_provider={}) set_zero_v_k(self.v_vert, offset_provider={}) - + print("rbf interpolation: start") # # 1. CALL rbf_vec_interpol_vertex mo_intp_rbf_rbf_vec_interpol_vertex( p_e_in=prognostic_state.vn, @@ -845,10 +834,11 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={"V2E": self.grid.get_v2e_connectivity(), "V2EDim": V2EDim}, ) - + print("rbf interpolation: end") # 2. HALO EXCHANGE -- CALL sync_patch_array_mult # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 + print("running calculate_nabla2_and_smag_coefficients_for_vn: start") calculate_nabla2_and_smag_coefficients_for_vn( diff_multfac_smag=self.diff_multfac_smag, tangent_orientation=tangent_orientation, @@ -866,16 +856,18 @@ def _do_diffusion_step( kh_smag_ec=self.kh_smag_ec, z_nabla2_e=self.z_nabla2_e, smag_offset=smag_offset, - horizontal_start=self.params.boundary_diffusion_start_index_edges, + horizontal_start=edge_start_lb_plus4, horizontal_end=edge_end_local_minus2, vertical_start=0, vertical_end=klevels, offset_provider={ "E2C2V": self.grid.get_e2c2v_connectivity(), - "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, 4), + "E2ECV": self.grid.get_e2ecv_connectivity(), + "E2C2VDim": E2C2VDim, }, ) - + print("running calculate_nabla2_and_smag_coefficients_for_vn: end") + print("running fused stencil fused stencil 02_03: start") fused_mo_nh_diffusion_stencil_02_03( kh_smag_ec=self.kh_smag_ec, vn=prognostic_state.vn, @@ -889,12 +881,18 @@ def _do_diffusion_step( horizontal_end=cell_end_local, vertical_start=0, vertical_end=klevels, - offset_provider={"C2E": self.grid.get_c2e_connectivity(), "C2EDim": C2EDim}, + offset_provider={ + "C2E": self.grid.get_c2e_connectivity(), + "C2EDim": C2EDim, + "Koff": KDim, + }, ) + print("running fused stencil fused stencil 02_03: end") # # # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH # # # 5. CALL rbf_vec_interpol_vertex_wp + print("rbf interpolation: start") mo_intp_rbf_rbf_vec_interpol_vertex( p_e_in=self.z_nabla2_e, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, @@ -907,42 +905,46 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={"V2E": self.grid.get_v2e_connectivity(), "V2EDim": V2EDim}, ) + print("rbf interpolation: end") # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult # # # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 # # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 # + + print("running fused stencil 04_05_06: start") fused_mo_nh_diffusion_stencil_04_05_06( - self.u_vert, - self.v_vert, - primal_normal_vert[0], - primal_normal_vert[1], - self.z_nabla2_e, - inverse_vertex_vertex_lengths, - inverse_primal_edge_lengths, - edge_areas, - self.kh_smag_e, - diff_multfac_vn, - self.interpolation_state.nudgecoeff_e, - prognostic_state.vn, - self.horizontal_edge_index, - self.nudgezone_diff, - self.fac_bdydiff_v, - edge_start_nudging_minus1, - prognostic_state.vn, + u_vert=self.u_vert, + v_vert=self.v_vert, + primal_normal_vert_v1=primal_normal_vert[0], + primal_normal_vert_v2=primal_normal_vert[1], + z_nabla2_e=self.z_nabla2_e, + inv_vert_vert_length=inverse_vertex_vertex_lengths, + inv_primal_edge_length=inverse_primal_edge_lengths, + area_edge=edge_areas, + kh_smag_e=self.kh_smag_e, + diff_multfac_vn=diff_multfac_vn, + nudgecoeff_e=self.interpolation_state.nudgecoeff_e, + vn=prognostic_state.vn, + horz_idx=self.horizontal_edge_index, + nudgezone_diff=self.nudgezone_diff, + fac_bdydiff_v=self.fac_bdydiff_v, + start_2nd_nudge_line_idx_e=int32(edge_start_nudging_minus1), horizontal_start=edge_start_nudging_plus_one, horizontal_end=edge_end_local, vertical_start=0, vertical_end=klevels, offset_provider={ "E2C2V": self.grid.get_e2c2v_connectivity(), - "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, 4), - "ECVDim": ECVDim, + "E2C2VDim": E2C2VDim, + "E2ECV": self.grid.get_e2ecv_connectivity(), }, ) # # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, # # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 + print("running fused stencil 04_05_06: end") + print("running fused stencil 07_08_09_10: start") fused_mo_nh_diffusion_stencil_07_08_09_10( area=cell_areas, geofac_n2s=self.interpolation_state.geofac_n2s, @@ -956,9 +958,13 @@ def _do_diffusion_step( diff_multfac_n2w=self.diff_multfac_n2w, vert_idx=self.vertical_index, horz_idx=self.horizontal_cell_index, - nrdmax=self.config.vertical_params._index_of_damping_height, - interior_idx=cell_start_interior, # h end index for stencil_09 and stencil_10 - halo_idx=cell_end_local, # h end index for stencil_09 and stencil_10, + nrdmax=self.vertical_params.index_of_damping_layer, + interior_idx=int32( + cell_start_interior + ), # h end index for stencil_09 and stencil_10 + halo_idx=int32( + cell_end_local + ), # h end index for stencil_09 and stencil_10, horizontal_start=cell_start_nudging, # h start index for stencil_07 and stencil_08 horizontal_end=cell_end_local_plus1, # h end index for stencil_07 and stencil_08 vertical_start=0, @@ -968,12 +974,14 @@ def _do_diffusion_step( "C2E2CODim": C2E2CODim, }, ) + print("running fused stencil 07_08_09_10: start") # # 8. HALO EXCHANGE: CALL sync_patch_array # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, # # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 # # # TODO check: kh_smag_e is an out field, should not be calculated in init? # + print("running fused stencil 11_12: start") fused_mo_nh_diffusion_stencil_11_12( theta_v=prognostic_state.theta_v, theta_ref_mc=self.metric_state.theta_ref_mc, @@ -987,8 +995,11 @@ def _do_diffusion_step( "E2C": self.grid.get_e2c_connectivity(), "C2E2C": self.grid.get_c2e2c_connectivity(), "C2E2CDim": C2E2CDim, + "E2CDim": E2CDim, }, ) + print("running fused stencil 11_12: end") + print("running fused stencil 13_14: start") fused_mo_nh_diffusion_stencil_13_14( kh_smag_e=self.kh_smag_e, inv_dual_edge_length=inverse_dual_edge_length, @@ -1002,28 +1013,33 @@ def _do_diffusion_step( offset_provider={ "C2E": self.grid.get_c2e_connectivity(), "E2C": self.grid.get_e2c_connectivity(), + "E2CDim": E2CDim, "C2EDim": C2EDim, }, ) - - mo_nh_diffusion_stencil_15_numpy( - mask_hdiff=self.metric_state.mask_hdiff, - zd_vertidx=self.metric_state.zd_vertidx, - vcoef=self.metric_state.zd_diffcoef, - zd_diffcoef=self.metric_state.zd_diffcoef, - geofac_n2s=self.interpolation_state.geofac_n2s, - theta_v=prognostic_state.theta_v, - z_temp=self.z_temp, - domain={ - CellDim: self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.halo(CellDim), - ), - }, - offset_provider={}, - ) - + print("running fused stencil 13_14: end") + print("running fused stencil 15: start") + # mo_nh_diffusion_stencil_15( + # self.metric_state.mask_hdiff, + # self.metric_state.zd_vertidx, + # self.metric_state.zd_diffcoef, + # self.interpolation_state.geofac_n2s_c, + # self.interpolation_state.geofac_n2s_nbh, + # self.metric_state.zd_intcoef, + # prognostic_state.theta_v, + # self.z_temp, + # cell_start_nudging, + # cell_end_local, + # 0, + # klevels, + # offset_provider={ + # "C2E2C": self.grid.get_c2e2c_connectivity(), + # "C2E2CDim": C2E2CDim, + # "Koff": KDim, + # }, + # ) + print("running fused stencil 15: end") + print("running fused stencil update_theta_and_exner: start") update_theta_and_exner( z_temp=self.z_temp, area=cell_areas, @@ -1036,4 +1052,5 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={}, ) + print("running fused stencil update_theta_and_exner: end") # 10. HALO EXCHANGE sync_patch_array diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py index 75fb57932..e624fe622 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +# flake8: noqa from gt4py.next.common import Field from gt4py.next.ffront.decorator import program @@ -77,10 +78,12 @@ def diffusion_run( interpolation_geofac_grg_y: Field[[CellDim, C2E2CODim], float], interpolation_nudgecoeff_e: Field[[EdgeDim], float], interpolation_geofac_n2s: Field[[CellDim, C2E2CODim], float], + interpolation_geofac_n2s_c: Field[[CellDim], float], + interpolation_geofac_n2s_nbh: Field[[CellDim, C2E2CDim], float], tangent_orientation: Field[[EdgeDim], float], inverse_primal_edge_lengths: Field[[EdgeDim], float], inverse_dual_edge_lengths: Field[[EdgeDim], float], - inverse_vertical_vertex_lengths: Field[[EdgeDim], float], + inverse_vert_vert_lengths: Field[[EdgeDim], float], primal_normal_vert_1: Field[[ECVDim], float], primal_normal_vert_2: Field[[ECVDim], float], dual_normal_vert_1: Field[[ECVDim], float], @@ -109,20 +112,21 @@ def diffusion_run( local_horizontal_cell_index: Field[[CellDim], int32], local_horizontal_edge_index: Field[[EdgeDim], int32], cell_startindex_interior: int32, - cell_startindex_nudging: int32, - cell_endindex_local_plus1: int32, - cell_endindex_local: int32, - edge_startindex_nudging_plus1: int32, + cell_halo_idx: int32, + cell_startindex_nudging: int, + cell_endindex_local_plus1: int, + cell_endindex_local: int, + edge_startindex_nudging_plus1: int, edge_startindex_nudging_minus1: int32, - edge_endindex_local: int32, - edge_endindex_local_minus2: int32, - vertex_startindex_lb_plus3: int32, - vertex_startindex_lb_plus1: int32, - vertex_endindex_local: int32, - vertex_endindex_local_minus1: int32, + edge_endindex_local: int, + edge_endindex_local_minus2: int, + vertex_startindex_lb_plus3: int, + vertex_startindex_lb_plus1: int, + vertex_endindex_local: int, + vertex_endindex_local_minus1: int, index_of_damping_height: int32, nlev: int, - boundary_diffusion_start_index_edges: int32, + boundary_diffusion_start_index_edges: int, ): _scale_k(local_enh_smag_fac, dtime, out=local_diff_multfac_smag) @@ -149,7 +153,7 @@ def diffusion_run( local_diff_multfac_smag, tangent_orientation, inverse_primal_edge_lengths, - inverse_vertical_vertex_lengths, + inverse_vert_vert_lengths, local_u_vert, local_v_vert, primal_normal_vert_1, @@ -161,6 +165,7 @@ def diffusion_run( local_smag_offset, out=(local_kh_smag_e, local_kh_smag_ec, local_z_nabla2_e), domain={ + # TODO wrong start index?? EdgeDim: (boundary_diffusion_start_index_edges, edge_endindex_local_minus2), KDim: (0, nlev), }, @@ -207,7 +212,7 @@ def diffusion_run( primal_normal_vert_1, primal_normal_vert_2, local_z_nabla2_e, - inverse_vertical_vertex_lengths, + inverse_vert_vert_lengths, inverse_primal_edge_lengths, edge_areas, local_kh_smag_e, @@ -242,7 +247,7 @@ def diffusion_run( local_horizontal_cell_index, index_of_damping_height, cell_startindex_interior, - cell_endindex_local, + cell_halo_idx, out=( prognostic_w, diagnostic_dwdx, @@ -283,7 +288,21 @@ def diffusion_run( }, ) - # MO_NH_DIFFUSION_STENCIL_15: needs index fields! + # MO_NH_DIFFUSION_STENCIL_15: as_offset index fields! + # _mo_nh_diffusion_stencil_15( + # metric_mask_hdiff, + # metric_zd_vertidx, + # metric_zd_diffcoef, + # interpolation_geofac_n2s_c, + # interpolation_geofac_n2s_nbh, + # metric_zd_intcoef, + # prognostic_theta_v, + # local_z_temp, + # domain = { + # CellDim:(cell_startindex_nudging, cell_endindex_local), + # KDim: (0,nlev) + # } + # ) _update_theta_and_exner( local_z_temp, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index 9e49d22f2..216593fa0 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -18,7 +18,7 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import NeighborTableOffsetProvider -from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim +from icon4py.common.dimension import CellDim, ECVDim, EdgeDim, KDim, VertexDim from icon4py.diffusion.horizontal import HorizontalMeshSize @@ -105,7 +105,7 @@ def with_start_end_indices( @builder def with_connectivities(self, connectivity: Dict[Dimension, np.ndarray]): self.connectivities.update( - {d.value.lower(): k for d, k in connectivity.items()} + {d.value.lower(): k.astype(int) for d, k in connectivity.items()} ) self.size.update({d: t.shape[1] for d, t in connectivity.items()}) @@ -164,15 +164,20 @@ def get_c2e2co_connectivity(self): return NeighborTableOffsetProvider(table, CellDim, CellDim, table.shape[1]) def get_e2c2v_connectivity(self): - return self.get_e2v_connectivity() - - def get_e2c2v_size(self): - self.connectivities["e2v"].shape[1] + table = self.connectivities["e2c2v"] + return NeighborTableOffsetProvider(table, EdgeDim, VertexDim, table.shape[1]) def get_v2e_connectivity(self): table = self.connectivities["v2e"] return NeighborTableOffsetProvider(table, VertexDim, EdgeDim, table.shape[1]) + def get_e2ecv_connectivity(self): + old_shape = self.connectivities["e2c2v"].shape + v2ecv_table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) + return NeighborTableOffsetProvider( + v2ecv_table, EdgeDim, ECVDim, v2ecv_table.shape[1] + ) + class VerticalModelParams: def __init__(self, vct_a: Field[[KDim], float], rayleigh_damping_height: float): @@ -183,10 +188,12 @@ def __init__(self, vct_a: Field[[KDim], float], rayleigh_damping_height: float): vct_a: field containing the physical heights of the k level rayleigh_damping_height: height of rayleigh damping in [m] mo_nonhydro_nml """ - self.rayleigh_damping_height = rayleigh_damping_height + self._rayleigh_damping_height = rayleigh_damping_height self._vct_a = vct_a self._index_of_damping_height = int32( - np.argmax(np.where(np.asarray(self._vct_a) >= self.rayleigh_damping_height)) + np.argmax( + np.where(np.asarray(self._vct_a) >= self._rayleigh_damping_height) + ) ) @property @@ -196,3 +203,7 @@ def index_of_damping_layer(self): @property def physical_heights(self) -> Field[[KDim], float]: return self._vct_a + + @property + def rayleigh_damping_height(self): + return self._rayleigh_damping_height diff --git a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py index 12620b7cf..2486708e0 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py @@ -13,9 +13,12 @@ from dataclasses import dataclass +import numpy as np from gt4py.next.common import Field +from gt4py.next.iterator.embedded import np_as_located_field from icon4py.common.dimension import ( + C2E2CDim, C2E2CODim, C2EDim, CellDim, @@ -53,5 +56,17 @@ class InterpolationState: geofac_grg_x: Field[ [CellDim, C2E2CODim], float ] # factor for green gauss gradient (nproma,4,nblks_c,2) - geofac_grg_y: Field[[CellDim, C2E2CODim], float] # TODO combine to tuple + geofac_grg_y: Field[ + [CellDim, C2E2CODim], float + ] # TODO combine geofac_grg_x and geofac_grg_y to tuple nudgecoeff_e: Field[[EdgeDim], float] # Nudgeing coeffients for edges + + @property + def geofac_n2s_c(self) -> Field[[CellDim], float]: + return np_as_located_field(CellDim)(np.asarray(self.geofac_n2s)[:, 0]) + + @property + def geofac_n2s_nbh(self) -> Field[[CellDim, C2E2CDim], float]: + return np_as_located_field(CellDim, C2E2CDim)( + np.asarray(self.geofac_n2s)[:, 1:] + ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py b/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py index 0f28b2ca1..069863583 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py @@ -24,7 +24,7 @@ class MetricState: wgtfac_c: Field[ [CellDim, KDim], float ] # weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_c) - mask_hdiff: Field[[CellDim, KDim], int] + mask_hdiff: Field[[CellDim, KDim], bool] zd_vertidx: Field[[CellDim, C2E2CDim, KDim], int] zd_diffcoef: Field[[CellDim, KDim], float] zd_intcoef: Field[[CellDim, C2E2CDim, KDim], float] diff --git a/atm_dyn_iconam/src/icon4py/diffusion/utils.py b/atm_dyn_iconam/src/icon4py/diffusion/utils.py index b3854d981..c7d9004fe 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/utils.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/utils.py @@ -15,8 +15,9 @@ import numpy as np from gt4py.next.common import Dimension, Field from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import broadcast, maximum, minimum +from gt4py.next.ffront.fbuiltins import broadcast, int32, maximum, minimum from gt4py.next.iterator.embedded import np_as_located_field +from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.common.dimension import KDim, Koff, VertexDim @@ -32,7 +33,7 @@ def _scale_k(field: Field[[KDim], float], factor: float) -> Field[[KDim], float] return field * factor -@program +@program(backend=gtfn_cpu.run_gtfn) def scale_k( field: Field[[KDim], float], factor: float, scaled_field: Field[[KDim], float] ): @@ -44,7 +45,7 @@ def _set_zero_v_k() -> Field[[VertexDim, KDim], float]: return broadcast(0.0, (VertexDim, KDim)) -@program +@program(backend=gtfn_cpu.run_gtfn) def set_zero_v_k(field: Field[[VertexDim, KDim], float]): _set_zero_v_k(out=field) @@ -80,7 +81,7 @@ def _setup_fields_for_initial_step( return diff_multfac_vn, smag_limit -@program +@program(backend=gtfn_cpu.run_gtfn) def setup_fields_for_initial_step( k4: float, hdiff_efdt_ratio: float, @@ -155,7 +156,7 @@ def _init_diffusion_local_fields_for_regular_timestemp( ) -@program +@program(backend=gtfn_cpu.run_gtfn) def init_diffusion_local_fields_for_regular_timestep( k4: float, dyn_substeps: float, @@ -193,7 +194,7 @@ def init_diffusion_local_fields_for_regular_timestep( def init_nabla2_factor_in_upper_damping_zone( - k_size: int, nrdmax: int, nshift: int, physical_heights: np.ndarray + k_size: int, nrdmax: int32, nshift: int, physical_heights: np.ndarray ) -> Field[[KDim], float]: """ Calculate diff_multfac_n2w. diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index fb880d0c1..99e6db30b 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -22,6 +22,7 @@ C2E2CODim, C2EDim, CellDim, + E2C2VDim, E2CDim, E2VDim, EdgeDim, @@ -159,7 +160,7 @@ def icon_grid(data_provider): .with_connectivities( {C2EDim: sp.c2e(), E2CDim: sp.e2c(), C2E2CDim: c2e2c, C2E2CODim: c2e2c0} ) - .with_connectivities({E2VDim: sp.e2v(), V2EDim: sp.v2e()}) + .with_connectivities({E2VDim: sp.e2v(), V2EDim: sp.v2e(), E2C2VDim: sp.e2c2v()}) ) return grid diff --git a/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index 967b2b153..29dd176d0 100644 --- a/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -12,7 +12,6 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( calculate_nabla2_and_smag_coefficients_for_vn, @@ -210,7 +209,8 @@ def test_calculate_nabla2_and_smag_coefficients_for_vn(): mesh.k_level, offset_provider={ "E2C2V": mesh.get_e2c2v_offset_provider(), - "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, mesh.n_e2c2v), + "E2ECV": mesh.get_e2ecv_offset_provider(), + "E2C2VDim": E2C2VDim, }, ) assert np.allclose(kh_smag_e_ref, kh_smag_e) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index a4b0c9f17..24316b5a9 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -255,7 +255,7 @@ def test_diffusion_init( theta_ref_mc=diffusion_savepoint_init.theta_ref_mc(), wgtfac_c=diffusion_savepoint_init.wgtfac_c(), zd_intcoef=diffusion_savepoint_init.zd_intcoef(), - zd_vertidx=diffusion_savepoint_init.zd_vertidx(), + zd_vertidx=diffusion_savepoint_init.zd_vertoffset(), zd_diffcoef=diffusion_savepoint_init.zd_diffcoef(), ) @@ -385,7 +385,7 @@ def test_verify_diffusion_init_against_first_regular_savepoint( theta_ref_mc=savepoint.theta_ref_mc(), wgtfac_c=savepoint.wgtfac_c(), zd_intcoef=savepoint.zd_intcoef(), - zd_vertidx=savepoint.zd_vertidx(), + zd_vertidx=savepoint.zd_vertoffset(), zd_diffcoef=savepoint.zd_diffcoef(), ) diffusion = Diffusion() @@ -430,7 +430,7 @@ def test_verify_diffusion_init_against_other_regular_savepoint( theta_ref_mc=diffusion_savepoint_init.theta_ref_mc(), wgtfac_c=diffusion_savepoint_init.wgtfac_c(), zd_intcoef=diffusion_savepoint_init.zd_intcoef(), - zd_vertidx=diffusion_savepoint_init.zd_vertidx(), + zd_vertidx=diffusion_savepoint_init.zd_vertoffset(), zd_diffcoef=diffusion_savepoint_init.zd_diffcoef(), ) @@ -447,45 +447,41 @@ def test_verify_diffusion_init_against_other_regular_savepoint( _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) +@pytest.mark.skip("fix: diffusion_stencil_15") +@pytest.mark.parametrize("run_with_program", [True, False]) @pytest.mark.datatest -@pytest.mark.skip def test_run_diffusion_single_step( + run_with_program, diffusion_savepoint_init, diffusion_savepoint_exit, + grid_savepoint, icon_grid, r04b09_diffusion_config, damping_height, ): - sp = diffusion_savepoint_init - vct_a = sp.vct_a() + ( + cell_areas, + diagnostic_state, + dtime, + dual_normal_vert, + edge_areas, + interpolation_state, + inverse_dual_edge_length, + inverse_primal_edge_lengths, + inverse_vertical_vertex_lengths, + metric_state, + orientation, + primal_normal_vert, + prognostic_state, + ) = _read_fields(diffusion_savepoint_init, grid_savepoint) + + vct_a = grid_savepoint.vct_a() vertical_params = VerticalModelParams( vct_a=vct_a, rayleigh_damping_height=damping_height ) - additional_parameters = DiffusionParams(r04b09_diffusion_config) - grg = sp.geofac_grg() - interpolation_state = InterpolationState( - e_bln_c_s=sp.e_bln_c_s(), - rbf_coeff_1=sp.rbf_vec_coeff_v1(), - rbf_coeff_2=sp.rbf_vec_coeff_v2(), - geofac_div=sp.geofac_div(), - geofac_n2s=sp.geofac_n2s(), - geofac_grg_x=grg[0], - geofac_grg_y=grg[1], - nudgecoeff_e=sp.nudgecoeff_e(), - ) - - metric_state = MetricState( - mask_hdiff=sp.mask_diff(), - theta_ref_mc=sp.theta_ref_mc(), - wgtfac_c=sp.wgtfac_c(), - zd_intcoef=sp.zd_intcoef(), - zd_vertidx=sp.zd_vertidx(), - zd_diffcoef=sp.zd_diffcoef(), - ) - - diffusion = Diffusion(run_program=True) + diffusion = Diffusion(run_program=run_with_program) diffusion.init( grid=icon_grid, config=r04b09_diffusion_config, @@ -494,34 +490,6 @@ def test_run_diffusion_single_step( metric_state=metric_state, interpolation_state=interpolation_state, ) - - diagnostic_state = DiagnosticState( - hdef_ic=sp.hdef_ic(), div_ic=sp.div_ic(), dwdx=sp.dwdx(), dwdy=sp.dwdy() - ) - prognostic_state = PrognosticState( - w=sp.w(), - vn=sp.vn(), - exner_pressure=sp.exner(), - theta_v=sp.theta_v(), - ) - - dtime = sp.get_metadata("dtime").get("dtime") - orientation = sp.tangent_orientation() - - inverse_primal_edge_lengths = sp.inverse_primal_edge_lengths() - inverse_vertical_vertex_lengths = sp.inv_vert_vert_length() - inverse_dual_edge_length = sp.inv_dual_edge_length() - primal_normal_vert: VectorTuple = ( - as_1D_sparse_field(sp.primal_normal_vert_x(), ECVDim), - as_1D_sparse_field(sp.primal_normal_vert_y(), ECVDim), - ) - dual_normal_vert: VectorTuple = ( - as_1D_sparse_field(sp.dual_normal_vert_x(), ECVDim), - as_1D_sparse_field(sp.dual_normal_vert_y(), ECVDim), - ) - edge_areas = sp.edge_areas() - cell_areas = sp.cell_areas() - diffusion.time_step( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, @@ -529,7 +497,7 @@ def test_run_diffusion_single_step( tangent_orientation=orientation, inverse_primal_edge_lengths=inverse_primal_edge_lengths, inverse_dual_edge_length=inverse_dual_edge_length, - inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, + inverse_vert_vert_lengths=inverse_vertical_vertex_lengths, primal_normal_vert=primal_normal_vert, dual_normal_vert=dual_normal_vert, edge_areas=edge_areas, @@ -551,66 +519,101 @@ def test_run_diffusion_single_step( ) -@pytest.mark.skip -@pytest.mark.datatest -def test_diffusion_five_steps( - r04b09_diffusion_config, - icon_grid, - diffusion_savepoint_init, - diffusion_savepoint_exit, - linit=True, - step_date_exit="2021-06-20T12:01:00.000", -): - sp = diffusion_savepoint_init - - diagnostic_state = DiagnosticState( - hdef_ic=sp.hdef_ic(), div_ic=sp.div_ic(), dwdx=sp.dwdx(), dwdy=sp.dwdy() - ) - prognostic_state = PrognosticState( - w=sp.w(), - vn=sp.vn(), - exner_pressure=sp.exner(), - theta_v=sp.theta_v(), - ) - grg = sp.geofac_grg() +def _read_fields(diffusion_savepoint_init, grid_savepoint): + grg = diffusion_savepoint_init.geofac_grg() interpolation_state = InterpolationState( - e_bln_c_s=sp.e_bln_c_s(), - rbf_coeff_1=sp.rbf_vec_coeff_v1(), - rbf_coeff_2=sp.rbf_vec_coeff_v2(), - geofac_div=sp.geofac_div(), - geofac_n2s=sp.geofac_n2s(), + e_bln_c_s=diffusion_savepoint_init.e_bln_c_s(), + rbf_coeff_1=diffusion_savepoint_init.rbf_vec_coeff_v1(), + rbf_coeff_2=diffusion_savepoint_init.rbf_vec_coeff_v2(), + geofac_div=diffusion_savepoint_init.geofac_div(), + geofac_n2s=diffusion_savepoint_init.geofac_n2s(), geofac_grg_x=grg[0], geofac_grg_y=grg[1], - nudgecoeff_e=sp.nudgecoeff_e(), + nudgecoeff_e=diffusion_savepoint_init.nudgecoeff_e(), ) - metric_state = MetricState( - mask_hdiff=sp.mask_diff(), - theta_ref_mc=sp.theta_ref_mc(), - wgtfac_c=sp.wgtfac_c(), - zd_intcoef=sp.zd_intcoef(), - zd_vertidx=sp.zd_vertidx(), - zd_diffcoef=sp.zd_diffcoef(), - ) - dtime = sp.get_metadata("dtime").get("dtime") - orientation = sp.tangent_orientation() - inverse_primal_edge_lengths = sp.inverse_primal_edge_lengths() - inverse_vertical_vertex_lengths = sp.inv_vert_vert_length() - inverse_dual_edge_length = sp.inv_dual_edge_length() + mask_hdiff=diffusion_savepoint_init.mask_diff(), + theta_ref_mc=diffusion_savepoint_init.theta_ref_mc(), + wgtfac_c=diffusion_savepoint_init.wgtfac_c(), + zd_intcoef=diffusion_savepoint_init.zd_intcoef(), + zd_vertidx=diffusion_savepoint_init.zd_vertoffset(), + zd_diffcoef=diffusion_savepoint_init.zd_diffcoef(), + ) + diagnostic_state = DiagnosticState( + hdef_ic=diffusion_savepoint_init.hdef_ic(), + div_ic=diffusion_savepoint_init.div_ic(), + dwdx=diffusion_savepoint_init.dwdx(), + dwdy=diffusion_savepoint_init.dwdy(), + ) + prognostic_state = PrognosticState( + w=diffusion_savepoint_init.w(), + vn=diffusion_savepoint_init.vn(), + exner_pressure=diffusion_savepoint_init.exner(), + theta_v=diffusion_savepoint_init.theta_v(), + ) + dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") + orientation = grid_savepoint.tangent_orientation() + inverse_primal_edge_lengths = grid_savepoint.inverse_primal_edge_lengths() + inverse_vert_vert_lengths = grid_savepoint.inv_vert_vert_length() + inverse_dual_edge_length = grid_savepoint.inv_dual_edge_length() primal_normal_vert: VectorTuple = ( - as_1D_sparse_field(sp.primal_normal_vert_x(), ECVDim), - as_1D_sparse_field(sp.primal_normal_vert_y(), ECVDim), + as_1D_sparse_field(grid_savepoint.primal_normal_vert_x(), ECVDim), + as_1D_sparse_field(grid_savepoint.primal_normal_vert_y(), ECVDim), ) dual_normal_vert: VectorTuple = ( - as_1D_sparse_field(sp.dual_normal_vert_x(), ECVDim), - as_1D_sparse_field(sp.dual_normal_vert_y(), ECVDim), + as_1D_sparse_field(grid_savepoint.dual_normal_vert_x(), ECVDim), + as_1D_sparse_field(grid_savepoint.dual_normal_vert_y(), ECVDim), + ) + edge_areas = grid_savepoint.edge_areas() + cell_areas = grid_savepoint.cell_areas() + return ( + cell_areas, + diagnostic_state, + dtime, + dual_normal_vert, + edge_areas, + interpolation_state, + inverse_dual_edge_length, + inverse_primal_edge_lengths, + inverse_vert_vert_lengths, + metric_state, + orientation, + primal_normal_vert, + prognostic_state, ) - edge_areas = sp.edge_areas() - cell_areas = sp.cell_areas() + + +@pytest.mark.skip("fix: diffusion_stencil_15") +@pytest.mark.datatest +def test_diffusion_five_steps( + damping_height, + r04b09_diffusion_config, + icon_grid, + grid_savepoint, + diffusion_savepoint_init, + diffusion_savepoint_exit, + linit=True, + step_date_exit="2021-06-20T12:01:00.000", +): + ( + cell_areas, + diagnostic_state, + dtime, + dual_normal_vert, + edge_areas, + interpolation_state, + inverse_dual_edge_length, + inverse_primal_edge_lengths, + inverse_vert_vert_lengths, + metric_state, + orientation, + primal_normal_vert, + prognostic_state, + ) = _read_fields(diffusion_savepoint_init, grid_savepoint) vertical_params = VerticalModelParams( - vct_a=(sp._vct_a()), rayleigh_damping_height=12500.0 + vct_a=grid_savepoint.vct_a(), rayleigh_damping_height=damping_height ) additional_parameters = DiffusionParams(r04b09_diffusion_config) @@ -631,7 +634,7 @@ def test_diffusion_five_steps( tangent_orientation=orientation, inverse_primal_edge_lengths=inverse_primal_edge_lengths, inverse_dual_edge_length=inverse_dual_edge_length, - inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, + inverse_vert_vert_lengths=inverse_vert_vert_lengths, primal_normal_vert=primal_normal_vert, dual_normal_vert=dual_normal_vert, edge_areas=edge_areas, @@ -645,7 +648,7 @@ def test_diffusion_five_steps( tangent_orientation=orientation, inverse_primal_edge_lengths=inverse_primal_edge_lengths, inverse_dual_edge_length=inverse_dual_edge_length, - inverse_vertical_vertex_lengths=inverse_vertical_vertex_lengths, + inverse_vert_vert_lengths=inverse_vert_vert_lengths, primal_normal_vert=primal_normal_vert, dual_normal_vert=dual_normal_vert, edge_areas=edge_areas, diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index 854d49f2b..256ded7aa 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -25,7 +25,8 @@ def test_horizontal_grid_cell_indices(icon_grid): ) == ( 20896, 20896, - ) # halo +1 + ) # halo + 1 + assert icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.local(CellDim), @@ -33,7 +34,7 @@ def test_horizontal_grid_cell_indices(icon_grid): ) == ( -1, 20896, - ) # halo in icon is (1,20896) why + ) # halo in icon is (1,20896) assert icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.interior(CellDim), @@ -45,7 +46,7 @@ def test_horizontal_grid_cell_indices(icon_grid): assert icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, - HorizontalMarkerIndex._INTERIOR_CELLS + 1, + HorizontalMarkerIndex.interior(CellDim) + 1, ) == ( 0, 850, diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 9f19e769a..9a037e60c 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -101,7 +101,7 @@ def cell_areas(self): return self._get_field("cell_areas", CellDim) def edge_areas(self): - return self._get_field("cell_areas", EdgeDim) + return self._get_field("edge_areas", EdgeDim) def inv_dual_edge_length(self): return self._get_field("inv_dual_edge_length", EdgeDim) @@ -126,6 +126,9 @@ def edge_end_index(self): # one off accounts for being exclusive [from:to) return self.serializer.read("e_end_index", self.savepoint) + def print_connectivity_info(name: str, ar: np.ndarray): + print(f" connectivity {name} {ar.shape}") + def refin_ctrl(self, dim: Dimension): if dim == CellDim: return self.serializer.read("c_refin_ctl", self.savepoint) @@ -137,24 +140,32 @@ def refin_ctrl(self, dim: Dimension): return None def c2e(self): - # subtract 1 to account for python being 0 based - return self.serializer.read("c2e", self.savepoint) - 1 + + return self._get_connectiviy_array("c2e") + + def _get_connectiviy_array(self, name: str): + connectivity = self.serializer.read(name, self.savepoint) - 1 + print(f" connectivity {name} : {connectivity.shape}") + return connectivity def c2e2c(self): - # subtract 1 to account for python being 0 based - return self.serializer.read("c2e2c", self.savepoint) - 1 + return self._get_connectiviy_array("c2e2c") def e2c(self): - # subtract 1 to account for python being 0 based - return self.serializer.read("e2c", self.savepoint) - 1 + return self._get_connectiviy_array("e2c") def e2v(self): - # subtract 1 to account for python being 0 based - return self.serializer.read("e2v", self.savepoint) - 1 + # array "e2v" is actually e2c2v + v_ = self._get_connectiviy_array("e2v")[:, 0:2] + print(f"real e2v {v_.shape}") + return v_ + + def e2c2v(self): + # array "e2v" is actually e2c2v + return self._get_connectiviy_array("e2v") def v2e(self): - # subtract 1 to account for python being 0 based - return self.serializer.read("v2e", self.savepoint) - 1 + return self._get_connectiviy_array("v2e") class IconDiffusionInitSavepoint(IconSavepoint): @@ -222,6 +233,9 @@ def zd_vertidx(self): # TODO fix this return self._get_field("zd_vertidx", CellDim, C2E2CDim, dtype=int) + def zd_vertoffset(self): + return self._get_field("zd_vertoffset", CellDim, C2E2CDim, KDim, dtype=int) + def rbf_vec_coeff_v1(self): return self._get_field("rbf_vec_coeff_v1", VertexDim, V2EDim) diff --git a/testutils/src/icon4py/testutils/simple_mesh.py b/testutils/src/icon4py/testutils/simple_mesh.py index 5b66aeed1..4bff39e00 100644 --- a/testutils/src/icon4py/testutils/simple_mesh.py +++ b/testutils/src/icon4py/testutils/simple_mesh.py @@ -425,3 +425,10 @@ def get_e2v_offset_provider(self) -> NeighborTableOffsetProvider: def get_e2c2v_offset_provider(self) -> NeighborTableOffsetProvider: return NeighborTableOffsetProvider(self.e2c2v, EdgeDim, VertexDim, self.n_e2c2v) + + def get_e2ecv_offset_provider(self): + old_shape = self.e2c2v.shape + v2ecv_table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) + return NeighborTableOffsetProvider( + v2ecv_table, EdgeDim, ECVDim, v2ecv_table.shape[1] + ) From 95fe3dab345a88683c75a374c70bc9c80ec60b1a Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 2 Mar 2023 16:37:05 +0100 Subject: [PATCH 078/263] fix renamings from upstream merge update typing-extensions version (used by icon-liskov) possibly availabe through --- .../fused_mo_nh_diffusion_stencil_02_03.py | 12 ++++++------ .../fused_mo_nh_diffusion_stencil_07_08_09_10.py | 6 +++--- base-requirements-dev.txt | 1 + 3 files changed, 10 insertions(+), 9 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py index 35cdf2127..dd75ec25a 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -15,11 +15,11 @@ from gt4py.next.ffront.fbuiltins import Field from gt4py.next.program_processors.runners import gtfn_cpu -from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulance import ( - _calculate_diagnostics_for_turbulance, +from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulence import ( + _calculate_diagnostics_for_turbulence, ) -from icon4py.atm_dyn_iconam.temporary_fields_for_turbulance_diagnostics import ( - _temporary_fields_for_turbulance_diagnostics, +from icon4py.atm_dyn_iconam.temporary_fields_for_turbulence_diagnostics import ( + _temporary_fields_for_turbulence_diagnostics, ) from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim @@ -33,10 +33,10 @@ def _fused_mo_nh_diffusion_stencil_02_03( diff_multfac_smag: Field[[KDim], float], wgtfac_c: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - kh_c, div = _temporary_fields_for_turbulance_diagnostics( + kh_c, div = _temporary_fields_for_turbulence_diagnostics( kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag ) - div_ic, hdef_ic = _calculate_diagnostics_for_turbulance(div, kh_c, wgtfac_c) + div_ic, hdef_ic = _calculate_diagnostics_for_turbulence(div, kh_c, wgtfac_c) return div_ic, hdef_ic diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py index 0c647c1ce..e1318c068 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -19,8 +19,8 @@ from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import ( _apply_nabla2_to_w_in_upper_damping_layer, ) -from icon4py.atm_dyn_iconam.calculate_horizontal_gradients_for_turbulance import ( - _calculate_horizontal_gradients_for_turbulance, +from icon4py.atm_dyn_iconam.calculate_horizontal_gradients_for_turbulence import ( + _calculate_horizontal_gradients_for_turbulence, ) from icon4py.atm_dyn_iconam.calculate_nabla2_for_w import _calculate_nabla2_for_w from icon4py.common.dimension import C2E2CODim, CellDim, KDim @@ -53,7 +53,7 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( dwdx, dwdy = where( vert_idx > int32(0), - _calculate_horizontal_gradients_for_turbulance( + _calculate_horizontal_gradients_for_turbulence( w_old, geofac_grg_x, geofac_grg_y ), (dwdx, dwdy), diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 29b446eac..db4d67bbf 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -15,6 +15,7 @@ flake8-mutable>=1.2.0 flake8-rst-docstrings>=0.0.14 isort~=5.10 mypy>=0.942 +typing-extensions>=4.5 pre-commit~=2.15 pytest>=6.1 pytest-cache>=1.0 From be5560d95b0c7e453b8ccfaf275e903edd2a48d3 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 2 Mar 2023 17:02:12 +0100 Subject: [PATCH 079/263] fix qa --- base-requirements-dev.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index db4d67bbf..754f55ea4 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -1,5 +1,4 @@ # VCS -#-e git+https://github.com/tehrengruber/gt4py.git@2a9dcb0367a1ec1ef6d78951358d700424867b07#egg=gt4py-functional -e git+https://github.com/tehrengruber/gt4py.git@reduce_foast_lowering_complexity#egg=gt4py git+https://github.com/GridTools/serialbox#egg=serialbox&subdirectory=src/serialbox-python From b0f142b2afd63db27f4cee4d6950877a4bf8bd77 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 23 Mar 2023 18:58:31 +0100 Subject: [PATCH 080/263] add first version of grid manager reading netcdf file - move joint fixtures to testutils - add dependencies for StrEnum (py11) and netcdf - add gridmanager --- .../src/icon4py/diffusion/icon_grid.py | 4 + atm_dyn_iconam/tests/conftest.py | 42 +--- atm_dyn_iconam/tests/test_diffusion.py | 1 + atm_dyn_iconam/tests/test_icon_grid.py | 1 + atm_dyn_iconam/tests/test_vertical.py | 1 + base-requirements-dev.txt | 2 + common/src/icon4py/grid/__init__.py | 12 ++ common/src/icon4py/grid/icon_grid.py | 181 ++++++++++++++++++ common/tests/conftest.py | 19 ++ common/tests/test_gridmanager.py | 175 +++++++++++++++++ testutils/src/icon4py/testutils/fixtures.py | 57 ++++++ 11 files changed, 454 insertions(+), 41 deletions(-) create mode 100644 common/src/icon4py/grid/__init__.py create mode 100644 common/src/icon4py/grid/icon_grid.py create mode 100644 common/tests/conftest.py create mode 100644 common/tests/test_gridmanager.py create mode 100644 testutils/src/icon4py/testutils/fixtures.py diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index 216593fa0..af2cead56 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -171,6 +171,10 @@ def get_v2e_connectivity(self): table = self.connectivities["v2e"] return NeighborTableOffsetProvider(table, VertexDim, EdgeDim, table.shape[1]) + def get_v2c_connectivity(self): + table = self.connectivities["v2c"] + return NeighborTableOffsetProvider(table, VertexDim, EdgeDim, table.shape[1]) + def get_e2ecv_connectivity(self): old_shape = self.connectivities["e2c2v"].shape v2ecv_table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 99e6db30b..65fbc89bf 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -10,12 +10,10 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -import tarfile -from pathlib import Path + import numpy as np import pytest -import wget from icon4py.common.dimension import ( C2E2CDim, @@ -32,39 +30,6 @@ from icon4py.diffusion.diffusion import DiffusionConfig from icon4py.diffusion.horizontal import HorizontalMeshSize from icon4py.diffusion.icon_grid import IconGrid, MeshConfig, VerticalMeshConfig -from icon4py.testutils.serialbox_utils import IconSerialDataProvider - - -data_uri = "https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download" -data_path = Path(__file__).parent.joinpath("ser_icondata") -extracted_path = data_path.joinpath("mch_ch_r04b09_dsl/ser_data") -data_file = data_path.joinpath("mch_ch_r04b09_dsl_v2.tar.gz").name - - -@pytest.fixture(scope="session") -def setup_icon_data(): - """ - Get the binary ICON data from a remote server. - - Session scoped fixture which is a prerequisite of all the other fixtures in this file. - """ - data_path.mkdir(parents=True, exist_ok=True) - if not any(data_path.iterdir()): - print( - f"directory {data_path} is empty: downloading data from {data_uri} and extracting" - ) - wget.download(data_uri, out=data_file) - # extract downloaded file - if not tarfile.is_tarfile(data_file): - raise NotImplementedError(f"{data_file} needs to be a valid tar file") - with tarfile.open(data_file, mode="r:*") as tf: - tf.extractall(path=data_path) - Path(data_file).unlink(missing_ok=True) - - -@pytest.fixture -def data_provider(setup_icon_data) -> IconSerialDataProvider: - return IconSerialDataProvider("icon_pydycore", str(extracted_path), True) @pytest.fixture @@ -165,11 +130,6 @@ def icon_grid(data_provider): return grid -@pytest.fixture -def grid_savepoint(data_provider): - return data_provider.from_savepoint_grid() - - @pytest.fixture def r04b09_diffusion_config(setup_icon_data) -> DiffusionConfig: """ diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 24316b5a9..7faff8355 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -32,6 +32,7 @@ from icon4py.testutils.serialbox_utils import IconDiffusionInitSavepoint from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import as_1D_sparse_field, random_field, zero_field +from icon4py.testutils.fixtures import setup_icon_data, data_provider, grid_savepoint def test_scale_k(): diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index 256ded7aa..672b5c51f 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -14,6 +14,7 @@ from icon4py.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.diffusion.horizontal import HorizontalMarkerIndex +from icon4py.testutils.fixtures import setup_icon_data, data_provider @pytest.mark.datatest diff --git a/atm_dyn_iconam/tests/test_vertical.py b/atm_dyn_iconam/tests/test_vertical.py index ccf0b7dbf..d553d102b 100644 --- a/atm_dyn_iconam/tests/test_vertical.py +++ b/atm_dyn_iconam/tests/test_vertical.py @@ -17,6 +17,7 @@ import pytest from icon4py.diffusion.icon_grid import VerticalModelParams +from icon4py.testutils.fixtures import setup_icon_data, data_provider, grid_savepoint @pytest.mark.parametrize( diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 754f55ea4..b0fa61844 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -25,3 +25,5 @@ setuptools>=40.8.0 wheel>=0.37.1 tox >= 3.25 wget >= 3.2 +netcdf4 >=1.6.3 +enum34 >= 1.1.10 diff --git a/common/src/icon4py/grid/__init__.py b/common/src/icon4py/grid/__init__.py new file mode 100644 index 000000000..15dfdb009 --- /dev/null +++ b/common/src/icon4py/grid/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/common/src/icon4py/grid/icon_grid.py b/common/src/icon4py/grid/icon_grid.py new file mode 100644 index 000000000..36eaa45cd --- /dev/null +++ b/common/src/icon4py/grid/icon_grid.py @@ -0,0 +1,181 @@ +# 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 logging +from typing import Optional +from uuid import UUID + +import numpy as np +from gt4py.eve import StrEnum +from gt4py.next.common import Dimension +from netCDF4 import Dataset + +from icon4py.common.dimension import ( + C2E2CDim, + C2EDim, + CellDim, + E2CDim, + E2VDim, + EdgeDim, + V2CDim, + V2EDim, + VertexDim, +) +from icon4py.diffusion.horizontal import HorizontalMeshSize +from icon4py.diffusion.icon_grid import IconGrid, MeshConfig, VerticalMeshConfig + + +class GridFileName(StrEnum): + pass + + +class GridFile: + class Property(GridFileName): + GRID_ID = "uuidOfHGrid" + PARENT_GRID_ID = "uuidOfParHGrid" + + class Offsets(GridFileName): + C2E2C = "neighbor_cell_index" + V2E2V = "vertices_of_vertex" + V2E = "edges_of_vertex" + V2C = "cells_of_vertex" + E2V = "edge_vertices" + E2C = "adjacent_cell_of_edge" + C2V = "vertex_of_cell" + C2E = "edge_of_cell" + + class Dimension(GridFileName): + VERTEX_NAME = "vertex" + EDGE_NAME = "edge" + CELL_NAME = "cell" + V2E_SIZE = "ne" # 6 + DIAMOND_EDGE_SIZE = "no" # 4 + C2E_SIZE = "nv" # 3 + E2V_SIZE = "nc" # 2 + + def __init__(self, dataset: Dataset): + self._dataset = dataset + self._log = logging.getLogger(__name__) + + def dimension(self, name: GridFileName) -> int: + return self._dataset.dimensions[name].size + + def int_field( + self, name: GridFileName, offset_value=-1, transpose=True + ) -> np.ndarray: + nc_variable = self._dataset.variables[name] + self._log.debug(f"{name}: {nc_variable}") + data = nc_variable[:] + offset_value + + data = np.array(data, dtype=np.int32) + return np.transpose(data) if transpose else data + + +class IconGridError(RuntimeError): + pass + + +class GridManager: + def __init__(self, grid_file: str): + self._log = logging.getLogger(__name__) + self._grid: Optional[IconGrid] = None + self._file_names = grid_file + + def init(self): + _, grid = self._read_from_gridfile(self._file_names) + self._grid = grid + + def _read_from_gridfile(self, fname: str): + try: + dataset = Dataset(fname, "r", format="NETCDF4") + self._log.debug(dataset) + grid_id = UUID(dataset.getncattr(GridFile.Property.GRID_ID)) + return grid_id, self.__class__.from_grid_dataset(dataset) + except: + self._log.error(f"gridfile {fname} not found, aborting") + exit(1) + + def get_c2e_connectivity(self): + return self._grid.get_c2e_connectivity() + + def get_e2v_connectivity(self): + return self._grid.get_e2v_connectivity() + + def get_e2c_connectivity(self): + return self._grid.get_e2c_connectivity() + + def get_c2e2c_connectivity(self): + return self._grid.get_c2e2c_connectivity() + + def get_v2c_connectivity(self): + return self._grid.get_v2e_connectivity() + + def get_c2e2co_connectivity(self): + return self._grid.get_c2e2co_connectivity() + + def get_e2c2v_connectivity(self): + return self._grid.get_e2c2v_connectivity() + + def get_v2e_connectivity(self): + return self._grid.get_v2e_connectivity() + + def get_e2ecv_connectivity(self): + return self._grid.get_e2ecv_connectivity() + + def get_size(self, dim: Dimension): + if dim == VertexDim: + return self._grid.config.num_vertices + elif dim == CellDim: + return self._grid.config.num_cells + elif dim == EdgeDim: + return self._grid.config.num_edges + else: + self._log.warning(f"cannot determine size of unknown dimension {dim}") + raise IconGridError(f"Unknown dimension {dim}") + + @staticmethod + def from_grid_dataset(dataset: Dataset) -> IconGrid: + reader = GridFile(dataset) + num_cells = reader.dimension(GridFile.Dimension.CELL_NAME) + num_edges = reader.dimension(GridFile.Dimension.EDGE_NAME) + num_vertices = reader.dimension(GridFile.Dimension.VERTEX_NAME) + grid_size = HorizontalMeshSize( + num_vertices=num_vertices, num_edges=num_edges, num_cells=num_cells + ) + + c2e = reader.int_field(GridFile.Offsets.C2E) + c2v = reader.int_field(GridFile.Offsets.C2V) + e2c = reader.int_field(GridFile.Offsets.E2C) + e2v = reader.int_field(GridFile.Offsets.E2V) + v2c = reader.int_field(GridFile.Offsets.V2C) + v2e = reader.int_field(GridFile.Offsets.V2E) + v2e2v = reader.int_field(GridFile.Offsets.V2E2V) + c2e2c = reader.int_field(GridFile.Offsets.C2E2C) + config = MeshConfig( + horizontal_config=grid_size, vertical_config=VerticalMeshConfig(0) + ) + icon_grid = ( + IconGrid() + .with_config(config) + .with_connectivities( + { + C2EDim: c2e, + E2CDim: e2c, + E2VDim: e2v, + V2EDim: v2e, + V2CDim: v2c, + C2E2CDim: c2e2c, + } + ) + ) + return icon_grid diff --git a/common/tests/conftest.py b/common/tests/conftest.py new file mode 100644 index 000000000..ca1a9e945 --- /dev/null +++ b/common/tests/conftest.py @@ -0,0 +1,19 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + + +from icon4py.testutils.fixtures import ( + data_provider, + grid_savepoint, + setup_icon_data, +) diff --git a/common/tests/test_gridmanager.py b/common/tests/test_gridmanager.py new file mode 100644 index 000000000..503353cb2 --- /dev/null +++ b/common/tests/test_gridmanager.py @@ -0,0 +1,175 @@ +# 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 logging +import pathlib +from uuid import uuid4 + +import numpy as np +import pytest +from netCDF4 import Dataset + +from icon4py.common.dimension import CellDim, EdgeDim, VertexDim +from icon4py.grid.icon_grid import GridFile, GridFileName, GridManager +from icon4py.testutils.simple_mesh import SimpleMesh + + +SIMPLE_MESH_NC = "./simple_mesh_grid.nc" + + +@pytest.fixture +def simple_mesh_path(): + return pathlib.Path(SIMPLE_MESH_NC).absolute() + + +@pytest.fixture(scope="session") +def simple_mesh_data(): + mesh = SimpleMesh() + dataset = Dataset(SIMPLE_MESH_NC, "w", format="NETCDF4") + dataset.setncattr(GridFile.Property.GRID_ID, str(uuid4())) + dataset.createDimension(GridFile.Dimension.VERTEX_NAME, size=mesh.n_vertices) + dataset.createDimension(GridFile.Dimension.EDGE_NAME, size=mesh.n_edges) + dataset.createDimension(GridFile.Dimension.CELL_NAME, size=mesh.n_cells) + dataset.createDimension(GridFile.Dimension.E2V_SIZE, size=mesh.n_e2v) + dataset.createDimension(GridFile.Dimension.DIAMOND_EDGE_SIZE, size=mesh.n_e2c2e) + dataset.createDimension(GridFile.Dimension.C2E_SIZE, size=mesh.n_c2e) + dataset.createDimension(GridFile.Dimension.V2E_SIZE, size=mesh.n_v2c) + + add_to_dataset( + dataset, + mesh.c2e, + GridFile.Offsets.C2E, + (GridFile.Dimension.C2E_SIZE, GridFile.Dimension.CELL_NAME), + ) + # add_to_dataset(data, mesh.c2v, GridFile.Offsets.C2V, (GridFile.Dimension.C2E_SIZE, GridFile.Dimension.CELL_NAME)) + add_to_dataset( + dataset, + mesh.e2c, + GridFile.Offsets.E2C, + (GridFile.Dimension.E2V_SIZE, GridFile.Dimension.EDGE_NAME), + ) + add_to_dataset( + dataset, + mesh.e2v, + GridFile.Offsets.E2V, + (GridFile.Dimension.E2V_SIZE, GridFile.Dimension.EDGE_NAME), + ) + + add_to_dataset( + dataset, + mesh.v2c, + GridFile.Offsets.V2C, + (GridFile.Dimension.V2E_SIZE, GridFile.Dimension.VERTEX_NAME), + ) + # TODO fix me there is no v2c in the simple mesh + add_to_dataset( + dataset, + np.zeros((mesh.n_cells, 3), dtype=np.int32), + GridFile.Offsets.C2V, + (GridFile.Dimension.C2E_SIZE, GridFile.Dimension.CELL_NAME), + ) + add_to_dataset( + dataset, + np.zeros((mesh.n_vertices, 4), dtype=np.int32), + GridFile.Offsets.V2E2V, + (GridFile.Dimension.DIAMOND_EDGE_SIZE, GridFile.Dimension.VERTEX_NAME), + ) + add_to_dataset( + dataset, + mesh.v2e, + GridFile.Offsets.V2E, + (GridFile.Dimension.V2E_SIZE, GridFile.Dimension.VERTEX_NAME), + ) + # add_to_dataset(data, mesh.v2e2v, GridFile.Offsets.V2E2V, (GridFile.Dimension.V2E_SIZE, GridFile.Dimension.VERTEX_NAME)) + add_to_dataset( + dataset, + mesh.c2e2c, + GridFile.Offsets.C2E2C, + (GridFile.Dimension.C2E_SIZE, GridFile.Dimension.CELL_NAME), + ) + dataset.close() + + +def add_to_dataset( + dataset: Dataset, + data: np.ndarray, + var_name: str, + dims: tuple[GridFileName, GridFileName], +): + var = dataset.createVariable(var_name, np.int32, dims) + var[:] = np.transpose(data)[:] + 1 + + +def test_gridparser_dimension(simple_mesh_data): + + data = Dataset(SIMPLE_MESH_NC, "r") + grid_parser = GridFile(data) + mesh = SimpleMesh() + assert grid_parser.dimension(GridFile.Dimension.CELL_NAME) == mesh.n_cells + assert grid_parser.dimension(GridFile.Dimension.VERTEX_NAME) == mesh.n_vertices + assert grid_parser.dimension(GridFile.Dimension.EDGE_NAME) == mesh.n_edges + + +def test_grid_parser_index_fields(simple_mesh_data, caplog): + caplog.set_level(logging.DEBUG) + data = Dataset(SIMPLE_MESH_NC, "r") + mesh = SimpleMesh() + grid_parser = GridFile(data) + + assert np.allclose(grid_parser.int_field(GridFile.Offsets.C2E), mesh.c2e) + assert np.allclose(grid_parser.int_field(GridFile.Offsets.E2C), mesh.e2c) + assert np.allclose(grid_parser.int_field(GridFile.Offsets.V2E), mesh.v2e) + assert np.allclose(grid_parser.int_field(GridFile.Offsets.V2C), mesh.v2c) + + +@pytest.mark.skip("TODO: how are -1 values handled?") +@pytest.mark.datatest +def test_gridmanager_read_gridfile(caplog, grid_savepoint): + caplog.set_level(logging.DEBUG) + fname = "/home/magdalena/data/exclaim/grids/mch_ch_r04b09_dsl/grid.nc" + gm = GridManager(fname) + gm.init() + num_vertex = gm.get_size(VertexDim) + num_edges = gm.get_size(EdgeDim) + num_cells = gm.get_size(CellDim) + assert np.allclose( + gm.get_v2e_connectivity().table, grid_savepoint.v2e()[0:num_vertex, :] + ) + assert np.allclose( + gm.get_v2c_connectivity().table, grid_savepoint.v2c()[0:num_vertex, :] + ) + assert np.allclose( + gm.get_e2c_connectivity().table, grid_savepoint.e2c()[0:num_edges, :] + ) + assert np.allclose( + gm.get_c2e_connectivity().table, grid_savepoint.c2e()[0:num_cells, :] + ) + assert np.allclose( + gm.get_e2v_connectivity().table, grid_savepoint.v2e()[0:num_vertex, :] + ) + + +@pytest.mark.parametrize("dim, size", [(CellDim, 18), (EdgeDim, 27), (VertexDim, 9)]) +def test_grid_manager_getsize(simple_mesh_data, simple_mesh_path, dim, size): + gm = GridManager(simple_mesh_path) + gm.init() + assert size == gm.get_size(dim) + + +def test_gridmanager_given_file_not_found_then_abort(): + fname = "./unknown_grid.nc" + with pytest.raises(SystemExit) as error: + gm = GridManager(fname) + gm.init() + assert error.type == SystemExit + assert error.value == 1 diff --git a/testutils/src/icon4py/testutils/fixtures.py b/testutils/src/icon4py/testutils/fixtures.py new file mode 100644 index 000000000..690feafc3 --- /dev/null +++ b/testutils/src/icon4py/testutils/fixtures.py @@ -0,0 +1,57 @@ +# 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 tarfile +from pathlib import Path + +import pytest +import wget + +from icon4py.testutils.serialbox_utils import IconSerialDataProvider + + +data_uri = "https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download" +data_path = Path(__file__).parent.joinpath("ser_icondata") +extracted_path = data_path.joinpath("mch_ch_r04b09_dsl/ser_data") +data_file = data_path.joinpath("mch_ch_r04b09_dsl_v2.tar.gz").name + + +@pytest.fixture(scope="session") +def setup_icon_data(): + """ + Get the binary ICON data from a remote server. + + Session scoped fixture which is a prerequisite of all the other fixtures in this file. + """ + data_path.mkdir(parents=True, exist_ok=True) + if not any(data_path.iterdir()): + print( + f"directory {data_path} is empty: downloading data from {data_uri} and extracting" + ) + wget.download(data_uri, out=data_file) + # extract downloaded file + if not tarfile.is_tarfile(data_file): + raise NotImplementedError(f"{data_file} needs to be a valid tar file") + with tarfile.open(data_file, mode="r:*") as tf: + tf.extractall(path=data_path) + Path(data_file).unlink(missing_ok=True) + + +@pytest.fixture +def data_provider(setup_icon_data) -> IconSerialDataProvider: + return IconSerialDataProvider("icon_pydycore", str(extracted_path), True) + + +@pytest.fixture +def grid_savepoint(data_provider): + return data_provider.from_savepoint_grid() From 0af0bd7ab5bee7cedd3276fac8ff101cda5d4b4b Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 24 Mar 2023 11:58:31 +0100 Subject: [PATCH 081/263] - isolate index transform in special class, - do not apply zero base offset to -1 values - get grid files from online storage --- atm_dyn_iconam/tests/test_diffusion.py | 18 +-- atm_dyn_iconam/tests/test_icon_grid.py | 2 +- atm_dyn_iconam/tests/test_vertical.py | 8 +- common/src/icon4py/grid/icon_grid.py | 66 +++++++--- common/tests/conftest.py | 2 +- common/tests/test_gridmanager.py | 113 ++++++++++++++---- .../src/icon4py/testutils/data_handling.py | 32 +++++ testutils/src/icon4py/testutils/fixtures.py | 16 +-- 8 files changed, 191 insertions(+), 66 deletions(-) create mode 100644 testutils/src/icon4py/testutils/data_handling.py diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 7faff8355..f16a96829 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -29,10 +29,14 @@ set_zero_v_k, setup_fields_for_initial_step, ) +from icon4py.testutils.fixtures import ( # noqa F401 + data_provider, + grid_savepoint, + setup_icon_data, +) from icon4py.testutils.serialbox_utils import IconDiffusionInitSavepoint from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import as_1D_sparse_field, random_field, zero_field -from icon4py.testutils.fixtures import setup_icon_data, data_provider, grid_savepoint def test_scale_k(): @@ -224,7 +228,7 @@ def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): @pytest.mark.datatest def test_diffusion_init( diffusion_savepoint_init, - grid_savepoint, + grid_savepoint, # noqa: F811 icon_grid, r04b09_diffusion_config, step_date_init, @@ -359,7 +363,7 @@ def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( @pytest.mark.datatest def test_verify_diffusion_init_against_first_regular_savepoint( diffusion_savepoint_init, - grid_savepoint, + grid_savepoint, # noqa: F811 r04b09_diffusion_config, icon_grid, damping_height, @@ -406,7 +410,7 @@ def test_verify_diffusion_init_against_first_regular_savepoint( @pytest.mark.parametrize("step_date_init", ["2021-06-20T12:00:50.000"]) def test_verify_diffusion_init_against_other_regular_savepoint( r04b09_diffusion_config, - grid_savepoint, + grid_savepoint, # noqa: F811 icon_grid, diffusion_savepoint_init, damping_height, @@ -455,7 +459,7 @@ def test_run_diffusion_single_step( run_with_program, diffusion_savepoint_init, diffusion_savepoint_exit, - grid_savepoint, + grid_savepoint, # noqa: F811 icon_grid, r04b09_diffusion_config, damping_height, @@ -520,7 +524,7 @@ def test_run_diffusion_single_step( ) -def _read_fields(diffusion_savepoint_init, grid_savepoint): +def _read_fields(diffusion_savepoint_init, grid_savepoint): # noqa: F811 grg = diffusion_savepoint_init.geofac_grg() interpolation_state = InterpolationState( @@ -591,7 +595,7 @@ def test_diffusion_five_steps( damping_height, r04b09_diffusion_config, icon_grid, - grid_savepoint, + grid_savepoint, # noqa: F811 diffusion_savepoint_init, diffusion_savepoint_exit, linit=True, diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index 672b5c51f..838a01cbb 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -14,7 +14,7 @@ from icon4py.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.diffusion.horizontal import HorizontalMarkerIndex -from icon4py.testutils.fixtures import setup_icon_data, data_provider +from icon4py.testutils.fixtures import data_provider, setup_icon_data # noqa @pytest.mark.datatest diff --git a/atm_dyn_iconam/tests/test_vertical.py b/atm_dyn_iconam/tests/test_vertical.py index d553d102b..550854ca4 100644 --- a/atm_dyn_iconam/tests/test_vertical.py +++ b/atm_dyn_iconam/tests/test_vertical.py @@ -17,7 +17,11 @@ import pytest from icon4py.diffusion.icon_grid import VerticalModelParams -from icon4py.testutils.fixtures import setup_icon_data, data_provider, grid_savepoint +from icon4py.testutils.fixtures import ( # noqa F401 + data_provider, + grid_savepoint, + setup_icon_data, +) @pytest.mark.parametrize( @@ -35,7 +39,7 @@ def test_nrdmax_calculation(max_h, damping, delta): @pytest.mark.datatest -def test_nrdmax_calculation_from_icon_input(icon_grid, grid_savepoint): +def test_nrdmax_calculation_from_icon_input(icon_grid, grid_savepoint): # noqa: F811 a = grid_savepoint.vct_a() damping_height = 12500 vertical_params = VerticalModelParams( diff --git a/common/src/icon4py/grid/icon_grid.py b/common/src/icon4py/grid/icon_grid.py index 36eaa45cd..2c1414440 100644 --- a/common/src/icon4py/grid/icon_grid.py +++ b/common/src/icon4py/grid/icon_grid.py @@ -12,6 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later import logging +from abc import ABC from typing import Optional from uuid import UUID @@ -40,6 +41,12 @@ class GridFileName(StrEnum): class GridFile: + """ + Represents and ICON netcdf grid file. + """ + + INVALID_INDEX = -1 + class Property(GridFileName): GRID_ID = "uuidOfHGrid" PARENT_GRID_ID = "uuidOfParHGrid" @@ -60,7 +67,7 @@ class Dimension(GridFileName): CELL_NAME = "cell" V2E_SIZE = "ne" # 6 DIAMOND_EDGE_SIZE = "no" # 4 - C2E_SIZE = "nv" # 3 + NEIGHBORING_EDGES_TO_CELL_SIZE = "nv" # 3 E2V_SIZE = "nc" # 2 def __init__(self, dataset: Dataset): @@ -70,12 +77,10 @@ def __init__(self, dataset: Dataset): def dimension(self, name: GridFileName) -> int: return self._dataset.dimensions[name].size - def int_field( - self, name: GridFileName, offset_value=-1, transpose=True - ) -> np.ndarray: + def int_field(self, name: GridFileName, transpose=True) -> np.ndarray: nc_variable = self._dataset.variables[name] self._log.debug(f"{name}: {nc_variable}") - data = nc_variable[:] + offset_value + data = nc_variable[:] data = np.array(data, dtype=np.int32) return np.transpose(data) if transpose else data @@ -85,9 +90,26 @@ class IconGridError(RuntimeError): pass +class GridTransformation(ABC): + def get_offset_for_field(self, array: np.ndarray): + return np.zeros(array.shape) + + +class ToGt4PyTransformation(GridTransformation): + def get_offset_for_field(self, array: np.ndarray): + """ + Calculate the index offset needed for usage with python. + + Fortran indices are 1-based, hence the offset is -1 for 0-based ness of python except for + INVALID values which are marked with -1 in the grid file and are kept such. + """ + return np.where(array < 0, 0, -1) + + class GridManager: - def __init__(self, grid_file: str): + def __init__(self, transformation: GridTransformation, grid_file: str): self._log = logging.getLogger(__name__) + self._transformation = transformation self._grid: Optional[IconGrid] = None self._file_names = grid_file @@ -95,13 +117,14 @@ def init(self): _, grid = self._read_from_gridfile(self._file_names) self._grid = grid - def _read_from_gridfile(self, fname: str): + def _read_from_gridfile(self, fname: str) -> tuple[UUID, IconGrid]: try: dataset = Dataset(fname, "r", format="NETCDF4") + self._log.debug(dataset) grid_id = UUID(dataset.getncattr(GridFile.Property.GRID_ID)) - return grid_id, self.__class__.from_grid_dataset(dataset) - except: + return grid_id, self.from_grid_dataset(dataset) + except FileNotFoundError: self._log.error(f"gridfile {fname} not found, aborting") exit(1) @@ -143,8 +166,12 @@ def get_size(self, dim: Dimension): self._log.warning(f"cannot determine size of unknown dimension {dim}") raise IconGridError(f"Unknown dimension {dim}") - @staticmethod - def from_grid_dataset(dataset: Dataset) -> IconGrid: + def _get_index_field(self, reader, name: GridFileName): + field = reader.int_field(name) + field = field + self._transformation.get_offset_for_field(field) + return field + + def from_grid_dataset(self, dataset: Dataset) -> IconGrid: reader = GridFile(dataset) num_cells = reader.dimension(GridFile.Dimension.CELL_NAME) num_edges = reader.dimension(GridFile.Dimension.EDGE_NAME) @@ -152,15 +179,16 @@ def from_grid_dataset(dataset: Dataset) -> IconGrid: grid_size = HorizontalMeshSize( num_vertices=num_vertices, num_edges=num_edges, num_cells=num_cells ) + c2e = self._get_index_field(reader, GridFile.Offsets.C2E) + c2v = self._get_index_field(reader, GridFile.Offsets.C2V) + + e2c = self._get_index_field(reader, GridFile.Offsets.E2C) - c2e = reader.int_field(GridFile.Offsets.C2E) - c2v = reader.int_field(GridFile.Offsets.C2V) - e2c = reader.int_field(GridFile.Offsets.E2C) - e2v = reader.int_field(GridFile.Offsets.E2V) - v2c = reader.int_field(GridFile.Offsets.V2C) - v2e = reader.int_field(GridFile.Offsets.V2E) - v2e2v = reader.int_field(GridFile.Offsets.V2E2V) - c2e2c = reader.int_field(GridFile.Offsets.C2E2C) + e2v = self._get_index_field(reader, GridFile.Offsets.E2V) + v2c = self._get_index_field(reader, GridFile.Offsets.V2C) + v2e = self._get_index_field(reader, GridFile.Offsets.V2E) + v2e2v = self._get_index_field(reader, GridFile.Offsets.V2E2V) + c2e2c = self._get_index_field(reader, GridFile.Offsets.C2E2C) config = MeshConfig( horizontal_config=grid_size, vertical_config=VerticalMeshConfig(0) ) diff --git a/common/tests/conftest.py b/common/tests/conftest.py index ca1a9e945..d79fbd985 100644 --- a/common/tests/conftest.py +++ b/common/tests/conftest.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later -from icon4py.testutils.fixtures import ( +from icon4py.testutils.fixtures import ( # noqa F401 data_provider, grid_savepoint, setup_icon_data, diff --git a/common/tests/test_gridmanager.py b/common/tests/test_gridmanager.py index 503353cb2..a7a213e7a 100644 --- a/common/tests/test_gridmanager.py +++ b/common/tests/test_gridmanager.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later import logging -import pathlib +from pathlib import Path from uuid import uuid4 import numpy as np @@ -20,16 +20,38 @@ from netCDF4 import Dataset from icon4py.common.dimension import CellDim, EdgeDim, VertexDim -from icon4py.grid.icon_grid import GridFile, GridFileName, GridManager +from icon4py.grid.icon_grid import ( + GridFile, + GridFileName, + GridManager, + GridTransformation, + ToGt4PyTransformation, +) +from icon4py.testutils.data_handling import download_and_extract from icon4py.testutils.simple_mesh import SimpleMesh SIMPLE_MESH_NC = "./simple_mesh_grid.nc" +grid_uri = "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download" +grids_path = Path(__file__).parent.joinpath("grids") +r04b09_dsl_grid_path = grids_path.joinpath("mch_ch_r04b09_dsl") +data_file = r04b09_dsl_grid_path.joinpath("mch_ch_r04b09_dsl_grids_v1.tar.gz").name + + +@pytest.fixture(scope="session") +def get_grid_files(): + """ + Get the grid files used for testing. + + Session scoped fixture which is a prerequisite of all the other fixtures in this file. + """ + download_and_extract(grid_uri, r04b09_dsl_grid_path, data_file) + @pytest.fixture def simple_mesh_path(): - return pathlib.Path(SIMPLE_MESH_NC).absolute() + return Path(SIMPLE_MESH_NC).absolute() @pytest.fixture(scope="session") @@ -42,14 +64,19 @@ def simple_mesh_data(): dataset.createDimension(GridFile.Dimension.CELL_NAME, size=mesh.n_cells) dataset.createDimension(GridFile.Dimension.E2V_SIZE, size=mesh.n_e2v) dataset.createDimension(GridFile.Dimension.DIAMOND_EDGE_SIZE, size=mesh.n_e2c2e) - dataset.createDimension(GridFile.Dimension.C2E_SIZE, size=mesh.n_c2e) + dataset.createDimension( + GridFile.Dimension.NEIGHBORING_EDGES_TO_CELL_SIZE, size=mesh.n_c2e + ) dataset.createDimension(GridFile.Dimension.V2E_SIZE, size=mesh.n_v2c) add_to_dataset( dataset, mesh.c2e, GridFile.Offsets.C2E, - (GridFile.Dimension.C2E_SIZE, GridFile.Dimension.CELL_NAME), + ( + GridFile.Dimension.NEIGHBORING_EDGES_TO_CELL_SIZE, + GridFile.Dimension.CELL_NAME, + ), ) # add_to_dataset(data, mesh.c2v, GridFile.Offsets.C2V, (GridFile.Dimension.C2E_SIZE, GridFile.Dimension.CELL_NAME)) add_to_dataset( @@ -76,7 +103,10 @@ def simple_mesh_data(): dataset, np.zeros((mesh.n_cells, 3), dtype=np.int32), GridFile.Offsets.C2V, - (GridFile.Dimension.C2E_SIZE, GridFile.Dimension.CELL_NAME), + ( + GridFile.Dimension.NEIGHBORING_EDGES_TO_CELL_SIZE, + GridFile.Dimension.CELL_NAME, + ), ) add_to_dataset( dataset, @@ -95,7 +125,10 @@ def simple_mesh_data(): dataset, mesh.c2e2c, GridFile.Offsets.C2E2C, - (GridFile.Dimension.C2E_SIZE, GridFile.Dimension.CELL_NAME), + ( + GridFile.Dimension.NEIGHBORING_EDGES_TO_CELL_SIZE, + GridFile.Dimension.CELL_NAME, + ), ) dataset.close() @@ -107,7 +140,7 @@ def add_to_dataset( dims: tuple[GridFileName, GridFileName], ): var = dataset.createVariable(var_name, np.int32, dims) - var[:] = np.transpose(data)[:] + 1 + var[:] = np.transpose(data)[:] def test_gridparser_dimension(simple_mesh_data): @@ -132,36 +165,62 @@ def test_grid_parser_index_fields(simple_mesh_data, caplog): assert np.allclose(grid_parser.int_field(GridFile.Offsets.V2C), mesh.v2c) -@pytest.mark.skip("TODO: how are -1 values handled?") +@pytest.mark.skip("TODO: handling of boundary values") @pytest.mark.datatest -def test_gridmanager_read_gridfile(caplog, grid_savepoint): +def test_gridmanager_eval_v2e(caplog, grid_savepoint, get_grid_files): caplog.set_level(logging.DEBUG) - fname = "/home/magdalena/data/exclaim/grids/mch_ch_r04b09_dsl/grid.nc" - gm = GridManager(fname) - gm.init() - num_vertex = gm.get_size(VertexDim) - num_edges = gm.get_size(EdgeDim) - num_cells = gm.get_size(CellDim) + fname = r04b09_dsl_grid_path.joinpath("grid.nc") + gm, num_cells, num_edges, num_vertex = _init_grid_manager(fname) assert np.allclose( gm.get_v2e_connectivity().table, grid_savepoint.v2e()[0:num_vertex, :] ) + + +@pytest.mark.skip("TODO: v2c not serialized??") +@pytest.mark.datatest +def test_gridmanager_eval_v2c(caplog, grid_savepoint, get_grid_files): + caplog.set_level(logging.DEBUG) + fname = r04b09_dsl_grid_path.joinpath("grid.nc") + gm, num_cells, num_edges, num_vertex = _init_grid_manager(fname) assert np.allclose( gm.get_v2c_connectivity().table, grid_savepoint.v2c()[0:num_vertex, :] ) + + +@pytest.mark.datatest +def test_gridmanager_eval_e2c(caplog, grid_savepoint, get_grid_files): + caplog.set_level(logging.DEBUG) + fname = r04b09_dsl_grid_path.joinpath("grid.nc") + gm, num_cells, num_edges, num_vertex = _init_grid_manager(fname) + assert np.allclose( gm.get_e2c_connectivity().table, grid_savepoint.e2c()[0:num_edges, :] ) + + +@pytest.mark.datatest +def test_gridmanager_eval_c2e(caplog, grid_savepoint, get_grid_files): + caplog.set_level(logging.DEBUG) + fname = r04b09_dsl_grid_path.joinpath("grid.nc") + gm, num_cells, num_edges, num_vertex = _init_grid_manager(fname) assert np.allclose( gm.get_c2e_connectivity().table, grid_savepoint.c2e()[0:num_cells, :] ) - assert np.allclose( - gm.get_e2v_connectivity().table, grid_savepoint.v2e()[0:num_vertex, :] - ) + + +def _init_grid_manager(fname): + gm = GridManager(ToGt4PyTransformation(), fname) + gm.init() + num_vertex = gm.get_size(VertexDim) + num_edges = gm.get_size(EdgeDim) + num_cells = gm.get_size(CellDim) + return gm, num_cells, num_edges, num_vertex @pytest.mark.parametrize("dim, size", [(CellDim, 18), (EdgeDim, 27), (VertexDim, 9)]) -def test_grid_manager_getsize(simple_mesh_data, simple_mesh_path, dim, size): - gm = GridManager(simple_mesh_path) +def test_grid_manager_getsize(simple_mesh_data, simple_mesh_path, dim, size, caplog): + caplog.set_level(logging.DEBUG) + gm = GridManager(GridTransformation(), simple_mesh_path) gm.init() assert size == gm.get_size(dim) @@ -169,7 +228,17 @@ def test_grid_manager_getsize(simple_mesh_data, simple_mesh_path, dim, size): def test_gridmanager_given_file_not_found_then_abort(): fname = "./unknown_grid.nc" with pytest.raises(SystemExit) as error: - gm = GridManager(fname) + gm = GridManager(GridTransformation(), fname) gm.init() assert error.type == SystemExit assert error.value == 1 + + +@pytest.mark.parametrize("size", [100, 1500, 20000]) +def test_gt4py_transform_offset_by_1_where_valid(size): + trafo = ToGt4PyTransformation() + input_field = np.random.randint(-1, size, (size,)) + offset = trafo.get_offset_for_field(input_field) + expected = np.where(input_field >= 0, -1, 0) + print(expected) + assert np.allclose(expected, offset) diff --git a/testutils/src/icon4py/testutils/data_handling.py b/testutils/src/icon4py/testutils/data_handling.py new file mode 100644 index 000000000..e069b24f6 --- /dev/null +++ b/testutils/src/icon4py/testutils/data_handling.py @@ -0,0 +1,32 @@ +# 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 tarfile +from pathlib import Path + +import wget + + +def download_and_extract(uri: str, local_path: Path, data_file: str): + local_path.mkdir(parents=True, exist_ok=True) + if not any(local_path.iterdir()): + print( + f"directory {local_path} is empty: downloading data from {uri} and extracting" + ) + wget.download(uri, out=data_file) + # extract downloaded file + if not tarfile.is_tarfile(data_file): + raise NotImplementedError(f"{data_file} needs to be a valid tar file") + with tarfile.open(data_file, mode="r:*") as tf: + tf.extractall(path=local_path) + Path(data_file).unlink(missing_ok=True) diff --git a/testutils/src/icon4py/testutils/fixtures.py b/testutils/src/icon4py/testutils/fixtures.py index 690feafc3..356121fa7 100644 --- a/testutils/src/icon4py/testutils/fixtures.py +++ b/testutils/src/icon4py/testutils/fixtures.py @@ -11,12 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import tarfile from pathlib import Path import pytest -import wget +from icon4py.testutils.data_handling import download_and_extract from icon4py.testutils.serialbox_utils import IconSerialDataProvider @@ -33,18 +32,7 @@ def setup_icon_data(): Session scoped fixture which is a prerequisite of all the other fixtures in this file. """ - data_path.mkdir(parents=True, exist_ok=True) - if not any(data_path.iterdir()): - print( - f"directory {data_path} is empty: downloading data from {data_uri} and extracting" - ) - wget.download(data_uri, out=data_file) - # extract downloaded file - if not tarfile.is_tarfile(data_file): - raise NotImplementedError(f"{data_file} needs to be a valid tar file") - with tarfile.open(data_file, mode="r:*") as tf: - tf.extractall(path=data_path) - Path(data_file).unlink(missing_ok=True) + download_and_extract(data_uri, data_path, data_file) @pytest.fixture From 0ae5571c24289c51701e7a2ac7b50866bfb6212e Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 28 Mar 2023 09:28:20 +0200 Subject: [PATCH 082/263] - boundary conditions handling WIP (1) --- common/src/icon4py/grid/icon_grid.py | 5 +- common/tests/test_gridmanager.py | 52 ++++++++++++++++++- testutils/src/icon4py/testutils/fixtures.py | 3 +- .../src/icon4py/testutils/serialbox_utils.py | 3 ++ 4 files changed, 59 insertions(+), 4 deletions(-) diff --git a/common/src/icon4py/grid/icon_grid.py b/common/src/icon4py/grid/icon_grid.py index 2c1414440..d5cfa1ec4 100644 --- a/common/src/icon4py/grid/icon_grid.py +++ b/common/src/icon4py/grid/icon_grid.py @@ -141,7 +141,7 @@ def get_c2e2c_connectivity(self): return self._grid.get_c2e2c_connectivity() def get_v2c_connectivity(self): - return self._grid.get_v2e_connectivity() + return self._grid.get_v2c_connectivity() def get_c2e2co_connectivity(self): return self._grid.get_c2e2co_connectivity() @@ -207,3 +207,6 @@ def from_grid_dataset(self, dataset: Dataset) -> IconGrid: ) ) return icon_grid + + def get_e2c2e_connectivity(self): + pass diff --git a/common/tests/test_gridmanager.py b/common/tests/test_gridmanager.py index a7a213e7a..e6da2f22d 100644 --- a/common/tests/test_gridmanager.py +++ b/common/tests/test_gridmanager.py @@ -165,6 +165,14 @@ def test_grid_parser_index_fields(simple_mesh_data, caplog): assert np.allclose(grid_parser.int_field(GridFile.Offsets.V2C), mesh.v2c) + + +#e2c2v - diamond: serial, simple + +#c2v: grid, ??? +#v2e2v: grid,??? + +# v2e: serial, simple, grid @pytest.mark.skip("TODO: handling of boundary values") @pytest.mark.datatest def test_gridmanager_eval_v2e(caplog, grid_savepoint, get_grid_files): @@ -175,8 +183,15 @@ def test_gridmanager_eval_v2e(caplog, grid_savepoint, get_grid_files): gm.get_v2e_connectivity().table, grid_savepoint.v2e()[0:num_vertex, :] ) +#v2c: serial, simple, grid +#@pytest.mark.skip("TODO: handling of boundary values") + +#mo_model_domimp_patches.f90 lines 2183ff +# ! Checks for the pentagon case and moves dummy cells to end. +# ! The dummy entry is either set to 0 or duplicated from the last one +# SUBROUTINE move_dummies_to_end(array, array_size, max_connectivity, duplicate) + -@pytest.mark.skip("TODO: v2c not serialized??") @pytest.mark.datatest def test_gridmanager_eval_v2c(caplog, grid_savepoint, get_grid_files): caplog.set_level(logging.DEBUG) @@ -186,7 +201,18 @@ def test_gridmanager_eval_v2c(caplog, grid_savepoint, get_grid_files): gm.get_v2c_connectivity().table, grid_savepoint.v2c()[0:num_vertex, :] ) +#e2v: serial, simple, grid +@pytest.mark.datatest +def test_gridmanager_eval_e2v(caplog, grid_savepoint, get_grid_files): + caplog.set_level(logging.DEBUG) + fname = r04b09_dsl_grid_path.joinpath("grid.nc") + gm, num_cells, num_edges, num_vertex = _init_grid_manager(fname) + assert np.allclose( + gm.get_e2v_connectivity().table, grid_savepoint.e2v()[0:num_edges, :] + ) + +# e2c :serial, simple, grid @pytest.mark.datatest def test_gridmanager_eval_e2c(caplog, grid_savepoint, get_grid_files): caplog.set_level(logging.DEBUG) @@ -198,6 +224,7 @@ def test_gridmanager_eval_e2c(caplog, grid_savepoint, get_grid_files): ) +#c2e: serial, simple, grid @pytest.mark.datatest def test_gridmanager_eval_c2e(caplog, grid_savepoint, get_grid_files): caplog.set_level(logging.DEBUG) @@ -207,6 +234,28 @@ def test_gridmanager_eval_c2e(caplog, grid_savepoint, get_grid_files): gm.get_c2e_connectivity().table, grid_savepoint.c2e()[0:num_cells, :] ) +# e2c2e (e2c2eo) - diamond: serial, simple +# @pytest.mark.datatest +# def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, get_grid_files): +# caplog.set_level(logging.DEBUG) +# fname = r04b09_dsl_grid_path.joinpath("grid.nc") +# gm, num_cells, num_edges, num_vertex = _init_grid_manager(fname) +# assert np.allclose( +# gm.get_e2c2e_connectivity().table, grid_savepoint.e2c2e()[0:num_cells, :] +# ) + +#c2e2c: serial, simple, grid +@pytest.mark.datatest +def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, get_grid_files): + caplog.set_level(logging.DEBUG) + fname = r04b09_dsl_grid_path.joinpath("grid.nc") + gm, num_cells, num_edges, num_vertex = _init_grid_manager(fname) + assert np.allclose( + gm.get_c2e2c_connectivity().table, grid_savepoint.c2e2c()[0:num_cells, :] + ) + + + def _init_grid_manager(fname): gm = GridManager(ToGt4PyTransformation(), fname) @@ -240,5 +289,4 @@ def test_gt4py_transform_offset_by_1_where_valid(size): input_field = np.random.randint(-1, size, (size,)) offset = trafo.get_offset_for_field(input_field) expected = np.where(input_field >= 0, -1, 0) - print(expected) assert np.allclose(expected, offset) diff --git a/testutils/src/icon4py/testutils/fixtures.py b/testutils/src/icon4py/testutils/fixtures.py index 356121fa7..f0c0e816e 100644 --- a/testutils/src/icon4py/testutils/fixtures.py +++ b/testutils/src/icon4py/testutils/fixtures.py @@ -19,7 +19,8 @@ from icon4py.testutils.serialbox_utils import IconSerialDataProvider -data_uri = "https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download" +#data_uri = "https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download" +data_uri= "https://polybox.ethz.ch/index.php/s/joglKLsbyIXOOp8/download" data_path = Path(__file__).parent.joinpath("ser_icondata") extracted_path = data_path.joinpath("mch_ch_r04b09_dsl/ser_data") data_file = data_path.joinpath("mch_ch_r04b09_dsl_v2.tar.gz").name diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 9a037e60c..8ea1b4a7c 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -167,6 +167,9 @@ def e2c2v(self): def v2e(self): return self._get_connectiviy_array("v2e") + def v2c(self): + return self._get_connectiviy_array("v2c") + class IconDiffusionInitSavepoint(IconSavepoint): def hdef_ic(self): From 72e6e45d3097b77f15b79799a2c2a37f24142c73 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 30 Mar 2023 14:59:03 +0200 Subject: [PATCH 083/263] add missing dimension --- common/src/icon4py/common/dimension.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/common/src/icon4py/common/dimension.py b/common/src/icon4py/common/dimension.py index 9193bb8b9..cc4bc5e81 100644 --- a/common/src/icon4py/common/dimension.py +++ b/common/src/icon4py/common/dimension.py @@ -27,6 +27,7 @@ E2VDim = Dimension("E2V", DimensionKind.LOCAL) C2EDim = Dimension("C2E", DimensionKind.LOCAL) V2CDim = Dimension("V2C", DimensionKind.LOCAL) +C2VDim = Dimension("C2V", DimensionKind.LOCAL) V2EDim = Dimension("V2E", DimensionKind.LOCAL) E2C2VDim = Dimension("E2C2V", DimensionKind.LOCAL) C2E2CODim = Dimension("C2E2CO", DimensionKind.LOCAL) @@ -36,6 +37,7 @@ E2C = FieldOffset("E2C", source=CellDim, target=(EdgeDim, E2CDim)) C2E = FieldOffset("C2E", source=EdgeDim, target=(CellDim, C2EDim)) V2C = FieldOffset("V2C", source=CellDim, target=(VertexDim, V2CDim)) +C2V = FieldOffset("C2V", source=VertexDim, target=(CellDim, C2VDim)) V2E = FieldOffset("V2E", source=EdgeDim, target=(VertexDim, V2EDim)) E2V = FieldOffset("E2V", source=VertexDim, target=(EdgeDim, E2VDim)) C2CE = FieldOffset("C2CE", source=CEDim, target=(CellDim, C2EDim)) From 8c92a131e7fb469f5bfb24bd70e439dcea99a80f Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 30 Mar 2023 14:59:56 +0200 Subject: [PATCH 084/263] wrong interpretation of icon short names: rename lb=lateral_boundary --- atm_dyn_iconam/src/icon4py/diffusion/horizontal.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py index 531926929..d29187be0 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py @@ -39,19 +39,19 @@ class HorizontalMarkerIndex: _MIN_RL_EDGE: Final[int] = _MIN_RL_EDGE_INT - (2 * NUM_GHOST_ROWS + 1) _MAX_RL_EDGE: Final[int] = 2 * _MAX_RL_CELL - _LOCAL_BOUNDARY_EDGES: Final[int] = 1 + _ICON_INDEX_OFFSET_EDGES + _LATERAL_BOUNDARY_EDGES: Final[int] = 1 + _ICON_INDEX_OFFSET_EDGES _INTERIOR_EDGES: Final[int] = _ICON_INDEX_OFFSET_EDGES _NUDGING_EDGES: Final[int] = _GRF_BOUNDARY_WIDTH_EDGES + _ICON_INDEX_OFFSET_EDGES _HALO_EDGES: Final[int] = _MIN_RL_EDGE_INT + _ICON_INDEX_OFFSET_EDGES _END_EDGES: Final[int] = 0 - _LOCAL_BOUNDARY_CELLS: Final[int] = 1 + _ICON_INDEX_OFFSET_CELLS + _LATERAL_BOUNDARY_CELLS: Final[int] = 1 + _ICON_INDEX_OFFSET_CELLS _INTERIOR_CELLS: Final[int] = _ICON_INDEX_OFFSET_CELLS _NUDGING_CELLS: Final[int] = _GRF_BOUNDARY_WIDTH_CELL + 1 + _ICON_INDEX_OFFSET_CELLS _HALO_CELLS: Final[int] = _MIN_RL_CELL_INT + _ICON_INDEX_OFFSET_CELLS _END_CELLS: Final[int] = 0 - _LOCAL_BOUNDARY_VERTICES = 1 + _ICON_INDEX_OFFSET_VERTEX + _LATERAL_BOUNDARY_VERTICES = 1 + _ICON_INDEX_OFFSET_VERTEX _INTERIOR_VERTICES: Final[int] = _ICON_INDEX_OFFSET_VERTEX _NUDGING_VERTICES: Final[int] = 0 _HALO_VERTICES: Final[int] = _MIN_RL_VERTEX_INT + _ICON_INDEX_OFFSET_VERTEX @@ -61,11 +61,11 @@ class HorizontalMarkerIndex: def local_boundary(cls, dim: Dimension) -> int: match (dim): case (dimension.CellDim): - return cls._LOCAL_BOUNDARY_CELLS + return cls._LATERAL_BOUNDARY_CELLS case (dimension.EdgeDim): - return cls._LOCAL_BOUNDARY_EDGES + return cls._LATERAL_BOUNDARY_EDGES case (dimension.VertexDim): - return cls._LOCAL_BOUNDARY_VERTICES + return cls._LATERAL_BOUNDARY_VERTICES @classmethod def local(cls, dim: Dimension) -> int: From f77dc9016cad6c545d6578242ff3a398b3ab197c Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 4 Apr 2023 14:26:50 +0200 Subject: [PATCH 085/263] wrong interpretation of icon short names: rename lb=lateral_boundary --- .../src/icon4py/diffusion/diffusion.py | 16 ++--- .../src/icon4py/diffusion/horizontal.py | 2 +- atm_dyn_iconam/tests/test_icon_grid.py | 66 ++++++++++--------- 3 files changed, 45 insertions(+), 39 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 2cba30b0e..b132213d3 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -607,7 +607,7 @@ def time_step( vertex_endindex_local, ) = self.grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 3, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, HorizontalMarkerIndex.local(VertexDim), ) @@ -616,13 +616,13 @@ def time_step( vertex_endindex_local_minus1, ) = self.grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 1, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, HorizontalMarkerIndex.local(VertexDim) - 1, ) edge_start_lb_plus4, _ = self.grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, ) diff_prog.diffusion_run( @@ -781,8 +781,8 @@ def _do_diffusion_step( edge_start_lb_plus4, _ = self.grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, ) ( @@ -799,7 +799,7 @@ def _do_diffusion_step( vertex_end_local, ) = self.grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 3, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, HorizontalMarkerIndex.local(VertexDim), ) ( @@ -807,7 +807,7 @@ def _do_diffusion_step( vertex_end_local_minus1, ) = self.grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 1, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, HorizontalMarkerIndex.local(VertexDim) - 1, ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py index d29187be0..f7a46ccbf 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py @@ -58,7 +58,7 @@ class HorizontalMarkerIndex: _END_VERTICES: Final[int] = 0 @classmethod - def local_boundary(cls, dim: Dimension) -> int: + def lateral_boundary(cls, dim: Dimension) -> int: match (dim): case (dimension.CellDim): return cls._LATERAL_BOUNDARY_CELLS diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index 838a01cbb..700158871 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -54,13 +54,13 @@ def test_horizontal_grid_cell_indices(icon_grid): ) # lb+1 assert icon_grid.get_indices_from_to( CellDim, - HorizontalMarkerIndex.local_boundary(CellDim) + 1, - HorizontalMarkerIndex.local_boundary(CellDim) + 1, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, ) == (850, 1688) assert icon_grid.get_indices_from_to( CellDim, - HorizontalMarkerIndex.local_boundary(CellDim) + 2, - HorizontalMarkerIndex.local_boundary(CellDim) + 2, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, ) == ( 1688, 2511, @@ -126,64 +126,64 @@ def test_horizontal_edge_indices(icon_grid): ) # nudging assert icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 7, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 7, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, ) == ( 4184, 4989, ) # lb +7 assert icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 6, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 6, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, ) == ( 3777, 4184, ) # lb +6 assert icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 5, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 5, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, ) == ( 2954, 3777, ) # lb +5 assert icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, ) == ( 2538, 2954, ) # lb +4 assert icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 3, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 3, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, ) == ( 1700, 2538, ) # lb +3 assert icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 2, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 2, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, ) == ( 1278, 1700, ) # lb +2 assert icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 1, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 1, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, ) == ( 428, 1278, ) # lb +1 assert icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim), - HorizontalMarkerIndex.local_boundary(EdgeDim), + HorizontalMarkerIndex.lateral_boundary(EdgeDim), + HorizontalMarkerIndex.lateral_boundary(EdgeDim), ) == ( 0, 428, @@ -210,28 +210,28 @@ def test_horizontal_vertex_indices(icon_grid): assert icon_grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim), - HorizontalMarkerIndex.local_boundary(VertexDim), + HorizontalMarkerIndex.lateral_boundary(VertexDim), + HorizontalMarkerIndex.lateral_boundary(VertexDim), ) == (0, 428) assert icon_grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 1, - HorizontalMarkerIndex.local_boundary(VertexDim) + 1, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, ) == (428, 850) assert icon_grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 2, - HorizontalMarkerIndex.local_boundary(VertexDim) + 2, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, ) == (850, 1266) assert icon_grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 3, - HorizontalMarkerIndex.local_boundary(VertexDim) + 3, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, ) == (1266, 1673) assert icon_grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 4, - HorizontalMarkerIndex.local_boundary(VertexDim) + 4, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, ) == (1673, 2071) assert icon_grid.get_indices_from_to( @@ -239,3 +239,9 @@ def test_horizontal_vertex_indices(icon_grid): HorizontalMarkerIndex.interior(VertexDim), HorizontalMarkerIndex.interior(VertexDim), ) == (2071, 10663) + + assert icon_grid.get_indices_from_to( + VertexDim, + HorizontalMarkerIndex.nudging(VertexDim), + HorizontalMarkerIndex.nudging(VertexDim), + ) == (10663, 10663) From d36b3e40e79987158d5fb8036903ae3c121c96ef Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 4 Apr 2023 18:43:59 +0200 Subject: [PATCH 086/263] fix index field tests (revert the last valid index workaround for serialized data in tests) --- .../grid/{icon_grid.py => grid_manager.py} | 154 +++++++++++-- common/tests/test_gridmanager.py | 214 ++++++++++++++---- 2 files changed, 295 insertions(+), 73 deletions(-) rename common/src/icon4py/grid/{icon_grid.py => grid_manager.py} (51%) diff --git a/common/src/icon4py/grid/icon_grid.py b/common/src/icon4py/grid/grid_manager.py similarity index 51% rename from common/src/icon4py/grid/icon_grid.py rename to common/src/icon4py/grid/grid_manager.py index d5cfa1ec4..edaec11fe 100644 --- a/common/src/icon4py/grid/icon_grid.py +++ b/common/src/icon4py/grid/grid_manager.py @@ -13,11 +13,11 @@ import logging from abc import ABC +from enum import Enum from typing import Optional from uuid import UUID import numpy as np -from gt4py.eve import StrEnum from gt4py.next.common import Dimension from netCDF4 import Dataset @@ -31,12 +31,13 @@ V2CDim, V2EDim, VertexDim, + C2VDim, ) -from icon4py.diffusion.horizontal import HorizontalMeshSize +from icon4py.diffusion.horizontal import HorizontalMeshSize, HorizontalMarkerIndex from icon4py.diffusion.icon_grid import IconGrid, MeshConfig, VerticalMeshConfig -class GridFileName(StrEnum): +class GridFileName(str, Enum): pass @@ -70,6 +71,24 @@ class Dimension(GridFileName): NEIGHBORING_EDGES_TO_CELL_SIZE = "nv" # 3 E2V_SIZE = "nc" # 2 + MAX_CHILD_DOMAINS = "max_chdom" + + # TODO: @magdalena what does the grf stand for 'grid_refinement'? + CELL_GRF = "cell_grf" + EDGE_GRF = "edge_grf" + VERTEX_GRF = "vert_grf" + + class GridRefinement(GridFileName): + CONTROL_CELLS = "refin_c_ctrl" + CONTROL_EDGES = "refin_e_ctrl" + CONTROL_VERTICES = "refin_v_ctrl" + START_INDEX_CELLS = "start_idx_c" + START_INDEX_EDGES = "start_idx_e" + START_INDEX_VERTICES = "start_idx_v" + END_INDEX_CELLS ="end_idx_c" + END_INDEX_EDGES = "end_idx_e" + END_INDEX_VERTICES = "end_idx_v" + def __init__(self, dataset: Dataset): self._dataset = dataset self._log = logging.getLogger(__name__) @@ -78,25 +97,39 @@ def dimension(self, name: GridFileName) -> int: return self._dataset.dimensions[name].size def int_field(self, name: GridFileName, transpose=True) -> np.ndarray: - nc_variable = self._dataset.variables[name] - self._log.debug(f"{name}: {nc_variable}") - data = nc_variable[:] + try: + nc_variable = self._dataset.variables[name] + self._log.debug(f"{name}: {nc_variable}") + data = nc_variable[:] - data = np.array(data, dtype=np.int32) - return np.transpose(data) if transpose else data + data = np.array(data, dtype=np.int32) + return np.transpose(data) if transpose else data + except KeyError as error: + msg = f"{name} does not exist in dataset" + self._log.warning(msg) + raise IconGridError(msg) class IconGridError(RuntimeError): pass +class IconDomainZone(str,Enum): + HALO = "halo" + LATERAL_BOUNDARY = "lb" + NUDGING = "nudging" + INTERIOR = "interior" # TODO @magdalena is this the same as "prognostic" if yes then change the naming + END = "end" -class GridTransformation(ABC): - def get_offset_for_field(self, array: np.ndarray): + +class IndexTransformation(ABC): + def get_offset_for_index_field(self, array: np.ndarray): return np.zeros(array.shape) -class ToGt4PyTransformation(GridTransformation): - def get_offset_for_field(self, array: np.ndarray): + + +class ToGt4PyTransformation(IndexTransformation): + def get_offset_for_index_field(self, array: np.ndarray): """ Calculate the index offset needed for usage with python. @@ -106,28 +139,92 @@ def get_offset_for_field(self, array: np.ndarray): return np.where(array < 0, 0, -1) +class GridDomainBoundaryTransformation(ABC): + def __init__(self): + pass + +class SingleNodeBoundaryTransform(GridDomainBoundaryTransformation): + pass + + class GridManager: - def __init__(self, transformation: GridTransformation, grid_file: str): + def __init__(self, transformation: IndexTransformation, grid_file: str): self._log = logging.getLogger(__name__) self._transformation = transformation self._grid: Optional[IconGrid] = None self._file_names = grid_file + self._start_indices = {} + self._end_indices = {} def init(self): - _, grid = self._read_from_gridfile(self._file_names) + dataset = self._read_gridfile(self._file_names) + start_indices, end_indices = self._read_boundaries(dataset) + self._start_indices.update(start_indices) + self._end_indices.update(end_indices) + _, grid= self._read_grid(dataset) self._grid = grid - def _read_from_gridfile(self, fname: str) -> tuple[UUID, IconGrid]: + def _read_gridfile(self, fname:str)->Dataset: try: - dataset = Dataset(fname, "r", format="NETCDF4") - + dataset = Dataset(self._file_names, "r", format="NETCDF4") self._log.debug(dataset) - grid_id = UUID(dataset.getncattr(GridFile.Property.GRID_ID)) - return grid_id, self.from_grid_dataset(dataset) + return dataset except FileNotFoundError: self._log.error(f"gridfile {fname} not found, aborting") exit(1) + def _read_boundaries(self, dataset): + reader = GridFile(dataset) + grf_vertices = reader.dimension(GridFile.Dimension.VERTEX_GRF) + grf_edges = reader.dimension(GridFile.Dimension.EDGE_GRF) + grf_cells = reader.dimension(GridFile.Dimension.CELL_GRF) + refin_c_ctl = reader.int_field(GridFile.GridRefinement.CONTROL_CELLS) + refin_v_ctl = reader.int_field(GridFile.GridRefinement.CONTROL_VERTICES) + refin_e_ctl = reader.int_field(GridFile.GridRefinement.CONTROL_EDGES) + start_indices = {} + end_indices = {} + start_indices[CellDim] = self._get_index_field(reader, GridFile.GridRefinement.START_INDEX_CELLS, transpose=False) + end_indices[CellDim] = self._get_index_field(reader, GridFile.GridRefinement.END_INDEX_CELLS, transpose=False, apply_offset=False) + start_indices[EdgeDim] = self._get_index_field(reader, GridFile.GridRefinement.START_INDEX_EDGES, transpose=False) + end_indices[EdgeDim] = self._get_index_field(reader, GridFile.GridRefinement.END_INDEX_EDGES, transpose=False, apply_offset=False) + start_indices[VertexDim] = self._get_index_field(reader, GridFile.GridRefinement.START_INDEX_VERTICES, transpose=False) + end_indices[VertexDim] = self._get_index_field(reader, GridFile.GridRefinement.END_INDEX_VERTICES, transpose=False, apply_offset=False) + + return start_indices, end_indices + + + + def get_domain_boundaries(self, dim: Dimension, zone: IconDomainZone) -> tuple[int, int]: + match zone: + case IconDomainZone.LATERAL_BOUNDARY: + index = HorizontalMarkerIndex.lateral_boundary(dim) + return self._get_start_end_indices(dim, index) + case IconDomainZone.INTERIOR: + index = HorizontalMarkerIndex.interior(dim) + return self._get_start_end_indices(dim, index) + case IconDomainZone.NUDGING: + index = HorizontalMarkerIndex.nudging(dim) + return self._get_start_end_indices(dim, index) + case IconDomainZone.HALO: + index = HorizontalMarkerIndex.local(dim) + return self._get_start_end_indices(dim, index) + case IconDomainZone.END: + index = HorizontalMarkerIndex.end(dim) + return self._get_start_end_indices(dim, index) + + def _get_start_end_indices(self, dim, index): + try: + return self._start_indices[dim][0][index], self._end_indices[dim][0][index] + except KeyError as error: + msg = f"start, end indices for dimension {dim} not present" + self._log.error(msg) + raise IconGridError(msg) + def _read_grid(self, dataset:Dataset) -> tuple[UUID, IconGrid]: + + grid_id = UUID(dataset.getncattr(GridFile.Property.GRID_ID)) + return grid_id, self.from_grid_dataset(dataset) + + def get_c2e_connectivity(self): return self._grid.get_c2e_connectivity() @@ -155,6 +252,10 @@ def get_v2e_connectivity(self): def get_e2ecv_connectivity(self): return self._grid.get_e2ecv_connectivity() + def get_e2c2e_connectivity(self): + return self._grid.get_e2c2e_connectivity() + + def get_size(self, dim: Dimension): if dim == VertexDim: return self._grid.config.num_vertices @@ -166,9 +267,12 @@ def get_size(self, dim: Dimension): self._log.warning(f"cannot determine size of unknown dimension {dim}") raise IconGridError(f"Unknown dimension {dim}") - def _get_index_field(self, reader, name: GridFileName): - field = reader.int_field(name) - field = field + self._transformation.get_offset_for_field(field) + + + def _get_index_field(self, reader, name: GridFileName, transpose=True, apply_offset = True): + field = reader.int_field(name, transpose=transpose) + if apply_offset: + field = field + self._transformation.get_offset_for_index_field(field) return field def from_grid_dataset(self, dataset: Dataset) -> IconGrid: @@ -176,6 +280,8 @@ def from_grid_dataset(self, dataset: Dataset) -> IconGrid: num_cells = reader.dimension(GridFile.Dimension.CELL_NAME) num_edges = reader.dimension(GridFile.Dimension.EDGE_NAME) num_vertices = reader.dimension(GridFile.Dimension.VERTEX_NAME) + + grid_size = HorizontalMeshSize( num_vertices=num_vertices, num_edges=num_edges, num_cells=num_cells ) @@ -203,10 +309,10 @@ def from_grid_dataset(self, dataset: Dataset) -> IconGrid: V2EDim: v2e, V2CDim: v2c, C2E2CDim: c2e2c, + C2VDim: c2v, } ) ) return icon_grid - def get_e2c2e_connectivity(self): - pass + diff --git a/common/tests/test_gridmanager.py b/common/tests/test_gridmanager.py index e6da2f22d..0badaed7b 100644 --- a/common/tests/test_gridmanager.py +++ b/common/tests/test_gridmanager.py @@ -20,12 +20,12 @@ from netCDF4 import Dataset from icon4py.common.dimension import CellDim, EdgeDim, VertexDim -from icon4py.grid.icon_grid import ( +from icon4py.grid.grid_manager import ( GridFile, GridFileName, GridManager, - GridTransformation, - ToGt4PyTransformation, + IndexTransformation, + ToGt4PyTransformation, IconDomainZone, ) from icon4py.testutils.data_handling import download_and_extract from icon4py.testutils.simple_mesh import SimpleMesh @@ -33,11 +33,13 @@ SIMPLE_MESH_NC = "./simple_mesh_grid.nc" -grid_uri = "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download" +mch_ch_r04b09_dsl_grid_uri = "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download" +r02b04_global_grid_uri = "https://polybox.ethz.ch/index.php/s/0EM8O8U53GKGsst/download" grids_path = Path(__file__).parent.joinpath("grids") r04b09_dsl_grid_path = grids_path.joinpath("mch_ch_r04b09_dsl") -data_file = r04b09_dsl_grid_path.joinpath("mch_ch_r04b09_dsl_grids_v1.tar.gz").name - +r04b09_dsl_data_file = r04b09_dsl_grid_path.joinpath("mch_ch_r04b09_dsl_grids_v1.tar.gz").name +r02b04_global_grid_path = grids_path.joinpath("icon_r02b04_global") +r02b04_global_data_file = r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_G.tar.gz").name @pytest.fixture(scope="session") def get_grid_files(): @@ -46,7 +48,8 @@ def get_grid_files(): Session scoped fixture which is a prerequisite of all the other fixtures in this file. """ - download_and_extract(grid_uri, r04b09_dsl_grid_path, data_file) + download_and_extract(mch_ch_r04b09_dsl_grid_uri, r04b09_dsl_grid_path, r04b09_dsl_data_file) + download_and_extract(r02b04_global_grid_uri, r02b04_global_grid_path, r02b04_global_data_file) @pytest.fixture @@ -64,6 +67,16 @@ def simple_mesh_data(): dataset.createDimension(GridFile.Dimension.CELL_NAME, size=mesh.n_cells) dataset.createDimension(GridFile.Dimension.E2V_SIZE, size=mesh.n_e2v) dataset.createDimension(GridFile.Dimension.DIAMOND_EDGE_SIZE, size=mesh.n_e2c2e) + dataset.createDimension(GridFile.Dimension.MAX_CHILD_DOMAINS, size=1) + # add dummy values for the grf dimensions + # TODO @magdalena fix to something more useful? + dataset.createDimension(GridFile.Dimension.CELL_GRF, size=14) + dataset.createDimension(GridFile.Dimension.EDGE_GRF, size=24) + dataset.createDimension(GridFile.Dimension.VERTEX_GRF, size=13) + add_to_dataset(dataset, np.zeros(mesh.n_edges), GridFile.GridRefinement.CONTROL_EDGES, (GridFile.Dimension.EDGE_NAME,)) + add_to_dataset(dataset, np.zeros(mesh.n_cells), GridFile.GridRefinement.CONTROL_CELLS, (GridFile.Dimension.CELL_NAME,)) + add_to_dataset(dataset, np.zeros(mesh.n_vertices), GridFile.GridRefinement.CONTROL_VERTICES,(GridFile.Dimension.VERTEX_NAME,)) + dataset.createDimension( GridFile.Dimension.NEIGHBORING_EDGES_TO_CELL_SIZE, size=mesh.n_c2e ) @@ -98,7 +111,7 @@ def simple_mesh_data(): GridFile.Offsets.V2C, (GridFile.Dimension.V2E_SIZE, GridFile.Dimension.VERTEX_NAME), ) - # TODO fix me there is no v2c in the simple mesh + # TODO @magdalena: there is no v2c in the simple mesh add_to_dataset( dataset, np.zeros((mesh.n_cells, 3), dtype=np.int32), @@ -130,6 +143,13 @@ def simple_mesh_data(): GridFile.Dimension.CELL_NAME, ), ) + + add_to_dataset(dataset, np.ones((1, 24), dtype=np.int32), GridFile.GridRefinement.START_INDEX_EDGES, (GridFile.Dimension.MAX_CHILD_DOMAINS, GridFile.Dimension.EDGE_GRF)) + add_to_dataset(dataset, np.ones((1, 14), dtype=np.int32), GridFile.GridRefinement.START_INDEX_CELLS, (GridFile.Dimension.MAX_CHILD_DOMAINS, GridFile.Dimension.CELL_GRF)) + add_to_dataset(dataset, np.ones((1, 13), dtype=np.int32), GridFile.GridRefinement.START_INDEX_VERTICES, (GridFile.Dimension.MAX_CHILD_DOMAINS, GridFile.Dimension.VERTEX_GRF)) + add_to_dataset(dataset, np.ones((1, 24), dtype=np.int32), GridFile.GridRefinement.END_INDEX_EDGES, (GridFile.Dimension.MAX_CHILD_DOMAINS, GridFile.Dimension.EDGE_GRF)) + add_to_dataset(dataset, np.ones((1, 14), dtype=np.int32), GridFile.GridRefinement.END_INDEX_CELLS, (GridFile.Dimension.MAX_CHILD_DOMAINS, GridFile.Dimension.CELL_GRF)) + add_to_dataset(dataset, np.ones((1, 13), dtype=np.int32), GridFile.GridRefinement.END_INDEX_VERTICES,(GridFile.Dimension.MAX_CHILD_DOMAINS, GridFile.Dimension.VERTEX_GRF)) dataset.close() @@ -166,110 +186,170 @@ def test_grid_parser_index_fields(simple_mesh_data, caplog): - +# TODO @magdalena add test cases for #e2c2v - diamond: serial, simple - #c2v: grid, ??? #v2e2v: grid,??? # v2e: serial, simple, grid -@pytest.mark.skip("TODO: handling of boundary values") @pytest.mark.datatest def test_gridmanager_eval_v2e(caplog, grid_savepoint, get_grid_files): caplog.set_level(logging.DEBUG) fname = r04b09_dsl_grid_path.joinpath("grid.nc") - gm, num_cells, num_edges, num_vertex = _init_grid_manager(fname) + gm = _init_grid_manager(fname) + num_vertex = gm.get_size(VertexDim) + seralized_v2e = grid_savepoint.v2e()[0:num_vertex, :] + # there are vertices at the boundary of a local domain or at a pentagon point that have less than + # 6 neighbors hence there are "Missing values" in the grid file + # they get substituted by the "last valid index" in preprocessing step in icon. + assert not has_invalid_index(seralized_v2e) + assert has_invalid_index(gm.get_v2e_connectivity().table) + reset_invalid_index(seralized_v2e) assert np.allclose( - gm.get_v2e_connectivity().table, grid_savepoint.v2e()[0:num_vertex, :] + gm.get_v2e_connectivity().table, seralized_v2e ) #v2c: serial, simple, grid -#@pytest.mark.skip("TODO: handling of boundary values") - -#mo_model_domimp_patches.f90 lines 2183ff -# ! Checks for the pentagon case and moves dummy cells to end. -# ! The dummy entry is either set to 0 or duplicated from the last one -# SUBROUTINE move_dummies_to_end(array, array_size, max_connectivity, duplicate) - - @pytest.mark.datatest def test_gridmanager_eval_v2c(caplog, grid_savepoint, get_grid_files): caplog.set_level(logging.DEBUG) fname = r04b09_dsl_grid_path.joinpath("grid.nc") - gm, num_cells, num_edges, num_vertex = _init_grid_manager(fname) + gm = _init_grid_manager(fname) + num_vertex = gm.get_size(VertexDim) + serialized_v2c = grid_savepoint.v2c()[0:num_vertex, :] + # there are vertices that have less than 6 neighboring cells: either pentagon points or + # vertices at the boundary of the domain for a limited area mode + # hence in the grid file there are "missing values" + # they get substituted by the "last valid index" in preprocessing step in icon. + assert not has_invalid_index(serialized_v2c) + assert has_invalid_index(gm.get_v2c_connectivity().table) + reset_invalid_index(serialized_v2c) + assert np.allclose( - gm.get_v2c_connectivity().table, grid_savepoint.v2c()[0:num_vertex, :] + gm.get_v2c_connectivity().table, serialized_v2c ) + +def reset_invalid_index(index_array: np.ndarray): + """ + helper function to revert mo_model_domimp_patches.f90: move_dummies_to_end_idxblk. + see: + # ! Checks for the pentagon case and moves dummy cells to end. + # ! The dummy entry is either set to 0 or duplicated from the last one + # SUBROUTINE move_dummies_to_end(array, array_size, max_connectivity, duplicate) + + After reading the grid file ICON moves all invalid indices (neighbors not existing in the + grid file) to the end of the neighbor list and replaces them with the "last valid neighbor index" + it is up to the user then to ensure that any coefficients in neighbor some multiplied with + these values are zero in order to "remove" them again from the sum. + + For testing we resubstitute those to the GridFile.INVALID_INDEX + Args: + index_array: the array where values the invalid values have to be reset + + Returns: an array where the spurious "last valid index" are replaced by GridFile.INVALID_INDEX + + """ + for i in range(0, index_array.shape[0]): + uq, index = np.unique(index_array[i, :], return_index=True) + index_array[i, max(index) + 1:] = GridFile.INVALID_INDEX + + #e2v: serial, simple, grid @pytest.mark.datatest def test_gridmanager_eval_e2v(caplog, grid_savepoint, get_grid_files): caplog.set_level(logging.DEBUG) fname = r04b09_dsl_grid_path.joinpath("grid.nc") - gm, num_cells, num_edges, num_vertex = _init_grid_manager(fname) + gm = _init_grid_manager(fname) + num_edges = gm.get_size(EdgeDim) + + serialized_e2v = grid_savepoint.e2v()[0:num_edges, :] + # all vertices in the system have to neighboring edges, there no edges that point nowhere + # hence this connectivity has no "missing values" in the grid file + assert not has_invalid_index(serialized_e2v) + assert not has_invalid_index(gm.get_e2v_connectivity().table) assert np.allclose( - gm.get_e2v_connectivity().table, grid_savepoint.e2v()[0:num_edges, :] + gm.get_e2v_connectivity().table, serialized_e2v ) +def has_invalid_index(ar :np.ndarray): + return np.any(np.where(ar == GridFile.INVALID_INDEX)) + + # e2c :serial, simple, grid @pytest.mark.datatest def test_gridmanager_eval_e2c(caplog, grid_savepoint, get_grid_files): caplog.set_level(logging.DEBUG) fname = r04b09_dsl_grid_path.joinpath("grid.nc") - gm, num_cells, num_edges, num_vertex = _init_grid_manager(fname) - + gm = _init_grid_manager(fname) + num_edges = gm.get_size(EdgeDim) + serialized_c2e = grid_savepoint.e2c()[0:num_edges, :] + # there are edges at the boundary that have only one + # neighboring cell, there are "missing values" in the grid file + # and here they do not get substituted in the ICON preprocessing + assert has_invalid_index(serialized_c2e) + assert has_invalid_index(gm.get_e2c_connectivity().table) assert np.allclose( - gm.get_e2c_connectivity().table, grid_savepoint.e2c()[0:num_edges, :] + gm.get_e2c_connectivity().table, serialized_c2e ) #c2e: serial, simple, grid @pytest.mark.datatest def test_gridmanager_eval_c2e(caplog, grid_savepoint, get_grid_files): + caplog.set_level(logging.DEBUG) + fname = r04b09_dsl_grid_path.joinpath("grid.nc") + gm = _init_grid_manager(fname) + num_cells = gm.get_size(CellDim) + + serialized_c2e = grid_savepoint.c2e()[0:num_cells, :] + # no cells with less than 3 neighboring edges exist, otherwise the cell is not there in the + # first place + # hence there are no "missing values" in the grid file + assert not has_invalid_index(serialized_c2e) + assert not has_invalid_index(gm.get_c2e_connectivity().table) + assert np.allclose( + gm.get_c2e_connectivity().table, serialized_c2e + ) + +#e2c2e (e2c2eo) - diamond: serial, simple_mesh +@pytest.mark.skip("does this array exist in grid file?") +@pytest.mark.datatest +def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, get_grid_files): caplog.set_level(logging.DEBUG) fname = r04b09_dsl_grid_path.joinpath("grid.nc") gm, num_cells, num_edges, num_vertex = _init_grid_manager(fname) + serialized_e2c2e = grid_savepoint.e2c2e()[0:num_cells, :] + assert has_invalid_index(serialized_e2c2e) + assert has_invalid_index(gm.get_e2c2e_connectivity().table) assert np.allclose( - gm.get_c2e_connectivity().table, grid_savepoint.c2e()[0:num_cells, :] + gm.get_e2c2e_connectivity().table, serialized_e2c2e ) -# e2c2e (e2c2eo) - diamond: serial, simple -# @pytest.mark.datatest -# def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, get_grid_files): -# caplog.set_level(logging.DEBUG) -# fname = r04b09_dsl_grid_path.joinpath("grid.nc") -# gm, num_cells, num_edges, num_vertex = _init_grid_manager(fname) -# assert np.allclose( -# gm.get_e2c2e_connectivity().table, grid_savepoint.e2c2e()[0:num_cells, :] -# ) - -#c2e2c: serial, simple, grid +#c2e2c: serial, simple_mesh, grid @pytest.mark.datatest def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, get_grid_files): caplog.set_level(logging.DEBUG) fname = r04b09_dsl_grid_path.joinpath("grid.nc") - gm, num_cells, num_edges, num_vertex = _init_grid_manager(fname) + gm = _init_grid_manager(fname) + num_cells = gm.get_size(CellDim) assert np.allclose( gm.get_c2e2c_connectivity().table, grid_savepoint.c2e2c()[0:num_cells, :] ) - def _init_grid_manager(fname): gm = GridManager(ToGt4PyTransformation(), fname) gm.init() - num_vertex = gm.get_size(VertexDim) - num_edges = gm.get_size(EdgeDim) - num_cells = gm.get_size(CellDim) - return gm, num_cells, num_edges, num_vertex + return gm @pytest.mark.parametrize("dim, size", [(CellDim, 18), (EdgeDim, 27), (VertexDim, 9)]) def test_grid_manager_getsize(simple_mesh_data, simple_mesh_path, dim, size, caplog): caplog.set_level(logging.DEBUG) - gm = GridManager(GridTransformation(), simple_mesh_path) + gm = GridManager(IndexTransformation(), simple_mesh_path) gm.init() assert size == gm.get_size(dim) @@ -277,7 +357,7 @@ def test_grid_manager_getsize(simple_mesh_data, simple_mesh_path, dim, size, cap def test_gridmanager_given_file_not_found_then_abort(): fname = "./unknown_grid.nc" with pytest.raises(SystemExit) as error: - gm = GridManager(GridTransformation(), fname) + gm = GridManager(IndexTransformation(), fname) gm.init() assert error.type == SystemExit assert error.value == 1 @@ -287,6 +367,42 @@ def test_gridmanager_given_file_not_found_then_abort(): def test_gt4py_transform_offset_by_1_where_valid(size): trafo = ToGt4PyTransformation() input_field = np.random.randint(-1, size, (size,)) - offset = trafo.get_offset_for_field(input_field) + offset = trafo.get_offset_for_index_field(input_field) expected = np.where(input_field >= 0, -1, 0) assert np.allclose(expected, offset) + +@pytest.mark.parametrize("zone, bounds", [ + (IconDomainZone.NUDGING,(3316, 4104)), + (IconDomainZone.HALO, (20896, 20896)), + (IconDomainZone.LATERAL_BOUNDARY, (0, 850)), + (IconDomainZone.INTERIOR, (4104, 20896)) +] ) +def test_start_end_indices_cells(get_grid_files, zone, bounds): + fname = r04b09_dsl_grid_path.joinpath("grid.nc") + gm = _init_grid_manager(fname) + assert gm.get_domain_boundaries(CellDim, zone) == bounds + + +@pytest.mark.parametrize("zone, bounds", [ + (IconDomainZone.NUDGING,(4989, 5387)), + (IconDomainZone.HALO, (31558, 31558)), + (IconDomainZone.LATERAL_BOUNDARY, (0, 428)), + (IconDomainZone.INTERIOR, (6176, 31558))]) +def test_start_end_indices_edges(get_grid_files, zone, bounds): + fname = r04b09_dsl_grid_path.joinpath("grid.nc") + gm = _init_grid_manager(fname) + assert gm.get_domain_boundaries(EdgeDim, zone) == bounds + + +@pytest.mark.parametrize("zone, bounds", [ + (IconDomainZone.NUDGING,(10663, 10663)), + (IconDomainZone.HALO, (10663, 10663)), + (IconDomainZone.LATERAL_BOUNDARY, (0, 428)), + (IconDomainZone.INTERIOR, (2071, 10663)), + (IconDomainZone.END, (10663, 10663)) +]) +def test_start_end_indices_vertices(get_grid_files, zone, bounds): + fname = r04b09_dsl_grid_path.joinpath("grid.nc") + gm = _init_grid_manager(fname) + assert gm.get_domain_boundaries(VertexDim, zone) == bounds + From f0c1fef6fa86ceceae086839c785ce99240083f3 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 4 Apr 2023 19:04:13 +0200 Subject: [PATCH 087/263] fix typo, add function in gridsavepoint to get number of fields --- .../src/icon4py/testutils/serialbox_utils.py | 30 ++++++++++++------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 8ea1b4a7c..c7a775390 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -10,6 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +from typing import Optional import numpy as np import serialbox as ser @@ -129,7 +130,7 @@ def edge_end_index(self): def print_connectivity_info(name: str, ar: np.ndarray): print(f" connectivity {name} {ar.shape}") - def refin_ctrl(self, dim: Dimension): + def refin_ctrl(self, dim: Dimension)-> Optional[np.ndarray]: if dim == CellDim: return self.serializer.read("c_refin_ctl", self.savepoint) elif dim == EdgeDim: @@ -139,36 +140,45 @@ def refin_ctrl(self, dim: Dimension): else: return None - def c2e(self): + def num(self, dim:Dimension)->Optional[int]: + if dim == CellDim: + return self.serializer.read("num_cells", self.savepoint)[0] + elif dim == EdgeDim: + return self.serializer.read("num_edges", self.savepoint)[0] + elif dim == VertexDim: + return self.serializer.read("num_vert", self.savepoint)[0] + else: + return None - return self._get_connectiviy_array("c2e") + def c2e(self): + return self._get_connectivity_array("c2e") - def _get_connectiviy_array(self, name: str): + def _get_connectivity_array(self, name: str): connectivity = self.serializer.read(name, self.savepoint) - 1 print(f" connectivity {name} : {connectivity.shape}") return connectivity def c2e2c(self): - return self._get_connectiviy_array("c2e2c") + return self._get_connectivity_array("c2e2c") def e2c(self): - return self._get_connectiviy_array("e2c") + return self._get_connectivity_array("e2c") def e2v(self): # array "e2v" is actually e2c2v - v_ = self._get_connectiviy_array("e2v")[:, 0:2] + v_ = self._get_connectivity_array("e2v")[:, 0:2] print(f"real e2v {v_.shape}") return v_ def e2c2v(self): # array "e2v" is actually e2c2v - return self._get_connectiviy_array("e2v") + return self._get_connectivity_array("e2v") def v2e(self): - return self._get_connectiviy_array("v2e") + return self._get_connectivity_array("v2e") def v2c(self): - return self._get_connectiviy_array("v2c") + return self._get_connectivity_array("v2c") class IconDiffusionInitSavepoint(IconSavepoint): From 38e74f544c2571e363c3bfe46940014bbcdf021c Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 4 Apr 2023 19:04:30 +0200 Subject: [PATCH 088/263] test dimensions against serialized data --- common/tests/test_gridmanager.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/common/tests/test_gridmanager.py b/common/tests/test_gridmanager.py index 0badaed7b..e3c33cd10 100644 --- a/common/tests/test_gridmanager.py +++ b/common/tests/test_gridmanager.py @@ -172,6 +172,18 @@ def test_gridparser_dimension(simple_mesh_data): assert grid_parser.dimension(GridFile.Dimension.VERTEX_NAME) == mesh.n_vertices assert grid_parser.dimension(GridFile.Dimension.EDGE_NAME) == mesh.n_edges +def test_gridfile_vertex_cell_edge_dimensions(grid_savepoint): + fname = r04b09_dsl_grid_path.joinpath("grid.nc") + data = Dataset(fname, "r") + grid_file = GridFile(data) + + assert grid_file.dimension(GridFile.Dimension.CELL_NAME) == grid_savepoint.num(CellDim) + assert grid_file.dimension(GridFile.Dimension.EDGE_NAME) == grid_savepoint.num(EdgeDim) + # TODO: @magdalena fix in serialized data. it returns the num_cells + assert grid_file.dimension(GridFile.Dimension.VERTEX_NAME) == grid_savepoint.num(VertexDim) + + + def test_grid_parser_index_fields(simple_mesh_data, caplog): caplog.set_level(logging.DEBUG) From ddd6b7699c79e73ed08c232872dbfcf05ff87854 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 11 Apr 2023 20:43:27 +0200 Subject: [PATCH 089/263] - move commonly used fixtures to testutils - move icon_grid.py to common/grid - move vertical stuff to vertical.py - use grid_manager to setup grid in test_diff --- .../src/icon4py/diffusion/diffusion.py | 8 +- atm_dyn_iconam/tests/conftest.py | 11 +- atm_dyn_iconam/tests/test_diffusion.py | 16 +- atm_dyn_iconam/tests/test_icon_grid.py | 2 +- base-requirements-dev.txt | 3 +- base-requirements.txt | 1 + common/src/icon4py/common/dimension.py | 4 +- common/src/icon4py/grid/grid_manager.py | 259 ++++++------ .../src/icon4py/grid}/horizontal.py | 58 ++- .../src/icon4py/grid}/icon_grid.py | 59 +-- common/src/icon4py/grid/vertical.py | 56 +++ common/tests/conftest.py | 2 + common/tests/test_gridmanager.py | 388 ++++++++++-------- common/tests/test_horizontal.py | 37 ++ .../tests/test_vertical.py | 10 +- .../py2f/wrappers/diffusion_wrapper.py | 4 +- testutils/src/icon4py/testutils/fixtures.py | 37 +- .../src/icon4py/testutils/serialbox_utils.py | 6 +- 18 files changed, 596 insertions(+), 365 deletions(-) rename {atm_dyn_iconam/src/icon4py/diffusion => common/src/icon4py/grid}/horizontal.py (79%) rename {atm_dyn_iconam/src/icon4py/diffusion => common/src/icon4py/grid}/icon_grid.py (77%) create mode 100644 common/src/icon4py/grid/vertical.py create mode 100644 common/tests/test_horizontal.py rename {atm_dyn_iconam => common}/tests/test_vertical.py (81%) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index b132213d3..d2890ba44 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -45,9 +45,6 @@ from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_15 import ( - mo_nh_diffusion_stencil_15, -) from icon4py.atm_dyn_iconam.update_theta_and_exner import update_theta_and_exner from icon4py.common.constants import ( CPD, @@ -68,8 +65,6 @@ VertexDim, ) from icon4py.diffusion.diagnostic_state import DiagnosticState -from icon4py.diffusion.horizontal import HorizontalMarkerIndex -from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.prognostic_state import PrognosticState @@ -81,6 +76,9 @@ setup_fields_for_initial_step, zero_field, ) +from icon4py.grid.horizontal import HorizontalMarkerIndex +from icon4py.grid.icon_grid import IconGrid +from icon4py.grid.vertical import VerticalGridConfig, VerticalModelParams VectorTuple = namedtuple("VectorTuple", "x y") diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 65fbc89bf..c16cf5dd1 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -28,8 +28,9 @@ VertexDim, ) from icon4py.diffusion.diffusion import DiffusionConfig -from icon4py.diffusion.horizontal import HorizontalMeshSize -from icon4py.diffusion.icon_grid import IconGrid, MeshConfig, VerticalMeshConfig +from icon4py.grid.horizontal import HorizontalGridSize +from icon4py.grid.icon_grid import GridConfig, IconGrid +from icon4py.grid.vertical import VerticalGridConfig @pytest.fixture @@ -105,13 +106,13 @@ def icon_grid(data_provider): edge_starts = sp.edge_start_index() edge_ends = sp.edge_end_index() - config = MeshConfig( - HorizontalMeshSize( + config = GridConfig( + HorizontalGridSize( num_vertices=sp_meta["nproma"], # or rather "num_vert" num_cells=sp_meta["nproma"], # or rather "num_cells" num_edges=sp_meta["nproma"], # or rather "num_edges" ), - VerticalMeshConfig(num_lev=sp_meta["nlev"]), + VerticalGridConfig(num_lev=sp_meta["nlev"]), ) c2e2c = sp.c2e2c() diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index f16a96829..99909850a 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -17,7 +17,6 @@ from icon4py.common.dimension import ECVDim, KDim, VertexDim from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.diffusion import Diffusion, DiffusionParams, VectorTuple -from icon4py.diffusion.icon_grid import VerticalModelParams from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.prognostic_state import PrognosticState @@ -29,9 +28,13 @@ set_zero_v_k, setup_fields_for_initial_step, ) +from icon4py.grid.grid_manager import GridManager, ToGt4PyTransformation +from icon4py.grid.vertical import VerticalGridConfig, VerticalModelParams from icon4py.testutils.fixtures import ( # noqa F401 data_provider, + get_grid_files, grid_savepoint, + r04b09_dsl_gridfile, setup_icon_data, ) from icon4py.testutils.serialbox_utils import IconDiffusionInitSavepoint @@ -225,16 +228,25 @@ def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor))) +def init_grid_manager(fname, config): + gm = GridManager(ToGt4PyTransformation(), fname, config) + gm.init() + return gm + + @pytest.mark.datatest def test_diffusion_init( diffusion_savepoint_init, grid_savepoint, # noqa: F811 - icon_grid, + r04b09_dsl_gridfile, # noqa: F811 r04b09_diffusion_config, step_date_init, damping_height, ): config = r04b09_diffusion_config + grid_config = VerticalGridConfig(grid_savepoint.num(KDim)) + gm = init_grid_manager(r04b09_dsl_gridfile, grid_config) + icon_grid = gm.get_grid() additional_parameters = DiffusionParams(config) vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index 700158871..1a7ba1396 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -13,7 +13,7 @@ import pytest from icon4py.common.dimension import CellDim, EdgeDim, VertexDim -from icon4py.diffusion.horizontal import HorizontalMarkerIndex +from icon4py.grid.horizontal import HorizontalMarkerIndex from icon4py.testutils.fixtures import data_provider, setup_icon_data # noqa diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index b0fa61844..2e7c83dbc 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -25,5 +25,4 @@ setuptools>=40.8.0 wheel>=0.37.1 tox >= 3.25 wget >= 3.2 -netcdf4 >=1.6.3 -enum34 >= 1.1.10 + diff --git a/base-requirements.txt b/base-requirements.txt index 4e204ea2f..fe57e7cbd 100644 --- a/base-requirements.txt +++ b/base-requirements.txt @@ -1,2 +1,3 @@ # VCS gt4py @ git+https://github.com/GridTools/gt4py.git@main +netcdf4 >=1.6.3 diff --git a/common/src/icon4py/common/dimension.py b/common/src/icon4py/common/dimension.py index cc4bc5e81..092871249 100644 --- a/common/src/icon4py/common/dimension.py +++ b/common/src/icon4py/common/dimension.py @@ -11,14 +11,14 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import DimensionKind +from gt4py.next.common import Dimension, DimensionKind from gt4py.next.ffront.fbuiltins import Dimension, FieldOffset KDim = Dimension("K", kind=DimensionKind.VERTICAL) KHalfDim = Dimension("KHalf", kind=DimensionKind.VERTICAL) EdgeDim = Dimension("Edge") -CellDim = Dimension("Cell") +CellDim: Dimension = Dimension("Cell") VertexDim = Dimension("Vertex") CEDim = Dimension("CE") ECDim = Dimension("EC") diff --git a/common/src/icon4py/grid/grid_manager.py b/common/src/icon4py/grid/grid_manager.py index edaec11fe..1d4b3176e 100644 --- a/common/src/icon4py/grid/grid_manager.py +++ b/common/src/icon4py/grid/grid_manager.py @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - +import dataclasses import logging from abc import ABC from enum import Enum @@ -18,12 +18,13 @@ from uuid import UUID import numpy as np -from gt4py.next.common import Dimension +from gt4py.next.common import Dimension, DimensionKind from netCDF4 import Dataset from icon4py.common.dimension import ( C2E2CDim, C2EDim, + C2VDim, CellDim, E2CDim, E2VDim, @@ -31,61 +32,74 @@ V2CDim, V2EDim, VertexDim, - C2VDim, ) -from icon4py.diffusion.horizontal import HorizontalMeshSize, HorizontalMarkerIndex -from icon4py.diffusion.icon_grid import IconGrid, MeshConfig, VerticalMeshConfig +from icon4py.grid.horizontal import HorizontalGridSize +from icon4py.grid.icon_grid import GridConfig, IconGrid +from icon4py.grid.vertical import VerticalGridConfig class GridFileName(str, Enum): pass +@dataclasses.dataclass +class GridFileField: + name: GridFileName + shape: tuple[int, ...] + + +def _validate_shape(data: np.array, field_definition: GridFileField): + if data.shape != field_definition.shape: + raise IconGridError( + f"invalid grid file field {field_definition.name} does not have dimension {field_definition.shape}" + ) + + class GridFile: - """ - Represents and ICON netcdf grid file. - """ + """Represent and ICON netcdf grid file.""" INVALID_INDEX = -1 - class Property(GridFileName): + class PropertyName(GridFileName): GRID_ID = "uuidOfHGrid" PARENT_GRID_ID = "uuidOfParHGrid" - class Offsets(GridFileName): - C2E2C = "neighbor_cell_index" - V2E2V = "vertices_of_vertex" - V2E = "edges_of_vertex" - V2C = "cells_of_vertex" - E2V = "edge_vertices" - E2C = "adjacent_cell_of_edge" - C2V = "vertex_of_cell" - C2E = "edge_of_cell" - - class Dimension(GridFileName): + class OffsetName(GridFileName): + # e2c2e/e2c2eO: diamond edges (including origin) -> calculate? + # e2c2v: diamond vertices + # + C2E2C = "neighbor_cell_index" # dims(nv=3, cell) + V2E2V = "vertices_of_vertex" # dims(ne=6, vertex) not in simple mesh, = v2c2v vertices in hexa/pentagon + V2E = "edges_of_vertex" # dims(ne=6, vertex) + V2C = "cells_of_vertex" # dims(ne=6, vertex) + E2V = "edge_vertices" # dims(nc=2, edge) + C2V = "vertex_of_cell" # dims(nv=3, cell) # not in simple mesh + E2C = "adjacent_cell_of_edge" # dims(nc=2, edge) + C2E = "edge_of_cell" # dims(nv=3, cell) + + class DimensionName(GridFileName): VERTEX_NAME = "vertex" EDGE_NAME = "edge" CELL_NAME = "cell" - V2E_SIZE = "ne" # 6 DIAMOND_EDGE_SIZE = "no" # 4 - NEIGHBORING_EDGES_TO_CELL_SIZE = "nv" # 3 - E2V_SIZE = "nc" # 2 - + NEIGHBORS_TO_VERTEX_SIZE = "ne" # 6 + NEIGHBORS_TO_CELL_SIZE = "nv" # 3 + NEIGHBORS_TO_EDGE_SIZE = "nc" # 2 MAX_CHILD_DOMAINS = "max_chdom" - # TODO: @magdalena what does the grf stand for 'grid_refinement'? + # TODO: @magdalena what does the grf abbrev. stand for 'grid_refinement'? CELL_GRF = "cell_grf" EDGE_GRF = "edge_grf" VERTEX_GRF = "vert_grf" - class GridRefinement(GridFileName): + class GridRefinementName(GridFileName): CONTROL_CELLS = "refin_c_ctrl" CONTROL_EDGES = "refin_e_ctrl" CONTROL_VERTICES = "refin_v_ctrl" START_INDEX_CELLS = "start_idx_c" START_INDEX_EDGES = "start_idx_e" START_INDEX_VERTICES = "start_idx_v" - END_INDEX_CELLS ="end_idx_c" + END_INDEX_CELLS = "end_idx_c" END_INDEX_EDGES = "end_idx_e" END_INDEX_VERTICES = "end_idx_v" @@ -99,12 +113,12 @@ def dimension(self, name: GridFileName) -> int: def int_field(self, name: GridFileName, transpose=True) -> np.ndarray: try: nc_variable = self._dataset.variables[name] - self._log.debug(f"{name}: {nc_variable}") - data = nc_variable[:] + self._log.debug(f"reading {name}: {nc_variable}") + data = nc_variable[:] data = np.array(data, dtype=np.int32) return np.transpose(data) if transpose else data - except KeyError as error: + except KeyError: msg = f"{name} does not exist in dataset" self._log.warning(msg) raise IconGridError(msg) @@ -113,21 +127,12 @@ def int_field(self, name: GridFileName, transpose=True) -> np.ndarray: class IconGridError(RuntimeError): pass -class IconDomainZone(str,Enum): - HALO = "halo" - LATERAL_BOUNDARY = "lb" - NUDGING = "nudging" - INTERIOR = "interior" # TODO @magdalena is this the same as "prognostic" if yes then change the naming - END = "end" - class IndexTransformation(ABC): def get_offset_for_index_field(self, array: np.ndarray): return np.zeros(array.shape) - - class ToGt4PyTransformation(IndexTransformation): def get_offset_for_index_field(self, array: np.ndarray): """ @@ -136,35 +141,29 @@ def get_offset_for_index_field(self, array: np.ndarray): Fortran indices are 1-based, hence the offset is -1 for 0-based ness of python except for INVALID values which are marked with -1 in the grid file and are kept such. """ - return np.where(array < 0, 0, -1) - - -class GridDomainBoundaryTransformation(ABC): - def __init__(self): - pass - -class SingleNodeBoundaryTransform(GridDomainBoundaryTransformation): - pass + return np.where(array == GridFile.INVALID_INDEX, 0, -1) class GridManager: - def __init__(self, transformation: IndexTransformation, grid_file: str): + def __init__( + self, + transformation: IndexTransformation, + grid_file: str, + config: VerticalGridConfig, + ): self._log = logging.getLogger(__name__) self._transformation = transformation + self._config = config self._grid: Optional[IconGrid] = None self._file_names = grid_file - self._start_indices = {} - self._end_indices = {} def init(self): dataset = self._read_gridfile(self._file_names) - start_indices, end_indices = self._read_boundaries(dataset) - self._start_indices.update(start_indices) - self._end_indices.update(end_indices) - _, grid= self._read_grid(dataset) + _, grid = self._read_grid(dataset) + self._grid = grid - def _read_gridfile(self, fname:str)->Dataset: + def _read_gridfile(self, fname: str) -> Dataset: try: dataset = Dataset(self._file_names, "r", format="NETCDF4") self._log.debug(dataset) @@ -173,57 +172,74 @@ def _read_gridfile(self, fname:str)->Dataset: self._log.error(f"gridfile {fname} not found, aborting") exit(1) - def _read_boundaries(self, dataset): + def _read_grid_refinement_information(self, dataset): reader = GridFile(dataset) - grf_vertices = reader.dimension(GridFile.Dimension.VERTEX_GRF) - grf_edges = reader.dimension(GridFile.Dimension.EDGE_GRF) - grf_cells = reader.dimension(GridFile.Dimension.CELL_GRF) - refin_c_ctl = reader.int_field(GridFile.GridRefinement.CONTROL_CELLS) - refin_v_ctl = reader.int_field(GridFile.GridRefinement.CONTROL_VERTICES) - refin_e_ctl = reader.int_field(GridFile.GridRefinement.CONTROL_EDGES) + + grf_vertices = reader.dimension(GridFile.DimensionName.VERTEX_GRF) + grf_edges = reader.dimension(GridFile.DimensionName.EDGE_GRF) + grf_cells = reader.dimension(GridFile.DimensionName.CELL_GRF) + refin_c_ctl = reader.int_field(GridFile.GridRefinementName.CONTROL_CELLS) + refin_v_ctl = reader.int_field(GridFile.GridRefinementName.CONTROL_VERTICES) + refin_e_ctl = reader.int_field(GridFile.GridRefinementName.CONTROL_EDGES) start_indices = {} end_indices = {} - start_indices[CellDim] = self._get_index_field(reader, GridFile.GridRefinement.START_INDEX_CELLS, transpose=False) - end_indices[CellDim] = self._get_index_field(reader, GridFile.GridRefinement.END_INDEX_CELLS, transpose=False, apply_offset=False) - start_indices[EdgeDim] = self._get_index_field(reader, GridFile.GridRefinement.START_INDEX_EDGES, transpose=False) - end_indices[EdgeDim] = self._get_index_field(reader, GridFile.GridRefinement.END_INDEX_EDGES, transpose=False, apply_offset=False) - start_indices[VertexDim] = self._get_index_field(reader, GridFile.GridRefinement.START_INDEX_VERTICES, transpose=False) - end_indices[VertexDim] = self._get_index_field(reader, GridFile.GridRefinement.END_INDEX_VERTICES, transpose=False, apply_offset=False) + start_indices[CellDim] = self._get_index_field( + reader, GridFile.GridRefinementName.START_INDEX_CELLS, transpose=False + ) + end_indices[CellDim] = self._get_index_field( + reader, + GridFile.GridRefinementName.END_INDEX_CELLS, + transpose=False, + apply_offset=False, + ) + start_indices[EdgeDim] = self._get_index_field( + reader, GridFile.GridRefinementName.START_INDEX_EDGES, transpose=False + ) + end_indices[EdgeDim] = self._get_index_field( + reader, + GridFile.GridRefinementName.END_INDEX_EDGES, + transpose=False, + apply_offset=False, + ) + start_indices[VertexDim] = self._get_index_field( + reader, GridFile.GridRefinementName.START_INDEX_VERTICES, transpose=False + ) + end_indices[VertexDim] = self._get_index_field( + reader, + GridFile.GridRefinementName.END_INDEX_VERTICES, + transpose=False, + apply_offset=False, + ) return start_indices, end_indices + # TODO @magdalena make HorizontalMarkerIndex a type that behaves and is compatible with an int + + def get_start_index(self, dim: Dimension, start_marker: int): + return self._get_index(dim, start_marker, self._grid.start_indices) + + def get_end_index(self, dim: Dimension, start_marker: int): + return self._get_index(dim, start_marker, self._grid.end_indices) + def get_grid(self): + return self._grid - def get_domain_boundaries(self, dim: Dimension, zone: IconDomainZone) -> tuple[int, int]: - match zone: - case IconDomainZone.LATERAL_BOUNDARY: - index = HorizontalMarkerIndex.lateral_boundary(dim) - return self._get_start_end_indices(dim, index) - case IconDomainZone.INTERIOR: - index = HorizontalMarkerIndex.interior(dim) - return self._get_start_end_indices(dim, index) - case IconDomainZone.NUDGING: - index = HorizontalMarkerIndex.nudging(dim) - return self._get_start_end_indices(dim, index) - case IconDomainZone.HALO: - index = HorizontalMarkerIndex.local(dim) - return self._get_start_end_indices(dim, index) - case IconDomainZone.END: - index = HorizontalMarkerIndex.end(dim) - return self._get_start_end_indices(dim, index) - - def _get_start_end_indices(self, dim, index): + def _get_index(self, dim: Dimension, start_marker: int, index_dict): + _CHILD_DOM = 0 + if dim.kind != DimensionKind.HORIZONTAL: + msg = f"getting start index in horizontal domain with non - horizontal dimension {dim}" + self._log.warning(msg) + raise IconGridError(msg) try: - return self._start_indices[dim][0][index], self._end_indices[dim][0][index] + return index_dict[dim][_CHILD_DOM][start_marker] except KeyError as error: msg = f"start, end indices for dimension {dim} not present" self._log.error(msg) raise IconGridError(msg) - def _read_grid(self, dataset:Dataset) -> tuple[UUID, IconGrid]: - - grid_id = UUID(dataset.getncattr(GridFile.Property.GRID_ID)) - return grid_id, self.from_grid_dataset(dataset) + def _read_grid(self, dataset: Dataset) -> tuple[UUID, IconGrid]: + grid_id = UUID(dataset.getncattr(GridFile.PropertyName.GRID_ID)) + return grid_id, self.from_grid_dataset(dataset) def get_c2e_connectivity(self): return self._grid.get_c2e_connectivity() @@ -255,7 +271,6 @@ def get_e2ecv_connectivity(self): def get_e2c2e_connectivity(self): return self._grid.get_e2c2e_connectivity() - def get_size(self, dim: Dimension): if dim == VertexDim: return self._grid.config.num_vertices @@ -267,36 +282,38 @@ def get_size(self, dim: Dimension): self._log.warning(f"cannot determine size of unknown dimension {dim}") raise IconGridError(f"Unknown dimension {dim}") - - - def _get_index_field(self, reader, name: GridFileName, transpose=True, apply_offset = True): - field = reader.int_field(name, transpose=transpose) + def _get_index_field( + self, reader, field: GridFileName, transpose=True, apply_offset=True + ): + field = reader.int_field(field, transpose=transpose) if apply_offset: field = field + self._transformation.get_offset_for_index_field(field) return field def from_grid_dataset(self, dataset: Dataset) -> IconGrid: reader = GridFile(dataset) - num_cells = reader.dimension(GridFile.Dimension.CELL_NAME) - num_edges = reader.dimension(GridFile.Dimension.EDGE_NAME) - num_vertices = reader.dimension(GridFile.Dimension.VERTEX_NAME) - + num_cells = reader.dimension(GridFile.DimensionName.CELL_NAME) + num_edges = reader.dimension(GridFile.DimensionName.EDGE_NAME) + num_vertices = reader.dimension(GridFile.DimensionName.VERTEX_NAME) - grid_size = HorizontalMeshSize( + grid_size = HorizontalGridSize( num_vertices=num_vertices, num_edges=num_edges, num_cells=num_cells ) - c2e = self._get_index_field(reader, GridFile.Offsets.C2E) - c2v = self._get_index_field(reader, GridFile.Offsets.C2V) - - e2c = self._get_index_field(reader, GridFile.Offsets.E2C) - - e2v = self._get_index_field(reader, GridFile.Offsets.E2V) - v2c = self._get_index_field(reader, GridFile.Offsets.V2C) - v2e = self._get_index_field(reader, GridFile.Offsets.V2E) - v2e2v = self._get_index_field(reader, GridFile.Offsets.V2E2V) - c2e2c = self._get_index_field(reader, GridFile.Offsets.C2E2C) - config = MeshConfig( - horizontal_config=grid_size, vertical_config=VerticalMeshConfig(0) + c2e = self._get_index_field(reader, GridFile.OffsetName.C2E) + c2v = self._get_index_field(reader, GridFile.OffsetName.C2V) + + e2c = self._get_index_field(reader, GridFile.OffsetName.E2C) + + e2v = self._get_index_field(reader, GridFile.OffsetName.E2V) + v2c = self._get_index_field(reader, GridFile.OffsetName.V2C) + v2e = self._get_index_field(reader, GridFile.OffsetName.V2E) + v2e2v = self._get_index_field(reader, GridFile.OffsetName.V2E2V) + c2e2c = self._get_index_field(reader, GridFile.OffsetName.C2E2C) + start_indices, end_indices = self._read_grid_refinement_information(dataset) + + config = GridConfig( + horizontal_config=grid_size, + vertical_config=self._config, ) icon_grid = ( IconGrid() @@ -312,7 +329,15 @@ def from_grid_dataset(self, dataset: Dataset) -> IconGrid: C2VDim: c2v, } ) + .with_start_end_indices( + CellDim, start_indices[CellDim], end_indices[CellDim] + ) + .with_start_end_indices( + EdgeDim, start_indices[EdgeDim], end_indices[EdgeDim] + ) + .with_start_end_indices( + VertexDim, start_indices[VertexDim], end_indices[VertexDim] + ) ) - return icon_grid - + return icon_grid diff --git a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py b/common/src/icon4py/grid/horizontal.py similarity index 79% rename from atm_dyn_iconam/src/icon4py/diffusion/horizontal.py rename to common/src/icon4py/grid/horizontal.py index f7a46ccbf..42e4d4e38 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py +++ b/common/src/icon4py/grid/horizontal.py @@ -10,6 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +from abc import ABC from dataclasses import dataclass from typing import Final @@ -19,7 +20,29 @@ from icon4py.common.dimension import CellDim, ECVDim, EdgeDim +# W605 + + class HorizontalMarkerIndex: + """ + Handles constants indexing into the start_index and end_index fields. + + ICON uses a double indexing scheme for field indices marking the start and end of special + grid zone: The constants defined here (from mo_impl_constants.f90 and mo_impl_constants_grf.f90) + are the indices that are used to index into the start_idx and end_idx arrays + provided by the grid file where for each dimension the start index of the horizontal + "zones" are defined: + f.ex. an inlined access of the field F: Field[[CellDim], double] at the starting point of the lateral boundary zone would be + + F[start_idx_c[_LATERAL_BOUNDARY_CELLS] + + + ICON uses a custom index range from [ICON_INDEX_OFFSET... ] such that the index 0 marks the + internal entities for _all_ dimensions (Cell, Edge, Vertex) that is why we define these + additional INDEX_OFFSETs here in order to swap back to a 0 base python array. + + """ + NUM_GHOST_ROWS: Final[int] = 2 # values from mo_impl_constants.f90 _ICON_INDEX_OFFSET_CELLS: Final[int] = 8 @@ -59,7 +82,8 @@ class HorizontalMarkerIndex: @classmethod def lateral_boundary(cls, dim: Dimension) -> int: - match (dim): + """TODO @magdalena""" + match dim: case (dimension.CellDim): return cls._LATERAL_BOUNDARY_CELLS case (dimension.EdgeDim): @@ -108,8 +132,38 @@ def end(cls, dim: Dimension) -> int: return cls._END_VERTICES +class IconHorizontalDomainZone(ABC): + def __init__(self, dim: Dimension): + self._marker = HorizontalMarkerIndex.nudging(dim) + + def __call__(self, *args, **kwargs): + return self._marker + + def __add__(self, other: int): + return self._marker + other + + +class Nudging(IconHorizontalDomainZone): + def __init__(self, dim: Dimension): + super().__init__(dim) + + +class LateralBoundary(IconHorizontalDomainZone): + def __init__(self, dim: Dimension): + super().__init__(dim) + + +class Interior(IconHorizontalDomainZone): + def __init__(self, dim: Dimension): + super().__init__(dim) + + +def nudging(dim: CellDim, offset=0): + return HorizontalMarkerIndex.nudging(dim) + offset + + @dataclass(frozen=True) -class HorizontalMeshSize: +class HorizontalGridSize: num_vertices: int num_edges: int num_cells: int diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/common/src/icon4py/grid/icon_grid.py similarity index 77% rename from atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py rename to common/src/icon4py/grid/icon_grid.py index af2cead56..9494fde69 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/common/src/icon4py/grid/icon_grid.py @@ -10,32 +10,22 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - from typing import Dict, Tuple import numpy as np -from gt4py.next.common import Dimension, DimensionKind, Field -from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.common import Dimension, DimensionKind from gt4py.next.iterator.embedded import NeighborTableOffsetProvider from icon4py.common.dimension import CellDim, ECVDim, EdgeDim, KDim, VertexDim -from icon4py.diffusion.horizontal import HorizontalMeshSize - - -class VerticalMeshConfig: - def __init__(self, num_lev: int): - self._num_lev = num_lev - - @property - def num_lev(self) -> int: - return self._num_lev +from icon4py.grid.horizontal import HorizontalGridSize +from icon4py.grid.vertical import VerticalGridConfig -class MeshConfig: +class GridConfig: def __init__( self, - horizontal_config: HorizontalMeshSize, - vertical_config: VerticalMeshConfig, + horizontal_config: HorizontalGridSize, + vertical_config: VerticalGridConfig, limited_area=True, ): self._vertical = vertical_config @@ -78,20 +68,21 @@ def wrapper(self, *args, **kwargs): class IconGrid: def __init__(self): - self.config: MeshConfig = None + self.config: GridConfig = None + self.start_indices = {} self.end_indices = {} self.connectivities: Dict[str, np.ndarray] = {} self.size: Dict[Dimension, int] = {} - def _update_size(self, config: MeshConfig): + def _update_size(self, config: GridConfig): self.size[VertexDim] = config.num_vertices self.size[CellDim] = config.num_cells self.size[EdgeDim] = config.num_edges self.size[KDim] = config.num_k_levels @builder - def with_config(self, config: MeshConfig): + def with_config(self, config: GridConfig): self.config = config self._update_size(config) @@ -181,33 +172,3 @@ def get_e2ecv_connectivity(self): return NeighborTableOffsetProvider( v2ecv_table, EdgeDim, ECVDim, v2ecv_table.shape[1] ) - - -class VerticalModelParams: - def __init__(self, vct_a: Field[[KDim], float], rayleigh_damping_height: float): - """ - Contains vertical physical parameters defined on the grid. - - Args: - vct_a: field containing the physical heights of the k level - rayleigh_damping_height: height of rayleigh damping in [m] mo_nonhydro_nml - """ - self._rayleigh_damping_height = rayleigh_damping_height - self._vct_a = vct_a - self._index_of_damping_height = int32( - np.argmax( - np.where(np.asarray(self._vct_a) >= self._rayleigh_damping_height) - ) - ) - - @property - def index_of_damping_layer(self): - return self._index_of_damping_height - - @property - def physical_heights(self) -> Field[[KDim], float]: - return self._vct_a - - @property - def rayleigh_damping_height(self): - return self._rayleigh_damping_height diff --git a/common/src/icon4py/grid/vertical.py b/common/src/icon4py/grid/vertical.py new file mode 100644 index 000000000..6c7513696 --- /dev/null +++ b/common/src/icon4py/grid/vertical.py @@ -0,0 +1,56 @@ +# 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.ffront.fbuiltins import Field, int32 + +from icon4py.common.dimension import KDim + + +class VerticalGridConfig: + def __init__(self, num_lev: int): + self._num_lev = num_lev + + @property + def num_lev(self) -> int: + return self._num_lev + + +class VerticalModelParams: + def __init__(self, vct_a: Field[[KDim], float], rayleigh_damping_height: float): + """ + Contains vertical physical parameters defined on the grid. + + Args: + vct_a: field containing the physical heights of the k level + rayleigh_damping_height: height of rayleigh damping in [m] mo_nonhydro_nml + """ + self._rayleigh_damping_height = rayleigh_damping_height + self._vct_a = vct_a + self._index_of_damping_height = int32( + np.argmax( + np.where(np.asarray(self._vct_a) >= self._rayleigh_damping_height) + ) + ) + + @property + def index_of_damping_layer(self) -> int32: + return self._index_of_damping_height + + @property + def physical_heights(self) -> Field[[KDim], float]: + return self._vct_a + + @property + def rayleigh_damping_height(self) -> float: + return self._rayleigh_damping_height diff --git a/common/tests/conftest.py b/common/tests/conftest.py index d79fbd985..6d4e393ba 100644 --- a/common/tests/conftest.py +++ b/common/tests/conftest.py @@ -14,6 +14,8 @@ from icon4py.testutils.fixtures import ( # noqa F401 data_provider, + get_grid_files, grid_savepoint, + r04b09_dsl_gridfile, setup_icon_data, ) diff --git a/common/tests/test_gridmanager.py b/common/tests/test_gridmanager.py index e3c33cd10..2dca01494 100644 --- a/common/tests/test_gridmanager.py +++ b/common/tests/test_gridmanager.py @@ -25,32 +25,14 @@ GridFileName, GridManager, IndexTransformation, - ToGt4PyTransformation, IconDomainZone, + ToGt4PyTransformation, ) -from icon4py.testutils.data_handling import download_and_extract +from icon4py.grid.horizontal import HorizontalMarkerIndex from icon4py.testutils.simple_mesh import SimpleMesh SIMPLE_MESH_NC = "./simple_mesh_grid.nc" -mch_ch_r04b09_dsl_grid_uri = "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download" -r02b04_global_grid_uri = "https://polybox.ethz.ch/index.php/s/0EM8O8U53GKGsst/download" -grids_path = Path(__file__).parent.joinpath("grids") -r04b09_dsl_grid_path = grids_path.joinpath("mch_ch_r04b09_dsl") -r04b09_dsl_data_file = r04b09_dsl_grid_path.joinpath("mch_ch_r04b09_dsl_grids_v1.tar.gz").name -r02b04_global_grid_path = grids_path.joinpath("icon_r02b04_global") -r02b04_global_data_file = r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_G.tar.gz").name - -@pytest.fixture(scope="session") -def get_grid_files(): - """ - Get the grid files used for testing. - - Session scoped fixture which is a prerequisite of all the other fixtures in this file. - """ - download_and_extract(mch_ch_r04b09_dsl_grid_uri, r04b09_dsl_grid_path, r04b09_dsl_data_file) - download_and_extract(r02b04_global_grid_uri, r02b04_global_grid_path, r02b04_global_data_file) - @pytest.fixture def simple_mesh_path(): @@ -61,99 +43,160 @@ def simple_mesh_path(): def simple_mesh_data(): mesh = SimpleMesh() dataset = Dataset(SIMPLE_MESH_NC, "w", format="NETCDF4") - dataset.setncattr(GridFile.Property.GRID_ID, str(uuid4())) - dataset.createDimension(GridFile.Dimension.VERTEX_NAME, size=mesh.n_vertices) - dataset.createDimension(GridFile.Dimension.EDGE_NAME, size=mesh.n_edges) - dataset.createDimension(GridFile.Dimension.CELL_NAME, size=mesh.n_cells) - dataset.createDimension(GridFile.Dimension.E2V_SIZE, size=mesh.n_e2v) - dataset.createDimension(GridFile.Dimension.DIAMOND_EDGE_SIZE, size=mesh.n_e2c2e) - dataset.createDimension(GridFile.Dimension.MAX_CHILD_DOMAINS, size=1) + dataset.setncattr(GridFile.PropertyName.GRID_ID, str(uuid4())) + dataset.createDimension(GridFile.DimensionName.VERTEX_NAME, size=mesh.n_vertices) + dataset.createDimension(GridFile.DimensionName.EDGE_NAME, size=mesh.n_edges) + dataset.createDimension(GridFile.DimensionName.CELL_NAME, size=mesh.n_cells) + dataset.createDimension( + GridFile.DimensionName.NEIGHBORS_TO_EDGE_SIZE, size=mesh.n_e2v + ) + dataset.createDimension(GridFile.DimensionName.DIAMOND_EDGE_SIZE, size=mesh.n_e2c2e) + dataset.createDimension(GridFile.DimensionName.MAX_CHILD_DOMAINS, size=1) # add dummy values for the grf dimensions # TODO @magdalena fix to something more useful? - dataset.createDimension(GridFile.Dimension.CELL_GRF, size=14) - dataset.createDimension(GridFile.Dimension.EDGE_GRF, size=24) - dataset.createDimension(GridFile.Dimension.VERTEX_GRF, size=13) - add_to_dataset(dataset, np.zeros(mesh.n_edges), GridFile.GridRefinement.CONTROL_EDGES, (GridFile.Dimension.EDGE_NAME,)) - add_to_dataset(dataset, np.zeros(mesh.n_cells), GridFile.GridRefinement.CONTROL_CELLS, (GridFile.Dimension.CELL_NAME,)) - add_to_dataset(dataset, np.zeros(mesh.n_vertices), GridFile.GridRefinement.CONTROL_VERTICES,(GridFile.Dimension.VERTEX_NAME,)) + dataset.createDimension(GridFile.DimensionName.CELL_GRF, size=14) + dataset.createDimension(GridFile.DimensionName.EDGE_GRF, size=24) + dataset.createDimension(GridFile.DimensionName.VERTEX_GRF, size=13) + _add_to_dataset( + dataset, + np.zeros(mesh.n_edges), + GridFile.GridRefinementName.CONTROL_EDGES, + (GridFile.DimensionName.EDGE_NAME,), + ) + _add_to_dataset( + dataset, + np.zeros(mesh.n_cells), + GridFile.GridRefinementName.CONTROL_CELLS, + (GridFile.DimensionName.CELL_NAME,), + ) + _add_to_dataset( + dataset, + np.zeros(mesh.n_vertices), + GridFile.GridRefinementName.CONTROL_VERTICES, + (GridFile.DimensionName.VERTEX_NAME,), + ) dataset.createDimension( - GridFile.Dimension.NEIGHBORING_EDGES_TO_CELL_SIZE, size=mesh.n_c2e + GridFile.DimensionName.NEIGHBORS_TO_CELL_SIZE, size=mesh.n_c2e + ) + dataset.createDimension( + GridFile.DimensionName.NEIGHBORS_TO_VERTEX_SIZE, size=mesh.n_v2c ) - dataset.createDimension(GridFile.Dimension.V2E_SIZE, size=mesh.n_v2c) - add_to_dataset( + _add_to_dataset( dataset, mesh.c2e, - GridFile.Offsets.C2E, + GridFile.OffsetName.C2E, ( - GridFile.Dimension.NEIGHBORING_EDGES_TO_CELL_SIZE, - GridFile.Dimension.CELL_NAME, + GridFile.DimensionName.NEIGHBORS_TO_CELL_SIZE, + GridFile.DimensionName.CELL_NAME, ), ) # add_to_dataset(data, mesh.c2v, GridFile.Offsets.C2V, (GridFile.Dimension.C2E_SIZE, GridFile.Dimension.CELL_NAME)) - add_to_dataset( + _add_to_dataset( dataset, mesh.e2c, - GridFile.Offsets.E2C, - (GridFile.Dimension.E2V_SIZE, GridFile.Dimension.EDGE_NAME), + GridFile.OffsetName.E2C, + ( + GridFile.DimensionName.NEIGHBORS_TO_EDGE_SIZE, + GridFile.DimensionName.EDGE_NAME, + ), ) - add_to_dataset( + _add_to_dataset( dataset, mesh.e2v, - GridFile.Offsets.E2V, - (GridFile.Dimension.E2V_SIZE, GridFile.Dimension.EDGE_NAME), + GridFile.OffsetName.E2V, + ( + GridFile.DimensionName.NEIGHBORS_TO_EDGE_SIZE, + GridFile.DimensionName.EDGE_NAME, + ), ) - add_to_dataset( + _add_to_dataset( dataset, mesh.v2c, - GridFile.Offsets.V2C, - (GridFile.Dimension.V2E_SIZE, GridFile.Dimension.VERTEX_NAME), + GridFile.OffsetName.V2C, + ( + GridFile.DimensionName.NEIGHBORS_TO_VERTEX_SIZE, + GridFile.DimensionName.VERTEX_NAME, + ), ) # TODO @magdalena: there is no v2c in the simple mesh - add_to_dataset( + _add_to_dataset( dataset, np.zeros((mesh.n_cells, 3), dtype=np.int32), - GridFile.Offsets.C2V, + GridFile.OffsetName.C2V, ( - GridFile.Dimension.NEIGHBORING_EDGES_TO_CELL_SIZE, - GridFile.Dimension.CELL_NAME, + GridFile.DimensionName.NEIGHBORS_TO_CELL_SIZE, + GridFile.DimensionName.CELL_NAME, ), ) - add_to_dataset( + _add_to_dataset( dataset, np.zeros((mesh.n_vertices, 4), dtype=np.int32), - GridFile.Offsets.V2E2V, - (GridFile.Dimension.DIAMOND_EDGE_SIZE, GridFile.Dimension.VERTEX_NAME), + GridFile.OffsetName.V2E2V, + (GridFile.DimensionName.DIAMOND_EDGE_SIZE, GridFile.DimensionName.VERTEX_NAME), ) - add_to_dataset( + _add_to_dataset( dataset, mesh.v2e, - GridFile.Offsets.V2E, - (GridFile.Dimension.V2E_SIZE, GridFile.Dimension.VERTEX_NAME), + GridFile.OffsetName.V2E, + ( + GridFile.DimensionName.NEIGHBORS_TO_VERTEX_SIZE, + GridFile.DimensionName.VERTEX_NAME, + ), ) # add_to_dataset(data, mesh.v2e2v, GridFile.Offsets.V2E2V, (GridFile.Dimension.V2E_SIZE, GridFile.Dimension.VERTEX_NAME)) - add_to_dataset( + _add_to_dataset( dataset, mesh.c2e2c, - GridFile.Offsets.C2E2C, + GridFile.OffsetName.C2E2C, ( - GridFile.Dimension.NEIGHBORING_EDGES_TO_CELL_SIZE, - GridFile.Dimension.CELL_NAME, + GridFile.DimensionName.NEIGHBORS_TO_CELL_SIZE, + GridFile.DimensionName.CELL_NAME, ), ) - add_to_dataset(dataset, np.ones((1, 24), dtype=np.int32), GridFile.GridRefinement.START_INDEX_EDGES, (GridFile.Dimension.MAX_CHILD_DOMAINS, GridFile.Dimension.EDGE_GRF)) - add_to_dataset(dataset, np.ones((1, 14), dtype=np.int32), GridFile.GridRefinement.START_INDEX_CELLS, (GridFile.Dimension.MAX_CHILD_DOMAINS, GridFile.Dimension.CELL_GRF)) - add_to_dataset(dataset, np.ones((1, 13), dtype=np.int32), GridFile.GridRefinement.START_INDEX_VERTICES, (GridFile.Dimension.MAX_CHILD_DOMAINS, GridFile.Dimension.VERTEX_GRF)) - add_to_dataset(dataset, np.ones((1, 24), dtype=np.int32), GridFile.GridRefinement.END_INDEX_EDGES, (GridFile.Dimension.MAX_CHILD_DOMAINS, GridFile.Dimension.EDGE_GRF)) - add_to_dataset(dataset, np.ones((1, 14), dtype=np.int32), GridFile.GridRefinement.END_INDEX_CELLS, (GridFile.Dimension.MAX_CHILD_DOMAINS, GridFile.Dimension.CELL_GRF)) - add_to_dataset(dataset, np.ones((1, 13), dtype=np.int32), GridFile.GridRefinement.END_INDEX_VERTICES,(GridFile.Dimension.MAX_CHILD_DOMAINS, GridFile.Dimension.VERTEX_GRF)) + _add_to_dataset( + dataset, + np.ones((1, 24), dtype=np.int32), + GridFile.GridRefinementName.START_INDEX_EDGES, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.EDGE_GRF), + ) + _add_to_dataset( + dataset, + np.ones((1, 14), dtype=np.int32), + GridFile.GridRefinementName.START_INDEX_CELLS, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.CELL_GRF), + ) + _add_to_dataset( + dataset, + np.ones((1, 13), dtype=np.int32), + GridFile.GridRefinementName.START_INDEX_VERTICES, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.VERTEX_GRF), + ) + _add_to_dataset( + dataset, + np.ones((1, 24), dtype=np.int32), + GridFile.GridRefinementName.END_INDEX_EDGES, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.EDGE_GRF), + ) + _add_to_dataset( + dataset, + np.ones((1, 14), dtype=np.int32), + GridFile.GridRefinementName.END_INDEX_CELLS, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.CELL_GRF), + ) + _add_to_dataset( + dataset, + np.ones((1, 13), dtype=np.int32), + GridFile.GridRefinementName.END_INDEX_VERTICES, + (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.VERTEX_GRF), + ) dataset.close() -def add_to_dataset( +def _add_to_dataset( dataset: Dataset, data: np.ndarray, var_name: str, @@ -168,21 +211,25 @@ def test_gridparser_dimension(simple_mesh_data): data = Dataset(SIMPLE_MESH_NC, "r") grid_parser = GridFile(data) mesh = SimpleMesh() - assert grid_parser.dimension(GridFile.Dimension.CELL_NAME) == mesh.n_cells - assert grid_parser.dimension(GridFile.Dimension.VERTEX_NAME) == mesh.n_vertices - assert grid_parser.dimension(GridFile.Dimension.EDGE_NAME) == mesh.n_edges + assert grid_parser.dimension(GridFile.DimensionName.CELL_NAME) == mesh.n_cells + assert grid_parser.dimension(GridFile.DimensionName.VERTEX_NAME) == mesh.n_vertices + assert grid_parser.dimension(GridFile.DimensionName.EDGE_NAME) == mesh.n_edges -def test_gridfile_vertex_cell_edge_dimensions(grid_savepoint): - fname = r04b09_dsl_grid_path.joinpath("grid.nc") - data = Dataset(fname, "r") + +def test_gridfile_vertex_cell_edge_dimensions(grid_savepoint, r04b09_dsl_gridfile): + data = Dataset(r04b09_dsl_gridfile, "r") grid_file = GridFile(data) - assert grid_file.dimension(GridFile.Dimension.CELL_NAME) == grid_savepoint.num(CellDim) - assert grid_file.dimension(GridFile.Dimension.EDGE_NAME) == grid_savepoint.num(EdgeDim) + assert grid_file.dimension(GridFile.DimensionName.CELL_NAME) == grid_savepoint.num( + CellDim + ) + assert grid_file.dimension(GridFile.DimensionName.EDGE_NAME) == grid_savepoint.num( + EdgeDim + ) # TODO: @magdalena fix in serialized data. it returns the num_cells - assert grid_file.dimension(GridFile.Dimension.VERTEX_NAME) == grid_savepoint.num(VertexDim) - - + assert grid_file.dimension( + GridFile.DimensionName.VERTEX_NAME + ) == grid_savepoint.num(VertexDim) def test_grid_parser_index_fields(simple_mesh_data, caplog): @@ -191,24 +238,22 @@ def test_grid_parser_index_fields(simple_mesh_data, caplog): mesh = SimpleMesh() grid_parser = GridFile(data) - assert np.allclose(grid_parser.int_field(GridFile.Offsets.C2E), mesh.c2e) - assert np.allclose(grid_parser.int_field(GridFile.Offsets.E2C), mesh.e2c) - assert np.allclose(grid_parser.int_field(GridFile.Offsets.V2E), mesh.v2e) - assert np.allclose(grid_parser.int_field(GridFile.Offsets.V2C), mesh.v2c) - + assert np.allclose(grid_parser.int_field(GridFile.OffsetName.C2E), mesh.c2e) + assert np.allclose(grid_parser.int_field(GridFile.OffsetName.E2C), mesh.e2c) + assert np.allclose(grid_parser.int_field(GridFile.OffsetName.V2E), mesh.v2e) + assert np.allclose(grid_parser.int_field(GridFile.OffsetName.V2C), mesh.v2c) # TODO @magdalena add test cases for -#e2c2v - diamond: serial, simple -#c2v: grid, ??? -#v2e2v: grid,??? +# e2c2v - diamond: serial, simple +# c2v: grid, ??? +# v2e2v: grid,??? # v2e: serial, simple, grid @pytest.mark.datatest -def test_gridmanager_eval_v2e(caplog, grid_savepoint, get_grid_files): +def test_gridmanager_eval_v2e(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) - fname = r04b09_dsl_grid_path.joinpath("grid.nc") - gm = _init_grid_manager(fname) + gm = init_grid_manager(r04b09_dsl_gridfile) num_vertex = gm.get_size(VertexDim) seralized_v2e = grid_savepoint.v2e()[0:num_vertex, :] # there are vertices at the boundary of a local domain or at a pentagon point that have less than @@ -217,16 +262,14 @@ def test_gridmanager_eval_v2e(caplog, grid_savepoint, get_grid_files): assert not has_invalid_index(seralized_v2e) assert has_invalid_index(gm.get_v2e_connectivity().table) reset_invalid_index(seralized_v2e) - assert np.allclose( - gm.get_v2e_connectivity().table, seralized_v2e - ) + assert np.allclose(gm.get_v2e_connectivity().table, seralized_v2e) + -#v2c: serial, simple, grid +# v2c: serial, simple, grid @pytest.mark.datatest -def test_gridmanager_eval_v2c(caplog, grid_savepoint, get_grid_files): +def test_gridmanager_eval_v2c(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) - fname = r04b09_dsl_grid_path.joinpath("grid.nc") - gm = _init_grid_manager(fname) + gm = init_grid_manager(r04b09_dsl_gridfile) num_vertex = gm.get_size(VertexDim) serialized_v2c = grid_savepoint.v2c()[0:num_vertex, :] # there are vertices that have less than 6 neighboring cells: either pentagon points or @@ -237,14 +280,12 @@ def test_gridmanager_eval_v2c(caplog, grid_savepoint, get_grid_files): assert has_invalid_index(gm.get_v2c_connectivity().table) reset_invalid_index(serialized_v2c) - assert np.allclose( - gm.get_v2c_connectivity().table, serialized_v2c - ) + assert np.allclose(gm.get_v2c_connectivity().table, serialized_v2c) def reset_invalid_index(index_array: np.ndarray): """ - helper function to revert mo_model_domimp_patches.f90: move_dummies_to_end_idxblk. + Helper function to revert mo_model_domimp_patches.f90: move_dummies_to_end_idxblk. see: # ! Checks for the pentagon case and moves dummy cells to end. # ! The dummy entry is either set to 0 or duplicated from the last one @@ -264,15 +305,14 @@ def reset_invalid_index(index_array: np.ndarray): """ for i in range(0, index_array.shape[0]): uq, index = np.unique(index_array[i, :], return_index=True) - index_array[i, max(index) + 1:] = GridFile.INVALID_INDEX + index_array[i, max(index) + 1 :] = GridFile.INVALID_INDEX -#e2v: serial, simple, grid +# e2v: serial, simple, grid @pytest.mark.datatest -def test_gridmanager_eval_e2v(caplog, grid_savepoint, get_grid_files): +def test_gridmanager_eval_e2v(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) - fname = r04b09_dsl_grid_path.joinpath("grid.nc") - gm = _init_grid_manager(fname) + gm = init_grid_manager(r04b09_dsl_gridfile) num_edges = gm.get_size(EdgeDim) serialized_e2v = grid_savepoint.e2v()[0:num_edges, :] @@ -280,21 +320,18 @@ def test_gridmanager_eval_e2v(caplog, grid_savepoint, get_grid_files): # hence this connectivity has no "missing values" in the grid file assert not has_invalid_index(serialized_e2v) assert not has_invalid_index(gm.get_e2v_connectivity().table) - assert np.allclose( - gm.get_e2v_connectivity().table, serialized_e2v - ) + assert np.allclose(gm.get_e2v_connectivity().table, serialized_e2v) -def has_invalid_index(ar :np.ndarray): +def has_invalid_index(ar: np.ndarray): return np.any(np.where(ar == GridFile.INVALID_INDEX)) # e2c :serial, simple, grid @pytest.mark.datatest -def test_gridmanager_eval_e2c(caplog, grid_savepoint, get_grid_files): +def test_gridmanager_eval_e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) - fname = r04b09_dsl_grid_path.joinpath("grid.nc") - gm = _init_grid_manager(fname) + gm = init_grid_manager(r04b09_dsl_gridfile) num_edges = gm.get_size(EdgeDim) serialized_c2e = grid_savepoint.e2c()[0:num_edges, :] # there are edges at the boundary that have only one @@ -302,17 +339,14 @@ def test_gridmanager_eval_e2c(caplog, grid_savepoint, get_grid_files): # and here they do not get substituted in the ICON preprocessing assert has_invalid_index(serialized_c2e) assert has_invalid_index(gm.get_e2c_connectivity().table) - assert np.allclose( - gm.get_e2c_connectivity().table, serialized_c2e - ) + assert np.allclose(gm.get_e2c_connectivity().table, serialized_c2e) -#c2e: serial, simple, grid +# c2e: serial, simple, grid @pytest.mark.datatest -def test_gridmanager_eval_c2e(caplog, grid_savepoint, get_grid_files): +def test_gridmanager_eval_c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) - fname = r04b09_dsl_grid_path.joinpath("grid.nc") - gm = _init_grid_manager(fname) + gm = init_grid_manager(r04b09_dsl_gridfile) num_cells = gm.get_size(CellDim) serialized_c2e = grid_savepoint.c2e()[0:num_cells, :] @@ -321,38 +355,33 @@ def test_gridmanager_eval_c2e(caplog, grid_savepoint, get_grid_files): # hence there are no "missing values" in the grid file assert not has_invalid_index(serialized_c2e) assert not has_invalid_index(gm.get_c2e_connectivity().table) - assert np.allclose( - gm.get_c2e_connectivity().table, serialized_c2e - ) + assert np.allclose(gm.get_c2e_connectivity().table, serialized_c2e) -#e2c2e (e2c2eo) - diamond: serial, simple_mesh + +# e2c2e (e2c2eo) - diamond: serial, simple_mesh @pytest.mark.skip("does this array exist in grid file?") @pytest.mark.datatest -def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, get_grid_files): +def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) - fname = r04b09_dsl_grid_path.joinpath("grid.nc") - gm, num_cells, num_edges, num_vertex = _init_grid_manager(fname) + gm, num_cells, num_edges, num_vertex = init_grid_manager(r04b09_dsl_gridfile) serialized_e2c2e = grid_savepoint.e2c2e()[0:num_cells, :] assert has_invalid_index(serialized_e2c2e) assert has_invalid_index(gm.get_e2c2e_connectivity().table) - assert np.allclose( - gm.get_e2c2e_connectivity().table, serialized_e2c2e - ) + assert np.allclose(gm.get_e2c2e_connectivity().table, serialized_e2c2e) + -#c2e2c: serial, simple_mesh, grid +# c2e2c: serial, simple_mesh, grid @pytest.mark.datatest -def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, get_grid_files): +def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) - fname = r04b09_dsl_grid_path.joinpath("grid.nc") - gm = _init_grid_manager(fname) + gm = init_grid_manager(r04b09_dsl_gridfile) num_cells = gm.get_size(CellDim) assert np.allclose( gm.get_c2e2c_connectivity().table, grid_savepoint.c2e2c()[0:num_cells, :] ) - -def _init_grid_manager(fname): +def init_grid_manager(fname): gm = GridManager(ToGt4PyTransformation(), fname) gm.init() return gm @@ -383,38 +412,55 @@ def test_gt4py_transform_offset_by_1_where_valid(size): expected = np.where(input_field >= 0, -1, 0) assert np.allclose(expected, offset) -@pytest.mark.parametrize("zone, bounds", [ - (IconDomainZone.NUDGING,(3316, 4104)), - (IconDomainZone.HALO, (20896, 20896)), - (IconDomainZone.LATERAL_BOUNDARY, (0, 850)), - (IconDomainZone.INTERIOR, (4104, 20896)) -] ) -def test_start_end_indices_cells(get_grid_files, zone, bounds): - fname = r04b09_dsl_grid_path.joinpath("grid.nc") - gm = _init_grid_manager(fname) - assert gm.get_domain_boundaries(CellDim, zone) == bounds - - -@pytest.mark.parametrize("zone, bounds", [ - (IconDomainZone.NUDGING,(4989, 5387)), - (IconDomainZone.HALO, (31558, 31558)), - (IconDomainZone.LATERAL_BOUNDARY, (0, 428)), - (IconDomainZone.INTERIOR, (6176, 31558))]) -def test_start_end_indices_edges(get_grid_files, zone, bounds): - fname = r04b09_dsl_grid_path.joinpath("grid.nc") - gm = _init_grid_manager(fname) - assert gm.get_domain_boundaries(EdgeDim, zone) == bounds - - -@pytest.mark.parametrize("zone, bounds", [ - (IconDomainZone.NUDGING,(10663, 10663)), - (IconDomainZone.HALO, (10663, 10663)), - (IconDomainZone.LATERAL_BOUNDARY, (0, 428)), - (IconDomainZone.INTERIOR, (2071, 10663)), - (IconDomainZone.END, (10663, 10663)) -]) -def test_start_end_indices_vertices(get_grid_files, zone, bounds): - fname = r04b09_dsl_grid_path.joinpath("grid.nc") - gm = _init_grid_manager(fname) - assert gm.get_domain_boundaries(VertexDim, zone) == bounds +@pytest.mark.parametrize( + "dim, marker, index", + [ + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim), 0), + (VertexDim, HorizontalMarkerIndex.local(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 2071), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim), 0), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 850), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 1688), + (CellDim, HorizontalMarkerIndex.local(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.interior(CellDim), 4104), + (CellDim, HorizontalMarkerIndex.nudging(CellDim), 3316), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim), 0), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2538), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 2954), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim), 31558), + (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 6176), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 4989), + ], +) +def test_get_start_indices(r04b09_dsl_gridfile, dim, marker, index): + gm = init_grid_manager(r04b09_dsl_gridfile) + assert gm.get_start_index(dim, marker) == index + + +@pytest.mark.parametrize( + "dim, marker, index", + [ + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim), 428), + (VertexDim, HorizontalMarkerIndex.local(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim), 850), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 1688), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 2511), + (CellDim, HorizontalMarkerIndex.local(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.local(CellDim) + 1, 20896), + (CellDim, HorizontalMarkerIndex.interior(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.nudging(CellDim), 4104), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim), 428), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2954), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 3777), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim), 31558), + (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 31558), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 5387), + ], +) +def test_get_start_indices(r04b09_dsl_gridfile, dim, marker, index): + gm = init_grid_manager(r04b09_dsl_gridfile) + assert gm.get_end_index(dim, marker) == index diff --git a/common/tests/test_horizontal.py b/common/tests/test_horizontal.py new file mode 100644 index 000000000..379e646d3 --- /dev/null +++ b/common/tests/test_horizontal.py @@ -0,0 +1,37 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from icon4py.common.dimension import CellDim +from icon4py.grid.horizontal import Nudging, nudging + + +def test_nudging_class_for_cells(): + nudging = Nudging(CellDim) + nudging_plus_one = Nudging(CellDim) + 1 + assert nudging() == 13 + assert nudging_plus_one == 14 + + +def test_nudging_for_cells(): + nudge = nudging(CellDim) + nudge_p1 = nudging(CellDim, 1) + assert nudge == 13 + assert nudge_p1 == 14 + + +def get_start_index(i): + pass + + +def test_start_index(): + get_start_index(nudging(CellDim, 1)) diff --git a/atm_dyn_iconam/tests/test_vertical.py b/common/tests/test_vertical.py similarity index 81% rename from atm_dyn_iconam/tests/test_vertical.py rename to common/tests/test_vertical.py index 550854ca4..42a560dee 100644 --- a/atm_dyn_iconam/tests/test_vertical.py +++ b/common/tests/test_vertical.py @@ -16,7 +16,7 @@ import numpy as np import pytest -from icon4py.diffusion.icon_grid import VerticalModelParams +from icon4py.grid.vertical import VerticalGridConfig, VerticalModelParams from icon4py.testutils.fixtures import ( # noqa F401 data_provider, grid_savepoint, @@ -31,7 +31,11 @@ def test_nrdmax_calculation(max_h, damping, delta): vct_a = np.arange(0, max_h, delta) vct_a = vct_a[::-1] - vertical_params = VerticalModelParams(rayleigh_damping_height=damping, vct_a=vct_a) + vertical_params = VerticalModelParams( + config=VerticalGridConfig(num_lev=10), + rayleigh_damping_height=damping, + vct_a=vct_a, + ) assert ( vertical_params.index_of_damping_layer == vct_a.shape[0] - math.ceil(damping / delta) - 1 @@ -43,7 +47,7 @@ def test_nrdmax_calculation_from_icon_input(icon_grid, grid_savepoint): # noqa: a = grid_savepoint.vct_a() damping_height = 12500 vertical_params = VerticalModelParams( - rayleigh_damping_height=damping_height, vct_a=a + VerticalGridConfig(num_lev=10), rayleigh_damping_height=damping_height, vct_a=a ) assert 9 == vertical_params.index_of_damping_layer a_array = np.asarray(a) diff --git a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py index 09e1dfce8..49019077a 100644 --- a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py +++ b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py @@ -31,11 +31,11 @@ DiffusionConfig, DiffusionParams, ) -from icon4py.diffusion.horizontal import CellParams, EdgeParams -from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.prognostic_state import PrognosticState +from icon4py.grid.horizontal import CellParams, EdgeParams +from icon4py.grid.icon_grid import IconGrid, VerticalModelParams from icon4py.py2f.cffi_utils import CffiMethod diff --git a/testutils/src/icon4py/testutils/fixtures.py b/testutils/src/icon4py/testutils/fixtures.py index f0c0e816e..e8c2270c4 100644 --- a/testutils/src/icon4py/testutils/fixtures.py +++ b/testutils/src/icon4py/testutils/fixtures.py @@ -19,11 +19,24 @@ from icon4py.testutils.serialbox_utils import IconSerialDataProvider -#data_uri = "https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download" -data_uri= "https://polybox.ethz.ch/index.php/s/joglKLsbyIXOOp8/download" +# data_uri = "https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download" +data_uri = "https://polybox.ethz.ch/index.php/s/joglKLsbyIXOOp8/download" data_path = Path(__file__).parent.joinpath("ser_icondata") extracted_path = data_path.joinpath("mch_ch_r04b09_dsl/ser_data") data_file = data_path.joinpath("mch_ch_r04b09_dsl_v2.tar.gz").name +mch_ch_r04b09_dsl_grid_uri = ( + "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download" +) +r02b04_global_grid_uri = "https://polybox.ethz.ch/index.php/s/0EM8O8U53GKGsst/download" +grids_path = Path(__file__).parent.joinpath("grids") +r04b09_dsl_grid_path = grids_path.joinpath("mch_ch_r04b09_dsl") +r04b09_dsl_data_file = r04b09_dsl_grid_path.joinpath( + "mch_ch_r04b09_dsl_grids_v1.tar.gz" +).name +r02b04_global_grid_path = grids_path.joinpath("icon_r02b04_global") +r02b04_global_data_file = r02b04_global_grid_path.joinpath( + "icon_grid_0013_R02B04_G.tar.gz" +).name @pytest.fixture(scope="session") @@ -44,3 +57,23 @@ def data_provider(setup_icon_data) -> IconSerialDataProvider: @pytest.fixture def grid_savepoint(data_provider): return data_provider.from_savepoint_grid() + + +@pytest.fixture(scope="session") +def get_grid_files(): + """ + Get the grid files used for testing. + + Session scoped fixture which is a prerequisite of all the other fixtures in this file. + """ + download_and_extract( + mch_ch_r04b09_dsl_grid_uri, r04b09_dsl_grid_path, r04b09_dsl_data_file + ) + download_and_extract( + r02b04_global_grid_uri, r02b04_global_grid_path, r02b04_global_data_file + ) + + +@pytest.fixture() +def r04b09_dsl_gridfile(get_grid_files): + return r04b09_dsl_grid_path.joinpath("grid.nc") diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index c7a775390..e46b512c6 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -130,7 +130,7 @@ def edge_end_index(self): def print_connectivity_info(name: str, ar: np.ndarray): print(f" connectivity {name} {ar.shape}") - def refin_ctrl(self, dim: Dimension)-> Optional[np.ndarray]: + def refin_ctrl(self, dim: Dimension) -> Optional[np.ndarray]: if dim == CellDim: return self.serializer.read("c_refin_ctl", self.savepoint) elif dim == EdgeDim: @@ -140,13 +140,15 @@ def refin_ctrl(self, dim: Dimension)-> Optional[np.ndarray]: else: return None - def num(self, dim:Dimension)->Optional[int]: + def num(self, dim: Dimension) -> Optional[int]: if dim == CellDim: return self.serializer.read("num_cells", self.savepoint)[0] elif dim == EdgeDim: return self.serializer.read("num_edges", self.savepoint)[0] elif dim == VertexDim: return self.serializer.read("num_vert", self.savepoint)[0] + elif dim == KDim: + return self.serializer.read("nlev", self.savepoint)[0] else: return None From 50464914d1e5117c6d82291e2074dcef89a34c3c Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 12 Apr 2023 09:17:13 +0200 Subject: [PATCH 090/263] use grid file everywhere in test_diffusion --- atm_dyn_iconam/tests/test_diffusion.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 99909850a..7afbeeb53 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -228,7 +228,8 @@ def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor))) -def init_grid_manager(fname, config): +def init_grid_manager(fname, nlev): + config = VerticalGridConfig(nlev) gm = GridManager(ToGt4PyTransformation(), fname, config) gm.init() return gm @@ -244,8 +245,8 @@ def test_diffusion_init( damping_height, ): config = r04b09_diffusion_config - grid_config = VerticalGridConfig(grid_savepoint.num(KDim)) - gm = init_grid_manager(r04b09_dsl_gridfile, grid_config) + + gm = init_grid_manager(r04b09_dsl_gridfile, grid_savepoint.num(KDim)) icon_grid = gm.get_grid() additional_parameters = DiffusionParams(config) vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) @@ -348,11 +349,13 @@ def _verify_init_values_against_savepoint( @pytest.mark.datatest @pytest.mark.parametrize("linit", [True]) def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( - diffusion_savepoint_init, r04b09_diffusion_config, icon_grid + diffusion_savepoint_init, r04b09_diffusion_config, r04b09_dsl_gridfile, grid_savepoint ): savepoint = diffusion_savepoint_init config = r04b09_diffusion_config + icon_grid = init_grid_manager(r04b09_dsl_gridfile, grid_savepoint.num(KDim)).get_grid() + params = DiffusionParams(config) expected_diff_multfac_vn = savepoint.diff_multfac_vn() expected_smag_limit = savepoint.smag_limit() @@ -377,7 +380,7 @@ def test_verify_diffusion_init_against_first_regular_savepoint( diffusion_savepoint_init, grid_savepoint, # noqa: F811 r04b09_diffusion_config, - icon_grid, + r04b09_dsl_gridfile, damping_height, ): config = r04b09_diffusion_config @@ -405,6 +408,7 @@ def test_verify_diffusion_init_against_first_regular_savepoint( zd_vertidx=savepoint.zd_vertoffset(), zd_diffcoef=savepoint.zd_diffcoef(), ) + icon_grid = init_grid_manager(r04b09_dsl_gridfile, grid_savepoint.num(KDim)).get_grid() diffusion = Diffusion() diffusion.init( config=config, @@ -423,7 +427,7 @@ def test_verify_diffusion_init_against_first_regular_savepoint( def test_verify_diffusion_init_against_other_regular_savepoint( r04b09_diffusion_config, grid_savepoint, # noqa: F811 - icon_grid, + r04b09_dsl_gridfile, diffusion_savepoint_init, damping_height, ): @@ -431,6 +435,7 @@ def test_verify_diffusion_init_against_other_regular_savepoint( additional_parameters = DiffusionParams(config) vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) + icon_grid = init_grid_manager(r04b09_dsl_gridfile, grid_savepoint.num(KDim)).get_grid() grg = diffusion_savepoint_init.geofac_grg() interpolation_state = InterpolationState( e_bln_c_s=diffusion_savepoint_init.e_bln_c_s(), @@ -464,7 +469,7 @@ def test_verify_diffusion_init_against_other_regular_savepoint( _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) -@pytest.mark.skip("fix: diffusion_stencil_15") +#@pytest.mark.skip("fix: diffusion_stencil_15") @pytest.mark.parametrize("run_with_program", [True, False]) @pytest.mark.datatest def test_run_diffusion_single_step( @@ -472,7 +477,7 @@ def test_run_diffusion_single_step( diffusion_savepoint_init, diffusion_savepoint_exit, grid_savepoint, # noqa: F811 - icon_grid, + r04b09_dsl_gridfile, r04b09_diffusion_config, damping_height, ): @@ -497,7 +502,7 @@ def test_run_diffusion_single_step( vct_a=vct_a, rayleigh_damping_height=damping_height ) additional_parameters = DiffusionParams(r04b09_diffusion_config) - + icon_grid = init_grid_manager(r04b09_dsl_gridfile, grid_savepoint.num(KDim)).get_grid() diffusion = Diffusion(run_program=run_with_program) diffusion.init( grid=icon_grid, From 4385eb509782afb6211ccc262ec9ff19b824d8c1 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 13 Apr 2023 09:53:10 +0200 Subject: [PATCH 091/263] - change interface of icon_grid.py: get_startindex, get_endindex instead of get_start_end_indices - calulcate e2c2v --- .../src/icon4py/diffusion/diffusion.py | 138 ++++---- atm_dyn_iconam/tests/test_icon_grid.py | 312 ++++++------------ common/src/icon4py/grid/grid_manager.py | 70 ++-- common/src/icon4py/grid/icon_grid.py | 32 +- common/tests/test_gridmanager.py | 49 ++- .../src/icon4py/testutils/serialbox_utils.py | 10 +- .../src/icon4py/testutils/simple_mesh.py | 12 + 7 files changed, 276 insertions(+), 347 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index d2890ba44..19d57cd4a 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -564,64 +564,59 @@ def time_step( ) else: print("run program") - ( - cell_startindex_nudging, - cell_endindex_local, - ) = self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.local(CellDim), - ) - ( - cell_startindex_interior, - cell_endindex_local_plus1, - ) = self.grid.get_indices_from_to( + cell_start_nudging_minus1 = self.grid.get_start_index(CellDim, + HorizontalMarkerIndex.nudging( + CellDim) - 1) + + cell_startindex_interior = self.grid.get_start_index(CellDim, + HorizontalMarkerIndex.interior(CellDim)) + cell_startindex_nudging = self.grid.get_start_index( CellDim, - HorizontalMarkerIndex.interior(CellDim), - HorizontalMarkerIndex.local(CellDim) - 1, - ) + HorizontalMarkerIndex.nudging(CellDim)) - ( - edge_startindex_nudging_plus1, - edge_endindex_local, - ) = self.grid.get_indices_from_to( + cell_endindex_local_plus1 = self.grid.get_end_index(CellDim, HorizontalMarkerIndex.local(CellDim) - 1) + cell_endindex_local = self.grid.get_end_index( + CellDim, HorizontalMarkerIndex.local(CellDim)) + + edge_startindex_nudging_plus1 = self.grid.get_start_index(EdgeDim, + HorizontalMarkerIndex.nudging( + EdgeDim) + 1) + edge_endindex_local = self.grid.get_end_index( EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim) + 1, HorizontalMarkerIndex.local(EdgeDim), ) - ( - edge_startindex_nudging_minus1, - edge_endindex_local_minus2, - ) = self.grid.get_indices_from_to( + edge_start_lb_plus4 = self.grid.get_start_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4 + ) + + + edge_startindex_nudging_minus1 = self.grid.get_start_index(EdgeDim, + HorizontalMarkerIndex.nudging( + EdgeDim) - 1) + + edge_endindex_local_minus2 = self.grid.get_end_index( EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim) - 1, HorizontalMarkerIndex.local(EdgeDim) - 2, ) - ( - vertex_startindex_lb_plus3, - vertex_endindex_local, - ) = self.grid.get_indices_from_to( + vertex_startindex_lb_plus3 = self.grid.get_start_index(VertexDim, + HorizontalMarkerIndex.lateral_boundary( + VertexDim) + 3) + vertex_endindex_local = self.grid.get_end_index( VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, HorizontalMarkerIndex.local(VertexDim), ) - ( - vertex_startindex_lb_plus1, - vertex_endindex_local_minus1, - ) = self.grid.get_indices_from_to( + vertex_startindex_lb_plus1 = self.grid.get_start_index(VertexDim, + HorizontalMarkerIndex.lateral_boundary( + VertexDim) + 1) + vertex_endindex_local_minus1 = self.grid.get_end_index( VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, HorizontalMarkerIndex.local(VertexDim) - 1, ) - edge_start_lb_plus4, _ = self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, - ) diff_prog.diffusion_run( diagnostic_hdef_ic=diagnostic_state.hdef_ic, @@ -751,61 +746,50 @@ def _do_diffusion_step( """ klevels = self.grid.n_lev() - k_start_end_minus2 = klevels - 2 + k_end_minus2 = klevels - 2 - cell_start_nudging_minus1, cell_end_local_plus1 = self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim) - 1, - HorizontalMarkerIndex.local(CellDim) - 1, - ) + cell_start_nudging_minus1 = self.grid.get_start_index(CellDim, HorizontalMarkerIndex.nudging(CellDim) - 1) - cell_start_interior, cell_end_local = self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.interior(CellDim), - HorizontalMarkerIndex.local(CellDim), - ) - cell_start_nudging, _ = self.grid.get_indices_from_to( + cell_startindex_interior = self.grid.get_start_index(CellDim,HorizontalMarkerIndex.interior(CellDim)) + cell_start_nudging = self.grid.get_start_index( CellDim, - HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.local(CellDim), - ) + HorizontalMarkerIndex.nudging(CellDim)) + + + cell_end_local_plus1 = self.grid.get_end_index(CellDim, HorizontalMarkerIndex.local(CellDim) - 1) + cell_end_local = self.grid.get_end_index( + CellDim, HorizontalMarkerIndex.local(CellDim)) + + edge_start_nudging_plus_one = self.grid.get_start_index(EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1) - edge_start_nudging_plus_one, edge_end_local = self.grid.get_indices_from_to( + edge_end_local = self.grid.get_end_index( EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim) + 1, HorizontalMarkerIndex.local(EdgeDim), ) - edge_start_lb_plus4, _ = self.grid.get_indices_from_to( + edge_start_lb_plus4 = self.grid.get_start_index( EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4 ) - ( - edge_start_nudging_minus1, - edge_end_local_minus2, - ) = self.grid.get_indices_from_to( + edge_start_nudging_minus1 = self.grid.get_start_index(EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) - 1) + + edge_end_local_minus2 = self.grid.get_end_index( EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim) - 1, HorizontalMarkerIndex.local(EdgeDim) - 2, ) - ( - vertex_start_local_boundary_plus3, - vertex_end_local, - ) = self.grid.get_indices_from_to( + + vertex_start_local_boundary_plus3 = self.grid.get_start_index(VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3) + vertex_end_local = self.grid.get_end_index( VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, HorizontalMarkerIndex.local(VertexDim), ) - ( - vertex_start_local_boundary_plus1, - vertex_end_local_minus1, - ) = self.grid.get_indices_from_to( + + vertex_start_local_boundary_plus1 = self.grid.get_start_index(VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1) + vertex_end_local_minus1 = self.grid.get_end_index( VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, HorizontalMarkerIndex.local(VertexDim) - 1, ) @@ -958,7 +942,7 @@ def _do_diffusion_step( horz_idx=self.horizontal_cell_index, nrdmax=self.vertical_params.index_of_damping_layer, interior_idx=int32( - cell_start_interior + cell_startindex_interior ), # h end index for stencil_09 and stencil_10 halo_idx=int32( cell_end_local @@ -987,7 +971,7 @@ def _do_diffusion_step( kh_smag_e=self.kh_smag_e, horizontal_start=edge_start_nudging_plus_one, horizontal_end=edge_end_local, - vertical_start=k_start_end_minus2, + vertical_start=k_end_minus2, vertical_end=klevels, offset_provider={ "E2C": self.grid.get_e2c_connectivity(), diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index 1a7ba1396..477804557 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -17,231 +17,109 @@ from icon4py.testutils.fixtures import data_provider, setup_icon_data # noqa + @pytest.mark.datatest -def test_horizontal_grid_cell_indices(icon_grid): - assert icon_grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.local(CellDim) - 1, - HorizontalMarkerIndex.local(CellDim) - 1, - ) == ( - 20896, - 20896, - ) # halo + 1 +@pytest.mark.parametrize( "marker, index",[ + ( HorizontalMarkerIndex.interior(CellDim), 4104), + (HorizontalMarkerIndex.interior(CellDim) + 1, 0), + (HorizontalMarkerIndex.local(CellDim) - 1, 20896), + (HorizontalMarkerIndex.local(CellDim), -1), # halo in icon is (1,20896) + (HorizontalMarkerIndex.nudging(CellDim), 3316), + (HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 2511), + (HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 1688), + (HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 850), + (HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 0) + ]) +def test_horizontal_cell_start_indices(icon_grid, marker, index): + assert index == icon_grid.get_start_index(CellDim, marker) - assert icon_grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.local(CellDim), - HorizontalMarkerIndex.local(CellDim), - ) == ( - -1, - 20896, - ) # halo in icon is (1,20896) - assert icon_grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.interior(CellDim), - HorizontalMarkerIndex.interior(CellDim), - ) == ( - 4104, - 20896, - ) # interior - assert icon_grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.interior(CellDim) + 1, - HorizontalMarkerIndex.interior(CellDim) + 1, - ) == ( - 0, - 850, - ) # lb+1 - assert icon_grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, - HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, - ) == (850, 1688) - assert icon_grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, - HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, - ) == ( - 1688, - 2511, - ) # lb+2 - assert icon_grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim) - 1, - HorizontalMarkerIndex.nudging(CellDim) - 1, - ) == ( - 2511, - 3316, - ) # lb+3 - assert icon_grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.nudging(CellDim), - ) == ( - 3316, - 4104, - ) # nudging +@pytest.mark.datatest +@pytest.mark.parametrize( "marker, index",[ + ( HorizontalMarkerIndex.interior(CellDim), 20896), + (HorizontalMarkerIndex.interior(CellDim) + 1, 850), + ( HorizontalMarkerIndex.local(CellDim) - 2,20896 ), + (HorizontalMarkerIndex.local(CellDim) - 1, 20896), + (HorizontalMarkerIndex.local(CellDim), 31558), + (HorizontalMarkerIndex.nudging(CellDim), 4104), + (HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 3316), + (HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 2511), + (HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 1688), + (HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 428),]) +def test_horizontal_edge_end_indices(icon_grid, marker, index): + assert index == icon_grid.get_end_index(CellDim, marker) +@pytest.mark.datatest +@pytest.mark.parametrize( "marker, index",[ + ( HorizontalMarkerIndex.interior(EdgeDim), 6176), + ( HorizontalMarkerIndex.local(EdgeDim) - 2,31558 ), + (HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), + (HorizontalMarkerIndex.local(EdgeDim), -1), + (HorizontalMarkerIndex.nudging(EdgeDim) + 1,5387), + (HorizontalMarkerIndex.nudging(EdgeDim), 4989), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7,4184 ), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 3777), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 2954), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2538), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 1700), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1278), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 428), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 0) + ]) +def test_horizontal_edge_start_indices(icon_grid, marker, index): + assert index == icon_grid.get_start_index(EdgeDim, marker) @pytest.mark.datatest -def test_horizontal_edge_indices(icon_grid): - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.interior(EdgeDim), - HorizontalMarkerIndex.interior(EdgeDim), - ) == (6176, 31558) - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.local(EdgeDim) - 2, - HorizontalMarkerIndex.local(EdgeDim) - 2, - ) == (31558, 31558) - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.local(EdgeDim) - 1, - HorizontalMarkerIndex.local(EdgeDim) - 1, - ) == (31558, 31558) - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.local(EdgeDim), - HorizontalMarkerIndex.local(EdgeDim), - ) == ( - -1, - 31558, - ) # halo in icon is (1, 31558) - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim) + 1, - HorizontalMarkerIndex.nudging(EdgeDim) + 1, - ) == ( - 5387, - 6176, - ) # nudging +1 - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim), - HorizontalMarkerIndex.nudging(EdgeDim), - ) == ( - 4989, - 5387, - ) # nudging - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, - ) == ( - 4184, - 4989, - ) # lb +7 - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, - ) == ( - 3777, - 4184, - ) # lb +6 - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, - ) == ( - 2954, - 3777, - ) # lb +5 - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, - ) == ( - 2538, - 2954, - ) # lb +4 - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, - ) == ( - 1700, - 2538, - ) # lb +3 - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, - ) == ( - 1278, - 1700, - ) # lb +2 - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, - ) == ( - 428, - 1278, - ) # lb +1 - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim), - HorizontalMarkerIndex.lateral_boundary(EdgeDim), - ) == ( - 0, - 428, - ) # lb +0 +@pytest.mark.parametrize( "marker, index",[ + ( HorizontalMarkerIndex.interior(EdgeDim), 31558), + ( HorizontalMarkerIndex.local(EdgeDim) - 2,31558 ), + (HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), + (HorizontalMarkerIndex.local(EdgeDim), 31558), + (HorizontalMarkerIndex.nudging(EdgeDim) + 1,6176), + (HorizontalMarkerIndex.nudging(EdgeDim), 5387), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4989), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 4184), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 3777), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2954), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 2538), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1700), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 1278), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 428),]) +def test_horizontal_edge_end_indices(icon_grid, marker, index): + assert index == icon_grid.get_end_index(EdgeDim, marker) @pytest.mark.datatest -def test_horizontal_vertex_indices(icon_grid): - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.end(VertexDim), - HorizontalMarkerIndex.end(VertexDim), - ) == (10663, 10663) - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.local(VertexDim), - HorizontalMarkerIndex.local(VertexDim), - ) == (-1, 10663) - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.local(VertexDim) - 1, - HorizontalMarkerIndex.local(VertexDim) - 1, - ) == (10663, 10663) +@pytest.mark.parametrize( "marker, index",[ + ( HorizontalMarkerIndex.interior(VertexDim), 2071), + (HorizontalMarkerIndex.local(VertexDim) - 1, 10663), + (HorizontalMarkerIndex.local(VertexDim), -1), + (HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), + (HorizontalMarkerIndex.nudging(VertexDim), 10663), + (HorizontalMarkerIndex.end(VertexDim), 10663), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 1673), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1266), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 850), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 428), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 0) + ]) +def test_horizontal_vertex_start_indices(icon_grid, marker, index): + assert index == icon_grid.get_start_index(VertexDim, marker) - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim), - HorizontalMarkerIndex.lateral_boundary(VertexDim), - ) == (0, 428) - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, - ) == (428, 850) - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, - ) == (850, 1266) - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, - ) == (1266, 1673) - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, - ) == (1673, 2071) +@pytest.mark.datatest +@pytest.mark.parametrize( "marker, index",[ + ( HorizontalMarkerIndex.interior(VertexDim), 10663), + ( HorizontalMarkerIndex.local(VertexDim) - 2,10663 ), + (HorizontalMarkerIndex.local(VertexDim) - 1, 10663), + (HorizontalMarkerIndex.local(VertexDim), 10663), + (HorizontalMarkerIndex.nudging(VertexDim) + 1,10663), + (HorizontalMarkerIndex.nudging(VertexDim), 10663), + (HorizontalMarkerIndex.end(VertexDim), 10663), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 2071), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1673), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 1266), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 850), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 428),]) +def test_horizontal_vertex_end_indices(icon_grid, marker, index): + assert index == icon_grid.get_end_index(VertexDim, marker) - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.interior(VertexDim), - HorizontalMarkerIndex.interior(VertexDim), - ) == (2071, 10663) - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.nudging(VertexDim), - HorizontalMarkerIndex.nudging(VertexDim), - ) == (10663, 10663) diff --git a/common/src/icon4py/grid/grid_manager.py b/common/src/icon4py/grid/grid_manager.py index 1d4b3176e..acfe56e33 100644 --- a/common/src/icon4py/grid/grid_manager.py +++ b/common/src/icon4py/grid/grid_manager.py @@ -31,7 +31,7 @@ EdgeDim, V2CDim, V2EDim, - VertexDim, + VertexDim, C2E2CODim, E2C2VDim, ) from icon4py.grid.horizontal import HorizontalGridSize from icon4py.grid.icon_grid import GridConfig, IconGrid @@ -66,8 +66,8 @@ class PropertyName(GridFileName): class OffsetName(GridFileName): # e2c2e/e2c2eO: diamond edges (including origin) -> calculate? - # e2c2v: diamond vertices - # + # e2c2v: diamond vertices: constructed from e2c and c2v + C2E2C = "neighbor_cell_index" # dims(nv=3, cell) V2E2V = "vertices_of_vertex" # dims(ne=6, vertex) not in simple mesh, = v2c2v vertices in hexa/pentagon V2E = "edges_of_vertex" # dims(ne=6, vertex) @@ -110,13 +110,13 @@ def __init__(self, dataset: Dataset): def dimension(self, name: GridFileName) -> int: return self._dataset.dimensions[name].size - def int_field(self, name: GridFileName, transpose=True) -> np.ndarray: + def int_field(self, name: GridFileName, transpose=True, dtype=np.int32) -> np.ndarray: try: nc_variable = self._dataset.variables[name] self._log.debug(f"reading {name}: {nc_variable}") data = nc_variable[:] - data = np.array(data, dtype=np.int32) + data = np.array(data, dtype=dtype) return np.transpose(data) if transpose else data except KeyError: msg = f"{name} does not exist in dataset" @@ -129,8 +129,8 @@ class IconGridError(RuntimeError): class IndexTransformation(ABC): - def get_offset_for_index_field(self, array: np.ndarray): - return np.zeros(array.shape) + def get_offset_for_index_field(self, array: np.ndarray,): + return np.zeros(array.shape, dtype=np.int32) class ToGt4PyTransformation(IndexTransformation): @@ -173,6 +173,7 @@ def _read_gridfile(self, fname: str) -> Dataset: exit(1) def _read_grid_refinement_information(self, dataset): + _CHILD_DOM = 0 reader = GridFile(dataset) grf_vertices = reader.dimension(GridFile.DimensionName.VERTEX_GRF) @@ -185,31 +186,31 @@ def _read_grid_refinement_information(self, dataset): end_indices = {} start_indices[CellDim] = self._get_index_field( reader, GridFile.GridRefinementName.START_INDEX_CELLS, transpose=False - ) + )[_CHILD_DOM] end_indices[CellDim] = self._get_index_field( reader, GridFile.GridRefinementName.END_INDEX_CELLS, transpose=False, - apply_offset=False, - ) + apply_offset=False, dtype=np.int64 + )[_CHILD_DOM] start_indices[EdgeDim] = self._get_index_field( - reader, GridFile.GridRefinementName.START_INDEX_EDGES, transpose=False - ) + reader, GridFile.GridRefinementName.START_INDEX_EDGES, transpose=False, dtype=np.int64 + )[_CHILD_DOM] end_indices[EdgeDim] = self._get_index_field( reader, GridFile.GridRefinementName.END_INDEX_EDGES, transpose=False, - apply_offset=False, - ) + apply_offset=False, dtype=np.int64 + )[_CHILD_DOM] start_indices[VertexDim] = self._get_index_field( - reader, GridFile.GridRefinementName.START_INDEX_VERTICES, transpose=False - ) + reader, GridFile.GridRefinementName.START_INDEX_VERTICES, transpose=False, dtype=np.int64 + )[_CHILD_DOM] end_indices[VertexDim] = self._get_index_field( reader, GridFile.GridRefinementName.END_INDEX_VERTICES, transpose=False, - apply_offset=False, - ) + apply_offset=False, dtype=np.int64 + )[_CHILD_DOM] return start_indices, end_indices @@ -225,13 +226,12 @@ def get_grid(self): return self._grid def _get_index(self, dim: Dimension, start_marker: int, index_dict): - _CHILD_DOM = 0 if dim.kind != DimensionKind.HORIZONTAL: msg = f"getting start index in horizontal domain with non - horizontal dimension {dim}" self._log.warning(msg) raise IconGridError(msg) try: - return index_dict[dim][_CHILD_DOM][start_marker] + return index_dict[dim][start_marker] except KeyError as error: msg = f"start, end indices for dimension {dim} not present" self._log.error(msg) @@ -256,6 +256,9 @@ def get_c2e2c_connectivity(self): def get_v2c_connectivity(self): return self._grid.get_v2c_connectivity() + def get_c2v_connectivity(self): + return self._grid.get_c2v_connectivity() + def get_c2e2co_connectivity(self): return self._grid.get_c2e2co_connectivity() @@ -283,9 +286,9 @@ def get_size(self, dim: Dimension): raise IconGridError(f"Unknown dimension {dim}") def _get_index_field( - self, reader, field: GridFileName, transpose=True, apply_offset=True + self, reader, field: GridFileName, transpose=True, apply_offset=True, dtype=np.int32 ): - field = reader.int_field(field, transpose=transpose) + field = reader.int_field(field, transpose=transpose, dtype=dtype) if apply_offset: field = field + self._transformation.get_offset_for_index_field(field) return field @@ -300,15 +303,19 @@ def from_grid_dataset(self, dataset: Dataset) -> IconGrid: num_vertices=num_vertices, num_edges=num_edges, num_cells=num_cells ) c2e = self._get_index_field(reader, GridFile.OffsetName.C2E) - c2v = self._get_index_field(reader, GridFile.OffsetName.C2V) - e2c = self._get_index_field(reader, GridFile.OffsetName.E2C) + e2c = self._get_index_field(reader, GridFile.OffsetName.E2C) + c2v = self._get_index_field(reader, GridFile.OffsetName.C2V) e2v = self._get_index_field(reader, GridFile.OffsetName.E2V) + + e2c2v = self.construct_diamond_array(c2v, e2c) + v2c = self._get_index_field(reader, GridFile.OffsetName.V2C) v2e = self._get_index_field(reader, GridFile.OffsetName.V2E) v2e2v = self._get_index_field(reader, GridFile.OffsetName.V2E2V) c2e2c = self._get_index_field(reader, GridFile.OffsetName.C2E2C) + c2e2c0 = np.column_stack((c2e2c, (np.asarray(range(c2e2c.shape[0]))))) start_indices, end_indices = self._read_grid_refinement_information(dataset) config = GridConfig( @@ -325,8 +332,10 @@ def from_grid_dataset(self, dataset: Dataset) -> IconGrid: E2VDim: e2v, V2EDim: v2e, V2CDim: v2c, - C2E2CDim: c2e2c, C2VDim: c2v, + C2E2CDim: c2e2c, + C2E2CODim: c2e2c0, + E2C2VDim: e2c2v, } ) .with_start_end_indices( @@ -341,3 +350,14 @@ def from_grid_dataset(self, dataset: Dataset) -> IconGrid: ) return icon_grid + + def construct_diamond_array(self, c2v:np.ndarray, e2c:np.ndarray): + dummy_c2v = np.append(c2v, GridFile.INVALID_INDEX * np.ones((1, c2v.shape[1]), dtype=np.int32), axis=0) + expanded = dummy_c2v[e2c[:, :], :] + sh = expanded.shape + flattened = expanded.reshape(sh[0], sh[1] * sh[2]) + return np.apply_along_axis(np.unique, 1, flattened) + + def get_c2v_connectivity(self): + return self._grid.get_c2v_connectivity() + diff --git a/common/src/icon4py/grid/icon_grid.py b/common/src/icon4py/grid/icon_grid.py index 9494fde69..fe9e89670 100644 --- a/common/src/icon4py/grid/icon_grid.py +++ b/common/src/icon4py/grid/icon_grid.py @@ -116,23 +116,23 @@ def num_vertices(self): def num_edges(self): return self.config.num_edges - def get_indices_from_to( - self, dim: Dimension, start_marker: int, end_marker: int - ) -> Tuple[int, int]: + def get_start_index(self, dim:Dimension, marker: int): """ - Use to specify domains of a field for field_operator. + Use to specify lower end of domains of a field for field_operators. - For a given dimension, returns the start and end index if a - horizontal region in a field given by the markers. + For a given dimension, returns the start index of the + horizontal region in a field given by the marker. + """ + return self.start_indices[dim][marker] - field operators will then run from start of the region given by the - start_marker to the end of the region given by the end_marker + def get_end_index(self, dim: Dimension, marker:int): """ - if dim.kind != DimensionKind.HORIZONTAL: - raise ValueError( - "only defined for {} dimension kind ", DimensionKind.HORIZONTAL - ) - return self.start_indices[dim][start_marker], self.end_indices[dim][end_marker] + Use to specify upper end of domains of a field for field_operators. + + For a given dimension, returns the end index of the + horizontal region in a field given by the marker. + """ + return self.end_indices[dim][marker] def get_c2e_connectivity(self): table = self.connectivities["c2e"] @@ -164,7 +164,11 @@ def get_v2e_connectivity(self): def get_v2c_connectivity(self): table = self.connectivities["v2c"] - return NeighborTableOffsetProvider(table, VertexDim, EdgeDim, table.shape[1]) + return NeighborTableOffsetProvider(table, VertexDim, CellDim, table.shape[1]) + + def get_c2v_connectivity(self): + table = self.connectivities["c2v"] + return NeighborTableOffsetProvider(table, VertexDim, CellDim, table.shape[1]) def get_e2ecv_connectivity(self): old_shape = self.connectivities["e2c2v"].shape diff --git a/common/tests/test_gridmanager.py b/common/tests/test_gridmanager.py index 2dca01494..0daa08200 100644 --- a/common/tests/test_gridmanager.py +++ b/common/tests/test_gridmanager.py @@ -28,6 +28,7 @@ ToGt4PyTransformation, ) from icon4py.grid.horizontal import HorizontalMarkerIndex +from icon4py.grid.vertical import VerticalGridConfig from icon4py.testutils.simple_mesh import SimpleMesh @@ -35,7 +36,7 @@ @pytest.fixture -def simple_mesh_path(): +def simple_mesh_path(simple_mesh_data): return Path(SIMPLE_MESH_NC).absolute() @@ -45,6 +46,7 @@ def simple_mesh_data(): dataset = Dataset(SIMPLE_MESH_NC, "w", format="NETCDF4") dataset.setncattr(GridFile.PropertyName.GRID_ID, str(uuid4())) dataset.createDimension(GridFile.DimensionName.VERTEX_NAME, size=mesh.n_vertices) + dataset.createDimension(GridFile.DimensionName.EDGE_NAME, size=mesh.n_edges) dataset.createDimension(GridFile.DimensionName.CELL_NAME, size=mesh.n_cells) dataset.createDimension( @@ -63,6 +65,7 @@ def simple_mesh_data(): GridFile.GridRefinementName.CONTROL_EDGES, (GridFile.DimensionName.EDGE_NAME,), ) + _add_to_dataset( dataset, np.zeros(mesh.n_cells), @@ -92,7 +95,7 @@ def simple_mesh_data(): GridFile.DimensionName.CELL_NAME, ), ) - # add_to_dataset(data, mesh.c2v, GridFile.Offsets.C2V, (GridFile.Dimension.C2E_SIZE, GridFile.Dimension.CELL_NAME)) + _add_to_dataset( dataset, mesh.e2c, @@ -121,10 +124,10 @@ def simple_mesh_data(): GridFile.DimensionName.VERTEX_NAME, ), ) - # TODO @magdalena: there is no v2c in the simple mesh + _add_to_dataset( dataset, - np.zeros((mesh.n_cells, 3), dtype=np.int32), + mesh.c2v, GridFile.OffsetName.C2V, ( GridFile.DimensionName.NEIGHBORS_TO_CELL_SIZE, @@ -206,6 +209,8 @@ def _add_to_dataset( var[:] = np.transpose(data)[:] + + def test_gridparser_dimension(simple_mesh_data): data = Dataset(SIMPLE_MESH_NC, "r") @@ -245,8 +250,6 @@ def test_grid_parser_index_fields(simple_mesh_data, caplog): # TODO @magdalena add test cases for -# e2c2v - diamond: serial, simple -# c2v: grid, ??? # v2e2v: grid,??? # v2e: serial, simple, grid @@ -265,6 +268,7 @@ def test_gridmanager_eval_v2e(caplog, grid_savepoint, r04b09_dsl_gridfile): assert np.allclose(gm.get_v2e_connectivity().table, seralized_v2e) + # v2c: serial, simple, grid @pytest.mark.datatest def test_gridmanager_eval_v2c(caplog, grid_savepoint, r04b09_dsl_gridfile): @@ -323,10 +327,13 @@ def test_gridmanager_eval_e2v(caplog, grid_savepoint, r04b09_dsl_gridfile): assert np.allclose(gm.get_e2v_connectivity().table, serialized_e2v) + def has_invalid_index(ar: np.ndarray): return np.any(np.where(ar == GridFile.INVALID_INDEX)) + + # e2c :serial, simple, grid @pytest.mark.datatest def test_gridmanager_eval_e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): @@ -358,7 +365,7 @@ def test_gridmanager_eval_c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): assert np.allclose(gm.get_c2e_connectivity().table, serialized_c2e) -# e2c2e (e2c2eo) - diamond: serial, simple_mesh +# e2c2e (e2c2eo) - diamond: serial, simple_mesh, grid file???? @pytest.mark.skip("does this array exist in grid file?") @pytest.mark.datatest def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): @@ -370,6 +377,8 @@ def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): assert np.allclose(gm.get_e2c2e_connectivity().table, serialized_e2c2e) + + # c2e2c: serial, simple_mesh, grid @pytest.mark.datatest def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): @@ -380,9 +389,25 @@ def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): gm.get_c2e2c_connectivity().table, grid_savepoint.c2e2c()[0:num_cells, :] ) - +def test_gridmanager_eval_e2c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + gm = init_grid_manager(r04b09_dsl_gridfile) + num_edges = gm.get_size(EdgeDim) + c2v = gm.get_c2v_connectivity().table + assert np.allclose( + gm.get_e2c2v_connectivity().table, grid_savepoint.e2c2v()[0:num_edges, :] + ) +@pytest.mark.skip("add to savepoint") +def test_gridmanager_eval_c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + gm = init_grid_manager(r04b09_dsl_gridfile) + num_cells = gm.get_size(CellDim) + c2v = gm.get_c2v_connectivity().table + assert np.allclose( + c2v, grid_savepoint.c2v()[0:num_cells, :] + ) def init_grid_manager(fname): - gm = GridManager(ToGt4PyTransformation(), fname) + gm = GridManager(ToGt4PyTransformation(), fname, VerticalGridConfig(65)) gm.init() return gm @@ -394,6 +419,12 @@ def test_grid_manager_getsize(simple_mesh_data, simple_mesh_path, dim, size, cap gm.init() assert size == gm.get_size(dim) +def test_grid_manager_diamond_offset(simple_mesh_path): + mesh = SimpleMesh() + gm = GridManager(IndexTransformation(), simple_mesh_path, VerticalGridConfig(num_lev=mesh.k_level)) + gm.init() + assert np.allclose(np.sort(gm.get_e2c2v_connectivity().table, 1), np.sort(mesh.diamond_arr, 1)) + def test_gridmanager_given_file_not_found_then_abort(): fname = "./unknown_grid.nc" diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index e46b512c6..751cdce04 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -142,13 +142,13 @@ def refin_ctrl(self, dim: Dimension) -> Optional[np.ndarray]: def num(self, dim: Dimension) -> Optional[int]: if dim == CellDim: - return self.serializer.read("num_cells", self.savepoint)[0] + return int(self.serializer.read("num_cells", self.savepoint)[0]) elif dim == EdgeDim: - return self.serializer.read("num_edges", self.savepoint)[0] + return int(self.serializer.read("num_edges", self.savepoint)[0]) elif dim == VertexDim: - return self.serializer.read("num_vert", self.savepoint)[0] + return int(self.serializer.read("num_vert", self.savepoint)[0]) elif dim == KDim: - return self.serializer.read("nlev", self.savepoint)[0] + return int(self.serializer.read("nlev", self.savepoint)[0]) else: return None @@ -173,7 +173,7 @@ def e2v(self): return v_ def e2c2v(self): - # array "e2v" is actually e2c2v + # array "e2v" is actually e2c2v, that is hexagon or pentagon return self._get_connectivity_array("e2v") def v2e(self): diff --git a/testutils/src/icon4py/testutils/simple_mesh.py b/testutils/src/icon4py/testutils/simple_mesh.py index 4bff39e00..47f0cdb4f 100644 --- a/testutils/src/icon4py/testutils/simple_mesh.py +++ b/testutils/src/icon4py/testutils/simple_mesh.py @@ -58,6 +58,14 @@ @dataclass class SimpleMeshData: + c2v_table = np.asarray([[0,1,4], [1,2,5], [2,0,3], + [0,3,4], [1,4,5], [2,5,3], + [3,4,7], [4,5,8], [5,3,6], + [3,6,7], [4,7,8], [5,8,6], + [6,7,1], [7,8,2], [8,6,0], + [6,0,1],[7,1,2], [8,2,0] + ]) + e2c2v_table = np.asarray( [ [0, 1, 4, 6], # 0 @@ -353,6 +361,7 @@ class SimpleMesh: def __init__(self, k_level: int = _DEFAULT_K_LEVEL): self.diamond_arr = SimpleMeshData.diamond_table + self.c2v = SimpleMeshData.c2v_table self.e2c = SimpleMeshData.e2c_table self.e2v = SimpleMeshData.e2v_table self.c2e = SimpleMeshData.c2e_table @@ -372,6 +381,7 @@ def __init__(self, k_level: int = _DEFAULT_K_LEVEL): self.n_e2c2e = self.e2c2e.shape[1] self.n_e2c2v = self.e2c2v.shape[1] self.n_v2c = self.v2c.shape[1] + self.n_c2v = self.c2v.shape[1] self.n_v2e = self.v2e.shape[1] self.n_cells = self.c2e.shape[0] self.n_edges = 27 @@ -396,6 +406,8 @@ def __init__(self, k_level: int = _DEFAULT_K_LEVEL): ECVDim: self.n_edges * self.n_e2c2v, } + def get_c2v_offset_provider(self)-> NeighborTableOffsetProvider: + return NeighborTableOffsetProvider(self.c2v, VertexDim, CellDim, self.n_c2v) def get_c2e_offset_provider(self) -> NeighborTableOffsetProvider: return NeighborTableOffsetProvider(self.c2e, CellDim, EdgeDim, self.n_c2e) From 3a2ee48fe9ce188728825eaff9e9f05018cb4a84 Mon Sep 17 00:00:00 2001 From: Magdalena Date: Thu, 27 Apr 2023 09:43:12 +0200 Subject: [PATCH 092/263] Dummy dycore driver (#174) Adds a python command line program that runs the diffusion module in a dummy time stepping scheme: Instead of time integration it just serves new slices of serialized data. other properties of the dummy driver: - configures simple logging - input data (grid, static grid properties, interpolation fields, initial data for prognostic and diagnostic variables, etc) are read from serialized test data. - configuration is programmatically initialized and not read from file. --- .../src/icon4py/diffusion/diffusion.py | 65 ++--- .../icon4py/diffusion/diffusion_program.py | 6 +- .../src/icon4py/diffusion/horizontal.py | 200 +++++++------ .../src/icon4py/diffusion/icon_grid.py | 3 +- .../icon4py/diffusion/interpolation_state.py | 6 +- atm_dyn_iconam/src/icon4py/diffusion/utils.py | 49 +++- atm_dyn_iconam/src/icon4py/driver/README.md | 20 ++ atm_dyn_iconam/src/icon4py/driver/__init__.py | 12 + .../src/icon4py/driver/dycore_driver.py | 268 ++++++++++++++++++ .../src/icon4py/driver/icon_configuration.py | 86 ++++++ atm_dyn_iconam/src/icon4py/driver/io_utils.py | 167 +++++++++++ atm_dyn_iconam/tests/conftest.py | 56 +--- atm_dyn_iconam/tests/test_diffusion.py | 151 ++++------ atm_dyn_iconam/tests/test_dycore_driver.py | 68 +++++ atm_dyn_iconam/tests/test_io_utils.py | 117 ++++++++ .../src/icon4py/testutils/serialbox_utils.py | 143 +++++++++- 16 files changed, 1123 insertions(+), 294 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/driver/README.md create mode 100644 atm_dyn_iconam/src/icon4py/driver/__init__.py create mode 100644 atm_dyn_iconam/src/icon4py/driver/dycore_driver.py create mode 100644 atm_dyn_iconam/src/icon4py/driver/icon_configuration.py create mode 100644 atm_dyn_iconam/src/icon4py/driver/io_utils.py create mode 100644 atm_dyn_iconam/tests/test_dycore_driver.py create mode 100644 atm_dyn_iconam/tests/test_io_utils.py diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 2cba30b0e..3ce0fe6c7 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -10,9 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - -# flake8: noqa - +import logging import math import sys from collections import namedtuple @@ -83,6 +81,9 @@ ) +# flake8: noqa +log = logging.getLogger(__name__) + VectorTuple = namedtuple("VectorTuple", "x y") @@ -93,8 +94,8 @@ class DiffusionConfig: Encapsulates namelist parameters and derived parameters. Values should be read from configuration. Default values are taken from the defaults in the corresponding ICON Fortran namelist files. - TODO: [ml] to be read from config - TODO: [ml] handle dependencies on other namelists (see below...) + TODO: @magdalena to be read from config + TODO: @magdalena handle dependencies on other namelists (see below...) """ def __init__( @@ -470,7 +471,7 @@ def _index_field(dim: Dimension, size=None): self.z_nabla2_e = _allocate(EdgeDim, KDim) self.z_temp = _allocate(CellDim, KDim) self.diff_multfac_smag = _allocate(KDim) - # TODO this is KHalfDim + # TODO @magdalena this is KHalfDim self.vertical_index = _index_field(KDim, self.grid.n_lev() + 1) self.horizontal_cell_index = _index_field(CellDim) self.horizontal_edge_index = _index_field(EdgeDim) @@ -565,7 +566,7 @@ def time_step( smag_offset=self.smag_offset, ) else: - print("run program") + log.info("running diffusion_program") ( cell_startindex_nudging, cell_endindex_local, @@ -624,7 +625,7 @@ def time_step( HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, ) - + log.info("diffusion program: start") diff_prog.diffusion_run( diagnostic_hdef_ic=diagnostic_state.hdef_ic, diagnostic_div_ic=diagnostic_state.div_ic, @@ -714,6 +715,7 @@ def time_step( "Koff": KDim, }, ) + log.info("diffusion program: end") def _do_diffusion_step( self, @@ -811,16 +813,15 @@ def _do_diffusion_step( HorizontalMarkerIndex.local(VertexDim) - 1, ) - # Oa TODO: logging - # 0b TODO: call timer start + # 0b call timer start # # 0c. dtime dependent stuff: enh_smag_factor, scale_k(self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={}) - # TODO: is this needed?, if not remove + # TODO: @magdalena is this needed?, if not remove set_zero_v_k(self.u_vert, offset_provider={}) set_zero_v_k(self.v_vert, offset_provider={}) - print("rbf interpolation: start") + log.debug("rbf interpolation: start") # # 1. CALL rbf_vec_interpol_vertex mo_intp_rbf_rbf_vec_interpol_vertex( p_e_in=prognostic_state.vn, @@ -834,11 +835,11 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={"V2E": self.grid.get_v2e_connectivity(), "V2EDim": V2EDim}, ) - print("rbf interpolation: end") + log.debug("rbf interpolation: end") # 2. HALO EXCHANGE -- CALL sync_patch_array_mult # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 - print("running calculate_nabla2_and_smag_coefficients_for_vn: start") + log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: start") calculate_nabla2_and_smag_coefficients_for_vn( diff_multfac_smag=self.diff_multfac_smag, tangent_orientation=tangent_orientation, @@ -866,8 +867,8 @@ def _do_diffusion_step( "E2C2VDim": E2C2VDim, }, ) - print("running calculate_nabla2_and_smag_coefficients_for_vn: end") - print("running fused stencil fused stencil 02_03: start") + log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: end") + log.debug("running fused stencil fused stencil 02_03: start") fused_mo_nh_diffusion_stencil_02_03( kh_smag_ec=self.kh_smag_ec, vn=prognostic_state.vn, @@ -887,12 +888,12 @@ def _do_diffusion_step( "Koff": KDim, }, ) - print("running fused stencil fused stencil 02_03: end") + log.debug("running fused stencil fused stencil 02_03: end") # # # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH # # # 5. CALL rbf_vec_interpol_vertex_wp - print("rbf interpolation: start") + log.debug("rbf interpolation: start") mo_intp_rbf_rbf_vec_interpol_vertex( p_e_in=self.z_nabla2_e, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, @@ -905,14 +906,14 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={"V2E": self.grid.get_v2e_connectivity(), "V2EDim": V2EDim}, ) - print("rbf interpolation: end") + log.debug("rbf interpolation: end") # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult # # # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 # # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 # - print("running fused stencil 04_05_06: start") + log.debug("running fused stencil 04_05_06: start") fused_mo_nh_diffusion_stencil_04_05_06( u_vert=self.u_vert, v_vert=self.v_vert, @@ -942,9 +943,9 @@ def _do_diffusion_step( ) # # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, # # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 - print("running fused stencil 04_05_06: end") + log.debug("running fused stencil 04_05_06: end") - print("running fused stencil 07_08_09_10: start") + log.debug("running fused stencil 07_08_09_10: start") fused_mo_nh_diffusion_stencil_07_08_09_10( area=cell_areas, geofac_n2s=self.interpolation_state.geofac_n2s, @@ -974,14 +975,14 @@ def _do_diffusion_step( "C2E2CODim": C2E2CODim, }, ) - print("running fused stencil 07_08_09_10: start") + log.debug("running fused stencil 07_08_09_10: start") # # 8. HALO EXCHANGE: CALL sync_patch_array # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, # # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 # - # # TODO check: kh_smag_e is an out field, should not be calculated in init? + # # TODO @magdalena check: kh_smag_e is an out field, should not be calculated in init? # - print("running fused stencil 11_12: start") + log.debug("running fused stencil 11_12: start") fused_mo_nh_diffusion_stencil_11_12( theta_v=prognostic_state.theta_v, theta_ref_mc=self.metric_state.theta_ref_mc, @@ -998,8 +999,8 @@ def _do_diffusion_step( "E2CDim": E2CDim, }, ) - print("running fused stencil 11_12: end") - print("running fused stencil 13_14: start") + log.debug("running fused stencil 11_12: end") + log.debug("running fused stencil 13_14: start") fused_mo_nh_diffusion_stencil_13_14( kh_smag_e=self.kh_smag_e, inv_dual_edge_length=inverse_dual_edge_length, @@ -1017,8 +1018,8 @@ def _do_diffusion_step( "C2EDim": C2EDim, }, ) - print("running fused stencil 13_14: end") - print("running fused stencil 15: start") + log.debug("running fused stencil 13_14: end") + log.debug("running fused stencil 15: start") # mo_nh_diffusion_stencil_15( # self.metric_state.mask_hdiff, # self.metric_state.zd_vertidx, @@ -1038,8 +1039,8 @@ def _do_diffusion_step( # "Koff": KDim, # }, # ) - print("running fused stencil 15: end") - print("running fused stencil update_theta_and_exner: start") + log.debug("running fused stencil 15: end") + log.debug("running fused stencil update_theta_and_exner: start") update_theta_and_exner( z_temp=self.z_temp, area=cell_areas, @@ -1052,5 +1053,5 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={}, ) - print("running fused stencil update_theta_and_exner: end") + log.debug("running fused stencil update_theta_and_exner: end") # 10. HALO EXCHANGE sync_patch_array diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py index e624fe622..755dab9ba 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py @@ -130,7 +130,7 @@ def diffusion_run( ): _scale_k(local_enh_smag_fac, dtime, out=local_diff_multfac_smag) - # TODO: is this needed?, if not remove + # TODO: @magdalena is this needed?, if not remove _set_zero_v_k(out=local_u_vert) _set_zero_v_k(out=local_v_vert) @@ -165,7 +165,7 @@ def diffusion_run( local_smag_offset, out=(local_kh_smag_e, local_kh_smag_ec, local_z_nabla2_e), domain={ - # TODO wrong start index?? + # TODO @magdalena wrong start index?? EdgeDim: (boundary_diffusion_start_index_edges, edge_endindex_local_minus2), KDim: (0, nlev), }, @@ -262,7 +262,7 @@ def diffusion_run( # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, # # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 # - # # TODO check: kh_smag_e is an out field, should not be calculated in init? + # # TODO @magdalena check: kh_smag_e is an out field, should not be calculated in init? # _fused_mo_nh_diffusion_stencil_11_12( diff --git a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py index 531926929..27f07105e 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py @@ -115,93 +115,121 @@ class HorizontalMeshSize: num_cells: int -@dataclass(frozen=True) +# TODO [@Magdalena] allow initialization with only partial values +# (becomes tedious for testing otherwise): hence this should +# that should not be a data class class EdgeParams: - tangent_orientation: Field[[EdgeDim], float] - """ - Orientation of vector product of the edge and the adjacent cell centers - v3 - / \ - / \ - / c1 \ - / | \ - v1----|--->v2 - \ | / - \ v / - \ c2 / - \ / - v4 - +1 or -1 depending on whether the vector product of - (v2-v1) x (c2-c1) points outside (+) or inside (-) the sphere - - defined in ICON in mo_model_domain.f90:t_grid_edges%tangent_orientation - """ - - primal_edge_lengths: Field[[EdgeDim], float] - """ - Length of the triangle edge. - - defined int ICON in mo_model_domain.f90:t_grid_edges%primal_edge_length - """ - - inverse_primal_edge_lengths: Field[[EdgeDim], float] - """ - Inverse of the triangle edge length: 1.0/primal_edge_length. - - defined int ICON in mo_model_domain.f90:t_grid_edges%inv_primal_edge_length - """ - - dual_edge_lengths: Field[[EdgeDim], float] - """ - Length of the hexagon/pentagon edge. - - defined int ICON in mo_model_domain.f90:t_grid_edges%dual_edge_length - """ - - inverse_dual_edge_lengths: Field[[EdgeDim], float] - """ - Inverse of hexagon/pentagon edge length: 1.0/dual_edge_length. - - defined int ICON in mo_model_domain.f90:t_grid_edges%inv_dual_edge_length - """ - - inverse_vertex_vertex_lengths: Field[[EdgeDim], float] - """ - Inverse distance between outer vertices of adjacent cells. - - v1-------- - | /| - | / | - | e | - | / | - |/ | - --------v2 - - inverse_vertex_vertex_length(e) = 1.0/|v2-v1| - - defined int ICON in mo_model_domain.f90:t_grid_edges%inv_vert_vert_length - """ - - primal_normal_vert: tuple[Field[[ECVDim], float], Field[[ECVDim], float]] - """ - Normal of the triangle edge, projected onto the location of the vertices - - defined int ICON in mo_model_domain.f90:t_grid_edges%primal_normal_vert - """ - - dual_normal_vert: tuple[Field[[ECVDim], float], Field[[ECVDim], float]] - """ - Tangent to the triangle edge, projected onto the location of vertices. - - defined int ICON in mo_model_domain.f90:t_grid_edges%dual_normal_vert - """ - - edge_areas: Field[[EdgeDim], float] - """ - Area of the quadrilateral (two triangles) adjacent to the edge. - - defined int ICON in mo_model_domain.f90:t_grid_edges%area_edge - """ + def __init__( + self, + tangent_orientation=None, + primal_edge_lengths=None, + inverse_primal_edge_lengths=None, + dual_edge_lengths=None, + inverse_dual_edge_lengths=None, + inverse_vertex_vertex_lengths=None, + primal_normal_vert_x=None, + primal_normal_vert_y=None, + dual_normal_vert_x=None, + dual_normal_vert_y=None, + edge_areas=None, + ): + + self.tangent_orientation: Field[[EdgeDim], float] = tangent_orientation + """ + Orientation of vector product of the edge and the adjacent cell centers + v3 + / \ + / \ + / c1 \ + / | \ + v1---|--->v2 + \ | / + \ v / + \ c2 / + \ / + v4 + +1 or -1 depending on whether the vector product of + (v2-v1) x (c2-c1) points outside (+) or inside (-) the sphere + + defined in ICON in mo_model_domain.f90:t_grid_edges%tangent_orientation + """ + + self.primal_edge_lengths: Field[[EdgeDim], float] = primal_edge_lengths + """ + Length of the triangle edge. + + defined int ICON in mo_model_domain.f90:t_grid_edges%primal_edge_length + """ + + self.inverse_primal_edge_lengths: Field[ + [EdgeDim], float + ] = inverse_primal_edge_lengths + """ + Inverse of the triangle edge length: 1.0/primal_edge_length. + + defined int ICON in mo_model_domain.f90:t_grid_edges%inv_primal_edge_length + """ + + self.dual_edge_lengths: Field[[EdgeDim], float] = dual_edge_lengths + """ + Length of the hexagon/pentagon edge. + + defined int ICON in mo_model_domain.f90:t_grid_edges%dual_edge_length + """ + + self.inverse_dual_edge_lengths: Field[ + [EdgeDim], float + ] = inverse_dual_edge_lengths + """ + Inverse of hexagon/pentagon edge length: 1.0/dual_edge_length. + + defined int ICON in mo_model_domain.f90:t_grid_edges%inv_dual_edge_length + """ + + self.inverse_vertex_vertex_lengths: Field[ + [EdgeDim], float + ] = inverse_vertex_vertex_lengths + """ + Inverse distance between outer vertices of adjacent cells. + + v1-------- + | /| + | / | + | e | + | / | + |/ | + --------v2 + + inverse_vertex_vertex_length(e) = 1.0/|v2-v1| + + defined int ICON in mo_model_domain.f90:t_grid_edges%inv_vert_vert_length + """ + + self.primal_normal_vert: tuple[ + Field[[ECVDim], float], Field[[ECVDim], float] + ] = (primal_normal_vert_x, primal_normal_vert_y) + """ + Normal of the triangle edge, projected onto the location of the vertices + + defined int ICON in mo_model_domain.f90:t_grid_edges%primal_normal_vert + """ + + self.dual_normal_vert: tuple[Field[[ECVDim], float], Field[[ECVDim], float]] = ( + dual_normal_vert_x, + dual_normal_vert_y, + ) + """ + Tangent to the triangle edge, projected onto the location of vertices. + + defined int ICON in mo_model_domain.f90:t_grid_edges%dual_normal_vert + """ + + self.edge_areas: Field[[EdgeDim], float] = edge_areas + """ + Area of the quadrilateral (two triangles) adjacent to the edge. + + defined int ICON in mo_model_domain.f90:t_grid_edges%area_edge + """ @dataclass(frozen=True) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index 216593fa0..07f6625e5 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -22,6 +22,7 @@ from icon4py.diffusion.horizontal import HorizontalMeshSize +# TODO @magdalena keep naming grid vs mesh consistent class VerticalMeshConfig: def __init__(self, num_lev: int): self._num_lev = num_lev @@ -110,7 +111,7 @@ def with_connectivities(self, connectivity: Dict[Dimension, np.ndarray]): self.size.update({d: t.shape[1] for d, t in connectivity.items()}) def limited_area(self): - # TODO defined in mo_grid_nml.f90 + # defined in mo_grid_nml.f90 return self.config.limited_area def n_lev(self): diff --git a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py index 2486708e0..2ac0ff4a7 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py @@ -33,7 +33,7 @@ class InterpolationState: """ represents the ICON interpolation state. - TODO [ml]: keep? does this state make sense at all? + TODO [magdalena]: keep? does this state make sense at all? """ e_bln_c_s: Field[ @@ -56,9 +56,7 @@ class InterpolationState: geofac_grg_x: Field[ [CellDim, C2E2CODim], float ] # factor for green gauss gradient (nproma,4,nblks_c,2) - geofac_grg_y: Field[ - [CellDim, C2E2CODim], float - ] # TODO combine geofac_grg_x and geofac_grg_y to tuple + geofac_grg_y: Field[[CellDim, C2E2CODim], float] nudgecoeff_e: Field[[EdgeDim], float] # Nudgeing coeffients for edges @property diff --git a/atm_dyn_iconam/src/icon4py/diffusion/utils.py b/atm_dyn_iconam/src/icon4py/diffusion/utils.py index c7d9004fe..3f6a2b706 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/utils.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/utils.py @@ -19,15 +19,58 @@ from gt4py.next.iterator.embedded import np_as_located_field from gt4py.next.program_processors.runners import gtfn_cpu -from icon4py.common.dimension import KDim, Koff, VertexDim +from icon4py.common.dimension import CellDim, EdgeDim, KDim, Koff, VertexDim -# TODO fix duplication: duplicated from test testutils/utils.py +# TODO [@Magdalena] fix duplication: duplicated from test testutils/utils.py def zero_field(mesh, *dims: Dimension, dtype=float): shapex = tuple(map(lambda x: mesh.size[x], dims)) return np_as_located_field(*dims)(np.zeros(shapex, dtype=dtype)) +@field_operator +def _identity_c_k( + field: Field[[CellDim, KDim], float] +) -> Field[[CellDim, KDim], float]: + return field + + +@field_operator +def _identity_e_k( + field: Field[[EdgeDim, KDim], float] +) -> Field[[EdgeDim, KDim], float]: + return field + + +@program(backend=gtfn_cpu.run_gtfn) +def copy_diagnostic_and_prognostics( + hdef_ic_new: Field[[CellDim, KDim], float], + hdef_ic: Field[[CellDim, KDim], float], + div_ic_new: Field[[CellDim, KDim], float], + div_ic: Field[[CellDim, KDim], float], + dwdx_new: Field[[CellDim, KDim], float], + dwdx: Field[[CellDim, KDim], float], + dwdy_new: Field[[CellDim, KDim], float], + dwdy: Field[[CellDim, KDim], float], + vn_new: Field[[EdgeDim, KDim], float], + vn: Field[[EdgeDim, KDim], float], + w_new: Field[[CellDim, KDim], float], + w: Field[[CellDim, KDim], float], + exner_new: Field[[CellDim, KDim], float], + exner: Field[[CellDim, KDim], float], + theta_v_new: Field[[CellDim, KDim], float], + theta_v: Field[[CellDim, KDim], float], +): + _identity_c_k(hdef_ic_new, out=hdef_ic) + _identity_c_k(div_ic_new, out=div_ic) + _identity_c_k(dwdx_new, out=dwdx) + _identity_c_k(dwdy_new, out=dwdy) + _identity_e_k(vn_new, out=vn) + _identity_c_k(w_new, out=w) + _identity_c_k(exner_new, out=exner) + _identity_c_k(theta_v_new, out=theta_v) + + @field_operator def _scale_k(field: Field[[KDim], float], factor: float) -> Field[[KDim], float]: return field * factor @@ -200,7 +243,7 @@ def init_nabla2_factor_in_upper_damping_zone( Calculate diff_multfac_n2w. numpy version, since gt4py does not allow non-constant indexing into fields - TODO: [ml] fix this once IndexedFields are implemented + TODO: @magdalena fix this once IndexedFields are implemented Args k_size: number of vertical levels diff --git a/atm_dyn_iconam/src/icon4py/driver/README.md b/atm_dyn_iconam/src/icon4py/driver/README.md new file mode 100644 index 000000000..a4839c882 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/driver/README.md @@ -0,0 +1,20 @@ +# dummy driver for python dycore code + +`dycore_driver.py` contains a simple python program to run the experimental ICON python port. The code mostly draws on serialized ICON data. It initializes the grid from serialized data from a `mch_ch_r04b09_dsl` run and configures a timeloop functionality based on that configuration. + +Currently, there is does no real timestepping instead it calls a dummy timestep that serves a batch of new serialized input fields from ICON. + +The code is meant to be changed and enlarged as we port new parts of the model. + +## usage + +``` +cd atm_dyn_iconam/src/icon4py +python driver/dycore_driver.py ../../tests/ser_icondata/mch_ch_r04b09_dsl/ser_data --n_steps=2 --run_path=/home/magdalena/temp/icon +``` + +#### remarks + +- first (required) arg is the folder where the serialized input data is stored. The same input data s for the unit tests is used. It can be obtained from [polybox](https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download) +- the serialized data used contains only 5 timesteps so `--n_steps > 2` will throw an exception. +- the code logs to file and to console. Debug logging is only going to file. The log directory can be changed with the --run_path option. diff --git a/atm_dyn_iconam/src/icon4py/driver/__init__.py b/atm_dyn_iconam/src/icon4py/driver/__init__.py new file mode 100644 index 000000000..15dfdb009 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/driver/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py new file mode 100644 index 000000000..57a8d378e --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -0,0 +1,268 @@ +# 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 logging +from datetime import datetime, timedelta +from pathlib import Path +from typing import Callable + +import click +import pytz +from devtools import Timer + +from icon4py.diffusion.diagnostic_state import DiagnosticState +from icon4py.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.diffusion.horizontal import CellParams, EdgeParams +from icon4py.diffusion.prognostic_state import PrognosticState +from icon4py.diffusion.utils import copy_diagnostic_and_prognostics +from icon4py.driver.icon_configuration import IconRunConfig, read_config +from icon4py.driver.io_utils import ( + SIMULATION_START_DATE, + configure_logging, + read_geometry_fields, + read_icon_grid, + read_initial_state, + read_static_fields, +) +from icon4py.testutils.serialbox_utils import IconSerialDataProvider + + +log = logging.getLogger(__name__) + + +class DummyAtmoNonHydro: + def __init__(self, data_provider: IconSerialDataProvider): + self.config = None + self.data_provider = data_provider + self.simulation_date = datetime.fromisoformat(SIMULATION_START_DATE) + + def init(self, config): + self.config = config + + def _next_physics_date(self, dtime: float): + dynamics_dtime = dtime / self.config.n_substeps + self.simulation_date += timedelta(seconds=dynamics_dtime) + + def _dynamics_timestep(self, dtime): + """Show structure with this dummy fucntion called inside substepping loop.""" + self._next_physics_date(dtime) + + def do_dynamics_substepping( + self, + dtime, + diagnostic_state: DiagnosticState, + prognostic_state: PrognosticState, + ): + for _ in range(self.config.n_substeps): + self._dynamics_timestep(dtime) + sp = self.data_provider.from_savepoint_diffusion_init( + linit=False, date=self.simulation_date.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + ) + new_p = sp.construct_prognostics() + new_d = sp.construct_diagnostics() + copy_diagnostic_and_prognostics( + new_d.hdef_ic, + diagnostic_state.hdef_ic, + new_d.div_ic, + diagnostic_state.div_ic, + new_d.dwdx, + diagnostic_state.dwdx, + new_d.dwdy, + diagnostic_state.dwdy, + new_p.vn, + prognostic_state.vn, + new_p.w, + prognostic_state.w, + new_p.exner_pressure, + prognostic_state.exner_pressure, + new_p.theta_v, + prognostic_state.theta_v, + offset_provider={}, + ) + + +class Timeloop: + @classmethod + def name(cls): + return cls.__name__ + + def __init__( + self, + config: IconRunConfig, + diffusion: Diffusion, + atmo_non_hydro: DummyAtmoNonHydro, + edge_geometry: EdgeParams, + cell_geometry: CellParams, + ): + self.config = config + self.diffusion = diffusion + self.atmo_non_hydro = atmo_non_hydro + self.edges = edge_geometry + self.cells = cell_geometry + + def _full_name(self, func: Callable): + return ":".join((self.__class__.__name__, func.__name__)) + + def _timestep( + self, + diagnostic_state: DiagnosticState, + prognostic_state: PrognosticState, + ): + + self.atmo_non_hydro.do_dynamics_substepping( + self.config.dtime, diagnostic_state, prognostic_state + ) + self.diffusion.time_step( + diagnostic_state, + prognostic_state, + self.config.dtime, + self.edges.tangent_orientation, + self.edges.inverse_primal_edge_lengths, + self.edges.inverse_dual_edge_lengths, + self.edges.inverse_vertex_vertex_lengths, + self.edges.primal_normal_vert, + self.edges.dual_normal_vert, + self.edges.edge_areas, + self.cells.area, + ) + + def __call__( + self, + diagnostic_state: DiagnosticState, + prognostic_state: PrognosticState, + ): + log.info( + f"starting time loop for dtime={self.config.dtime} n_timesteps={self.config.n_time_steps}" + ) + log.info("running initial step to diffuse fields before timeloop starts") + self.diffusion.initial_step( + diagnostic_state, + prognostic_state, + self.config.dtime, + self.edges.tangent_orientation, + self.edges.inverse_primal_edge_lengths, + self.edges.inverse_dual_edge_lengths, + self.edges.inverse_vertex_vertex_lengths, + self.edges.primal_normal_vert, + self.edges.dual_normal_vert, + self.edges.edge_areas, + self.cells.area, + ) + log.info( + f"starting real time loop for dtime={self.config.dtime} n_timesteps={self.config.n_time_steps}" + ) + timer = Timer(self._full_name(self._timestep)) + for t in range(self.config.n_time_steps): + log.info(f"run timestep : {t}") + timer.start() + self._timestep(diagnostic_state, prognostic_state) + timer.capture() + timer.summary(True) + + +def initialize(n_time_steps, file_path: Path): + """ + Inititalize the driver run. + + "reads" in + - configuration + + - grid information + + - (serialized) input fields, initial + + Returns: + tl: configured timeloop, + prognostic_state: initial state fro prognostic and diagnostic variables + diagnostic_state: + """ + experiment_name = "mch_ch_r04b09_dsl" + log.info(f"reading configuration: experiment {experiment_name}") + config = read_config(experiment_name, n_time_steps=n_time_steps) + log.info("initializing the grid") + icon_grid = read_icon_grid(file_path) + log.info("reading input fields") + (edge_geometry, cell_geometry, vertical_geometry) = read_geometry_fields(file_path) + (metric_state, interpolation_state) = read_static_fields(file_path) + + log.info("initializing dycore") + diffusion_params = DiffusionParams(config.diffusion_config) + diffusion = Diffusion(run_program=False) + diffusion.init( + icon_grid, + config.diffusion_config, + diffusion_params, + vertical_geometry, + metric_state, + interpolation_state, + ) + + data_provider, diagnostic_state, prognostic_state = read_initial_state(file_path) + + atmo_non_hydro = DummyAtmoNonHydro(data_provider) + atmo_non_hydro.init(config=config.dycore_config) + + tl = Timeloop( + config=config.run_config, + diffusion=diffusion, + atmo_non_hydro=atmo_non_hydro, + edge_geometry=edge_geometry, + cell_geometry=cell_geometry, + ) + return tl, diagnostic_state, prognostic_state + + +@click.command() +@click.argument("input_path") +@click.option("--run_path", default="", help="folder for output") +@click.option( + "--n_steps", default=5, help="number of time steps to run, max 5 is supported" +) +def run( + input_path, + run_path, + n_steps, +): + """ + Run the driver. + + usage: python driver/dycore_driver.py ../../tests/ser_icondata/mch_ch_r04b09_dsl/ser_data + + steps: + 1. initialize model: + + a) load config + + b) initialize grid + + c) initialize/configure components ie "granules" + + d) setup the time loop + + 2. run time loop + + """ + start_time = datetime.now().astimezone(pytz.UTC) + configure_logging(run_path, start_time) + log.info(f"Starting ICON dycore run: {datetime.isoformat(start_time)}") + log.info(f"input args: input_path={input_path}, n_time_steps={n_steps}") + timeloop, diagnostic_state, prognostic_state = initialize(n_steps, Path(input_path)) + log.info("dycore configuring: DONE") + log.info("timeloop: START") + + timeloop(diagnostic_state, prognostic_state) + + log.info("timeloop: DONE") + + +if __name__ == "__main__": + run() diff --git a/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py b/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py new file mode 100644 index 000000000..0c55ae88f --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py @@ -0,0 +1,86 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from dataclasses import dataclass +from typing import Optional + +from icon4py.diffusion.diffusion import DiffusionConfig + + +@dataclass +class IconRunConfig: + n_time_steps: int = 5 + dtime: float = 600.0 + + +@dataclass +class AtmoNonHydroConfig: + n_substeps: int = 5 + + +@dataclass +class IconConfig: + run_config: IconRunConfig + diffusion_config: DiffusionConfig + dycore_config: AtmoNonHydroConfig + + +# TODO @magdalena move to io_utils? +def read_config(experiment: Optional[str], n_time_steps: int) -> IconConfig: + def _default_run_config(n_steps: int): + if n_steps > 5: + raise NotImplementedError("only five dummy timesteps available") + return IconRunConfig(n_time_steps=n_steps) + + def mch_ch_r04b09_diffusion_config(): + return DiffusionConfig( + diffusion_type=5, + hdiff_w=True, + hdiff_vn=True, + type_t_diffu=2, + type_vn_diffu=1, + hdiff_efdt_ratio=24.0, + hdiff_w_efdt_ratio=15.0, + smagorinski_scaling_factor=0.025, + zdiffu_t=True, + velocity_boundary_diffusion_denom=150.0, + max_nudging_coeff=0.075, + ) + + def _default_diffusion_config(): + return DiffusionConfig() + + def _default_config(n_steps): + run_config = _default_run_config(n_steps) + return run_config, _default_diffusion_config(), AtmoNonHydroConfig() + + def _mch_ch_r04b09_config(n_steps): + return ( + IconRunConfig(n_time_steps=n_steps, dtime=10.0), + mch_ch_r04b09_diffusion_config(), + AtmoNonHydroConfig(), + ) + + if experiment == "mch_ch_r04b09_dsl": + (model_run_config, diffusion_config, dycore_config) = _mch_ch_r04b09_config( + n_time_steps + ) + else: + (model_run_config, diffusion_config, dycore_config) = _default_config( + n_time_steps + ) + return IconConfig( + run_config=model_run_config, + diffusion_config=diffusion_config, + dycore_config=dycore_config, + ) diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py new file mode 100644 index 000000000..858f9c9e7 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -0,0 +1,167 @@ +# 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 logging +from datetime import datetime +from enum import Enum +from pathlib import Path + +from icon4py.diffusion.diagnostic_state import DiagnosticState +from icon4py.diffusion.horizontal import CellParams, EdgeParams +from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams +from icon4py.diffusion.interpolation_state import InterpolationState +from icon4py.diffusion.metric_state import MetricState +from icon4py.diffusion.prognostic_state import PrognosticState +from icon4py.testutils import serialbox_utils +from icon4py.testutils.serialbox_utils import IconSerialDataProvider + + +SIMULATION_START_DATE = "2021-06-20T12:00:10.000" +log = logging.getLogger(__name__) + + +class SerializationType(str, Enum): + SB = "serialbox" + NC = "netcdf" + + +def read_icon_grid(path: Path, ser_type=SerializationType.SB) -> IconGrid: + """ + Return IconGrid parsed from a given input type. + + Factory method that returns an icon grid dependeing on the ser_type. + + Args: + path: str - path where to find the input data + ser_type: str - type of input data. Currently only 'sb (serialbox)' is supported. It reads from ppser serialized test data + """ + if ser_type == SerializationType.SB: + return ( + serialbox_utils.IconSerialDataProvider( + "icon_pydycore", str(path.absolute()), False + ) + .from_savepoint_grid() + .construct_icon_grid() + ) + else: + raise NotImplementedError("Only ser_type='sb' is implemented so far.") + + +def read_initial_state( + gridfile_path: Path, +) -> tuple[IconSerialDataProvider, DiagnosticState, PrognosticState]: + """ + Read prognostic and diagnostic state from serialized data. + + Args: + gridfile_path: path the the serialized input data + + Returns: a tuple containing the data_provider, the initial diagnostic and prognostic state. + The data_provider is returned such that further timesteps of diagnostics and prognostics can be + read from within the dummy timeloop + + """ + data_provider = serialbox_utils.IconSerialDataProvider( + "icon_pydycore", str(gridfile_path), False + ) + init_savepoint = data_provider.from_savepoint_diffusion_init( + linit=True, date=SIMULATION_START_DATE + ) + prognostic_state = init_savepoint.construct_prognostics() + diagnostic_state = init_savepoint.construct_diagnostics() + return data_provider, diagnostic_state, prognostic_state + + +def read_geometry_fields( + path: Path, ser_type=SerializationType.SB +) -> tuple[EdgeParams, CellParams, VerticalModelParams]: + """ + Read fields containing grid properties. + + Args: + path: path to the serialized input data + ser_type: (optional) defualt so SB=serialbox, type of input data to be read + + Returns: a tuple containing fields describing edges, cells, vertical properties of the model + the data is originally obtained from the grid file (horizontal fields) or some special input files. + """ + if ser_type == SerializationType.SB: + sp = serialbox_utils.IconSerialDataProvider( + "icon_pydycore", str(path.absolute()), False + ).from_savepoint_grid() + edge_geometry = sp.construct_edge_geometry() + cell_geometry = sp.construct_cell_geometry() + vertical_geometry = VerticalModelParams( + vct_a=sp.vct_a(), rayleigh_damping_height=12500 + ) + return edge_geometry, cell_geometry, vertical_geometry + else: + raise NotImplementedError("Only ser_type='sb' is implemented so far.") + + +def read_static_fields( + path: Path, ser_type=SerializationType.SB +) -> tuple[MetricState, InterpolationState]: + """ + Read fields for metric and interpolation state. + + Args: + path: path to the serialized input data + ser_type: (optional) defualt so SB=serialbox, type of input data to be read + + Returns: a tuple containing the metric_state and interpolation state, + the fields are precalculated in the icon setup. + + """ + if ser_type == SerializationType.SB: + sp = serialbox_utils.IconSerialDataProvider( + "icon_pydycore", str(path.absolute()), False + ).from_savepoint_diffusion_init(linit=True, date=SIMULATION_START_DATE) + metric_state = sp.construct_metric_state() + interpolation_state = sp.construct_interpolation_state() + return metric_state, interpolation_state + else: + raise NotImplementedError("Only ser_type='sb' is implemented so far.") + + +def configure_logging(run_path: str, start_time): + """ + Configure logging. + + Log output is sent to console and to a file. + + Args: + run_path: path to the output folder where the logfile should be stored + start_time: start time of the model run + + """ + run_dir = ( + Path(run_path).absolute() if run_path else Path(__file__).absolute().parent + ) + run_dir.mkdir(exist_ok=True) + logfile = run_dir.joinpath( + f"dummy_dycore_driver_{datetime.isoformat(start_time)}.log" + ) + logfile.touch(exist_ok=True) + logging.basicConfig( + level=logging.DEBUG, + format="%(asctime)s %(filename)-20s (%(lineno)-4d) : %(funcName)-20s: %(levelname)-8s %(message)s", + filemode="w", + filename=logfile, + ) + console_handler = logging.StreamHandler() + formatter = logging.Formatter( + "%(asctime)s %(filename)-20s : %(funcName)-20s: %(levelname)-8s %(message)s" + ) + console_handler.setFormatter(formatter) + console_handler.setLevel(logging.INFO) + logging.getLogger("").addHandler(console_handler) diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 99e6db30b..ba7d2b336 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -13,25 +13,10 @@ import tarfile from pathlib import Path -import numpy as np import pytest import wget -from icon4py.common.dimension import ( - C2E2CDim, - C2E2CODim, - C2EDim, - CellDim, - E2C2VDim, - E2CDim, - E2VDim, - EdgeDim, - V2EDim, - VertexDim, -) from icon4py.diffusion.diffusion import DiffusionConfig -from icon4py.diffusion.horizontal import HorizontalMeshSize -from icon4py.diffusion.icon_grid import IconGrid, MeshConfig, VerticalMeshConfig from icon4py.testutils.serialbox_utils import IconSerialDataProvider @@ -123,46 +108,13 @@ def diffusion_savepoint_exit(data_provider, step_date_exit): @pytest.fixture -def icon_grid(data_provider): +def icon_grid(grid_savepoint): """ Load the icon grid from an ICON savepoint. - Uses the default save_point from 'savepoint_init' fixture, however these data don't change for - different time steps. - """ - sp = data_provider.from_savepoint_grid() - sp_meta = sp.get_metadata("nproma", "nlev", "num_vert", "num_cells", "num_edges") - - cell_starts = sp.cells_start_index() - cell_ends = sp.cells_end_index() - vertex_starts = sp.vertex_start_index() - vertex_ends = sp.vertex_end_index() - edge_starts = sp.edge_start_index() - edge_ends = sp.edge_end_index() - - config = MeshConfig( - HorizontalMeshSize( - num_vertices=sp_meta["nproma"], # or rather "num_vert" - num_cells=sp_meta["nproma"], # or rather "num_cells" - num_edges=sp_meta["nproma"], # or rather "num_edges" - ), - VerticalMeshConfig(num_lev=sp_meta["nlev"]), - ) - - c2e2c = sp.c2e2c() - c2e2c0 = np.column_stack((c2e2c, (np.asarray(range(c2e2c.shape[0]))))) - grid = ( - IconGrid() - .with_config(config) - .with_start_end_indices(VertexDim, vertex_starts, vertex_ends) - .with_start_end_indices(EdgeDim, edge_starts, edge_ends) - .with_start_end_indices(CellDim, cell_starts, cell_ends) - .with_connectivities( - {C2EDim: sp.c2e(), E2CDim: sp.e2c(), C2E2CDim: c2e2c, C2E2CODim: c2e2c0} - ) - .with_connectivities({E2VDim: sp.e2v(), V2EDim: sp.v2e(), E2C2VDim: sp.e2c2v()}) - ) - return grid + Uses the special grid_savepoint that contains data from p_patch + """ + return grid_savepoint.construct_icon_grid() @pytest.fixture diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 24316b5a9..1ee598828 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -14,13 +14,12 @@ import numpy as np import pytest -from icon4py.common.dimension import ECVDim, KDim, VertexDim -from icon4py.diffusion.diagnostic_state import DiagnosticState -from icon4py.diffusion.diffusion import Diffusion, DiffusionParams, VectorTuple +from icon4py.common.dimension import KDim, VertexDim +from icon4py.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.diffusion.horizontal import CellParams, EdgeParams from icon4py.diffusion.icon_grid import VerticalModelParams from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState -from icon4py.diffusion.prognostic_state import PrognosticState from icon4py.diffusion.utils import ( _en_smag_fac_for_zero_nshift, _setup_runtime_diff_multfac_vn, @@ -31,7 +30,7 @@ ) from icon4py.testutils.serialbox_utils import IconDiffusionInitSavepoint from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import as_1D_sparse_field, random_field, zero_field +from icon4py.testutils.utils import random_field, zero_field def test_scale_k(): @@ -460,18 +459,12 @@ def test_run_diffusion_single_step( damping_height, ): ( - cell_areas, - diagnostic_state, dtime, - dual_normal_vert, - edge_areas, + cell_geometry, + edge_geometry, + diagnostic_state, interpolation_state, - inverse_dual_edge_length, - inverse_primal_edge_lengths, - inverse_vertical_vertex_lengths, metric_state, - orientation, - primal_normal_vert, prognostic_state, ) = _read_fields(diffusion_savepoint_init, grid_savepoint) @@ -494,14 +487,14 @@ def test_run_diffusion_single_step( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime, - tangent_orientation=orientation, - inverse_primal_edge_lengths=inverse_primal_edge_lengths, - inverse_dual_edge_length=inverse_dual_edge_length, - inverse_vert_vert_lengths=inverse_vertical_vertex_lengths, - primal_normal_vert=primal_normal_vert, - dual_normal_vert=dual_normal_vert, - edge_areas=edge_areas, - cell_areas=cell_areas, + tangent_orientation=edge_geometry.tangent_orientation, + inverse_primal_edge_lengths=edge_geometry.inverse_primal_edge_lengths, + inverse_dual_edge_length=edge_geometry.inverse_dual_edge_lengths, + inverse_vert_vert_lengths=edge_geometry.inverse_vertex_vertex_lengths, + primal_normal_vert=edge_geometry.primal_normal_vert, + dual_normal_vert=edge_geometry.dual_normal_vert, + edge_areas=edge_geometry.edge_areas, + cell_areas=cell_geometry.area, ) icon_result_exner = diffusion_savepoint_exit.exner() @@ -520,70 +513,30 @@ def test_run_diffusion_single_step( def _read_fields(diffusion_savepoint_init, grid_savepoint): - - grg = diffusion_savepoint_init.geofac_grg() - interpolation_state = InterpolationState( - e_bln_c_s=diffusion_savepoint_init.e_bln_c_s(), - rbf_coeff_1=diffusion_savepoint_init.rbf_vec_coeff_v1(), - rbf_coeff_2=diffusion_savepoint_init.rbf_vec_coeff_v2(), - geofac_div=diffusion_savepoint_init.geofac_div(), - geofac_n2s=diffusion_savepoint_init.geofac_n2s(), - geofac_grg_x=grg[0], - geofac_grg_y=grg[1], - nudgecoeff_e=diffusion_savepoint_init.nudgecoeff_e(), - ) - metric_state = MetricState( - mask_hdiff=diffusion_savepoint_init.mask_diff(), - theta_ref_mc=diffusion_savepoint_init.theta_ref_mc(), - wgtfac_c=diffusion_savepoint_init.wgtfac_c(), - zd_intcoef=diffusion_savepoint_init.zd_intcoef(), - zd_vertidx=diffusion_savepoint_init.zd_vertoffset(), - zd_diffcoef=diffusion_savepoint_init.zd_diffcoef(), - ) - diagnostic_state = DiagnosticState( - hdef_ic=diffusion_savepoint_init.hdef_ic(), - div_ic=diffusion_savepoint_init.div_ic(), - dwdx=diffusion_savepoint_init.dwdx(), - dwdy=diffusion_savepoint_init.dwdy(), - ) - prognostic_state = PrognosticState( - w=diffusion_savepoint_init.w(), - vn=diffusion_savepoint_init.vn(), - exner_pressure=diffusion_savepoint_init.exner(), - theta_v=diffusion_savepoint_init.theta_v(), - ) dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") - orientation = grid_savepoint.tangent_orientation() - inverse_primal_edge_lengths = grid_savepoint.inverse_primal_edge_lengths() - inverse_vert_vert_lengths = grid_savepoint.inv_vert_vert_length() - inverse_dual_edge_length = grid_savepoint.inv_dual_edge_length() - primal_normal_vert: VectorTuple = ( - as_1D_sparse_field(grid_savepoint.primal_normal_vert_x(), ECVDim), - as_1D_sparse_field(grid_savepoint.primal_normal_vert_y(), ECVDim), - ) - dual_normal_vert: VectorTuple = ( - as_1D_sparse_field(grid_savepoint.dual_normal_vert_x(), ECVDim), - as_1D_sparse_field(grid_savepoint.dual_normal_vert_y(), ECVDim), - ) - edge_areas = grid_savepoint.edge_areas() - cell_areas = grid_savepoint.cell_areas() + edge_geometry, cell_geometry = _read_geometry_fields(grid_savepoint) + + interpolation_state = diffusion_savepoint_init.construct_interpolation_state() + metric_state = diffusion_savepoint_init.construct_metric_state() + diagnostic_state = diffusion_savepoint_init.construct_diagnostics() + prognostic_state = diffusion_savepoint_init.construct_prognostics() return ( - cell_areas, - diagnostic_state, dtime, - dual_normal_vert, - edge_areas, + cell_geometry, + edge_geometry, + diagnostic_state, interpolation_state, - inverse_dual_edge_length, - inverse_primal_edge_lengths, - inverse_vert_vert_lengths, metric_state, - orientation, - primal_normal_vert, prognostic_state, ) +def _read_geometry_fields(grid_savepoint): + edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() + cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() + return edge_geometry, cell_geometry + + @pytest.mark.skip("fix: diffusion_stencil_15") @pytest.mark.datatest def test_diffusion_five_steps( @@ -597,18 +550,12 @@ def test_diffusion_five_steps( step_date_exit="2021-06-20T12:01:00.000", ): ( - cell_areas, - diagnostic_state, dtime, - dual_normal_vert, - edge_areas, + cell_geometry, + edge_geometry, + diagnostic_state, interpolation_state, - inverse_dual_edge_length, - inverse_primal_edge_lengths, - inverse_vert_vert_lengths, metric_state, - orientation, - primal_normal_vert, prognostic_state, ) = _read_fields(diffusion_savepoint_init, grid_savepoint) @@ -631,28 +578,28 @@ def test_diffusion_five_steps( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime, - tangent_orientation=orientation, - inverse_primal_edge_lengths=inverse_primal_edge_lengths, - inverse_dual_edge_length=inverse_dual_edge_length, - inverse_vert_vert_lengths=inverse_vert_vert_lengths, - primal_normal_vert=primal_normal_vert, - dual_normal_vert=dual_normal_vert, - edge_areas=edge_areas, - cell_areas=cell_areas, + tangent_orientation=edge_geometry.tangent_orientation, + inverse_primal_edge_lengths=edge_geometry.inverse_primal_edge_lengths, + inverse_dual_edge_length=edge_geometry.inverse_dual_edge_lengths, + inverse_vert_vert_lengths=edge_geometry.inverse_vertex_vertex_lengths, + primal_normal_vert=edge_geometry.primal_normal_vert, + dual_normal_vert=edge_geometry.dual_normal_vert, + edge_areas=edge_geometry.edge_areas, + cell_areas=cell_geometry.area, ) for _ in range(4): diffusion.time_step( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime, - tangent_orientation=orientation, - inverse_primal_edge_lengths=inverse_primal_edge_lengths, - inverse_dual_edge_length=inverse_dual_edge_length, - inverse_vert_vert_lengths=inverse_vert_vert_lengths, - primal_normal_vert=primal_normal_vert, - dual_normal_vert=dual_normal_vert, - edge_areas=edge_areas, - cell_areas=cell_areas, + tangent_orientation=edge_geometry.tangent_orientation, + inverse_primal_edge_lengths=edge_geometry.inverse_primal_edge_lengths, + inverse_dual_edge_length=edge_geometry.inverse_dual_edge_lengths, + inverse_vert_vert_lengths=edge_geometry.inverse_vertex_vertex_lengths, + primal_normal_vert=edge_geometry.primal_normal_vert, + dual_normal_vert=edge_geometry.dual_normal_vert, + edge_areas=edge_geometry.edge_areas, + cell_areas=cell_geometry.area, ) icon_result_exner = diffusion_savepoint_exit.exner() diff --git a/atm_dyn_iconam/tests/test_dycore_driver.py b/atm_dyn_iconam/tests/test_dycore_driver.py new file mode 100644 index 000000000..089dd8e0f --- /dev/null +++ b/atm_dyn_iconam/tests/test_dycore_driver.py @@ -0,0 +1,68 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.common.dimension import CellDim, EdgeDim, KDim +from icon4py.diffusion.utils import copy_diagnostic_and_prognostics +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + + +def test_copy_diagnostic_and_prognostics(): + mesh = SimpleMesh() + f1_in = random_field(mesh, CellDim, KDim) + f1_out = zero_field(mesh, CellDim, KDim) + f2_in = random_field(mesh, CellDim, KDim) + f2_out = zero_field(mesh, CellDim, KDim) + f3_in = random_field(mesh, CellDim, KDim) + f3_out = zero_field(mesh, CellDim, KDim) + f4_in = random_field(mesh, CellDim, KDim) + f4_out = zero_field(mesh, CellDim, KDim) + f5_in = random_field(mesh, EdgeDim, KDim) + f5_out = zero_field(mesh, EdgeDim, KDim) + f6_in = random_field(mesh, CellDim, KDim) + f6_out = zero_field(mesh, CellDim, KDim) + f7_in = random_field(mesh, CellDim, KDim) + f7_out = zero_field(mesh, CellDim, KDim) + f8_in = random_field(mesh, CellDim, KDim) + f8_out = zero_field(mesh, CellDim, KDim) + + copy_diagnostic_and_prognostics( + f1_in, + f1_out, + f2_in, + f2_out, + f3_in, + f3_out, + f4_in, + f4_out, + f5_in, + f5_out, + f6_in, + f6_out, + f7_in, + f7_out, + f8_in, + f8_out, + offset_provider={}, + ) + + assert np.allclose(f1_in, f1_out) + assert np.allclose(f2_in, f2_out) + assert np.allclose(f3_in, f3_out) + assert np.allclose(f4_in, f4_out) + assert np.allclose(f5_in, f5_out) + assert np.allclose(f6_in, f6_out) + assert np.allclose(f7_in, f7_out) + assert np.allclose(f8_in, f8_out) diff --git a/atm_dyn_iconam/tests/test_io_utils.py b/atm_dyn_iconam/tests/test_io_utils.py new file mode 100644 index 000000000..a03d810db --- /dev/null +++ b/atm_dyn_iconam/tests/test_io_utils.py @@ -0,0 +1,117 @@ +# 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 pathlib + +import pytest + +from icon4py.diffusion.horizontal import CellParams, EdgeParams +from icon4py.diffusion.icon_grid import VerticalModelParams +from icon4py.driver.io_utils import ( + SerializationType, + read_geometry_fields, + read_icon_grid, + read_static_fields, +) + + +test_data_path = pathlib.Path(__file__).parent.joinpath( + "./ser_icondata/mch_ch_r04b09_dsl/ser_data" +) + + +@pytest.mark.parametrize( + "read_fun", (read_geometry_fields, read_static_fields, read_icon_grid) +) +def test_read_geometry_fields_not_implemented_type(read_fun): + with pytest.raises(NotImplementedError, match=r"Only ser_type='sb'"): + read_fun(path=test_data_path, ser_type=SerializationType.NC) + + +def assert_grid_size_and_connectivities(grid): + assert grid.num_edges() > 0 + assert grid.num_cells() > 0 + assert grid.num_vertices() > 0 + assert grid.get_e2v_connectivity() + assert grid.get_v2e_connectivity() + assert grid.get_c2e_connectivity() + assert grid.get_e2c_connectivity() + assert grid.get_e2c2v_connectivity() + assert grid.get_c2e2c_connectivity() + assert grid.get_e2ecv_connectivity() + + +@pytest.mark.datatest +def test_read_icon_grid_for_type_sb(): + grid = read_icon_grid(test_data_path, ser_type=SerializationType.SB) + assert_grid_size_and_connectivities(grid) + + +@pytest.mark.datatest +def test_read_static_fields_for_type_sb(): + metric_state, interpolation_state = read_static_fields( + test_data_path, ser_type=SerializationType.SB + ) + assert_metric_state_fields(metric_state) + assert_interpolation_state_fields(interpolation_state) + + +@pytest.mark.datatest +def test_read_geometry_fields_for_type_sb(): + edge_geometry, cell_geometry, vertical_geometry = read_geometry_fields( + test_data_path, ser_type=SerializationType.SB + ) + assert_edge_geometry_fields(edge_geometry) + assert_cell_geometry_fields(cell_geometry) + assert_vertical_params(vertical_geometry) + + +def assert_vertical_params(vertical_geometry: VerticalModelParams): + assert vertical_geometry.physical_heights + assert vertical_geometry.index_of_damping_layer > 0 + assert vertical_geometry.rayleigh_damping_height > 0 + + +def assert_cell_geometry_fields(cell_geometry: CellParams): + assert cell_geometry.area + + +def assert_edge_geometry_fields(edge_geometry: EdgeParams): + assert edge_geometry.edge_areas + assert edge_geometry.primal_normal_vert + assert edge_geometry.inverse_primal_edge_lengths + assert edge_geometry.tangent_orientation + assert edge_geometry.inverse_dual_edge_lengths + assert edge_geometry.dual_normal_vert + + +def assert_metric_state_fields(metric_state): + assert metric_state.wgtfac_c + assert metric_state.zd_intcoef + assert metric_state.zd_diffcoef + assert metric_state.theta_ref_mc + assert metric_state.mask_hdiff + assert metric_state.zd_vertidx + + +def assert_interpolation_state_fields(interpolation_state): + assert interpolation_state.geofac_n2s + assert interpolation_state.e_bln_c_s + assert interpolation_state.nudgecoeff_e + assert interpolation_state.geofac_n2s_nbh + assert interpolation_state.geofac_div + assert interpolation_state.geofac_grg_y + assert interpolation_state.geofac_grg_x + assert interpolation_state.rbf_coeff_2 + assert interpolation_state.rbf_coeff_1 + assert interpolation_state.geofac_n2s_c diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 9a037e60c..ad312d679 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -10,6 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import logging import numpy as np import serialbox as ser @@ -23,24 +24,40 @@ C2EDim, CellDim, E2C2VDim, + E2CDim, + E2VDim, + ECVDim, EdgeDim, KDim, V2EDim, VertexDim, ) +from icon4py.diffusion.diagnostic_state import DiagnosticState +from icon4py.diffusion.diffusion import VectorTuple +from icon4py.diffusion.horizontal import ( + CellParams, + EdgeParams, + HorizontalMeshSize, +) +from icon4py.diffusion.icon_grid import IconGrid, MeshConfig, VerticalMeshConfig +from icon4py.diffusion.interpolation_state import InterpolationState +from icon4py.diffusion.metric_state import MetricState +from icon4py.diffusion.prognostic_state import PrognosticState +from icon4py.testutils.utils import as_1D_sparse_field class IconSavepoint: def __init__(self, sp: ser.Savepoint, ser: ser.Serializer): self.savepoint = sp self.serializer = ser + self.log = logging.getLogger((__name__)) - def print_meta_info(self): - print(self.savepoint.metainfo) + def log_meta_info(self): + self.log.info(self.savepoint.metainfo) def _get_field(self, name, *dimensions, dtype=float): buffer = np.squeeze(self.serializer.read(name, self.savepoint).astype(dtype)) - print(f"{name} {buffer.shape}") + self.log.debug(f"{name} {buffer.shape}") return np_as_located_field(*dimensions)(buffer) def get_metadata(self, *names): @@ -68,7 +85,7 @@ def _read_int32(self, name: str): def read_int(self, name: str): buffer = self.serializer.read(name, self.savepoint).astype(int) - print(f"{name} {buffer.shape}") + self.log.debug(f"{name} {buffer.shape}") return buffer @@ -126,8 +143,8 @@ def edge_end_index(self): # one off accounts for being exclusive [from:to) return self.serializer.read("e_end_index", self.savepoint) - def print_connectivity_info(name: str, ar: np.ndarray): - print(f" connectivity {name} {ar.shape}") + def print_connectivity_info(self, name: str, ar: np.ndarray): + self.log.debug(f" connectivity {name} {ar.shape}") def refin_ctrl(self, dim: Dimension): if dim == CellDim: @@ -145,7 +162,7 @@ def c2e(self): def _get_connectiviy_array(self, name: str): connectivity = self.serializer.read(name, self.savepoint) - 1 - print(f" connectivity {name} : {connectivity.shape}") + self.log.debug(f" connectivity {name} : {connectivity.shape}") return connectivity def c2e2c(self): @@ -157,7 +174,7 @@ def e2c(self): def e2v(self): # array "e2v" is actually e2c2v v_ = self._get_connectiviy_array("e2v")[:, 0:2] - print(f"real e2v {v_.shape}") + self.log.debug(f"real e2v {v_.shape}") return v_ def e2c2v(self): @@ -167,6 +184,70 @@ def e2c2v(self): def v2e(self): return self._get_connectiviy_array("v2e") + def construct_icon_grid(self) -> IconGrid: + sp_meta = self.get_metadata( + "nproma", "nlev", "num_vert", "num_cells", "num_edges" + ) + cell_starts = self.cells_start_index() + cell_ends = self.cells_end_index() + vertex_starts = self.vertex_start_index() + vertex_ends = self.vertex_end_index() + edge_starts = self.edge_start_index() + edge_ends = self.edge_end_index() + config = MeshConfig( + HorizontalMeshSize( + num_vertices=sp_meta["nproma"], # or rather "num_vert" + num_cells=sp_meta["nproma"], # or rather "num_cells" + num_edges=sp_meta["nproma"], # or rather "num_edges" + ), + VerticalMeshConfig(num_lev=sp_meta["nlev"]), + ) + c2e2c = self.c2e2c() + c2e2c0 = np.column_stack((c2e2c, (np.asarray(range(c2e2c.shape[0]))))) + grid = ( + IconGrid() + .with_config(config) + .with_start_end_indices(VertexDim, vertex_starts, vertex_ends) + .with_start_end_indices(EdgeDim, edge_starts, edge_ends) + .with_start_end_indices(CellDim, cell_starts, cell_ends) + .with_connectivities( + { + C2EDim: self.c2e(), + E2CDim: self.e2c(), + C2E2CDim: c2e2c, + C2E2CODim: c2e2c0, + } + ) + .with_connectivities( + {E2VDim: self.e2v(), V2EDim: self.v2e(), E2C2VDim: self.e2c2v()} + ) + ) + return grid + + def construct_edge_geometry(self) -> EdgeParams: + primal_normal_vert: VectorTuple = ( + as_1D_sparse_field(self.primal_normal_vert_x(), ECVDim), + as_1D_sparse_field(self.primal_normal_vert_y(), ECVDim), + ) + dual_normal_vert: VectorTuple = ( + as_1D_sparse_field(self.dual_normal_vert_x(), ECVDim), + as_1D_sparse_field(self.dual_normal_vert_y(), ECVDim), + ) + return EdgeParams( + tangent_orientation=self.tangent_orientation(), + inverse_primal_edge_lengths=self.inverse_primal_edge_lengths(), + inverse_dual_edge_lengths=self.inv_dual_edge_length(), + inverse_vertex_vertex_lengths=self.inv_vert_vert_length(), + primal_normal_vert_x=primal_normal_vert[0], + primal_normal_vert_y=primal_normal_vert[1], + dual_normal_vert_x=dual_normal_vert[0], + dual_normal_vert_y=dual_normal_vert[1], + edge_areas=self.edge_areas(), + ) + + def construct_cell_geometry(self) -> CellParams: + return CellParams(area=self.cell_areas()) + class IconDiffusionInitSavepoint(IconSavepoint): def hdef_ic(self): @@ -269,6 +350,45 @@ def diff_multfac_w(self): def diff_multfac_vn(self): return self.serializer.read("diff_multfac_vn", self.savepoint) + def construct_interpolation_state(self) -> InterpolationState: + grg = self.geofac_grg() + return InterpolationState( + e_bln_c_s=self.e_bln_c_s(), + rbf_coeff_1=self.rbf_vec_coeff_v1(), + rbf_coeff_2=self.rbf_vec_coeff_v2(), + geofac_div=self.geofac_div(), + geofac_n2s=self.geofac_n2s(), + geofac_grg_x=grg[0], + geofac_grg_y=grg[1], + nudgecoeff_e=self.nudgecoeff_e(), + ) + + def construct_metric_state(self) -> MetricState: + return MetricState( + mask_hdiff=self.mask_diff(), + theta_ref_mc=self.theta_ref_mc(), + wgtfac_c=self.wgtfac_c(), + zd_intcoef=self.zd_intcoef(), + zd_vertidx=self.zd_vertoffset(), + zd_diffcoef=self.zd_diffcoef(), + ) + + def construct_prognostics(self) -> PrognosticState: + return PrognosticState( + w=self.w(), + vn=self.vn(), + exner_pressure=self.exner(), + theta_v=self.theta_v(), + ) + + def construct_diagnostics(self) -> DiagnosticState: + return DiagnosticState( + hdef_ic=self.hdef_ic(), + div_ic=self.div_ic(), + dwdx=self.dwdx(), + dwdy=self.dwdy(), + ) + class IconDiffusionExitSavepoint(IconSavepoint): def vn(self): @@ -290,11 +410,12 @@ def __init__(self, fname_prefix, path=".", do_print=False): self.serializer: ser.Serializer = None self.file_path: str = path self.fname = f"{fname_prefix}_rank{str(self.rank)}" + self.log = logging.getLogger(__name__) self._init_serializer(do_print) def _init_serializer(self, do_print: bool): if not self.fname: - print(" WARNING: no filename! closing serializer") + self.log.warning(" WARNING: no filename! closing serializer") self.serializer = ser.Serializer( ser.OpenModeKind.Read, self.file_path, self.fname ) @@ -302,8 +423,8 @@ def _init_serializer(self, do_print: bool): self.print_info() def print_info(self): - print(f"SAVEPOINTS: {self.serializer.savepoint_list()}") - print(f"FIELDNAMES: {self.serializer.fieldnames()}") + self.log.info(f"SAVEPOINTS: {self.serializer.savepoint_list()}") + self.log.info(f"FIELDNAMES: {self.serializer.fieldnames()}") def from_savepoint_grid(self) -> IconGridSavePoint: savepoint = self.serializer.savepoint["icon-grid"].id[1].as_savepoint() From 583100cbe9b395ba4d4e3cc54cc8836f76fc4b33 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Mon, 1 May 2023 19:10:48 +0200 Subject: [PATCH 093/263] - fix test after moving: vertical_config out of vertical model parameters - expected fails of e2c2v test --- .../src/icon4py/diffusion/diffusion.py | 91 +++++---- atm_dyn_iconam/tests/test_diffusion.py | 29 ++- atm_dyn_iconam/tests/test_icon_grid.py | 188 ++++++++++-------- common/src/icon4py/common/dimension.py | 2 +- common/src/icon4py/grid/grid_manager.py | 51 +++-- common/src/icon4py/grid/horizontal.py | 2 +- common/src/icon4py/grid/icon_grid.py | 22 +- common/tests/test_gridmanager.py | 48 +++-- common/tests/test_vertical.py | 7 +- .../py2f/wrappers/diffusion_wrapper.py | 5 +- testutils/src/icon4py/testutils/fixtures.py | 2 +- .../src/icon4py/testutils/simple_mesh.py | 32 ++- 12 files changed, 287 insertions(+), 192 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 19d57cd4a..f9669dd6b 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -565,54 +565,56 @@ def time_step( else: print("run program") - cell_start_nudging_minus1 = self.grid.get_start_index(CellDim, - HorizontalMarkerIndex.nudging( - CellDim) - 1) + cell_start_nudging_minus1 = self.grid.get_start_index( + CellDim, HorizontalMarkerIndex.nudging(CellDim) - 1 + ) - cell_startindex_interior = self.grid.get_start_index(CellDim, - HorizontalMarkerIndex.interior(CellDim)) + cell_startindex_interior = self.grid.get_start_index( + CellDim, HorizontalMarkerIndex.interior(CellDim) + ) cell_startindex_nudging = self.grid.get_start_index( - CellDim, - HorizontalMarkerIndex.nudging(CellDim)) + CellDim, HorizontalMarkerIndex.nudging(CellDim) + ) - cell_endindex_local_plus1 = self.grid.get_end_index(CellDim, HorizontalMarkerIndex.local(CellDim) - 1) + cell_endindex_local_plus1 = self.grid.get_end_index( + CellDim, HorizontalMarkerIndex.local(CellDim) - 1 + ) cell_endindex_local = self.grid.get_end_index( - CellDim, HorizontalMarkerIndex.local(CellDim)) + CellDim, HorizontalMarkerIndex.local(CellDim) + ) - edge_startindex_nudging_plus1 = self.grid.get_start_index(EdgeDim, - HorizontalMarkerIndex.nudging( - EdgeDim) + 1) + edge_startindex_nudging_plus1 = self.grid.get_start_index( + EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1 + ) edge_endindex_local = self.grid.get_end_index( EdgeDim, HorizontalMarkerIndex.local(EdgeDim), ) edge_start_lb_plus4 = self.grid.get_start_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4 + EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4 ) - - edge_startindex_nudging_minus1 = self.grid.get_start_index(EdgeDim, - HorizontalMarkerIndex.nudging( - EdgeDim) - 1) + edge_startindex_nudging_minus1 = self.grid.get_start_index( + EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) - 1 + ) edge_endindex_local_minus2 = self.grid.get_end_index( EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2, ) - vertex_startindex_lb_plus3 = self.grid.get_start_index(VertexDim, - HorizontalMarkerIndex.lateral_boundary( - VertexDim) + 3) + vertex_startindex_lb_plus3 = self.grid.get_start_index( + VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3 + ) vertex_endindex_local = self.grid.get_end_index( VertexDim, HorizontalMarkerIndex.local(VertexDim), ) - vertex_startindex_lb_plus1 = self.grid.get_start_index(VertexDim, - HorizontalMarkerIndex.lateral_boundary( - VertexDim) + 1) + vertex_startindex_lb_plus1 = self.grid.get_start_index( + VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1 + ) vertex_endindex_local_minus1 = self.grid.get_end_index( VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, @@ -748,20 +750,27 @@ def _do_diffusion_step( klevels = self.grid.n_lev() k_end_minus2 = klevels - 2 - cell_start_nudging_minus1 = self.grid.get_start_index(CellDim, HorizontalMarkerIndex.nudging(CellDim) - 1) - + cell_start_nudging_minus1 = self.grid.get_start_index( + CellDim, HorizontalMarkerIndex.nudging(CellDim) - 1 + ) - cell_startindex_interior = self.grid.get_start_index(CellDim,HorizontalMarkerIndex.interior(CellDim)) + cell_startindex_interior = self.grid.get_start_index( + CellDim, HorizontalMarkerIndex.interior(CellDim) + ) cell_start_nudging = self.grid.get_start_index( - CellDim, - HorizontalMarkerIndex.nudging(CellDim)) - + CellDim, HorizontalMarkerIndex.nudging(CellDim) + ) - cell_end_local_plus1 = self.grid.get_end_index(CellDim, HorizontalMarkerIndex.local(CellDim) - 1) + cell_end_local_plus1 = self.grid.get_end_index( + CellDim, HorizontalMarkerIndex.local(CellDim) - 1 + ) cell_end_local = self.grid.get_end_index( - CellDim, HorizontalMarkerIndex.local(CellDim)) + CellDim, HorizontalMarkerIndex.local(CellDim) + ) - edge_start_nudging_plus_one = self.grid.get_start_index(EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1) + edge_start_nudging_plus_one = self.grid.get_start_index( + EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1 + ) edge_end_local = self.grid.get_end_index( EdgeDim, @@ -769,25 +778,29 @@ def _do_diffusion_step( ) edge_start_lb_plus4 = self.grid.get_start_index( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4 + EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4 ) - edge_start_nudging_minus1 = self.grid.get_start_index(EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) - 1) + edge_start_nudging_minus1 = self.grid.get_start_index( + EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) - 1 + ) edge_end_local_minus2 = self.grid.get_end_index( EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2, ) - - vertex_start_local_boundary_plus3 = self.grid.get_start_index(VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3) + vertex_start_local_boundary_plus3 = self.grid.get_start_index( + VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3 + ) vertex_end_local = self.grid.get_end_index( VertexDim, HorizontalMarkerIndex.local(VertexDim), ) - vertex_start_local_boundary_plus1 = self.grid.get_start_index(VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1) + vertex_start_local_boundary_plus1 = self.grid.get_start_index( + VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1 + ) vertex_end_local_minus1 = self.grid.get_end_index( VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 7afbeeb53..5c93eafee 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -349,12 +349,17 @@ def _verify_init_values_against_savepoint( @pytest.mark.datatest @pytest.mark.parametrize("linit", [True]) def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( - diffusion_savepoint_init, r04b09_diffusion_config, r04b09_dsl_gridfile, grid_savepoint + diffusion_savepoint_init, + r04b09_diffusion_config, + r04b09_dsl_gridfile, # noqa: F811 + grid_savepoint, # noqa: F811 ): savepoint = diffusion_savepoint_init config = r04b09_diffusion_config - icon_grid = init_grid_manager(r04b09_dsl_gridfile, grid_savepoint.num(KDim)).get_grid() + icon_grid = init_grid_manager( + r04b09_dsl_gridfile, grid_savepoint.num(KDim) + ).get_grid() params = DiffusionParams(config) expected_diff_multfac_vn = savepoint.diff_multfac_vn() @@ -380,7 +385,7 @@ def test_verify_diffusion_init_against_first_regular_savepoint( diffusion_savepoint_init, grid_savepoint, # noqa: F811 r04b09_diffusion_config, - r04b09_dsl_gridfile, + r04b09_dsl_gridfile, # noqa: F811 damping_height, ): config = r04b09_diffusion_config @@ -408,7 +413,9 @@ def test_verify_diffusion_init_against_first_regular_savepoint( zd_vertidx=savepoint.zd_vertoffset(), zd_diffcoef=savepoint.zd_diffcoef(), ) - icon_grid = init_grid_manager(r04b09_dsl_gridfile, grid_savepoint.num(KDim)).get_grid() + icon_grid = init_grid_manager( + r04b09_dsl_gridfile, grid_savepoint.num(KDim) + ).get_grid() diffusion = Diffusion() diffusion.init( config=config, @@ -427,7 +434,7 @@ def test_verify_diffusion_init_against_first_regular_savepoint( def test_verify_diffusion_init_against_other_regular_savepoint( r04b09_diffusion_config, grid_savepoint, # noqa: F811 - r04b09_dsl_gridfile, + r04b09_dsl_gridfile, # noqa: F811 diffusion_savepoint_init, damping_height, ): @@ -435,7 +442,9 @@ def test_verify_diffusion_init_against_other_regular_savepoint( additional_parameters = DiffusionParams(config) vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) - icon_grid = init_grid_manager(r04b09_dsl_gridfile, grid_savepoint.num(KDim)).get_grid() + icon_grid = init_grid_manager( + r04b09_dsl_gridfile, grid_savepoint.num(KDim) + ).get_grid() grg = diffusion_savepoint_init.geofac_grg() interpolation_state = InterpolationState( e_bln_c_s=diffusion_savepoint_init.e_bln_c_s(), @@ -469,15 +478,15 @@ def test_verify_diffusion_init_against_other_regular_savepoint( _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) -#@pytest.mark.skip("fix: diffusion_stencil_15") @pytest.mark.parametrize("run_with_program", [True, False]) @pytest.mark.datatest +@pytest.mark.xfail() def test_run_diffusion_single_step( run_with_program, diffusion_savepoint_init, diffusion_savepoint_exit, grid_savepoint, # noqa: F811 - r04b09_dsl_gridfile, + r04b09_dsl_gridfile, # noqa: F811 r04b09_diffusion_config, damping_height, ): @@ -502,7 +511,9 @@ def test_run_diffusion_single_step( vct_a=vct_a, rayleigh_damping_height=damping_height ) additional_parameters = DiffusionParams(r04b09_diffusion_config) - icon_grid = init_grid_manager(r04b09_dsl_gridfile, grid_savepoint.num(KDim)).get_grid() + icon_grid = init_grid_manager( + r04b09_dsl_gridfile, grid_savepoint.num(KDim) + ).get_grid() diffusion = Diffusion(run_program=run_with_program) diffusion.init( grid=icon_grid, diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index 477804557..bdd59f722 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -17,109 +17,131 @@ from icon4py.testutils.fixtures import data_provider, setup_icon_data # noqa - @pytest.mark.datatest -@pytest.mark.parametrize( "marker, index",[ - ( HorizontalMarkerIndex.interior(CellDim), 4104), - (HorizontalMarkerIndex.interior(CellDim) + 1, 0), - (HorizontalMarkerIndex.local(CellDim) - 1, 20896), - (HorizontalMarkerIndex.local(CellDim), -1), # halo in icon is (1,20896) - (HorizontalMarkerIndex.nudging(CellDim), 3316), - (HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 2511), - (HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 1688), - (HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 850), - (HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 0) - ]) +@pytest.mark.parametrize( + "marker, index", + [ + (HorizontalMarkerIndex.interior(CellDim), 4104), + (HorizontalMarkerIndex.interior(CellDim) + 1, 0), + (HorizontalMarkerIndex.local(CellDim) - 1, 20896), + (HorizontalMarkerIndex.local(CellDim), -1), # halo in icon is (1,20896) + (HorizontalMarkerIndex.nudging(CellDim), 3316), + (HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 2511), + (HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 1688), + (HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 850), + (HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 0), + ], +) def test_horizontal_cell_start_indices(icon_grid, marker, index): assert index == icon_grid.get_start_index(CellDim, marker) + @pytest.mark.datatest -@pytest.mark.parametrize( "marker, index",[ - ( HorizontalMarkerIndex.interior(CellDim), 20896), - (HorizontalMarkerIndex.interior(CellDim) + 1, 850), - ( HorizontalMarkerIndex.local(CellDim) - 2,20896 ), - (HorizontalMarkerIndex.local(CellDim) - 1, 20896), - (HorizontalMarkerIndex.local(CellDim), 31558), - (HorizontalMarkerIndex.nudging(CellDim), 4104), - (HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 3316), - (HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 2511), - (HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 1688), - (HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 428),]) -def test_horizontal_edge_end_indices(icon_grid, marker, index): +@pytest.mark.parametrize( + "marker, index", + [ + (HorizontalMarkerIndex.interior(CellDim), 20896), + (HorizontalMarkerIndex.interior(CellDim) + 1, 850), + (HorizontalMarkerIndex.local(CellDim) - 2, 20896), + (HorizontalMarkerIndex.local(CellDim) - 1, 20896), + (HorizontalMarkerIndex.local(CellDim), 20896), + (HorizontalMarkerIndex.nudging(CellDim), 4104), + (HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 3316), + (HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 2511), + (HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 1688), + (HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 850), + ], +) +def test_horizontal_cell_end_indices(icon_grid, marker, index): assert index == icon_grid.get_end_index(CellDim, marker) + @pytest.mark.datatest -@pytest.mark.parametrize( "marker, index",[ - ( HorizontalMarkerIndex.interior(EdgeDim), 6176), - ( HorizontalMarkerIndex.local(EdgeDim) - 2,31558 ), - (HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), - (HorizontalMarkerIndex.local(EdgeDim), -1), - (HorizontalMarkerIndex.nudging(EdgeDim) + 1,5387), - (HorizontalMarkerIndex.nudging(EdgeDim), 4989), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7,4184 ), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 3777), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 2954), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2538), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 1700), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1278), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 428), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 0) - ]) +@pytest.mark.parametrize( + "marker, index", + [ + (HorizontalMarkerIndex.interior(EdgeDim), 6176), + (HorizontalMarkerIndex.local(EdgeDim) - 2, 31558), + (HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), + (HorizontalMarkerIndex.local(EdgeDim), -1), + (HorizontalMarkerIndex.nudging(EdgeDim) + 1, 5387), + (HorizontalMarkerIndex.nudging(EdgeDim), 4989), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4184), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 3777), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 2954), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2538), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 1700), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1278), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 428), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 0), + ], +) def test_horizontal_edge_start_indices(icon_grid, marker, index): assert index == icon_grid.get_start_index(EdgeDim, marker) + @pytest.mark.datatest -@pytest.mark.parametrize( "marker, index",[ - ( HorizontalMarkerIndex.interior(EdgeDim), 31558), - ( HorizontalMarkerIndex.local(EdgeDim) - 2,31558 ), - (HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), - (HorizontalMarkerIndex.local(EdgeDim), 31558), - (HorizontalMarkerIndex.nudging(EdgeDim) + 1,6176), - (HorizontalMarkerIndex.nudging(EdgeDim), 5387), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4989), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 4184), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 3777), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2954), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 2538), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1700), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 1278), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 428),]) +@pytest.mark.parametrize( + "marker, index", + [ + (HorizontalMarkerIndex.interior(EdgeDim), 31558), + (HorizontalMarkerIndex.local(EdgeDim) - 2, 31558), + (HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), + (HorizontalMarkerIndex.local(EdgeDim), 31558), + (HorizontalMarkerIndex.nudging(EdgeDim) + 1, 6176), + (HorizontalMarkerIndex.nudging(EdgeDim), 5387), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4989), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 4184), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 3777), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2954), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 2538), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1700), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 1278), + (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 428), + ], +) def test_horizontal_edge_end_indices(icon_grid, marker, index): assert index == icon_grid.get_end_index(EdgeDim, marker) @pytest.mark.datatest -@pytest.mark.parametrize( "marker, index",[ - ( HorizontalMarkerIndex.interior(VertexDim), 2071), - (HorizontalMarkerIndex.local(VertexDim) - 1, 10663), - (HorizontalMarkerIndex.local(VertexDim), -1), - (HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), - (HorizontalMarkerIndex.nudging(VertexDim), 10663), - (HorizontalMarkerIndex.end(VertexDim), 10663), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 1673), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1266), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 850), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 428), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 0) - ]) +@pytest.mark.parametrize( + "marker, index", + [ + (HorizontalMarkerIndex.interior(VertexDim), 2071), + (HorizontalMarkerIndex.local(VertexDim) - 1, 10663), + (HorizontalMarkerIndex.local(VertexDim), -1), + (HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), + (HorizontalMarkerIndex.nudging(VertexDim), 10663), + (HorizontalMarkerIndex.end(VertexDim), 10663), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 1673), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1266), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 850), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 428), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 0), + ], +) def test_horizontal_vertex_start_indices(icon_grid, marker, index): assert index == icon_grid.get_start_index(VertexDim, marker) + @pytest.mark.datatest -@pytest.mark.parametrize( "marker, index",[ - ( HorizontalMarkerIndex.interior(VertexDim), 10663), - ( HorizontalMarkerIndex.local(VertexDim) - 2,10663 ), - (HorizontalMarkerIndex.local(VertexDim) - 1, 10663), - (HorizontalMarkerIndex.local(VertexDim), 10663), - (HorizontalMarkerIndex.nudging(VertexDim) + 1,10663), - (HorizontalMarkerIndex.nudging(VertexDim), 10663), - (HorizontalMarkerIndex.end(VertexDim), 10663), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 2071), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1673), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 1266), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 850), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 428),]) +@pytest.mark.parametrize( + "marker, index", + [ + (HorizontalMarkerIndex.interior(VertexDim), 10663), + (HorizontalMarkerIndex.local(VertexDim) - 2, 10663), + (HorizontalMarkerIndex.local(VertexDim) - 1, 10663), + (HorizontalMarkerIndex.local(VertexDim), 10663), + (HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), + (HorizontalMarkerIndex.nudging(VertexDim), 10663), + (HorizontalMarkerIndex.end(VertexDim), 10663), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 2071), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1673), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 1266), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 850), + (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 428), + ], +) def test_horizontal_vertex_end_indices(icon_grid, marker, index): assert index == icon_grid.get_end_index(VertexDim, marker) - - diff --git a/common/src/icon4py/common/dimension.py b/common/src/icon4py/common/dimension.py index 092871249..37f9d6655 100644 --- a/common/src/icon4py/common/dimension.py +++ b/common/src/icon4py/common/dimension.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import Dimension, DimensionKind +from gt4py.next.common import DimensionKind from gt4py.next.ffront.fbuiltins import Dimension, FieldOffset diff --git a/common/src/icon4py/grid/grid_manager.py b/common/src/icon4py/grid/grid_manager.py index acfe56e33..ad881697a 100644 --- a/common/src/icon4py/grid/grid_manager.py +++ b/common/src/icon4py/grid/grid_manager.py @@ -23,15 +23,17 @@ from icon4py.common.dimension import ( C2E2CDim, + C2E2CODim, C2EDim, C2VDim, CellDim, + E2C2VDim, E2CDim, E2VDim, EdgeDim, V2CDim, V2EDim, - VertexDim, C2E2CODim, E2C2VDim, + VertexDim, ) from icon4py.grid.horizontal import HorizontalGridSize from icon4py.grid.icon_grid import GridConfig, IconGrid @@ -110,7 +112,9 @@ def __init__(self, dataset: Dataset): def dimension(self, name: GridFileName) -> int: return self._dataset.dimensions[name].size - def int_field(self, name: GridFileName, transpose=True, dtype=np.int32) -> np.ndarray: + def int_field( + self, name: GridFileName, transpose=True, dtype=np.int32 + ) -> np.ndarray: try: nc_variable = self._dataset.variables[name] @@ -129,7 +133,10 @@ class IconGridError(RuntimeError): class IndexTransformation(ABC): - def get_offset_for_index_field(self, array: np.ndarray,): + def get_offset_for_index_field( + self, + array: np.ndarray, + ): return np.zeros(array.shape, dtype=np.int32) @@ -191,25 +198,34 @@ def _read_grid_refinement_information(self, dataset): reader, GridFile.GridRefinementName.END_INDEX_CELLS, transpose=False, - apply_offset=False, dtype=np.int64 + apply_offset=False, + dtype=np.int64, )[_CHILD_DOM] start_indices[EdgeDim] = self._get_index_field( - reader, GridFile.GridRefinementName.START_INDEX_EDGES, transpose=False, dtype=np.int64 + reader, + GridFile.GridRefinementName.START_INDEX_EDGES, + transpose=False, + dtype=np.int64, )[_CHILD_DOM] end_indices[EdgeDim] = self._get_index_field( reader, GridFile.GridRefinementName.END_INDEX_EDGES, transpose=False, - apply_offset=False, dtype=np.int64 + apply_offset=False, + dtype=np.int64, )[_CHILD_DOM] start_indices[VertexDim] = self._get_index_field( - reader, GridFile.GridRefinementName.START_INDEX_VERTICES, transpose=False, dtype=np.int64 + reader, + GridFile.GridRefinementName.START_INDEX_VERTICES, + transpose=False, + dtype=np.int64, )[_CHILD_DOM] end_indices[VertexDim] = self._get_index_field( reader, GridFile.GridRefinementName.END_INDEX_VERTICES, transpose=False, - apply_offset=False, dtype=np.int64 + apply_offset=False, + dtype=np.int64, )[_CHILD_DOM] return start_indices, end_indices @@ -232,7 +248,7 @@ def _get_index(self, dim: Dimension, start_marker: int, index_dict): raise IconGridError(msg) try: return index_dict[dim][start_marker] - except KeyError as error: + except KeyError: msg = f"start, end indices for dimension {dim} not present" self._log.error(msg) raise IconGridError(msg) @@ -286,7 +302,12 @@ def get_size(self, dim: Dimension): raise IconGridError(f"Unknown dimension {dim}") def _get_index_field( - self, reader, field: GridFileName, transpose=True, apply_offset=True, dtype=np.int32 + self, + reader, + field: GridFileName, + transpose=True, + apply_offset=True, + dtype=np.int32, ): field = reader.int_field(field, transpose=transpose, dtype=dtype) if apply_offset: @@ -304,7 +325,6 @@ def from_grid_dataset(self, dataset: Dataset) -> IconGrid: ) c2e = self._get_index_field(reader, GridFile.OffsetName.C2E) - e2c = self._get_index_field(reader, GridFile.OffsetName.E2C) c2v = self._get_index_field(reader, GridFile.OffsetName.C2V) e2v = self._get_index_field(reader, GridFile.OffsetName.E2V) @@ -351,8 +371,12 @@ def from_grid_dataset(self, dataset: Dataset) -> IconGrid: return icon_grid - def construct_diamond_array(self, c2v:np.ndarray, e2c:np.ndarray): - dummy_c2v = np.append(c2v, GridFile.INVALID_INDEX * np.ones((1, c2v.shape[1]), dtype=np.int32), axis=0) + def construct_diamond_array(self, c2v: np.ndarray, e2c: np.ndarray): + dummy_c2v = np.append( + c2v, + GridFile.INVALID_INDEX * np.ones((1, c2v.shape[1]), dtype=np.int32), + axis=0, + ) expanded = dummy_c2v[e2c[:, :], :] sh = expanded.shape flattened = expanded.reshape(sh[0], sh[1] * sh[2]) @@ -360,4 +384,3 @@ def construct_diamond_array(self, c2v:np.ndarray, e2c:np.ndarray): def get_c2v_connectivity(self): return self._grid.get_c2v_connectivity() - diff --git a/common/src/icon4py/grid/horizontal.py b/common/src/icon4py/grid/horizontal.py index 42e4d4e38..a56f74e4c 100644 --- a/common/src/icon4py/grid/horizontal.py +++ b/common/src/icon4py/grid/horizontal.py @@ -82,7 +82,7 @@ class HorizontalMarkerIndex: @classmethod def lateral_boundary(cls, dim: Dimension) -> int: - """TODO @magdalena""" + """Return the ICON constants for lateral boundary markers.""" match dim: case (dimension.CellDim): return cls._LATERAL_BOUNDARY_CELLS diff --git a/common/src/icon4py/grid/icon_grid.py b/common/src/icon4py/grid/icon_grid.py index fe9e89670..c91542b2a 100644 --- a/common/src/icon4py/grid/icon_grid.py +++ b/common/src/icon4py/grid/icon_grid.py @@ -10,10 +10,10 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -from typing import Dict, Tuple +from typing import Dict import numpy as np -from gt4py.next.common import Dimension, DimensionKind +from gt4py.next.common import Dimension from gt4py.next.iterator.embedded import NeighborTableOffsetProvider from icon4py.common.dimension import CellDim, ECVDim, EdgeDim, KDim, VertexDim @@ -116,22 +116,22 @@ def num_vertices(self): def num_edges(self): return self.config.num_edges - def get_start_index(self, dim:Dimension, marker: int): + def get_start_index(self, dim: Dimension, marker: int): """ - Use to specify lower end of domains of a field for field_operators. + Use to specify lower end of domains of a field for field_operators. - For a given dimension, returns the start index of the - horizontal region in a field given by the marker. + For a given dimension, returns the start index of the + horizontal region in a field given by the marker. """ return self.start_indices[dim][marker] - def get_end_index(self, dim: Dimension, marker:int): + def get_end_index(self, dim: Dimension, marker: int): """ - Use to specify upper end of domains of a field for field_operators. + Use to specify upper end of domains of a field for field_operators. - For a given dimension, returns the end index of the - horizontal region in a field given by the marker. - """ + For a given dimension, returns the end index of the + horizontal region in a field given by the marker. + """ return self.end_indices[dim][marker] def get_c2e_connectivity(self): diff --git a/common/tests/test_gridmanager.py b/common/tests/test_gridmanager.py index 0daa08200..ed3070402 100644 --- a/common/tests/test_gridmanager.py +++ b/common/tests/test_gridmanager.py @@ -209,8 +209,6 @@ def _add_to_dataset( var[:] = np.transpose(data)[:] - - def test_gridparser_dimension(simple_mesh_data): data = Dataset(SIMPLE_MESH_NC, "r") @@ -221,6 +219,7 @@ def test_gridparser_dimension(simple_mesh_data): assert grid_parser.dimension(GridFile.DimensionName.EDGE_NAME) == mesh.n_edges +@pytest.mark.xfail def test_gridfile_vertex_cell_edge_dimensions(grid_savepoint, r04b09_dsl_gridfile): data = Dataset(r04b09_dsl_gridfile, "r") grid_file = GridFile(data) @@ -268,7 +267,6 @@ def test_gridmanager_eval_v2e(caplog, grid_savepoint, r04b09_dsl_gridfile): assert np.allclose(gm.get_v2e_connectivity().table, seralized_v2e) - # v2c: serial, simple, grid @pytest.mark.datatest def test_gridmanager_eval_v2c(caplog, grid_savepoint, r04b09_dsl_gridfile): @@ -289,6 +287,8 @@ def test_gridmanager_eval_v2c(caplog, grid_savepoint, r04b09_dsl_gridfile): def reset_invalid_index(index_array: np.ndarray): """ + Revert changes from mo_model_domimp_patches. + Helper function to revert mo_model_domimp_patches.f90: move_dummies_to_end_idxblk. see: # ! Checks for the pentagon case and moves dummy cells to end. @@ -327,13 +327,10 @@ def test_gridmanager_eval_e2v(caplog, grid_savepoint, r04b09_dsl_gridfile): assert np.allclose(gm.get_e2v_connectivity().table, serialized_e2v) - def has_invalid_index(ar: np.ndarray): return np.any(np.where(ar == GridFile.INVALID_INDEX)) - - # e2c :serial, simple, grid @pytest.mark.datatest def test_gridmanager_eval_e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): @@ -366,7 +363,7 @@ def test_gridmanager_eval_c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): # e2c2e (e2c2eo) - diamond: serial, simple_mesh, grid file???? -@pytest.mark.skip("does this array exist in grid file?") +@pytest.mark.skip(" TODO @magdalena: does this array exist in grid file?") @pytest.mark.datatest def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) @@ -377,8 +374,6 @@ def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): assert np.allclose(gm.get_e2c2e_connectivity().table, serialized_e2c2e) - - # c2e2c: serial, simple_mesh, grid @pytest.mark.datatest def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): @@ -389,23 +384,29 @@ def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): gm.get_c2e2c_connectivity().table, grid_savepoint.c2e2c()[0:num_cells, :] ) + +@pytest.mark.xfail def test_gridmanager_eval_e2c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) gm = init_grid_manager(r04b09_dsl_gridfile) num_edges = gm.get_size(EdgeDim) - c2v = gm.get_c2v_connectivity().table + # TODO: for boundary edges there are only 2 valid vertices in the serialized data, not 3 three + # the "far" (adjacent to edge normal ) is not there. why? + # despite that ordering is different assert np.allclose( gm.get_e2c2v_connectivity().table, grid_savepoint.e2c2v()[0:num_edges, :] ) -@pytest.mark.skip("add to savepoint") + + +@pytest.mark.skip("TODO @magdalena: add this connectivity to savepoint") def test_gridmanager_eval_c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) gm = init_grid_manager(r04b09_dsl_gridfile) num_cells = gm.get_size(CellDim) c2v = gm.get_c2v_connectivity().table - assert np.allclose( - c2v, grid_savepoint.c2v()[0:num_cells, :] - ) + assert np.allclose(c2v, grid_savepoint.c2v()[0:num_cells, :]) + + def init_grid_manager(fname): gm = GridManager(ToGt4PyTransformation(), fname, VerticalGridConfig(65)) gm.init() @@ -415,21 +416,30 @@ def init_grid_manager(fname): @pytest.mark.parametrize("dim, size", [(CellDim, 18), (EdgeDim, 27), (VertexDim, 9)]) def test_grid_manager_getsize(simple_mesh_data, simple_mesh_path, dim, size, caplog): caplog.set_level(logging.DEBUG) - gm = GridManager(IndexTransformation(), simple_mesh_path) + gm = GridManager( + IndexTransformation(), simple_mesh_path, VerticalGridConfig(num_lev=80) + ) gm.init() assert size == gm.get_size(dim) + def test_grid_manager_diamond_offset(simple_mesh_path): mesh = SimpleMesh() - gm = GridManager(IndexTransformation(), simple_mesh_path, VerticalGridConfig(num_lev=mesh.k_level)) + gm = GridManager( + IndexTransformation(), + simple_mesh_path, + VerticalGridConfig(num_lev=mesh.k_level), + ) gm.init() - assert np.allclose(np.sort(gm.get_e2c2v_connectivity().table, 1), np.sort(mesh.diamond_arr, 1)) + assert np.allclose( + np.sort(gm.get_e2c2v_connectivity().table, 1), np.sort(mesh.diamond_arr, 1) + ) def test_gridmanager_given_file_not_found_then_abort(): fname = "./unknown_grid.nc" with pytest.raises(SystemExit) as error: - gm = GridManager(IndexTransformation(), fname) + gm = GridManager(IndexTransformation(), fname, VerticalGridConfig(num_lev=80)) gm.init() assert error.type == SystemExit assert error.value == 1 @@ -492,6 +502,6 @@ def test_get_start_indices(r04b09_dsl_gridfile, dim, marker, index): (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 5387), ], ) -def test_get_start_indices(r04b09_dsl_gridfile, dim, marker, index): +def test_get_start_indices(r04b09_dsl_gridfile, dim, marker, index): # noqa: F811 gm = init_grid_manager(r04b09_dsl_gridfile) assert gm.get_end_index(dim, marker) == index diff --git a/common/tests/test_vertical.py b/common/tests/test_vertical.py index 42a560dee..296c2802b 100644 --- a/common/tests/test_vertical.py +++ b/common/tests/test_vertical.py @@ -16,7 +16,7 @@ import numpy as np import pytest -from icon4py.grid.vertical import VerticalGridConfig, VerticalModelParams +from icon4py.grid.vertical import VerticalModelParams from icon4py.testutils.fixtures import ( # noqa F401 data_provider, grid_savepoint, @@ -32,7 +32,6 @@ def test_nrdmax_calculation(max_h, damping, delta): vct_a = np.arange(0, max_h, delta) vct_a = vct_a[::-1] vertical_params = VerticalModelParams( - config=VerticalGridConfig(num_lev=10), rayleigh_damping_height=damping, vct_a=vct_a, ) @@ -43,11 +42,11 @@ def test_nrdmax_calculation(max_h, damping, delta): @pytest.mark.datatest -def test_nrdmax_calculation_from_icon_input(icon_grid, grid_savepoint): # noqa: F811 +def test_nrdmax_calculation_from_icon_input(grid_savepoint): # noqa: F811 a = grid_savepoint.vct_a() damping_height = 12500 vertical_params = VerticalModelParams( - VerticalGridConfig(num_lev=10), rayleigh_damping_height=damping_height, vct_a=a + rayleigh_damping_height=damping_height, vct_a=a ) assert 9 == vertical_params.index_of_damping_layer a_array = np.asarray(a) diff --git a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py index 49019077a..7f4f9d75d 100644 --- a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py +++ b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py @@ -35,7 +35,8 @@ from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.prognostic_state import PrognosticState from icon4py.grid.horizontal import CellParams, EdgeParams -from icon4py.grid.icon_grid import IconGrid, VerticalModelParams +from icon4py.grid.icon_grid import IconGrid +from icon4py.grid.vertical import VerticalModelParams from icon4py.py2f.cffi_utils import CffiMethod @@ -80,7 +81,7 @@ def diffusion_init( Fortran ICON Diffusion component (aka Diffusion granule) """ - grid = IconGrid() # TODO where to get this from + grid = IconGrid() # TODO @magdalena: where to get this from? edges_params = EdgeParams( tangent_orientation=tangent_orientation, primal_edge_lengths=primal_edge_lengths, diff --git a/testutils/src/icon4py/testutils/fixtures.py b/testutils/src/icon4py/testutils/fixtures.py index e8c2270c4..0c2365e9a 100644 --- a/testutils/src/icon4py/testutils/fixtures.py +++ b/testutils/src/icon4py/testutils/fixtures.py @@ -19,7 +19,7 @@ from icon4py.testutils.serialbox_utils import IconSerialDataProvider -# data_uri = "https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download" +# "https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download" data_uri = "https://polybox.ethz.ch/index.php/s/joglKLsbyIXOOp8/download" data_path = Path(__file__).parent.joinpath("ser_icondata") extracted_path = data_path.joinpath("mch_ch_r04b09_dsl/ser_data") diff --git a/testutils/src/icon4py/testutils/simple_mesh.py b/testutils/src/icon4py/testutils/simple_mesh.py index 47f0cdb4f..7f742423c 100644 --- a/testutils/src/icon4py/testutils/simple_mesh.py +++ b/testutils/src/icon4py/testutils/simple_mesh.py @@ -58,13 +58,28 @@ @dataclass class SimpleMeshData: - c2v_table = np.asarray([[0,1,4], [1,2,5], [2,0,3], - [0,3,4], [1,4,5], [2,5,3], - [3,4,7], [4,5,8], [5,3,6], - [3,6,7], [4,7,8], [5,8,6], - [6,7,1], [7,8,2], [8,6,0], - [6,0,1],[7,1,2], [8,2,0] - ]) + c2v_table = np.asarray( + [ + [0, 1, 4], + [1, 2, 5], + [2, 0, 3], + [0, 3, 4], + [1, 4, 5], + [2, 5, 3], + [3, 4, 7], + [4, 5, 8], + [5, 3, 6], + [3, 6, 7], + [4, 7, 8], + [5, 8, 6], + [6, 7, 1], + [7, 8, 2], + [8, 6, 0], + [6, 0, 1], + [7, 1, 2], + [8, 2, 0], + ] + ) e2c2v_table = np.asarray( [ @@ -406,8 +421,9 @@ def __init__(self, k_level: int = _DEFAULT_K_LEVEL): ECVDim: self.n_edges * self.n_e2c2v, } - def get_c2v_offset_provider(self)-> NeighborTableOffsetProvider: + def get_c2v_offset_provider(self) -> NeighborTableOffsetProvider: return NeighborTableOffsetProvider(self.c2v, VertexDim, CellDim, self.n_c2v) + def get_c2e_offset_provider(self) -> NeighborTableOffsetProvider: return NeighborTableOffsetProvider(self.c2e, CellDim, EdgeDim, self.n_c2e) From cbaa3507adce45dd2e5b2d62a201cdca74e776b5 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 3 May 2023 10:11:18 +0200 Subject: [PATCH 094/263] add return value type to get_start/end_index --- common/src/icon4py/grid/icon_grid.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/common/src/icon4py/grid/icon_grid.py b/common/src/icon4py/grid/icon_grid.py index c91542b2a..e03705bd3 100644 --- a/common/src/icon4py/grid/icon_grid.py +++ b/common/src/icon4py/grid/icon_grid.py @@ -116,7 +116,7 @@ def num_vertices(self): def num_edges(self): return self.config.num_edges - def get_start_index(self, dim: Dimension, marker: int): + def get_start_index(self, dim: Dimension, marker: int)->int: """ Use to specify lower end of domains of a field for field_operators. @@ -125,7 +125,7 @@ def get_start_index(self, dim: Dimension, marker: int): """ return self.start_indices[dim][marker] - def get_end_index(self, dim: Dimension, marker: int): + def get_end_index(self, dim: Dimension, marker: int)->int: """ Use to specify upper end of domains of a field for field_operators. From 6d26b85c79e897b14fc9305ef8b747e2e33d2307 Mon Sep 17 00:00:00 2001 From: Magdalena Date: Wed, 3 May 2023 16:34:33 +0200 Subject: [PATCH 095/263] Merge main into greenline dycore (#201) merge icon4py main back into greenline-dycore --- .flake8 | 2 +- .pre-commit-config.yaml | 2 +- README.md | 2 +- .../mo_nh_diffusion_stencil_15.py | 10 +-- .../mo_solve_nonhydro_stencil_20.py | 7 +- .../mo_solve_nonhydro_stencil_21.py | 7 +- jenkins/spack-daily | 33 +++++++++ liskov/README.md | 6 +- liskov/src/icon4py/liskov/cli.py | 20 ++++- .../src/icon4py/liskov/codegen/exceptions.py | 16 ++++ liskov/src/icon4py/liskov/codegen/f90.py | 40 +++++++--- liskov/src/icon4py/liskov/codegen/generate.py | 67 +++++++++++------ .../src/icon4py/liskov/codegen/interface.py | 2 +- .../src/icon4py/liskov/external/exceptions.py | 28 +++++++ liskov/src/icon4py/liskov/external/gt4py.py | 9 +-- .../src/icon4py/liskov/external/metadata.py | 73 +++++++++++++++++++ .../src/icon4py/liskov/parsing/deserialise.py | 33 ++++++--- .../src/icon4py/liskov/parsing/exceptions.py | 8 -- liskov/src/icon4py/liskov/parsing/parse.py | 18 +++-- liskov/src/icon4py/liskov/parsing/scan.py | 12 ++- liskov/src/icon4py/liskov/parsing/types.py | 15 +++- liskov/src/icon4py/liskov/pipeline.py | 16 ++-- liskov/tests/test_code_metadata.py | 68 +++++++++++++++++ liskov/tests/test_deserialiser.py | 43 +++++++++-- liskov/tests/test_external.py | 4 +- liskov/tests/test_generation.py | 33 ++++----- liskov/tests/test_parser.py | 13 ++-- liskov/tests/test_validation.py | 12 ++- pyutils/setup.cfg | 2 +- .../{bindings => icon4pygen}/__init__.py | 0 .../{pyutils => icon4pygen}/backend.py | 6 +- .../bindings}/__init__.py | 0 .../bindings/codegen}/__init__.py | 0 .../{ => icon4pygen}/bindings/codegen/cpp.py | 10 ++- .../{ => icon4pygen}/bindings/codegen/f90.py | 4 +- .../bindings/codegen/header.py | 8 +- .../bindings/codegen/render}/__init__.py | 0 .../bindings/codegen/render/field.py | 6 +- .../bindings/codegen/render/location.py | 2 +- .../bindings/codegen/render/offset.py | 4 +- .../bindings/codegen/type_conversion.py | 0 .../bindings/codegen/types.py | 2 +- .../{ => icon4pygen}/bindings/entities.py | 20 +++-- .../{ => icon4pygen}/bindings/exceptions.py | 0 .../{ => icon4pygen}/bindings/locations.py | 2 +- .../{ => icon4pygen}/bindings/py.typed | 0 .../{ => icon4pygen}/bindings/utils.py | 22 +++++- .../{ => icon4pygen}/bindings/workflow.py | 12 +-- .../icon4pygen.py => icon4pygen/cli.py} | 6 +- .../{pyutils => icon4pygen}/exceptions.py | 0 .../{pyutils => icon4pygen}/icochainsize.py | 0 .../{pyutils => icon4pygen}/metadata.py | 5 +- .../icon4py/{pyutils => icon4pygen}/py.typed | 0 .../tests/{ => icon4pygen}/test_backend.py | 4 +- .../tests/{ => icon4pygen}/test_codegen.py | 2 +- .../tests/{ => icon4pygen}/test_exceptions.py | 12 +-- .../{ => icon4pygen}/test_field_rendering.py | 4 +- .../{ => icon4pygen}/test_icochainsize.py | 4 +- .../tests/{ => icon4pygen}/test_metadata.py | 2 +- 59 files changed, 550 insertions(+), 188 deletions(-) create mode 100644 jenkins/spack-daily create mode 100644 liskov/src/icon4py/liskov/codegen/exceptions.py create mode 100644 liskov/src/icon4py/liskov/external/exceptions.py create mode 100644 liskov/src/icon4py/liskov/external/metadata.py create mode 100644 liskov/tests/test_code_metadata.py rename pyutils/src/icon4py/{bindings => icon4pygen}/__init__.py (100%) rename pyutils/src/icon4py/{pyutils => icon4pygen}/backend.py (95%) rename pyutils/src/icon4py/{bindings/codegen => icon4pygen/bindings}/__init__.py (100%) rename pyutils/src/icon4py/{bindings/codegen/render => icon4pygen/bindings/codegen}/__init__.py (100%) rename pyutils/src/icon4py/{ => icon4pygen}/bindings/codegen/cpp.py (98%) rename pyutils/src/icon4py/{ => icon4pygen}/bindings/codegen/f90.py (99%) rename pyutils/src/icon4py/{ => icon4pygen}/bindings/codegen/header.py (95%) rename pyutils/src/icon4py/{pyutils => icon4pygen/bindings/codegen/render}/__init__.py (100%) rename pyutils/src/icon4py/{ => icon4pygen}/bindings/codegen/render/field.py (94%) rename pyutils/src/icon4py/{ => icon4pygen}/bindings/codegen/render/location.py (92%) rename pyutils/src/icon4py/{ => icon4pygen}/bindings/codegen/render/offset.py (96%) rename pyutils/src/icon4py/{ => icon4pygen}/bindings/codegen/type_conversion.py (100%) rename pyutils/src/icon4py/{ => icon4pygen}/bindings/codegen/types.py (97%) rename pyutils/src/icon4py/{ => icon4pygen}/bindings/entities.py (93%) rename pyutils/src/icon4py/{ => icon4pygen}/bindings/exceptions.py (100%) rename pyutils/src/icon4py/{ => icon4pygen}/bindings/locations.py (96%) rename pyutils/src/icon4py/{ => icon4pygen}/bindings/py.typed (100%) rename pyutils/src/icon4py/{ => icon4pygen}/bindings/utils.py (66%) rename pyutils/src/icon4py/{ => icon4pygen}/bindings/workflow.py (85%) rename pyutils/src/icon4py/{pyutils/icon4pygen.py => icon4pygen/cli.py} (93%) rename pyutils/src/icon4py/{pyutils => icon4pygen}/exceptions.py (100%) rename pyutils/src/icon4py/{pyutils => icon4pygen}/icochainsize.py (100%) rename pyutils/src/icon4py/{pyutils => icon4pygen}/metadata.py (98%) rename pyutils/src/icon4py/{pyutils => icon4pygen}/py.typed (100%) rename pyutils/tests/{ => icon4pygen}/test_backend.py (93%) rename pyutils/tests/{ => icon4pygen}/test_codegen.py (99%) rename pyutils/tests/{ => icon4pygen}/test_exceptions.py (92%) rename pyutils/tests/{ => icon4pygen}/test_field_rendering.py (96%) rename pyutils/tests/{ => icon4pygen}/test_icochainsize.py (92%) rename pyutils/tests/{ => icon4pygen}/test_metadata.py (97%) diff --git a/.flake8 b/.flake8 index 4213954b8..ce35d1d1c 100644 --- a/.flake8 +++ b/.flake8 @@ -42,5 +42,5 @@ rst-roles = py:obj, obj, per-file-ignores = - pyutils/src/icon4py/pyutils/icochainsize.py:E800 + pyutils/src/icon4py/icon4pygen/icochainsize.py:E800 atm_dyn_iconam/src/icon4py/diffusion/horizontal.py:W605 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d38e84f07..480995b3a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -126,7 +126,7 @@ repos: - id: mypy name: mypy static type checker entry: | - mypy --install-types --non-interactive -p icon4py.common -p icon4py.pyutils -p icon4py.testutils + mypy --install-types --non-interactive -p icon4py.common -p icon4py.icon4pygen -p icon4py.testutils language: system types_or: [python, pyi] always_run: true diff --git a/README.md b/README.md index b5636096f..07613e4f3 100644 --- a/README.md +++ b/README.md @@ -70,7 +70,7 @@ source .venv/bin/activate pip install --upgrade wheel # Install a specific ICON4Py subpackage and its dependencies -cd _SUBPACKAGE_ # where _SUBPACKAGE_ in atm_dyn_iconam | common | pyutils | testutils | ... +cd _SUBPACKAGE_ # where _SUBPACKAGE_ in atm_dyn_iconam | common | icon4pygen | testutils | ... pip install -r requirements-dev.txt ``` diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py index d4d8ba529..a0d3fca5b 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py @@ -16,6 +16,7 @@ from gt4py.next.iterator.builtins import ( deref, if_, + list_get, named_range, shift, unstructured_domain, @@ -24,14 +25,13 @@ from gt4py.next.type_system import type_specifications as ts from icon4py.common.dimension import C2E2C, C2E2CDim, CellDim, KDim, Koff -from icon4py.pyutils.metadata import FieldInfo +from icon4py.icon4pygen.metadata import FieldInfo -@fundef def step(i, geofac_n2s_nbh, vcoef, theta_v, zd_vertoffset): - d_vcoef = deref(shift(i)(vcoef)) - s_theta_v = shift(C2E2C, i, Koff, deref(shift(i)(zd_vertoffset)))(theta_v) - return deref(shift(i)(geofac_n2s_nbh)) * ( + d_vcoef = list_get(i, deref(vcoef)) + s_theta_v = shift(C2E2C, i, Koff, list_get(i, deref(zd_vertoffset)))(theta_v) + return list_get(i, deref(geofac_n2s_nbh)) * ( d_vcoef * deref(s_theta_v) + (1.0 - d_vcoef) * deref(shift(Koff, 1)(s_theta_v)) ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py index e107aafdb..06a3829b2 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py @@ -15,6 +15,7 @@ from gt4py.next.ffront import program_ast as past from gt4py.next.iterator.builtins import ( deref, + list_get, named_range, shift, unstructured_domain, @@ -23,7 +24,7 @@ from gt4py.next.type_system import type_specifications as ts from icon4py.common.dimension import E2C, CellDim, E2CDim, EdgeDim, KDim, Koff -from icon4py.pyutils.metadata import FieldInfo +from icon4py.icon4pygen.metadata import FieldInfo @fundef @@ -35,11 +36,11 @@ def step( z_dexner_dz_c_1, z_dexner_dz_c_2, ): - d_ikoffset = deref(shift(i)(ikoffset)) + d_ikoffset = list_get(i, deref(ikoffset)) d_z_exner_exp_pr = deref(shift(Koff, d_ikoffset, E2C, i)(z_exner_ex_pr)) d_z_dexner_dz_c_1 = deref(shift(Koff, d_ikoffset, E2C, i)(z_dexner_dz_c_1)) d_z_dexner_dz_c_2 = deref(shift(Koff, d_ikoffset, E2C, i)(z_dexner_dz_c_2)) - d_zdiff_gradp = deref(shift(i)(zdiff_gradp)) + d_zdiff_gradp = list_get(i, deref(zdiff_gradp)) return d_z_exner_exp_pr + d_zdiff_gradp * ( d_z_dexner_dz_c_1 + d_zdiff_gradp * d_z_dexner_dz_c_2 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py index 036422b11..4fcd5229a 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py @@ -15,6 +15,7 @@ from gt4py.next.ffront import program_ast as past from gt4py.next.iterator.builtins import ( deref, + list_get, named_range, power, shift, @@ -24,19 +25,19 @@ from gt4py.next.type_system import type_specifications as ts from icon4py.common.dimension import E2C, CellDim, E2CDim, EdgeDim, KDim, Koff -from icon4py.pyutils.metadata import FieldInfo +from icon4py.icon4pygen.metadata import FieldInfo @fundef def step(i, theta_v, ikoffset, zdiff_gradp, theta_v_ic, inv_ddqz_z_full): - d_ikoffset = deref(shift(i)(ikoffset)) + d_ikoffset = list_get(i, deref(ikoffset)) d_theta_v = deref(shift(Koff, d_ikoffset, E2C, i)(theta_v)) s_theta_v_ic = shift(Koff, d_ikoffset, E2C, i)(theta_v_ic) d_theta_v_ic = deref(s_theta_v_ic) d_theta_v_ic_p1 = deref(shift(Koff, 1)(s_theta_v_ic)) d_inv_ddqz_z_full = deref(shift(Koff, d_ikoffset, E2C, i)(inv_ddqz_z_full)) - d_zdiff_gradp = deref(shift(i)(zdiff_gradp)) + d_zdiff_gradp = list_get(i, deref(zdiff_gradp)) return ( d_theta_v + d_zdiff_gradp * (d_theta_v_ic - d_theta_v_ic_p1) * d_inv_ddqz_z_full diff --git a/jenkins/spack-daily b/jenkins/spack-daily new file mode 100644 index 000000000..852724ad3 --- /dev/null +++ b/jenkins/spack-daily @@ -0,0 +1,33 @@ +pipeline { + agent none + stages { + stage('Tests') { + matrix { + agent { label "${NODENAME}" } + axes { + axis { + name 'NODENAME' + values 'daint', 'balfrin' + } + } + post { + always { + echo 'Cleaning up workspace' + deleteDir() + } + } + stages { + stage('Install and Test') { + steps { + sh """ + git clone --depth 1 --recurse-submodules --shallow-submodules https://github.com/C2SM/spack-c2sm.git + . ./spack-c2sm/setup-env.sh + spack install -v --test=root py-icon4py@main%gcc ^py-gt4py@main%gcc + """ + } + } + } + } + } + } +} diff --git a/liskov/README.md b/liskov/README.md index ef517a117..ffa4a5580 100644 --- a/liskov/README.md +++ b/liskov/README.md @@ -8,14 +8,14 @@ To install the icon4py-liskov package, follow the instructions in the `README.md ## Description -The icon4py-liskov package includes the `icon_liskov` CLI tool which takes a fortran file as input and processes it with the ICON-Liskov DSL Preprocessor. This preprocessor adds the necessary `USE` statements and generates OpenACC `DATA CREATE` statements and declares DSL input/output fields based on directives in the input file. The preprocessor also processes stencils defined in the input file using the `START STENCIL` and `END STENCIL` directives, inserting the necessary code to run the stencils and adding nvtx profile statements if specified with the `--profile` flag. +The icon4py-liskov package includes the `icon_liskov` CLI tool which takes a fortran file as input and processes it with the ICON-Liskov DSL Preprocessor. This preprocessor adds the necessary `USE` statements and generates OpenACC `DATA CREATE` statements and declares DSL input/output fields based on directives in the input file. The preprocessor also processes stencils defined in the input file using the `START STENCIL` and `END STENCIL` directives, inserting the necessary code to run the stencils and adding nvtx profile statements if specified with the `--profile` or `-p` flag. Additionally, specifying the `--metadatagen` or `-m` flag will result in the generation of runtime metadata at the top of the generated file. Note that this requires `git` to be available in the shell from which `icon_liskov` is executed. ### Usage To use the `icon_liskov` tool, run the following command: ```bash -icon_liskov [--profile] +icon_liskov [--profile] [--metadatagen] ``` Where `input_filepath` is the path to the input file to be processed, and `output_filepath` is the path to the output file. The optional `--profile` flag adds nvtx profile statements to the stencils. @@ -30,7 +30,7 @@ This directive generates the necessary `USE` statements to import the Fortran to #### `!$DSL START CREATE()` -This directive generates an OpenACC `DATA CREATE` statement for all output fields used in each DSL (icon4py) stencil. +This directive generates an OpenACC `DATA CREATE` statement for all output fields used in each DSL (icon4py) stencil. The directive also takes an **optional** keyword argument to specify extra fields to include in the `DATA CREATE` statement called `extra_fields`. Here you can specify a comma-separated list of strings which should be added to the `DATA CREATE` statement as follows `extra_fields=foo,bar`. #### `!$DSL END CREATE()` diff --git a/liskov/src/icon4py/liskov/cli.py b/liskov/src/icon4py/liskov/cli.py index a27798d70..8e0fe2743 100644 --- a/liskov/src/icon4py/liskov/cli.py +++ b/liskov/src/icon4py/liskov/cli.py @@ -40,24 +40,36 @@ @click.option( "--profile", "-p", is_flag=True, help="Add nvtx profile statements to stencils." ) +@click.option( + "--metadatagen", + "-m", + is_flag=True, + help="Add metadata header with information about program (requires git).", +) def main( - input_filepath: pathlib.Path, output_filepath: pathlib.Path, profile: bool + input_filepath: pathlib.Path, + output_filepath: pathlib.Path, + profile: bool, + metadatagen: bool, ) -> None: """Command line interface for interacting with the ICON-Liskov DSL Preprocessor. Usage: - icon_liskov [--profile] + icon_liskov [-p] [-m] Options: -p --profile Add nvtx profile statements to stencils. + -m --metadatagen Add metadata header with information about program (requires git). Arguments: input_filepath Path to the input file to process. output_filepath Path to the output file to generate. """ - parsed = parse_fortran_file(input_filepath) + parsed = parse_fortran_file(input_filepath, output_filepath) parsed_checked = load_gt4py_stencils(parsed) - run_code_generation(parsed_checked, input_filepath, output_filepath, profile) + run_code_generation( + parsed_checked, input_filepath, output_filepath, profile, metadatagen + ) if __name__ == "__main__": diff --git a/liskov/src/icon4py/liskov/codegen/exceptions.py b/liskov/src/icon4py/liskov/codegen/exceptions.py new file mode 100644 index 000000000..cd7ca6839 --- /dev/null +++ b/liskov/src/icon4py/liskov/codegen/exceptions.py @@ -0,0 +1,16 @@ +# 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 + + +class UndeclaredFieldError(Exception): + pass diff --git a/liskov/src/icon4py/liskov/codegen/f90.py b/liskov/src/icon4py/liskov/codegen/f90.py index bdc62df51..0d2b0ed2d 100644 --- a/liskov/src/icon4py/liskov/codegen/f90.py +++ b/liskov/src/icon4py/liskov/codegen/f90.py @@ -19,12 +19,14 @@ from gt4py.eve.codegen import JinjaTemplate as as_jinja from gt4py.eve.codegen import TemplatedGenerator -from icon4py.bindings.utils import format_fortran_code +from icon4py.icon4pygen.bindings.utils import format_fortran_code +from icon4py.liskov.codegen.exceptions import UndeclaredFieldError from icon4py.liskov.codegen.interface import ( CodeGenInput, DeclareData, StartStencilData, ) +from icon4py.liskov.external.metadata import CodeMetadata def enclose_in_parentheses(string: str) -> str: @@ -107,6 +109,26 @@ def get_array_dims(association: str) -> str: return "".join(list(dims)) +class MetadataStatement(eve.Node): + metadata: CodeMetadata + + +class MetadataStatementGenerator(TemplatedGenerator): + MetadataStatement = as_jinja( + """\ + !+-+-+-+-+-+-+-+-+-+ +-+-+-+-+ +-+-+-+-+-+-+-+-+-+-+-+ + ! GENERATED WITH ICON-LISKOV + !+-+-+-+-+-+-+-+-+-+ +-+-+-+-+ +-+-+-+-+-+-+-+-+-+-+-+ + ! Generated on: {{ _this_node.metadata.generated_on }} + ! Input filepath: {{ _this_node.metadata.cli_params['input_filepath'] }} + ! Profiling active: {{ _this_node.metadata.cli_params['profile'] }} + ! Git version tag: {{ _this_node.metadata.tag }} + ! Git commit hash: {{ _this_node.metadata.commit_hash }} + !+-+-+-+-+-+-+-+-+-+ +-+-+-+-+ +-+-+-+-+-+-+-+-+-+-+-+ + """ + ) + + class EndStencilStatement(eve.Node): stencil_data: StartStencilData profile: bool @@ -251,7 +273,7 @@ def __post_init__(self) -> None: # type: ignore @staticmethod def make_copy_declaration(f: Field) -> CopyDeclaration: if f.dims is None: - raise Exception(f"{f.variable} not declared!") + raise UndeclaredFieldError(f"{f.variable} was not declared!") lh_idx = render_index(f.dims) @@ -330,6 +352,7 @@ class ImportsStatementGenerator(TemplatedGenerator): class StartCreateStatement(eve.Node): stencils: list[StartStencilData] + extra_fields: Optional[list[str]] out_field_names: list[str] = eve.datamodels.field(init=False) def __post_init__(self) -> None: # type: ignore @@ -348,18 +371,17 @@ def __post_init__(self) -> None: # type: ignore class StartCreateStatementGenerator(TemplatedGenerator): StartCreateStatement = as_jinja( """ - #ifdef __DSL_VERIFY - dsl_verify = .TRUE. - #else - dsl_verify = .FALSE. - #endif - !$ACC DATA CREATE( & + {%- if _this_node.extra_fields -%} + {%- for name in extra_fields %} + !$ACC {{ name }} {%- if not loop.last -%}, & {% else %}, & {%- endif -%} + {%- endfor %} + {%- endif -%} {%- for name in out_field_names %} !$ACC {{ name }}_before {%- if not loop.last -%}, & {% else %} & {%- endif -%} {%- endfor %} !$ACC ), & - !$ACC IF ( i_am_accel_node .AND. dsl_verify) + !$ACC IF ( i_am_accel_node ) """ ) diff --git a/liskov/src/icon4py/liskov/codegen/generate.py b/liskov/src/icon4py/liskov/codegen/generate.py index 6cbb057cc..26db2e626 100644 --- a/liskov/src/icon4py/liskov/codegen/generate.py +++ b/liskov/src/icon4py/liskov/codegen/generate.py @@ -12,11 +12,10 @@ # SPDX-License-Identifier: GPL-3.0-or-later from dataclasses import dataclass -from typing import Optional, Sequence, Type +from typing import Any, Optional, Sequence, Type import gt4py.eve as eve from gt4py.eve.codegen import TemplatedGenerator -from typing_extensions import Any from icon4py.liskov.codegen.f90 import ( DeclareStatement, @@ -33,6 +32,8 @@ ImportsStatementGenerator, InsertStatement, InsertStatementGenerator, + MetadataStatement, + MetadataStatementGenerator, StartCreateStatement, StartCreateStatementGenerator, StartProfileStatement, @@ -48,6 +49,7 @@ UnusedDirective, ) from icon4py.liskov.common import Step +from icon4py.liskov.external.metadata import CodeMetadata from icon4py.liskov.logger import setup_logger @@ -64,10 +66,16 @@ class GeneratedCode: class IntegrationGenerator(Step): - def __init__(self, directives: DeserialisedDirectives, profile: bool): + def __init__( + self, + directives: DeserialisedDirectives, + profile: bool, + metadata_gen: bool, + ): self.profile = profile self.directives = directives self.generated: list[GeneratedCode] = [] + self.metadata_gen = metadata_gen def __call__(self, data: Any = None) -> list[GeneratedCode]: """Generate all f90 code for integration. @@ -75,6 +83,7 @@ def __call__(self, data: Any = None) -> list[GeneratedCode]: Args: profile: A boolean indicating whether to include profiling calls in the generated code. """ + self._generate_metadata() self._generate_create() self._generate_imports() self._generate_declare() @@ -106,6 +115,18 @@ def _generate( code = GeneratedCode(source=source, startln=startln, endln=endln) self.generated.append(code) + def _generate_metadata(self) -> None: + """Generate metadata about the current liskov execution.""" + if self.metadata_gen: + logger.info("Generating icon-liskov metadata.") + self._generate( + MetadataStatement, + MetadataStatementGenerator, + startln=0, + endln=0, + metadata=CodeMetadata(), + ) + def _generate_declare(self) -> None: """Generate f90 code for declaration statements.""" for i, declare in enumerate(self.directives.Declare): @@ -206,6 +227,7 @@ def _generate_create(self) -> None: self.directives.StartCreate.startln, self.directives.StartCreate.endln, stencils=self.directives.StartStencil, + extra_fields=self.directives.StartCreate.extra_fields, ) self._generate( @@ -229,26 +251,27 @@ def _generate_endif(self) -> None: def _generate_profile(self) -> None: """Generate additional nvtx profiling statements.""" - if self.directives.StartProfile != UnusedDirective: - for start in self.directives.StartProfile: # type: ignore - logger.info("Generating nvtx start statement.") - self._generate( - StartProfileStatement, - StartProfileStatementGenerator, - start.startln, - start.endln, - name=start.name, - ) + if self.profile: + if self.directives.StartProfile != UnusedDirective: + for start in self.directives.StartProfile: # type: ignore + logger.info("Generating nvtx start statement.") + self._generate( + StartProfileStatement, + StartProfileStatementGenerator, + start.startln, + start.endln, + name=start.name, + ) - if self.directives.EndProfile != UnusedDirective: - for end in self.directives.EndProfile: # type: ignore - logger.info("Generating nvtx end statement.") - self._generate( - EndProfileStatement, - EndProfileStatementGenerator, - end.startln, - end.endln, - ) + if self.directives.EndProfile != UnusedDirective: + for end in self.directives.EndProfile: # type: ignore + logger.info("Generating nvtx end statement.") + self._generate( + EndProfileStatement, + EndProfileStatementGenerator, + end.startln, + end.endln, + ) def _generate_insert(self) -> None: """Generate free form statement from insert directive.""" diff --git a/liskov/src/icon4py/liskov/codegen/interface.py b/liskov/src/icon4py/liskov/codegen/interface.py index f89fccd9e..dfa879f0c 100644 --- a/liskov/src/icon4py/liskov/codegen/interface.py +++ b/liskov/src/icon4py/liskov/codegen/interface.py @@ -58,7 +58,7 @@ class ImportsData(CodeGenInput): @dataclass class StartCreateData(CodeGenInput): - ... + extra_fields: Optional[list[str]] @dataclass diff --git a/liskov/src/icon4py/liskov/external/exceptions.py b/liskov/src/icon4py/liskov/external/exceptions.py new file mode 100644 index 000000000..2ae393192 --- /dev/null +++ b/liskov/src/icon4py/liskov/external/exceptions.py @@ -0,0 +1,28 @@ +# 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 + + +class MissingClickContextError(Exception): + pass + + +class MissingGitError(Exception): + pass + + +class UnknownStencilError(Exception): + pass + + +class IncompatibleFieldError(Exception): + pass diff --git a/liskov/src/icon4py/liskov/external/gt4py.py b/liskov/src/icon4py/liskov/external/gt4py.py index 618c18115..c0096a92a 100644 --- a/liskov/src/icon4py/liskov/external/gt4py.py +++ b/liskov/src/icon4py/liskov/external/gt4py.py @@ -13,18 +13,18 @@ import importlib from inspect import getmembers +from typing import Any from gt4py.next.ffront.decorator import Program -from typing_extensions import Any +from icon4py.icon4pygen.metadata import get_stencil_info from icon4py.liskov.codegen.interface import DeserialisedDirectives from icon4py.liskov.common import Step -from icon4py.liskov.logger import setup_logger -from icon4py.liskov.parsing.exceptions import ( +from icon4py.liskov.external.exceptions import ( IncompatibleFieldError, UnknownStencilError, ) -from icon4py.pyutils.metadata import get_stencil_info +from icon4py.liskov.logger import setup_logger logger = setup_logger(__name__) @@ -58,7 +58,6 @@ def _collect_icon4py_stencil(self, stencil_name: str) -> Program: """Collect and return the ICON4PY stencil program with the given name.""" err_counter = 0 for pkg in self._STENCIL_PACKAGES: - try: module_name = f"icon4py.{pkg}.{stencil_name}" module = importlib.import_module(module_name) diff --git a/liskov/src/icon4py/liskov/external/metadata.py b/liskov/src/icon4py/liskov/external/metadata.py new file mode 100644 index 000000000..c47fedcc3 --- /dev/null +++ b/liskov/src/icon4py/liskov/external/metadata.py @@ -0,0 +1,73 @@ +# 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 datetime +import subprocess +from pathlib import Path +from typing import Any + +import click + +from icon4py.liskov.external.exceptions import ( + MissingClickContextError, + MissingGitError, +) + + +class CodeMetadata: + """Class that handles retrieval of icon-liskov runtime metadata.""" + + def __init__(self) -> None: + self.generated_on = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") + self.parent_dir = Path(__file__).parent + + @property + def cli_params(self) -> dict[str, Any]: + try: + ctx = click.get_current_context() + return ctx.params + except Exception as e: + raise MissingClickContextError( + f"Cannot fetch click context in this thread as no click command has been executed.\n {e}" + ) + + @property + def commit_hash(self) -> str: + """Get the latest git commit hash.""" + try: + return ( + subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=self.parent_dir + ) + .decode() + .strip() + ) + except Exception as e: + raise MissingGitError( + f"Git is not available or there is no commit or tag.\n {e}" + ) + + @property + def tag(self) -> str: + """Get the latest git tag.""" + try: + return ( + subprocess.check_output( + ["git", "describe", "--tags", "--abbrev=0"], cwd=self.parent_dir + ) + .decode() + .strip() + ) + except Exception as e: + raise MissingGitError( + f"Git is not available or there is no commit or tag.\n {e}" + ) diff --git a/liskov/src/icon4py/liskov/parsing/deserialise.py b/liskov/src/icon4py/liskov/parsing/deserialise.py index 5b8850b71..fa88ae048 100644 --- a/liskov/src/icon4py/liskov/parsing/deserialise.py +++ b/liskov/src/icon4py/liskov/parsing/deserialise.py @@ -78,6 +78,10 @@ def _extract_boolean_kwarg( return False +def pop_item_from_dict(dictionary: dict, key: str, default_value: str) -> str: + return dictionary.pop(key, default_value) + + class DirectiveInputFactory(Protocol): def __call__( self, parsed: ts.ParsedDict @@ -117,12 +121,6 @@ def __call__(self, parsed: ts.ParsedDict) -> CodeGenInput: return self.dtype(startln=extracted.startln, endln=extracted.endln) -@dataclass -class StartCreateDataFactory(RequiredSingleUseDataFactory): - directive_cls: Type[ts.ParsedDirective] = ts.StartCreate - dtype: Type[StartCreateData] = StartCreateData - - @dataclass class EndCreateDataFactory(RequiredSingleUseDataFactory): directive_cls: Type[ts.ParsedDirective] = ts.EndCreate @@ -147,8 +145,23 @@ class EndProfileDataFactory(OptionalMultiUseDataFactory): dtype: Type[EndProfileData] = EndProfileData -def pop_item_from_dict(dictionary: dict, key: str, default_value: str) -> str: - return dictionary.pop(key, default_value) +@dataclass +class StartCreateDataFactory(DataFactoryBase): + directive_cls: Type[ts.ParsedDirective] = ts.StartCreate + dtype: Type[StartCreateData] = StartCreateData + + def __call__(self, parsed: ts.ParsedDict) -> StartCreateData: + directive = extract_directive(parsed["directives"], self.directive_cls)[0] + + named_args = parsed["content"]["StartCreate"][0] + + extra_fields = None + if named_args: + extra_fields = named_args["extra_fields"].split(",") + + return self.dtype( + startln=directive.startln, endln=directive.endln, extra_fields=extra_fields + ) @dataclass @@ -252,7 +265,7 @@ def __call__(self, parsed: ts.ParsedDict) -> list[StartStencilData]: for i, directive in enumerate(directives): named_args = parsed["content"]["StartStencil"][i] acc_present = string_to_bool( - pop_item_from_dict(named_args, "accpresent", "false") + pop_item_from_dict(named_args, "accpresent", "true") ) mergecopy = string_to_bool( pop_item_from_dict(named_args, "mergecopy", "false") @@ -334,7 +347,6 @@ def _make_field_associations( """Create a list of FieldAssociation objects.""" fields = [] for field_name, association in field_args.items(): - # skipped as handled by _update_field_tolerances if any([field_name.endswith(tol) for tol in TOLERANCE_ARGS]): continue @@ -354,7 +366,6 @@ def _update_tolerances( """Set relative and absolute tolerance for a given field if set in the directives.""" for field_name, association in named_args.items(): for tol in TOLERANCE_ARGS: - _tol = f"_{tol}" if field_name.endswith(_tol): diff --git a/liskov/src/icon4py/liskov/parsing/exceptions.py b/liskov/src/icon4py/liskov/parsing/exceptions.py index 89b7b8b6b..9ddacb94a 100644 --- a/liskov/src/icon4py/liskov/parsing/exceptions.py +++ b/liskov/src/icon4py/liskov/parsing/exceptions.py @@ -38,11 +38,3 @@ class MissingBoundsError(Exception): class MissingDirectiveArgumentError(Exception): pass - - -class IncompatibleFieldError(Exception): - pass - - -class UnknownStencilError(Exception): - pass diff --git a/liskov/src/icon4py/liskov/parsing/parse.py b/liskov/src/icon4py/liskov/parsing/parse.py index a5b96c87b..93b42d0cb 100644 --- a/liskov/src/icon4py/liskov/parsing/parse.py +++ b/liskov/src/icon4py/liskov/parsing/parse.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import collections +import shutil import sys from pathlib import Path from typing import Sequence @@ -28,16 +29,18 @@ class DirectivesParser(Step): - def __init__(self, filepath: Path) -> None: + def __init__(self, input_filepath: Path, output_filepath: Path) -> None: """Initialize a DirectivesParser instance. This class parses a Sequence of RawDirective objects and returns a dictionary of parsed directives and their associated content. Args: directives: Sequence of RawDirective objects to parse. - filepath: Path to file being parsed. + input_filepath Path to the input file to process. + output_filepath Path to the output file to generate. """ - self.filepath = filepath + self.input_filepath = input_filepath + self.output_filepath = output_filepath def __call__(self, directives: list[ts.RawDirective]) -> ts.ParsedDict: """Parse the directives and return a dictionary of parsed directives and their associated content. @@ -45,13 +48,16 @@ def __call__(self, directives: list[ts.RawDirective]) -> ts.ParsedDict: Returns: ParsedType: Dictionary of parsed directives and their associated content. """ - logger.info(f"Parsing DSL Preprocessor directives at {self.filepath}") + logger.info(f"Parsing DSL Preprocessor directives at {self.input_filepath}") if len(directives) != 0: typed = self._determine_type(directives) preprocessed = self._preprocess(typed) self._run_validation_passes(preprocessed) return dict(directives=preprocessed, content=self._parse(preprocessed)) - logger.warning(f"No DSL Preprocessor directives found in {self.filepath}") + logger.warning( + f"No DSL Preprocessor directives found in {self.input_filepath}, copying to {self.output_filepath}" + ) + shutil.copyfile(self.input_filepath, self.output_filepath) sys.exit() @staticmethod @@ -87,7 +93,7 @@ def _run_validation_passes( ) -> None: """Run validation passes on the directives.""" for validator in VALIDATORS: - validator(self.filepath).validate(preprocessed) + validator(self.input_filepath).validate(preprocessed) @staticmethod def _clean_string(string: str) -> str: diff --git a/liskov/src/icon4py/liskov/parsing/scan.py b/liskov/src/icon4py/liskov/parsing/scan.py index c5e8dd942..84c44955e 100644 --- a/liskov/src/icon4py/liskov/parsing/scan.py +++ b/liskov/src/icon4py/liskov/parsing/scan.py @@ -30,7 +30,7 @@ class Scanned: class DirectivesScanner(Step): - def __init__(self, filepath: Path) -> None: + def __init__(self, input_filepath: Path) -> None: r"""Class for scanning a file for ICON-Liskov DSL directives. A directive must start with !$DSL ( with the @@ -47,9 +47,9 @@ def __init__(self, filepath: Path) -> None: !$DSL b=test) Args: - filepath: Path to file to scan for directives. + input_filepath: Path to file to scan for directives. """ - self.filepath = filepath + self.input_filepath = input_filepath def __call__(self, data: Any = None) -> list[ts.RawDirective]: """Scan filepath for directives and return them along with their line numbers. @@ -58,12 +58,10 @@ def __call__(self, data: Any = None) -> list[ts.RawDirective]: A list of RawDirective objects containing the scanned directives and their line numbers. """ directives = [] - with self.filepath.open() as f: - + with self.input_filepath.open() as f: scanned_directives = [] lines = f.readlines() for lnumber, string in enumerate(lines): - if ts.DIRECTIVE_IDENT in string: stripped = string.strip() eol = stripped[-1] @@ -86,7 +84,7 @@ def __call__(self, data: Any = None) -> list[ts.RawDirective]: raise DirectiveSyntaxError( f"Error in directive on line number: {lnumber + 1}\n Used invalid end of line character." ) - logger.info(f"Scanning for directives at {self.filepath}") + logger.info(f"Scanning for directives at {self.input_filepath}") return directives @staticmethod diff --git a/liskov/src/icon4py/liskov/parsing/types.py b/liskov/src/icon4py/liskov/parsing/types.py index b49664f02..9fb5ee29d 100644 --- a/liskov/src/icon4py/liskov/parsing/types.py +++ b/liskov/src/icon4py/liskov/parsing/types.py @@ -13,6 +13,7 @@ from dataclasses import dataclass, field from typing import ( Any, + Optional, Protocol, Sequence, Type, @@ -80,6 +81,18 @@ def get_content(self) -> dict[str, str]: return content +@dataclass(eq=False) +class WithOptionalArguments(TypedDirective): + regex: str = field(default=r"(?:.+?=.+?|)", init=False) + + def get_content(self) -> Optional[dict[str, str]]: + args = self.string.replace(f"{self.pattern}", "")[1:-1] + if len(args) > 0: + content = dict([args.split("=")]) + return content + return None + + @dataclass(eq=False) class WithoutArguments(TypedDirective): # matches an empty string at the beginning of a line @@ -115,7 +128,7 @@ class Imports(WithoutArguments): pattern = "IMPORTS" -class StartCreate(WithoutArguments): +class StartCreate(WithOptionalArguments): pattern = "START CREATE" diff --git a/liskov/src/icon4py/liskov/pipeline.py b/liskov/src/icon4py/liskov/pipeline.py index bc6f57a2a..d452ce6f7 100644 --- a/liskov/src/icon4py/liskov/pipeline.py +++ b/liskov/src/icon4py/liskov/pipeline.py @@ -24,7 +24,7 @@ @linear_pipeline -def parse_fortran_file(filepath: Path) -> list[Step]: +def parse_fortran_file(input_filepath: Path, output_filepath: Path) -> list[Step]: """Execute a pipeline to parse and deserialize directives from a file. The pipeline consists of three steps: DirectivesScanner, DirectivesParser, and @@ -34,14 +34,15 @@ def parse_fortran_file(filepath: Path) -> list[Step]: DeserialisedDirectives object. Args: - filepath (Path): The file path of the directives file. + input_filepath: Path to the input file to process. + output_filepath: Path to the output file to generate. Returns: DeserialisedDirectives: The deserialized directives object. """ return [ - DirectivesScanner(filepath), - DirectivesParser(filepath), + DirectivesScanner(input_filepath), + DirectivesParser(input_filepath, output_filepath), DirectiveDeserialiser(), ] @@ -65,6 +66,7 @@ def run_code_generation( input_filepath: Path, output_filepath: Path, profile: bool, + metadatagen: bool, ) -> list[Step]: """Execute a pipeline to generate and write code for a set of directives. @@ -74,10 +76,12 @@ def run_code_generation( Args: parsed: The deserialized directives object. - filepath: The file path to write the generated code to. + input_filepath: The original file containing the DSL preprocessor directives. + output_filepath: The file path to write the generated code to. profile: A flag to indicate if profiling information should be included in the generated code. + metadatagen: A flag to indicate if a metadata header should be included in the generated code. """ return [ - IntegrationGenerator(parsed, profile), + IntegrationGenerator(parsed, profile, metadatagen), IntegrationWriter(input_filepath, output_filepath), ] diff --git a/liskov/tests/test_code_metadata.py b/liskov/tests/test_code_metadata.py new file mode 100644 index 000000000..6dfb12800 --- /dev/null +++ b/liskov/tests/test_code_metadata.py @@ -0,0 +1,68 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +from datetime import datetime +from pathlib import Path +from unittest import mock + +import pytest + +from icon4py.liskov.external.exceptions import MissingGitError +from icon4py.liskov.external.metadata import CodeMetadata + + +@pytest.fixture +def module_parent(): + import icon4py.liskov.external.metadata as meta + + return Path(meta.__file__).parent + + +def test_generated_on(): + with mock.patch("datetime.datetime") as mock_datetime: + mock_datetime.now.return_value = datetime(2023, 3, 1, 12, 0, 0) + metadata = CodeMetadata() + assert metadata.generated_on == "2023-03-01 12:00:00" + + +def test_tag(module_parent): + with mock.patch("subprocess.check_output", side_effect=[b"v1.0.0"]) as mock_git: + metadata = CodeMetadata() + assert metadata.tag == "v1.0.0" + mock_git.assert_called_once_with( + ["git", "describe", "--tags", "--abbrev=0"], cwd=module_parent + ) + + +def test_commit_hash(module_parent): + with mock.patch( + "subprocess.check_output", side_effect=[b"abcdef123456"] + ) as mock_git: + metadata = CodeMetadata() + assert metadata.commit_hash == "abcdef123456" + mock_git.assert_called_once_with( + ["git", "rev-parse", "HEAD"], cwd=module_parent + ) + + +def test_no_git(): + with mock.patch("subprocess.check_output", side_effect=MissingGitError()): + metadata = CodeMetadata() + with pytest.raises(MissingGitError): + metadata.tag + + +def test_click_context(): + with mock.patch("click.get_current_context") as mock_click: + mock_click.return_value.params = {"foo": "bar"} + metadata = CodeMetadata() + assert metadata.cli_params == {"foo": "bar"} diff --git a/liskov/tests/test_deserialiser.py b/liskov/tests/test_deserialiser.py index c474dd2cf..f32c4e2c0 100644 --- a/liskov/tests/test_deserialiser.py +++ b/liskov/tests/test_deserialiser.py @@ -51,14 +51,6 @@ @pytest.mark.parametrize( "factory_class, directive_type, startln, endln, string, expected", [ - ( - StartCreateDataFactory, - ts.StartCreate, - "START CREATE()", - 1, - 1, - StartCreateData, - ), (EndCreateDataFactory, ts.EndCreate, "END CREATE", 2, 2, EndCreateData), (ImportsDataFactory, ts.Imports, "IMPORTS", 3, 3, ImportsData), (EndIfDataFactory, ts.EndIf, "ENDIF", 4, 4, EndIfData), @@ -172,6 +164,41 @@ def test_data_factories_with_args(factory, target, mock_data): assert all([isinstance(r, target) for r in result]) +@pytest.mark.parametrize( + "mock_data, extra_fields", + [ + ( + { + "directives": [ts.StartCreate("START CREATE(extra_fields=foo)", 5, 5)], + "content": {"StartCreate": [{"extra_fields": "foo"}]}, + }, + ["foo"], + ), + ( + { + "directives": [ + ts.StartCreate("START CREATE(extra_fields=foo,xyz)", 5, 5) + ], + "content": {"StartCreate": [{"extra_fields": "foo,xyz"}]}, + }, + ["foo", "xyz"], + ), + ( + { + "directives": [ts.StartCreate("START CREATE()", 5, 5)], + "content": {"StartCreate": [None]}, + }, + None, + ), + ], +) +def test_start_create_factory(mock_data, extra_fields): + factory = StartCreateDataFactory() + result = factory(mock_data) + assert isinstance(result, StartCreateData) + assert result.extra_fields == extra_fields + + @pytest.mark.parametrize( "factory,target,mock_data", [ diff --git a/liskov/tests/test_external.py b/liskov/tests/test_external.py index 0d11730e9..c97afed2d 100644 --- a/liskov/tests/test_external.py +++ b/liskov/tests/test_external.py @@ -22,11 +22,11 @@ FieldAssociationData, StartStencilData, ) -from icon4py.liskov.external.gt4py import UpdateFieldsWithGt4PyStencils -from icon4py.liskov.parsing.exceptions import ( +from icon4py.liskov.external.exceptions import ( IncompatibleFieldError, UnknownStencilError, ) +from icon4py.liskov.external.gt4py import UpdateFieldsWithGt4PyStencils def test_stencil_collector(): diff --git a/liskov/tests/test_generation.py b/liskov/tests/test_generation.py index 2c8facda8..a83f66b8e 100644 --- a/liskov/tests/test_generation.py +++ b/liskov/tests/test_generation.py @@ -11,7 +11,6 @@ # # SPDX-License-Identifier: GPL-3.0-or-later - import pytest from icon4py.liskov.codegen.generate import IntegrationGenerator @@ -82,7 +81,9 @@ def serialised_directives(): suffix="before", ) imports_data = ImportsData(startln=7, endln=8) - start_create_data = StartCreateData(startln=9, endln=10) + start_create_data = StartCreateData( + extra_fields=["foo", "bar"], startln=9, endln=10 + ) end_create_data = EndCreateData(startln=11, endln=11) endif_data = EndIfData(startln=12, endln=12) start_profile_data = StartProfileData(startln=13, endln=13, name="test_stencil") @@ -106,21 +107,17 @@ def serialised_directives(): @pytest.fixture def expected_start_create_source(): return """ -#ifdef __DSL_VERIFY - dsl_verify = .TRUE. -#else - dsl_verify = .FALSE. -#endif - - !$ACC DATA CREATE( & - !$ACC out1_before, & - !$ACC out2_before, & - !$ACC out3_before, & - !$ACC out4_before, & - !$ACC out5_before, & - !$ACC out6_before & - !$ACC ), & - !$ACC IF ( i_am_accel_node .AND. dsl_verify)""" +!$ACC DATA CREATE( & +!$ACC foo, & +!$ACC bar, & +!$ACC out1_before, & +!$ACC out2_before, & +!$ACC out3_before, & +!$ACC out4_before, & +!$ACC out5_before, & +!$ACC out6_before & +!$ACC ), & +!$ACC IF ( i_am_accel_node )""" @pytest.fixture @@ -205,7 +202,7 @@ def expected_insert_source(): @pytest.fixture def generator(serialised_directives): - return IntegrationGenerator(serialised_directives, profile=True) + return IntegrationGenerator(serialised_directives, profile=True, metadata_gen=False) def test_generate( diff --git a/liskov/tests/test_parser.py b/liskov/tests/test_parser.py index b4f9069e3..8101988c5 100644 --- a/liskov/tests/test_parser.py +++ b/liskov/tests/test_parser.py @@ -44,11 +44,11 @@ def test_parse_no_input(): defaultdict(list, {"Imports": [{}]}), ), ( - ts.StartCreate("START CREATE()", 2, 2), + ts.StartCreate("START CREATE(extra_fields=foo)", 2, 2), "START CREATE()", 2, 2, - defaultdict(list, {"StartCreate": [{}]}), + defaultdict(list, {"StartCreate": [{"extra_fields": "foo"}]}), ), ( ts.StartStencil( @@ -83,8 +83,9 @@ def test_parse_single_directive(directive, string, startln, endln, expected_cont ) def test_file_parsing(make_f90_tmpfile, stencil, num_directives, num_content): fpath = make_f90_tmpfile(content=stencil) + opath = fpath.with_suffix(".gen") directives = scan_for_directives(fpath) - parser = DirectivesParser(fpath) + parser = DirectivesParser(fpath, opath) parsed = parser(directives) directives = parsed["directives"] @@ -99,8 +100,9 @@ def test_file_parsing(make_f90_tmpfile, stencil, num_directives, num_content): def test_directive_parser_no_directives_found(make_f90_tmpfile): fpath = make_f90_tmpfile(content=NO_DIRECTIVES_STENCIL) + opath = fpath.with_suffix(".gen") directives = scan_for_directives(fpath) - parser = DirectivesParser(fpath) + parser = DirectivesParser(fpath, opath) with pytest.raises(SystemExit): parser(directives) @@ -118,9 +120,10 @@ def test_unsupported_directives( directive, ): fpath = make_f90_tmpfile(content=stencil) + opath = fpath.with_suffix(".gen") insert_new_lines(fpath, [directive]) directives = scan_for_directives(fpath) - parser = DirectivesParser(fpath) + parser = DirectivesParser(fpath, opath) with pytest.raises( UnsupportedDirectiveError, diff --git a/liskov/tests/test_validation.py b/liskov/tests/test_validation.py index 0e484b63f..03bd6e5b5 100644 --- a/liskov/tests/test_validation.py +++ b/liskov/tests/test_validation.py @@ -46,8 +46,9 @@ def test_directive_semantics_validation_unbalanced_stencil_directives( make_f90_tmpfile, stencil, directive ): fpath = make_f90_tmpfile(stencil + directive) + opath = fpath.with_suffix(".gen") directives = scan_for_directives(fpath) - parser = DirectivesParser(fpath) + parser = DirectivesParser(fpath, opath) with pytest.raises(UnbalancedStencilDirectiveError): parser(directives) @@ -82,9 +83,10 @@ def test_directive_semantics_validation_repeated_directives( make_f90_tmpfile, directive ): fpath = make_f90_tmpfile(content=SINGLE_STENCIL) + opath = fpath.with_suffix(".gen") insert_new_lines(fpath, [directive]) directives = scan_for_directives(fpath) - parser = DirectivesParser(fpath) + parser = DirectivesParser(fpath, opath) with pytest.raises( RepeatedDirectiveError, @@ -101,9 +103,10 @@ def test_directive_semantics_validation_repeated_directives( ) def test_directive_semantics_validation_repeated_stencil(make_f90_tmpfile, directive): fpath = make_f90_tmpfile(content=SINGLE_STENCIL) + opath = fpath.with_suffix(".gen") insert_new_lines(fpath, [directive]) directives = scan_for_directives(fpath) - parser = DirectivesParser(fpath) + parser = DirectivesParser(fpath, opath) parser(directives) @@ -120,8 +123,9 @@ def test_directive_semantics_validation_required_directives( ): new = SINGLE_STENCIL.replace(directive, "") fpath = make_f90_tmpfile(content=new) + opath = fpath.with_suffix(".gen") directives = scan_for_directives(fpath) - parser = DirectivesParser(fpath) + parser = DirectivesParser(fpath, opath) with pytest.raises( RequiredDirectivesError, diff --git a/pyutils/setup.cfg b/pyutils/setup.cfg index 2209a5a16..de7bddd91 100644 --- a/pyutils/setup.cfg +++ b/pyutils/setup.cfg @@ -53,4 +53,4 @@ exclude = [options.entry_points] console_scripts = - icon4pygen = icon4py.pyutils.icon4pygen:main + icon4pygen = icon4py.icon4pygen.cli:main diff --git a/pyutils/src/icon4py/bindings/__init__.py b/pyutils/src/icon4py/icon4pygen/__init__.py similarity index 100% rename from pyutils/src/icon4py/bindings/__init__.py rename to pyutils/src/icon4py/icon4pygen/__init__.py diff --git a/pyutils/src/icon4py/pyutils/backend.py b/pyutils/src/icon4py/icon4pygen/backend.py similarity index 95% rename from pyutils/src/icon4py/pyutils/backend.py rename to pyutils/src/icon4py/icon4pygen/backend.py index 3540a03dc..caa138d5d 100644 --- a/pyutils/src/icon4py/pyutils/backend.py +++ b/pyutils/src/icon4py/icon4pygen/backend.py @@ -16,10 +16,10 @@ from gt4py.next.iterator import ir as itir from gt4py.next.program_processors.codegens.gtfn.gtfn_backend import generate -from icon4py.bindings.utils import write_string from icon4py.common.dimension import Koff -from icon4py.pyutils.exceptions import MultipleFieldOperatorException -from icon4py.pyutils.metadata import StencilInfo +from icon4py.icon4pygen.bindings.utils import write_string +from icon4py.icon4pygen.exceptions import MultipleFieldOperatorException +from icon4py.icon4pygen.metadata import StencilInfo H_START = "horizontal_start" diff --git a/pyutils/src/icon4py/bindings/codegen/__init__.py b/pyutils/src/icon4py/icon4pygen/bindings/__init__.py similarity index 100% rename from pyutils/src/icon4py/bindings/codegen/__init__.py rename to pyutils/src/icon4py/icon4pygen/bindings/__init__.py diff --git a/pyutils/src/icon4py/bindings/codegen/render/__init__.py b/pyutils/src/icon4py/icon4pygen/bindings/codegen/__init__.py similarity index 100% rename from pyutils/src/icon4py/bindings/codegen/render/__init__.py rename to pyutils/src/icon4py/icon4pygen/bindings/codegen/__init__.py diff --git a/pyutils/src/icon4py/bindings/codegen/cpp.py b/pyutils/src/icon4py/icon4pygen/bindings/codegen/cpp.py similarity index 98% rename from pyutils/src/icon4py/bindings/codegen/cpp.py rename to pyutils/src/icon4py/icon4pygen/bindings/codegen/cpp.py index 9faa31f8b..b7d65be70 100644 --- a/pyutils/src/icon4py/bindings/codegen/cpp.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/codegen/cpp.py @@ -18,7 +18,7 @@ from gt4py.eve.codegen import JinjaTemplate as as_jinja from gt4py.eve.codegen import Node, TemplatedGenerator, format_source -from icon4py.bindings.codegen.header import ( +from icon4py.icon4pygen.bindings.codegen.header import ( CppFreeFunc, CppRunAndVerifyFuncDeclaration, CppRunFuncDeclaration, @@ -27,9 +27,11 @@ run_func_declaration, run_verify_func_declaration, ) -from icon4py.bindings.codegen.render.offset import GpuTriMeshOffsetRenderer -from icon4py.bindings.entities import Field, Offset -from icon4py.bindings.utils import write_string +from icon4py.icon4pygen.bindings.codegen.render.offset import ( + GpuTriMeshOffsetRenderer, +) +from icon4py.icon4pygen.bindings.entities import Field, Offset +from icon4py.icon4pygen.bindings.utils import write_string class CppDefGenerator(TemplatedGenerator): diff --git a/pyutils/src/icon4py/bindings/codegen/f90.py b/pyutils/src/icon4py/icon4pygen/bindings/codegen/f90.py similarity index 99% rename from pyutils/src/icon4py/bindings/codegen/f90.py rename to pyutils/src/icon4py/icon4pygen/bindings/codegen/f90.py index 2b75a8460..c4d22d753 100644 --- a/pyutils/src/icon4py/bindings/codegen/f90.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/codegen/f90.py @@ -19,8 +19,8 @@ from gt4py.eve.codegen import JinjaTemplate as as_jinja from gt4py.eve.codegen import TemplatedGenerator -from icon4py.bindings.entities import Field, Offset -from icon4py.bindings.utils import format_fortran_code, write_string +from icon4py.icon4pygen.bindings.entities import Field, Offset +from icon4py.icon4pygen.bindings.utils import format_fortran_code, write_string _DOMAIN_ARGS = [ diff --git a/pyutils/src/icon4py/bindings/codegen/header.py b/pyutils/src/icon4py/icon4pygen/bindings/codegen/header.py similarity index 95% rename from pyutils/src/icon4py/bindings/codegen/header.py rename to pyutils/src/icon4py/icon4pygen/bindings/codegen/header.py index 548ab5bd8..321dae71c 100644 --- a/pyutils/src/icon4py/bindings/codegen/header.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/codegen/header.py @@ -19,8 +19,8 @@ from gt4py.eve.codegen import JinjaTemplate as as_jinja from gt4py.eve.codegen import TemplatedGenerator, format_source -from icon4py.bindings.entities import Field -from icon4py.bindings.utils import write_string +from icon4py.icon4pygen.bindings.entities import Field +from icon4py.icon4pygen.bindings.utils import write_string run_func_declaration = as_jinja( @@ -42,7 +42,11 @@ {%- for field in _this_node.out_fields -%} {{ field.renderer.render_ctype('c++') }} {{ field.renderer.render_pointer() }} {{ field.name }}_{{ suffix }}, {%- endfor -%} + {%- if _this_node.tol_fields -%} const int verticalStart, const int verticalEnd, const int horizontalStart, const int horizontalEnd, + {%- else -%} + const int verticalStart, const int verticalEnd, const int horizontalStart, const int horizontalEnd + {%- endif -%} {%- for field in _this_node.tol_fields -%} const double {{ field.name }}_rel_tol, const double {{ field.name }}_abs_tol diff --git a/pyutils/src/icon4py/pyutils/__init__.py b/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/__init__.py similarity index 100% rename from pyutils/src/icon4py/pyutils/__init__.py rename to pyutils/src/icon4py/icon4pygen/bindings/codegen/render/__init__.py diff --git a/pyutils/src/icon4py/bindings/codegen/render/field.py b/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/field.py similarity index 94% rename from pyutils/src/icon4py/bindings/codegen/render/field.py rename to pyutils/src/icon4py/icon4pygen/bindings/codegen/render/field.py index d93f8a245..6d9a2d5fa 100644 --- a/pyutils/src/icon4py/bindings/codegen/render/field.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/field.py @@ -13,12 +13,12 @@ from dataclasses import dataclass -from icon4py.bindings.codegen.type_conversion import ( +from icon4py.icon4pygen.bindings.codegen.type_conversion import ( BUILTIN_TO_CPP_TYPE, BUILTIN_TO_ISO_C_TYPE, ) -from icon4py.bindings.codegen.types import FieldEntity -from icon4py.bindings.exceptions import BindingsRenderingException +from icon4py.icon4pygen.bindings.codegen.types import FieldEntity +from icon4py.icon4pygen.bindings.exceptions import BindingsRenderingException @dataclass(frozen=True) diff --git a/pyutils/src/icon4py/bindings/codegen/render/location.py b/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/location.py similarity index 92% rename from pyutils/src/icon4py/bindings/codegen/render/location.py rename to pyutils/src/icon4py/icon4pygen/bindings/codegen/render/location.py index e75e5fd2d..b90f32d84 100644 --- a/pyutils/src/icon4py/bindings/codegen/render/location.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/location.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from icon4py.bindings.exceptions import BindingsRenderingException +from icon4py.icon4pygen.bindings.exceptions import BindingsRenderingException class LocationRenderer: diff --git a/pyutils/src/icon4py/bindings/codegen/render/offset.py b/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/offset.py similarity index 96% rename from pyutils/src/icon4py/bindings/codegen/render/offset.py rename to pyutils/src/icon4py/icon4pygen/bindings/codegen/render/offset.py index b1c90f489..ee496d3ad 100644 --- a/pyutils/src/icon4py/bindings/codegen/render/offset.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/codegen/render/offset.py @@ -14,8 +14,8 @@ from dataclasses import dataclass from typing import Sequence -from icon4py.bindings.codegen.types import OffsetEntity -from icon4py.bindings.locations import BasicLocation, ChainedLocation +from icon4py.icon4pygen.bindings.codegen.types import OffsetEntity +from icon4py.icon4pygen.bindings.locations import BasicLocation, ChainedLocation @dataclass(frozen=True) diff --git a/pyutils/src/icon4py/bindings/codegen/type_conversion.py b/pyutils/src/icon4py/icon4pygen/bindings/codegen/type_conversion.py similarity index 100% rename from pyutils/src/icon4py/bindings/codegen/type_conversion.py rename to pyutils/src/icon4py/icon4pygen/bindings/codegen/type_conversion.py diff --git a/pyutils/src/icon4py/bindings/codegen/types.py b/pyutils/src/icon4py/icon4pygen/bindings/codegen/types.py similarity index 97% rename from pyutils/src/icon4py/bindings/codegen/types.py rename to pyutils/src/icon4py/icon4pygen/bindings/codegen/types.py index 166414f92..385467fb8 100644 --- a/pyutils/src/icon4py/bindings/codegen/types.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/codegen/types.py @@ -17,7 +17,7 @@ from gt4py.next.type_system import type_specifications as ts -from icon4py.bindings.locations import ( +from icon4py.icon4pygen.bindings.locations import ( BasicLocation, ChainedLocation, CompoundLocation, diff --git a/pyutils/src/icon4py/bindings/entities.py b/pyutils/src/icon4py/icon4pygen/bindings/entities.py similarity index 93% rename from pyutils/src/icon4py/bindings/entities.py rename to pyutils/src/icon4py/icon4pygen/bindings/entities.py index 6c38a713f..16ae8f6b0 100644 --- a/pyutils/src/icon4py/bindings/entities.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/entities.py @@ -17,11 +17,17 @@ from gt4py.next.ffront import program_ast as past from gt4py.next.type_system import type_specifications as ts -from icon4py.bindings.codegen.render.field import FieldRenderer -from icon4py.bindings.codegen.render.offset import OffsetRenderer -from icon4py.bindings.codegen.types import FieldEntity, FieldIntent, OffsetEntity -from icon4py.bindings.exceptions import BindingsTypeConsistencyException -from icon4py.bindings.locations import ( +from icon4py.icon4pygen.bindings.codegen.render.field import FieldRenderer +from icon4py.icon4pygen.bindings.codegen.render.offset import OffsetRenderer +from icon4py.icon4pygen.bindings.codegen.types import ( + FieldEntity, + FieldIntent, + OffsetEntity, +) +from icon4py.icon4pygen.bindings.exceptions import ( + BindingsTypeConsistencyException, +) +from icon4py.icon4pygen.bindings.locations import ( BASIC_LOCATIONS, BasicLocation, Cell, @@ -30,8 +36,8 @@ Edge, Vertex, ) -from icon4py.bindings.utils import calc_num_neighbors -from icon4py.pyutils.metadata import FieldInfo +from icon4py.icon4pygen.bindings.utils import calc_num_neighbors +from icon4py.icon4pygen.metadata import FieldInfo def chain_from_str(chain: list[str] | str) -> list[BasicLocation]: diff --git a/pyutils/src/icon4py/bindings/exceptions.py b/pyutils/src/icon4py/icon4pygen/bindings/exceptions.py similarity index 100% rename from pyutils/src/icon4py/bindings/exceptions.py rename to pyutils/src/icon4py/icon4pygen/bindings/exceptions.py diff --git a/pyutils/src/icon4py/bindings/locations.py b/pyutils/src/icon4py/icon4pygen/bindings/locations.py similarity index 96% rename from pyutils/src/icon4py/bindings/locations.py rename to pyutils/src/icon4py/icon4pygen/bindings/locations.py index 1220f712a..303d77655 100644 --- a/pyutils/src/icon4py/bindings/locations.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/locations.py @@ -15,8 +15,8 @@ from gt4py.next.ffront.fbuiltins import Dimension -from icon4py.bindings.codegen.render.location import LocationRenderer from icon4py.common.dimension import CellDim, EdgeDim, VertexDim +from icon4py.icon4pygen.bindings.codegen.render.location import LocationRenderer class BasicLocation: diff --git a/pyutils/src/icon4py/bindings/py.typed b/pyutils/src/icon4py/icon4pygen/bindings/py.typed similarity index 100% rename from pyutils/src/icon4py/bindings/py.typed rename to pyutils/src/icon4py/icon4pygen/bindings/py.typed diff --git a/pyutils/src/icon4py/bindings/utils.py b/pyutils/src/icon4py/icon4pygen/bindings/utils.py similarity index 66% rename from pyutils/src/icon4py/bindings/utils.py rename to pyutils/src/icon4py/icon4pygen/bindings/utils.py index b91b621d8..d33c56a98 100644 --- a/pyutils/src/icon4py/bindings/utils.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/utils.py @@ -11,12 +11,17 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +import os import subprocess +import sys from pathlib import Path from gt4py.next.common import Dimension -from icon4py.pyutils.icochainsize import IcoChainSize +from icon4py.icon4pygen.icochainsize import IcoChainSize + + +PYTHON_PATH = sys.executable def check_dir_exists(dirpath: Path) -> None: @@ -34,6 +39,19 @@ def calc_num_neighbors(dim_list: list[Dimension], includes_center: bool) -> int: def format_fortran_code(source: str) -> str: - args = ["fprettify"] + """Format fortran code using fprettify. + + The path to fprettify needs to be set explicitly for the + non-Spack build process as Liskov does not activate a virtual environment. + + Variable SPACK_ROOT is always set in a spack build, used to assume that fprettify + is in PATH + """ + if os.getenv("SPACK_ROOT") is not None: + fprettify_path = "fprettify" + else: + bin_path = Path(PYTHON_PATH).parent + fprettify_path = bin_path / "fprettify" + args = [str(fprettify_path)] p1 = subprocess.Popen(args, stdout=subprocess.PIPE, stdin=subprocess.PIPE) return p1.communicate(source.encode("UTF-8"))[0].decode("UTF-8").rstrip() diff --git a/pyutils/src/icon4py/bindings/workflow.py b/pyutils/src/icon4py/icon4pygen/bindings/workflow.py similarity index 85% rename from pyutils/src/icon4py/bindings/workflow.py rename to pyutils/src/icon4py/icon4pygen/bindings/workflow.py index 1b76c3ca5..ef7b88529 100644 --- a/pyutils/src/icon4py/bindings/workflow.py +++ b/pyutils/src/icon4py/icon4pygen/bindings/workflow.py @@ -13,12 +13,12 @@ from pathlib import Path -from icon4py.bindings.codegen.cpp import generate_cpp_definition -from icon4py.bindings.codegen.f90 import generate_f90_file -from icon4py.bindings.codegen.header import generate_cpp_header -from icon4py.bindings.entities import Field, Offset -from icon4py.bindings.utils import check_dir_exists -from icon4py.pyutils.metadata import StencilInfo +from icon4py.icon4pygen.bindings.codegen.cpp import generate_cpp_definition +from icon4py.icon4pygen.bindings.codegen.f90 import generate_f90_file +from icon4py.icon4pygen.bindings.codegen.header import generate_cpp_header +from icon4py.icon4pygen.bindings.entities import Field, Offset +from icon4py.icon4pygen.bindings.utils import check_dir_exists +from icon4py.icon4pygen.metadata import StencilInfo class PyBindGen: diff --git a/pyutils/src/icon4py/pyutils/icon4pygen.py b/pyutils/src/icon4py/icon4pygen/cli.py similarity index 93% rename from pyutils/src/icon4py/pyutils/icon4pygen.py rename to pyutils/src/icon4py/icon4pygen/cli.py index 2217a8529..382726abf 100644 --- a/pyutils/src/icon4py/pyutils/icon4pygen.py +++ b/pyutils/src/icon4py/icon4pygen/cli.py @@ -73,9 +73,9 @@ def main( outpath: represents a path to the folder in which to write all generated code. """ - from icon4py.bindings.workflow import PyBindGen - from icon4py.pyutils.backend import GTHeader - from icon4py.pyutils.metadata import get_stencil_info, import_definition + from icon4py.icon4pygen.backend import GTHeader + from icon4py.icon4pygen.bindings.workflow import PyBindGen + from icon4py.icon4pygen.metadata import get_stencil_info, import_definition fencil_def = import_definition(fencil) stencil_info = get_stencil_info(fencil_def, is_global) diff --git a/pyutils/src/icon4py/pyutils/exceptions.py b/pyutils/src/icon4py/icon4pygen/exceptions.py similarity index 100% rename from pyutils/src/icon4py/pyutils/exceptions.py rename to pyutils/src/icon4py/icon4pygen/exceptions.py diff --git a/pyutils/src/icon4py/pyutils/icochainsize.py b/pyutils/src/icon4py/icon4pygen/icochainsize.py similarity index 100% rename from pyutils/src/icon4py/pyutils/icochainsize.py rename to pyutils/src/icon4py/icon4pygen/icochainsize.py diff --git a/pyutils/src/icon4py/pyutils/metadata.py b/pyutils/src/icon4py/icon4pygen/metadata.py similarity index 98% rename from pyutils/src/icon4py/pyutils/metadata.py rename to pyutils/src/icon4py/icon4pygen/metadata.py index 4e53d9154..01e8a0afd 100644 --- a/pyutils/src/icon4py/pyutils/metadata.py +++ b/pyutils/src/icon4py/icon4pygen/metadata.py @@ -27,11 +27,11 @@ from gt4py.next.type_system import type_specifications as ts from icon4py.common.dimension import CellDim, EdgeDim, Koff, VertexDim -from icon4py.pyutils.exceptions import ( +from icon4py.icon4pygen.exceptions import ( InvalidConnectivityException, MultipleFieldOperatorException, ) -from icon4py.pyutils.icochainsize import IcoChainSize +from icon4py.icon4pygen.icochainsize import IcoChainSize @dataclass(frozen=True) @@ -212,6 +212,7 @@ def scan_for_offsets(fvprog: Program) -> list[eve.concepts.SymbolRef]: fvprog.itir.pre_walk_values() .if_isinstance(itir.OffsetLiteral) .getattr("value") + .if_isinstance(str) .to_list() ) all_dim_labels = [dim.value for dim in all_dims if dim.kind == DimensionKind.LOCAL] diff --git a/pyutils/src/icon4py/pyutils/py.typed b/pyutils/src/icon4py/icon4pygen/py.typed similarity index 100% rename from pyutils/src/icon4py/pyutils/py.typed rename to pyutils/src/icon4py/icon4pygen/py.typed diff --git a/pyutils/tests/test_backend.py b/pyutils/tests/icon4pygen/test_backend.py similarity index 93% rename from pyutils/tests/test_backend.py rename to pyutils/tests/icon4pygen/test_backend.py index 53b06118f..8fe160051 100644 --- a/pyutils/tests/test_backend.py +++ b/pyutils/tests/icon4pygen/test_backend.py @@ -13,8 +13,8 @@ import pytest from gt4py.next.iterator import ir as itir -from icon4py.pyutils import backend -from icon4py.pyutils.backend import GTHeader +from icon4py.icon4pygen import backend +from icon4py.icon4pygen.backend import GTHeader @pytest.mark.parametrize( diff --git a/pyutils/tests/test_codegen.py b/pyutils/tests/icon4pygen/test_codegen.py similarity index 99% rename from pyutils/tests/test_codegen.py rename to pyutils/tests/icon4pygen/test_codegen.py index 9502bd9aa..d262d893f 100644 --- a/pyutils/tests/test_codegen.py +++ b/pyutils/tests/icon4pygen/test_codegen.py @@ -19,7 +19,7 @@ from click.testing import CliRunner import icon4py.atm_dyn_iconam -from icon4py.pyutils.icon4pygen import main +from icon4py.icon4pygen.cli import main from icon4py.testutils.utils import get_stencil_module_path diff --git a/pyutils/tests/test_exceptions.py b/pyutils/tests/icon4pygen/test_exceptions.py similarity index 92% rename from pyutils/tests/test_exceptions.py rename to pyutils/tests/icon4pygen/test_exceptions.py index 0ac8b746c..d9dd65824 100644 --- a/pyutils/tests/test_exceptions.py +++ b/pyutils/tests/icon4pygen/test_exceptions.py @@ -15,15 +15,15 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Dimension, Field -from icon4py.bindings.codegen.render.location import LocationRenderer -from icon4py.bindings.entities import Offset, chain_from_str -from icon4py.bindings.exceptions import ( +from icon4py.common.dimension import EdgeDim, KDim +from icon4py.icon4pygen.bindings.codegen.render.location import LocationRenderer +from icon4py.icon4pygen.bindings.entities import Offset, chain_from_str +from icon4py.icon4pygen.bindings.exceptions import ( BindingsRenderingException, BindingsTypeConsistencyException, ) -from icon4py.bindings.workflow import PyBindGen -from icon4py.common.dimension import EdgeDim, KDim -from icon4py.pyutils.metadata import get_stencil_info +from icon4py.icon4pygen.bindings.workflow import PyBindGen +from icon4py.icon4pygen.metadata import get_stencil_info def test_invalid_offset(): diff --git a/pyutils/tests/test_field_rendering.py b/pyutils/tests/icon4pygen/test_field_rendering.py similarity index 96% rename from pyutils/tests/test_field_rendering.py rename to pyutils/tests/icon4pygen/test_field_rendering.py index a826597ce..79234ecd4 100644 --- a/pyutils/tests/test_field_rendering.py +++ b/pyutils/tests/icon4pygen/test_field_rendering.py @@ -14,9 +14,9 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.bindings.workflow import PyBindGen from icon4py.common.dimension import E2CDim, EdgeDim, KDim -from icon4py.pyutils.metadata import get_stencil_info +from icon4py.icon4pygen.bindings.workflow import PyBindGen +from icon4py.icon4pygen.metadata import get_stencil_info def test_horizontal_field_sid_rendering(): diff --git a/pyutils/tests/test_icochainsize.py b/pyutils/tests/icon4pygen/test_icochainsize.py similarity index 92% rename from pyutils/tests/test_icochainsize.py rename to pyutils/tests/icon4pygen/test_icochainsize.py index 49578bcd6..0d7f8fb94 100644 --- a/pyutils/tests/test_icochainsize.py +++ b/pyutils/tests/icon4pygen/test_icochainsize.py @@ -13,8 +13,8 @@ import pytest -from icon4py.pyutils.exceptions import InvalidConnectivityException -from icon4py.pyutils.metadata import provide_offset +from icon4py.icon4pygen.exceptions import InvalidConnectivityException +from icon4py.icon4pygen.metadata import provide_offset @pytest.mark.parametrize( diff --git a/pyutils/tests/test_metadata.py b/pyutils/tests/icon4pygen/test_metadata.py similarity index 97% rename from pyutils/tests/test_metadata.py rename to pyutils/tests/icon4pygen/test_metadata.py index 5cb5bae26..39359acb1 100644 --- a/pyutils/tests/test_metadata.py +++ b/pyutils/tests/icon4pygen/test_metadata.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.decorator import field_operator, program from icon4py.common.dimension import CellDim, KDim -from icon4py.pyutils.metadata import _get_field_infos, provide_neighbor_table +from icon4py.icon4pygen.metadata import _get_field_infos, provide_neighbor_table chain_false_skipvalues = [ From b6950de86542d757a1bb46b225c69e7d9b4bb0c6 Mon Sep 17 00:00:00 2001 From: Magdalena Date: Fri, 5 May 2023 09:37:44 +0200 Subject: [PATCH 096/263] Back to gt4py main (#202) get back to gt4py main branch and remove workarounds. --- .../src/icon4py/diffusion/diffusion.py | 19 ++----------------- ...ate_nabla2_and_smag_coefficients_for_vn.py | 1 - ...est_mo_intp_rbf_rbf_vec_interpol_vertex.py | 2 +- .../tests/test_mo_nh_diffusion_stencil_15.py | 1 - .../test_mo_solve_nonhydro_stencil_20.py | 2 -- .../test_mo_solve_nonhydro_stencil_21.py | 2 -- base-requirements-dev.txt | 2 +- py2f/src/icon4py/py2f/codegen.py | 4 ++-- 8 files changed, 6 insertions(+), 27 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 3ce0fe6c7..c4f0d093a 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -700,17 +700,11 @@ def time_step( boundary_diffusion_start_index_edges=edge_start_lb_plus4, offset_provider={ "V2E": self.grid.get_v2e_connectivity(), - "V2EDim": V2EDim, "E2C2V": self.grid.get_e2c2v_connectivity(), - "E2C2VDim": E2C2VDim, - "E2CDim": E2CDim, "E2ECV": self.grid.get_e2ecv_connectivity(), "C2E": self.grid.get_c2e_connectivity(), - "C2EDim": C2EDim, "E2C": self.grid.get_e2c_connectivity(), "C2E2C": self.grid.get_c2e2c_connectivity(), - "C2E2CDim": C2E2CDim, - "C2E2CODim": C2E2CODim, "C2E2CO": self.grid.get_c2e2co_connectivity(), "Koff": KDim, }, @@ -833,7 +827,7 @@ def _do_diffusion_step( horizontal_end=vertex_end_local_minus1, vertical_start=0, vertical_end=klevels, - offset_provider={"V2E": self.grid.get_v2e_connectivity(), "V2EDim": V2EDim}, + offset_provider={"V2E": self.grid.get_v2e_connectivity()}, ) log.debug("rbf interpolation: end") # 2. HALO EXCHANGE -- CALL sync_patch_array_mult @@ -864,7 +858,6 @@ def _do_diffusion_step( offset_provider={ "E2C2V": self.grid.get_e2c2v_connectivity(), "E2ECV": self.grid.get_e2ecv_connectivity(), - "E2C2VDim": E2C2VDim, }, ) log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: end") @@ -884,7 +877,6 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={ "C2E": self.grid.get_c2e_connectivity(), - "C2EDim": C2EDim, "Koff": KDim, }, ) @@ -904,7 +896,7 @@ def _do_diffusion_step( horizontal_end=vertex_end_local, vertical_start=0, vertical_end=klevels, - offset_provider={"V2E": self.grid.get_v2e_connectivity(), "V2EDim": V2EDim}, + offset_provider={"V2E": self.grid.get_v2e_connectivity()}, ) log.debug("rbf interpolation: end") # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult @@ -937,7 +929,6 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={ "E2C2V": self.grid.get_e2c2v_connectivity(), - "E2C2VDim": E2C2VDim, "E2ECV": self.grid.get_e2ecv_connectivity(), }, ) @@ -972,7 +963,6 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={ "C2E2CO": self.grid.get_c2e2co_connectivity(), - "C2E2CODim": C2E2CODim, }, ) log.debug("running fused stencil 07_08_09_10: start") @@ -995,8 +985,6 @@ def _do_diffusion_step( offset_provider={ "E2C": self.grid.get_e2c_connectivity(), "C2E2C": self.grid.get_c2e2c_connectivity(), - "C2E2CDim": C2E2CDim, - "E2CDim": E2CDim, }, ) log.debug("running fused stencil 11_12: end") @@ -1014,8 +1002,6 @@ def _do_diffusion_step( offset_provider={ "C2E": self.grid.get_c2e_connectivity(), "E2C": self.grid.get_e2c_connectivity(), - "E2CDim": E2CDim, - "C2EDim": C2EDim, }, ) log.debug("running fused stencil 13_14: end") @@ -1035,7 +1021,6 @@ def _do_diffusion_step( # klevels, # offset_provider={ # "C2E2C": self.grid.get_c2e2c_connectivity(), - # "C2E2CDim": C2E2CDim, # "Koff": KDim, # }, # ) diff --git a/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index 29dd176d0..6482e5952 100644 --- a/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/atm_dyn_iconam/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -210,7 +210,6 @@ def test_calculate_nabla2_and_smag_coefficients_for_vn(): offset_provider={ "E2C2V": mesh.get_e2c2v_offset_provider(), "E2ECV": mesh.get_e2ecv_offset_provider(), - "E2C2VDim": E2C2VDim, }, ) assert np.allclose(kh_smag_e_ref, kh_smag_e) diff --git a/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index 2eb90a21a..24b3a5df9 100644 --- a/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/atm_dyn_iconam/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -52,7 +52,7 @@ def test_mo_intp_rbf_rbf_vec_interpol_vertex(): mesh.n_vertices, 0, mesh.k_level, - offset_provider={"V2E": mesh.get_v2e_offset_provider(), "V2EDim": V2EDim}, + offset_provider={"V2E": mesh.get_v2e_offset_provider()}, ) p_u_out_ref, p_v_out_ref = mo_intp_rbf_rbf_vec_interpol_vertex_numpy( mesh.v2e, np.asarray(p_e_in), np.asarray(ptr_coeff_1), np.asarray(ptr_coeff_2) diff --git a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py index 4ae2fb0b1..b1bd0b63c 100644 --- a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py @@ -61,7 +61,6 @@ def mo_nh_diffusion_stencil_15_numpy( ) -@pytest.mark.skip("new lowering: dims in offset provider") def test_mo_nh_diffusion_stencil_15(): mesh = SimpleMesh() diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py index 4ee09f0e7..040a88ba2 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_20.py @@ -12,7 +12,6 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -import pytest from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_20 import ( mo_solve_nonhydro_stencil_20, @@ -65,7 +64,6 @@ def at_neighbor(i): return z_gradh_exner -@pytest.mark.skip("new lowering: dims in offset provider") def test_mo_solve_nonhydro_stencil_20(): mesh = SimpleMesh() diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py index 9e76b51ec..7e77db024 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_21.py @@ -12,7 +12,6 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -import pytest from gt4py.next.iterator.embedded import constant_field from icon4py.atm_dyn_iconam.mo_solve_nonhydro_stencil_21 import ( @@ -86,7 +85,6 @@ def _apply_index_field(shape, to_index, neighbor_table, offset_field): return z_hydro_corr -@pytest.mark.skip("new lowering: dims in offset provider") def test_mo_solve_nonhydro_stencil_21(): mesh = SimpleMesh() diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 754f55ea4..03c4d2b85 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -1,5 +1,5 @@ # VCS --e git+https://github.com/tehrengruber/gt4py.git@reduce_foast_lowering_complexity#egg=gt4py +-e git+https://github.com/GridTools/gt4py.git@main#egg=gt4py git+https://github.com/GridTools/serialbox#egg=serialbox&subdirectory=src/serialbox-python diff --git a/py2f/src/icon4py/py2f/codegen.py b/py2f/src/icon4py/py2f/codegen.py index f85592c3f..32fab0ddd 100644 --- a/py2f/src/icon4py/py2f/codegen.py +++ b/py2f/src/icon4py/py2f/codegen.py @@ -18,11 +18,11 @@ from gt4py.eve.codegen import TemplatedGenerator from gt4py.next.type_system.type_specifications import ScalarKind -from icon4py.bindings.codegen.type_conversion import ( +from icon4py.icon4pygen.bindings.codegen.type_conversion import ( BUILTIN_TO_CPP_TYPE, BUILTIN_TO_ISO_C_TYPE, ) -from icon4py.bindings.utils import write_string +from icon4py.icon4pygen.bindings.utils import write_string class DimensionType(Node): From 99931bc69b22fb10f2c60dac05d1f2c724271cc9 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 16 May 2023 14:52:51 +0200 Subject: [PATCH 097/263] fix typo in naming --- common/tests/test_gridmanager.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/common/tests/test_gridmanager.py b/common/tests/test_gridmanager.py index ed3070402..55d7cd381 100644 --- a/common/tests/test_gridmanager.py +++ b/common/tests/test_gridmanager.py @@ -337,13 +337,13 @@ def test_gridmanager_eval_e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) gm = init_grid_manager(r04b09_dsl_gridfile) num_edges = gm.get_size(EdgeDim) - serialized_c2e = grid_savepoint.e2c()[0:num_edges, :] + serialized_e2c = grid_savepoint.e2c()[0:num_edges, :] # there are edges at the boundary that have only one # neighboring cell, there are "missing values" in the grid file # and here they do not get substituted in the ICON preprocessing - assert has_invalid_index(serialized_c2e) + assert has_invalid_index(serialized_e2c) assert has_invalid_index(gm.get_e2c_connectivity().table) - assert np.allclose(gm.get_e2c_connectivity().table, serialized_c2e) + assert np.allclose(gm.get_e2c_connectivity().table, serialized_e2c) # c2e: serial, simple, grid @@ -390,7 +390,6 @@ def test_gridmanager_eval_e2c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) gm = init_grid_manager(r04b09_dsl_gridfile) num_edges = gm.get_size(EdgeDim) - # TODO: for boundary edges there are only 2 valid vertices in the serialized data, not 3 three # the "far" (adjacent to edge normal ) is not there. why? # despite that ordering is different assert np.allclose( From 6eee9a2998ade7a81c02baf183237e74b51d0eb4 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 23 May 2023 11:18:36 +0200 Subject: [PATCH 098/263] rename diffusion.timestep to run --- atm_dyn_iconam/src/icon4py/driver/dycore_driver.py | 4 ++-- atm_dyn_iconam/tests/test_diffusion.py | 6 +++--- py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index 57a8d378e..d015271ae 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -121,7 +121,7 @@ def _timestep( self.atmo_non_hydro.do_dynamics_substepping( self.config.dtime, diagnostic_state, prognostic_state ) - self.diffusion.time_step( + self.diffusion.run( diagnostic_state, prognostic_state, self.config.dtime, @@ -144,7 +144,7 @@ def __call__( f"starting time loop for dtime={self.config.dtime} n_timesteps={self.config.n_time_steps}" ) log.info("running initial step to diffuse fields before timeloop starts") - self.diffusion.initial_step( + self.diffusion.initial_run( diagnostic_state, prognostic_state, self.config.dtime, diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 1ee598828..f9d68c924 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -483,7 +483,7 @@ def test_run_diffusion_single_step( metric_state=metric_state, interpolation_state=interpolation_state, ) - diffusion.time_step( + diffusion.run( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime, @@ -574,7 +574,7 @@ def test_diffusion_five_steps( metric_state=metric_state, interpolation_state=interpolation_state, ) - diffusion.initial_step( + diffusion.initial_run( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime, @@ -588,7 +588,7 @@ def test_diffusion_five_steps( cell_areas=cell_geometry.area, ) for _ in range(4): - diffusion.time_step( + diffusion.run( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime, diff --git a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py index 09e1dfce8..fb55f4411 100644 --- a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py +++ b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py @@ -143,13 +143,13 @@ def diffusion_run( theta_v=theta_v, ) if linit: - diffusion.initial_step( + diffusion.initial_run( diagnostic_state, prognostic_state, dtime, ) else: - diffusion.time_step(diagnostic_state, prognostic_state, dtime) + diffusion.run(diagnostic_state, prognostic_state, dtime) class DuplicateInitializationException(Exception): From 9f0aebef6e6cd5d517a753bd2cab73493728458e Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 24 May 2023 16:17:44 +0200 Subject: [PATCH 099/263] add mo_nh_diffusion_stencil_15 to the diffusion.py and diffusion_program.py --- .../src/icon4py/diffusion/diffusion.py | 51 ++++++++++--------- .../icon4py/diffusion/diffusion_program.py | 41 ++++++++------- .../src/icon4py/diffusion/icon_grid.py | 16 +++++- .../icon4py/diffusion/interpolation_state.py | 10 ++-- .../src/icon4py/diffusion/metric_state.py | 6 +-- atm_dyn_iconam/src/icon4py/diffusion/utils.py | 1 - 6 files changed, 74 insertions(+), 51 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index c4f0d093a..ef3010996 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -528,7 +528,7 @@ def initial_step( 0.0, ) - def time_step( + def run( self, diagnostic_state: DiagnosticState, prognostic_state: PrognosticState, @@ -637,10 +637,10 @@ def time_step( prognostic_theta_v=prognostic_state.theta_v, metric_theta_ref_mc=self.metric_state.theta_ref_mc, metric_wgtfac_c=self.metric_state.wgtfac_c, - metric_mask_hdiff=self.metric_state.mask_hdiff, - metric_zd_vertidx=self.metric_state.zd_vertidx, - metric_zd_diffcoef=self.metric_state.zd_diffcoef, - metric_zd_intcoef=self.metric_state.zd_intcoef, + metric_mask_hdiff=self.metric_state.mask_hdiff, # + metric_zd_vertidx=self.metric_state.zd_vertidx, # + metric_zd_diffcoef=self.metric_state.zd_diffcoef, # + metric_zd_intcoef=self.metric_state.zd_intcoef, # interpolation_e_bln_c_s=self.interpolation_state.e_bln_c_s, interpolation_rbf_coeff_1=self.interpolation_state.rbf_coeff_1, interpolation_rbf_coeff_2=self.interpolation_state.rbf_coeff_2, @@ -650,7 +650,7 @@ def time_step( interpolation_nudgecoeff_e=self.interpolation_state.nudgecoeff_e, interpolation_geofac_n2s=self.interpolation_state.geofac_n2s, interpolation_geofac_n2s_c=self.interpolation_state.geofac_n2s_c, - interpolation_geofac_n2s_nbh=self.interpolation_state.geofac_n2s_nbh, + interpolation_geofac_n2s_nbh=self.interpolation_state.geofac_n2s_nbh, # tangent_orientation=tangent_orientation, inverse_primal_edge_lengths=inverse_primal_edge_lengths, inverse_dual_edge_lengths=inverse_dual_edge_length, @@ -705,6 +705,7 @@ def time_step( "C2E": self.grid.get_c2e_connectivity(), "E2C": self.grid.get_e2c_connectivity(), "C2E2C": self.grid.get_c2e2c_connectivity(), + "C2CEC": self.grid.get_c2cec_connectivity(), "C2E2CO": self.grid.get_c2e2co_connectivity(), "Koff": KDim, }, @@ -1006,24 +1007,26 @@ def _do_diffusion_step( ) log.debug("running fused stencil 13_14: end") log.debug("running fused stencil 15: start") - # mo_nh_diffusion_stencil_15( - # self.metric_state.mask_hdiff, - # self.metric_state.zd_vertidx, - # self.metric_state.zd_diffcoef, - # self.interpolation_state.geofac_n2s_c, - # self.interpolation_state.geofac_n2s_nbh, - # self.metric_state.zd_intcoef, - # prognostic_state.theta_v, - # self.z_temp, - # cell_start_nudging, - # cell_end_local, - # 0, - # klevels, - # offset_provider={ - # "C2E2C": self.grid.get_c2e2c_connectivity(), - # "Koff": KDim, - # }, - # ) + mo_nh_diffusion_stencil_15( + mask=self.metric_state.mask_hdiff, + zd_vertoffset=self.metric_state.zd_vertidx, + zd_diffcoef=self.metric_state.zd_diffcoef, + geofac_n2s_c=self.interpolation_state.geofac_n2s_c, + geofac_n2s_nbh=self.interpolation_state.geofac_n2s_nbh, + vcoef=self.metric_state.zd_intcoef, + theta_v=prognostic_state.theta_v, + z_temp=self.z_temp, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={ + "C2CEC": self.grid.get_c2cec_connectivity(), + "C2E2C": self.grid.get_c2e2c_connectivity(), + "Koff": KDim, + }, + ) + log.debug("running fused stencil 15: end") log.debug("running fused stencil update_theta_and_exner: start") update_theta_and_exner( diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py index 755dab9ba..430685af6 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py @@ -39,11 +39,15 @@ from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( _mo_intp_rbf_rbf_vec_interpol_vertex, ) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_15 import ( + _mo_nh_diffusion_stencil_15, +) from icon4py.atm_dyn_iconam.update_theta_and_exner import _update_theta_and_exner from icon4py.common.dimension import ( C2E2CDim, C2E2CODim, C2EDim, + CECDim, CellDim, ECVDim, EdgeDim, @@ -66,10 +70,10 @@ def diffusion_run( prognostic_theta_v: Field[[CellDim, KDim], float], metric_theta_ref_mc: Field[[CellDim, KDim], float], metric_wgtfac_c: Field[[CellDim, KDim], float], - metric_mask_hdiff: Field[[CellDim, KDim], int], - metric_zd_vertidx: Field[[CellDim, C2E2CDim, KDim], int], + metric_mask_hdiff: Field[[CellDim, KDim], bool], + metric_zd_vertidx: Field[[CECDim, KDim], int32], metric_zd_diffcoef: Field[[CellDim, KDim], float], - metric_zd_intcoef: Field[[CellDim, C2E2CDim, KDim], float], + metric_zd_intcoef: Field[[CECDim, KDim], float], interpolation_e_bln_c_s: Field[[CellDim, C2EDim], float], interpolation_rbf_coeff_1: Field[[VertexDim, V2EDim], float], interpolation_rbf_coeff_2: Field[[VertexDim, V2EDim], float], @@ -79,7 +83,7 @@ def diffusion_run( interpolation_nudgecoeff_e: Field[[EdgeDim], float], interpolation_geofac_n2s: Field[[CellDim, C2E2CODim], float], interpolation_geofac_n2s_c: Field[[CellDim], float], - interpolation_geofac_n2s_nbh: Field[[CellDim, C2E2CDim], float], + interpolation_geofac_n2s_nbh: Field[[CECDim], float], tangent_orientation: Field[[EdgeDim], float], inverse_primal_edge_lengths: Field[[EdgeDim], float], inverse_dual_edge_lengths: Field[[EdgeDim], float], @@ -289,20 +293,21 @@ def diffusion_run( ) # MO_NH_DIFFUSION_STENCIL_15: as_offset index fields! - # _mo_nh_diffusion_stencil_15( - # metric_mask_hdiff, - # metric_zd_vertidx, - # metric_zd_diffcoef, - # interpolation_geofac_n2s_c, - # interpolation_geofac_n2s_nbh, - # metric_zd_intcoef, - # prognostic_theta_v, - # local_z_temp, - # domain = { - # CellDim:(cell_startindex_nudging, cell_endindex_local), - # KDim: (0,nlev) - # } - # ) + _mo_nh_diffusion_stencil_15( + metric_mask_hdiff, + metric_zd_vertidx, + metric_zd_diffcoef, + interpolation_geofac_n2s_c, + interpolation_geofac_n2s_nbh, + metric_zd_intcoef, + prognostic_theta_v, + local_z_temp, + domain={ + CellDim: (cell_startindex_nudging, cell_endindex_local), + KDim: (0, nlev), + }, + out=local_z_temp, + ) _update_theta_and_exner( local_z_temp, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index 07f6625e5..670380bef 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -18,7 +18,14 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import NeighborTableOffsetProvider -from icon4py.common.dimension import CellDim, ECVDim, EdgeDim, KDim, VertexDim +from icon4py.common.dimension import ( + CECDim, + CellDim, + ECVDim, + EdgeDim, + KDim, + VertexDim, +) from icon4py.diffusion.horizontal import HorizontalMeshSize @@ -179,6 +186,13 @@ def get_e2ecv_connectivity(self): v2ecv_table, EdgeDim, ECVDim, v2ecv_table.shape[1] ) + def get_c2cec_connectivity(self): + old_shape = self.connectivities["c2e2c"].shape + c2cec_table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) + return NeighborTableOffsetProvider( + c2cec_table, CellDim, CECDim, c2cec_table.shape[1] + ) + class VerticalModelParams: def __init__(self, vct_a: Field[[KDim], float], rayleigh_damping_height: float): diff --git a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py index 2ac0ff4a7..4a864cb72 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py @@ -18,9 +18,9 @@ from gt4py.next.iterator.embedded import np_as_located_field from icon4py.common.dimension import ( - C2E2CDim, C2E2CODim, C2EDim, + CECDim, CellDim, EdgeDim, V2EDim, @@ -64,7 +64,9 @@ def geofac_n2s_c(self) -> Field[[CellDim], float]: return np_as_located_field(CellDim)(np.asarray(self.geofac_n2s)[:, 0]) @property - def geofac_n2s_nbh(self) -> Field[[CellDim, C2E2CDim], float]: - return np_as_located_field(CellDim, C2E2CDim)( - np.asarray(self.geofac_n2s)[:, 1:] + def geofac_n2s_nbh(self) -> Field[[CECDim], float]: + np_array = np.asarray(self.geofac_n2s)[:, 1:] + old_shape = np_array.shape + return np_as_located_field(CECDim)( + np_array.reshape(old_shape[0] * old_shape[1]) ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py b/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py index 069863583..33d6672f9 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py @@ -15,7 +15,7 @@ from gt4py.next.common import Field -from icon4py.common.dimension import C2E2CDim, CellDim, KDim +from icon4py.common.dimension import CECDim, CellDim, KDim @dataclass @@ -25,6 +25,6 @@ class MetricState: [CellDim, KDim], float ] # weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_c) mask_hdiff: Field[[CellDim, KDim], bool] - zd_vertidx: Field[[CellDim, C2E2CDim, KDim], int] + zd_vertidx: Field[[CECDim, KDim], int] zd_diffcoef: Field[[CellDim, KDim], float] - zd_intcoef: Field[[CellDim, C2E2CDim, KDim], float] + zd_intcoef: Field[[CECDim, KDim], float] diff --git a/atm_dyn_iconam/src/icon4py/diffusion/utils.py b/atm_dyn_iconam/src/icon4py/diffusion/utils.py index 3f6a2b706..b03f01051 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/utils.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/utils.py @@ -107,7 +107,6 @@ def _setup_runtime_diff_multfac_vn( return broadcast(minimum(con, dyn), (KDim,)) -# @field_operator(backend=run_gtfn) @field_operator def _setup_initial_diff_multfac_vn( k4: float, hdiff_efdt_ratio: float From 77eafbb1cc63758051d454872f532c6eba93e919 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 24 May 2023 16:18:04 +0200 Subject: [PATCH 100/263] downgrade typing-extensions --- base-requirements-dev.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 03c4d2b85..b9fe27a5f 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -14,7 +14,7 @@ flake8-mutable>=1.2.0 flake8-rst-docstrings>=0.0.14 isort~=5.10 mypy>=0.942 -typing-extensions>=4.5 +typing-extensions==4.5.0 pre-commit~=2.15 pytest>=6.1 pytest-cache>=1.0 From 6740747af518470d5c5bf84b551c0c25bbcf7e16 Mon Sep 17 00:00:00 2001 From: Magdalena Date: Thu, 25 May 2023 08:43:15 +0200 Subject: [PATCH 101/263] (refactoring): rename local_boundary to lateral_boundary --- .../src/icon4py/diffusion/diffusion.py | 16 ++--- .../src/icon4py/diffusion/horizontal.py | 2 +- atm_dyn_iconam/tests/test_icon_grid.py | 60 +++++++++---------- .../tests/test_mo_nh_diffusion_stencil_15.py | 1 - 4 files changed, 39 insertions(+), 40 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index c4f0d093a..813463602 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -608,7 +608,7 @@ def time_step( vertex_endindex_local, ) = self.grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 3, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, HorizontalMarkerIndex.local(VertexDim), ) @@ -617,13 +617,13 @@ def time_step( vertex_endindex_local_minus1, ) = self.grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 1, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, HorizontalMarkerIndex.local(VertexDim) - 1, ) edge_start_lb_plus4, _ = self.grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, ) log.info("diffusion program: start") diff_prog.diffusion_run( @@ -777,8 +777,8 @@ def _do_diffusion_step( edge_start_lb_plus4, _ = self.grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, ) ( @@ -795,7 +795,7 @@ def _do_diffusion_step( vertex_end_local, ) = self.grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 3, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, HorizontalMarkerIndex.local(VertexDim), ) ( @@ -803,7 +803,7 @@ def _do_diffusion_step( vertex_end_local_minus1, ) = self.grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 1, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, HorizontalMarkerIndex.local(VertexDim) - 1, ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py index 27f07105e..82c3a206e 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py @@ -58,7 +58,7 @@ class HorizontalMarkerIndex: _END_VERTICES: Final[int] = 0 @classmethod - def local_boundary(cls, dim: Dimension) -> int: + def lateral_boundary(cls, dim: Dimension) -> int: match (dim): case (dimension.CellDim): return cls._LOCAL_BOUNDARY_CELLS diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index 256ded7aa..c3de4fbb1 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -53,13 +53,13 @@ def test_horizontal_grid_cell_indices(icon_grid): ) # lb+1 assert icon_grid.get_indices_from_to( CellDim, - HorizontalMarkerIndex.local_boundary(CellDim) + 1, - HorizontalMarkerIndex.local_boundary(CellDim) + 1, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, ) == (850, 1688) assert icon_grid.get_indices_from_to( CellDim, - HorizontalMarkerIndex.local_boundary(CellDim) + 2, - HorizontalMarkerIndex.local_boundary(CellDim) + 2, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, ) == ( 1688, 2511, @@ -125,64 +125,64 @@ def test_horizontal_edge_indices(icon_grid): ) # nudging assert icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 7, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 7, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, ) == ( 4184, 4989, ) # lb +7 assert icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 6, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 6, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, ) == ( 3777, 4184, ) # lb +6 assert icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 5, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 5, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, ) == ( 2954, 3777, ) # lb +5 assert icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, ) == ( 2538, 2954, ) # lb +4 assert icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 3, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 3, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, ) == ( 1700, 2538, ) # lb +3 assert icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 2, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 2, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, ) == ( 1278, 1700, ) # lb +2 assert icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 1, - HorizontalMarkerIndex.local_boundary(EdgeDim) + 1, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, ) == ( 428, 1278, ) # lb +1 assert icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local_boundary(EdgeDim), - HorizontalMarkerIndex.local_boundary(EdgeDim), + HorizontalMarkerIndex.lateral_boundary(EdgeDim), + HorizontalMarkerIndex.lateral_boundary(EdgeDim), ) == ( 0, 428, @@ -209,28 +209,28 @@ def test_horizontal_vertex_indices(icon_grid): assert icon_grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim), - HorizontalMarkerIndex.local_boundary(VertexDim), + HorizontalMarkerIndex.lateral_boundary(VertexDim), + HorizontalMarkerIndex.lateral_boundary(VertexDim), ) == (0, 428) assert icon_grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 1, - HorizontalMarkerIndex.local_boundary(VertexDim) + 1, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, ) == (428, 850) assert icon_grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 2, - HorizontalMarkerIndex.local_boundary(VertexDim) + 2, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, ) == (850, 1266) assert icon_grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 3, - HorizontalMarkerIndex.local_boundary(VertexDim) + 3, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, ) == (1266, 1673) assert icon_grid.get_indices_from_to( VertexDim, - HorizontalMarkerIndex.local_boundary(VertexDim) + 4, - HorizontalMarkerIndex.local_boundary(VertexDim) + 4, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, ) == (1673, 2071) assert icon_grid.get_indices_from_to( diff --git a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py index b1bd0b63c..829fb9391 100644 --- a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py @@ -12,7 +12,6 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np -import pytest from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_15 import ( mo_nh_diffusion_stencil_15, From 88dbc7481bfa29dae3c8a1a8d9d23aa9a6510517 Mon Sep 17 00:00:00 2001 From: Magdalena Date: Thu, 25 May 2023 08:53:11 +0200 Subject: [PATCH 102/263] (fix) renable test in icon4pygen codegen tests were set to `skip` because they didn't run while we were running on a special `gt4py` branch --- pyutils/tests/icon4pygen/test_codegen.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/pyutils/tests/icon4pygen/test_codegen.py b/pyutils/tests/icon4pygen/test_codegen.py index d262d893f..a63dc2ebf 100644 --- a/pyutils/tests/icon4pygen/test_codegen.py +++ b/pyutils/tests/icon4pygen/test_codegen.py @@ -109,7 +109,6 @@ def check_code_was_generated(stencil_name: str) -> None: check_cpp_codegen(f"{stencil_name}.cpp") -@pytest.mark.skip("raises exception due to dims in offset provider") @pytest.mark.parametrize(("stencil_module", "stencil_name"), atm_dyn_iconam_fencils()) def test_codegen_atm_dyn_iconam(cli, stencil_module, stencil_name) -> None: module_path = get_stencil_module_path(stencil_module, stencil_name) @@ -124,13 +123,3 @@ def test_invalid_module_path(cli) -> None: result = cli.invoke(main, [module_path, BLOCK_SIZE, LEVELS_PER_THREAD, OUTPATH]) assert result.exit_code == 1 assert isinstance(result.exception, ModuleNotFoundError) - - -@pytest.mark.skip("raises exception due to dims in offset provider") -def test_codegen_mo_nh_diffusion_stencil_14(cli) -> None: - stencil_name = "mo_nh_diffusion_stencil_14" - module_path = get_stencil_module_path("atm_dyn_iconam", stencil_name) - with cli.isolated_filesystem(): - result = cli.invoke(main, [module_path, BLOCK_SIZE, LEVELS_PER_THREAD, OUTPATH]) - assert result.exit_code == 0 - check_code_was_generated(stencil_name) From 9bcd4e97d8a59682a50e269cc607b70b8115a450 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 26 May 2023 10:48:14 +0200 Subject: [PATCH 103/263] - add mpi4py dependency - init mpi and get procs info --- .../src/icon4py/decomposition/decomposed.py | 18 +++++++++++ .../src/icon4py/driver/dycore_driver.py | 2 +- .../src/icon4py/driver/parallel_setup.py | 32 +++++++++++++++++++ .../tests/mpi_tests/test_parallel_setup.py | 27 ++++++++++++++++ base-requirements-dev.txt | 2 ++ 5 files changed, 80 insertions(+), 1 deletion(-) create mode 100644 atm_dyn_iconam/src/icon4py/decomposition/decomposed.py create mode 100644 atm_dyn_iconam/src/icon4py/driver/parallel_setup.py create mode 100644 atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py new file mode 100644 index 000000000..0eedd496e --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -0,0 +1,18 @@ +import mpi4py.MPI + + +class ProcessProperties(): + def __init__(self, name='', rank = 0): + self._communicator_name: str = name + self._rank: int = rank + + def rank(self): + return self._rank + + def comm_name(self): + return self._communicator_name + @classmethod + def from_mpi_comm(cls, comm: mpi4py.MPI.Comm): + return ProcessProperties(comm.Get_name(), comm.Get_rank()) + + diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index 57a8d378e..2ceada3f6 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -194,7 +194,7 @@ def initialize(n_time_steps, file_path: Path): (edge_geometry, cell_geometry, vertical_geometry) = read_geometry_fields(file_path) (metric_state, interpolation_state) = read_static_fields(file_path) - log.info("initializing dycore") + log.info("initializing diffusion") diffusion_params = DiffusionParams(config.diffusion_config) diffusion = Diffusion(run_program=False) diffusion.init( diff --git a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py new file mode 100644 index 000000000..a63837cda --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py @@ -0,0 +1,32 @@ +from typing import Optional, Union + +import mpi4py +from mpi4py.MPI import Comm + + +from icon4py.decomposition.decomposed import ProcessProperties + +mpi4py.rc.initialize = False + +CommId = Union[int, Comm, None] + +def get_processor_properties(comm_id: CommId = None): + init_mpi() + + def _get_current_comm_or_comm_world(comm_id: CommId)->Comm: + if isinstance(comm_id, int): + comm = Comm.f2py(comm_id) + elif isinstance(comm_id, Comm): + comm = comm_id + else: + comm = mpi4py.MPI.COMM_WORLD + return comm + + current_comm = _get_current_comm_or_comm_world(comm_id) + return ProcessProperties.from_mpi_comm(current_comm) + + +def init_mpi(): + from mpi4py import MPI + if not MPI.Is_initialized(): + MPI.Init() diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py new file mode 100644 index 000000000..c1145b4f0 --- /dev/null +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -0,0 +1,27 @@ +import mpi4py +import pytest + +from icon4py.decomposition.decomposed import ProcessProperties +from icon4py.driver.parallel_setup import get_processor_properties + +""" +running tests with mpi: + +mpirun -np 2 python -m pytest --with-mpi tests/test_parallel_setup.py + +mpirun -np 2 pytest -v --with-mpi tests/mpi_tests/ + + +""" + +@pytest.fixture +def mpi(): + from mpi4py import MPI + mpi4py.rc.initialize = False + return MPI + +@pytest.mark.mpi() +def test_processor_properties_from_comm_world(mpi): + props = get_processor_properties() + assert props.rank() < mpi.COMM_WORLD.Get_size() + assert props.comm_name() == mpi.COMM_WORLD.Get_name() diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 03c4d2b85..b34f5b856 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -25,3 +25,5 @@ setuptools>=40.8.0 wheel>=0.37.1 tox >= 3.25 wget >= 3.2 +mpi4py<=3.1.4 +pytest-mpi>=0.6 From 3e786a43d93e9c649a8c1b98bed90d00ad21a519 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 26 May 2023 18:12:38 +0200 Subject: [PATCH 104/263] read in decomposition data (WIP 1) --- .../src/icon4py/decomposition/decomposed.py | 3 +- .../src/icon4py/diffusion/diffusion.py | 6 +- .../src/icon4py/diffusion/icon_grid.py | 7 +- atm_dyn_iconam/src/icon4py/diffusion/utils.py | 8 ++ atm_dyn_iconam/src/icon4py/driver/io_utils.py | 31 ++++- .../src/icon4py/driver/parallel_setup.py | 35 +++++- atm_dyn_iconam/tests/conftest.py | 16 ++- .../tests/mpi_tests/test_parallel_setup.py | 29 +++-- atm_dyn_iconam/tests/test_io_utils.py | 5 +- .../src/icon4py/testutils/serialbox_utils.py | 108 +++++++++++++----- 10 files changed, 191 insertions(+), 57 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index 0eedd496e..00db3eaeb 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -5,10 +5,11 @@ class ProcessProperties(): def __init__(self, name='', rank = 0): self._communicator_name: str = name self._rank: int = rank - + @property def rank(self): return self._rank + @property def comm_name(self): return self._communicator_name @classmethod diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 813463602..f4282fe72 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -763,11 +763,7 @@ def _do_diffusion_step( HorizontalMarkerIndex.local(CellDim), ) - cell_start_nudging, _ = self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.local(CellDim), - ) + edge_start_nudging_plus_one, edge_end_local = self.grid.get_indices_from_to( EdgeDim, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index 07f6625e5..a7a0b2528 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -20,6 +20,7 @@ from icon4py.common.dimension import CellDim, ECVDim, EdgeDim, KDim, VertexDim from icon4py.diffusion.horizontal import HorizontalMeshSize +from icon4py.diffusion.utils import builder # TODO @magdalena keep naming grid vs mesh consistent @@ -69,12 +70,6 @@ def num_cells(self): return self._horizontal.num_cells -def builder(func): - def wrapper(self, *args, **kwargs): - func(self, *args, **kwargs) - return self - - return wrapper class IconGrid: diff --git a/atm_dyn_iconam/src/icon4py/diffusion/utils.py b/atm_dyn_iconam/src/icon4py/diffusion/utils.py index 3f6a2b706..1ed208af5 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/utils.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/utils.py @@ -265,3 +265,11 @@ def init_nabla2_factor_in_upper_damping_zone( ** 4 ) return np_as_located_field(KDim)(buffer) + + +def builder(func): + def wrapper(self, *args, **kwargs): + func(self, *args, **kwargs) + return self + + return wrapper diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index 858f9c9e7..647964bd9 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -15,15 +15,18 @@ from enum import Enum from pathlib import Path +from icon4py.decomposition.decomposed import ProcessProperties from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.horizontal import CellParams, EdgeParams from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.prognostic_state import PrognosticState +from icon4py.driver.parallel_setup import DecompositionInfo from icon4py.testutils import serialbox_utils from icon4py.testutils.serialbox_utils import IconSerialDataProvider +SERIALBOX_ONLY_MSG = "Only ser_type='sb (Serialbox)' is implemented so far." SIMULATION_START_DATE = "2021-06-20T12:00:10.000" log = logging.getLogger(__name__) @@ -105,7 +108,31 @@ def read_geometry_fields( ) return edge_geometry, cell_geometry, vertical_geometry else: - raise NotImplementedError("Only ser_type='sb' is implemented so far.") + raise NotImplementedError(SERIALBOX_ONLY_MSG) + + +# /home/magdalena/data/exclaim/dycore/mch_ch_r04b09_dsl/node2/mch_ch_r04b09_dsl/icon_grid +def read_decomp_info( + path: Path, procs_props: ProcessProperties, ser_type = SerializationType.SB, +)->DecompositionInfo: + if ser_type == SerializationType.SB: + sp = serialbox_utils.IconSerialDataProvider("icon_grid", str(path.absolute()), True, + procs_props.rank) + return \ + sp.from_savepoint_grid().construct_decomposition_info() + else: + raise NotImplementedError(SERIALBOX_ONLY_MSG) + +def read_grid( + path: Path, procs_props: ProcessProperties, ser_type = SerializationType.SB, +)->IconGrid: + if ser_type == SerializationType.SB: + sp = serialbox_utils.IconSerialDataProvider("icon_grid", str(path.absolute()), True, + procs_props.rank) + return \ + sp.from_savepoint_grid().construct_icon_grid() + else: + raise NotImplementedError(SERIALBOX_ONLY_MSG) def read_static_fields( @@ -130,7 +157,7 @@ def read_static_fields( interpolation_state = sp.construct_interpolation_state() return metric_state, interpolation_state else: - raise NotImplementedError("Only ser_type='sb' is implemented so far.") + raise NotImplementedError(SERIALBOX_ONLY_MSG) def configure_logging(run_path: str, start_time): diff --git a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py index a63837cda..c86b98705 100644 --- a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py @@ -1,11 +1,15 @@ -from typing import Optional, Union +from enum import Enum +from typing import Union import mpi4py +import numpy as np +import numpy.ma as ma +from gt4py.next.common import Dimension from mpi4py.MPI import Comm from icon4py.decomposition.decomposed import ProcessProperties - +from icon4py.diffusion.utils import builder mpi4py.rc.initialize = False CommId = Union[int, Comm, None] @@ -30,3 +34,30 @@ def init_mpi(): from mpi4py import MPI if not MPI.Is_initialized(): MPI.Init() + + + + +class DecompositionInfo: + class EntryType(int, Enum): + ALL = 0, + OWNED = 1, + HALO = 2 + @builder + def with_dimension(self, dim:Dimension, global_index:np.ndarray, owner_mask:np.ndarray): + masked_global_index = ma.array(global_index, mask=owner_mask) + self._global_index[dim] = masked_global_index + + def __init__(self): + self._global_index = {} + + + + def global_index(self, dim:Dimension, type:EntryType = EntryType.ALL): + match(type): + case DecompositionInfo.EntryType.ALL: + return ma.getdata(self._global_index[dim], subok=False) + case _: + raise NotImplementedError() + + diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index ba7d2b336..e9cd9fd58 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -10,21 +10,30 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import pathlib import tarfile from pathlib import Path +import mpi4py import pytest import wget from icon4py.diffusion.diffusion import DiffusionConfig -from icon4py.testutils.serialbox_utils import IconSerialDataProvider - +from icon4py.driver.io_utils import read_decomp_info, SerializationType +from icon4py.driver.parallel_setup import get_processor_properties, DecompositionInfo +from icon4py.testutils.serialbox_utils import IconSerialDataProvider, IconGridSavePoint data_uri = "https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download" data_path = Path(__file__).parent.joinpath("ser_icondata") extracted_path = data_path.joinpath("mch_ch_r04b09_dsl/ser_data") data_file = data_path.joinpath("mch_ch_r04b09_dsl_v2.tar.gz").name +@pytest.fixture +def mpi(): + from mpi4py import MPI + mpi4py.rc.initialize = False + return MPI + @pytest.fixture(scope="session") def setup_icon_data(): @@ -52,6 +61,9 @@ def data_provider(setup_icon_data) -> IconSerialDataProvider: return IconSerialDataProvider("icon_pydycore", str(extracted_path), True) + + + @pytest.fixture def linit(): """ diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index c1145b4f0..4fcdd04c3 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -1,8 +1,12 @@ +import pathlib + import mpi4py import pytest +from icon4py.common.dimension import CellDim from icon4py.decomposition.decomposed import ProcessProperties -from icon4py.driver.parallel_setup import get_processor_properties +from icon4py.driver.io_utils import SerializationType, read_decomp_info, read_grid +from icon4py.driver.parallel_setup import get_processor_properties, DecompositionInfo """ running tests with mpi: @@ -14,14 +18,21 @@ """ -@pytest.fixture -def mpi(): - from mpi4py import MPI - mpi4py.rc.initialize = False - return MPI -@pytest.mark.mpi() + + +@pytest.mark.mpi def test_processor_properties_from_comm_world(mpi): props = get_processor_properties() - assert props.rank() < mpi.COMM_WORLD.Get_size() - assert props.comm_name() == mpi.COMM_WORLD.Get_name() + assert props.rank < mpi.COMM_WORLD.Get_size() + assert props.comm_name == mpi.COMM_WORLD.Get_name() + + +def test_decomposition_info(): + path = pathlib.Path( + "/home/magdalena/data/exclaim/dycore/mch_ch_r04b09_dsl/node2/mch_ch_r04b09_dsl/icon_grid") + props = get_processor_properties() + decomposition_info = read_decomp_info(path, props, SerializationType.SB) + icon_grid = read_grid(path, props, SerializationType.SB) + global_index = decomposition_info.global_index(CellDim, DecompositionInfo.EntryType.ALL) + assert global_index.shape[0] == icon_grid.num_cells() diff --git a/atm_dyn_iconam/tests/test_io_utils.py b/atm_dyn_iconam/tests/test_io_utils.py index a03d810db..e98578628 100644 --- a/atm_dyn_iconam/tests/test_io_utils.py +++ b/atm_dyn_iconam/tests/test_io_utils.py @@ -23,7 +23,7 @@ read_icon_grid, read_static_fields, ) - +from icon4py.driver.parallel_setup import get_processor_properties test_data_path = pathlib.Path(__file__).parent.joinpath( "./ser_icondata/mch_ch_r04b09_dsl/ser_data" @@ -115,3 +115,6 @@ def assert_interpolation_state_fields(interpolation_state): assert interpolation_state.rbf_coeff_2 assert interpolation_state.rbf_coeff_1 assert interpolation_state.geofac_n2s_c + + + diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index ad312d679..336867fad 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -18,6 +18,7 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import np_as_located_field +from icon4py.common import dimension from icon4py.common.dimension import ( C2E2CDim, C2E2CODim, @@ -43,6 +44,7 @@ from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.prognostic_state import PrognosticState +from icon4py.driver.parallel_setup import DecompositionInfo from icon4py.testutils.utils import as_1D_sparse_field @@ -68,7 +70,7 @@ def _read_int32_shift1(self, name: str): """ Read a index field and shift it by -1. - use for start indeces: the shift accounts for the zero based python + use for start indices: the shift accounts for the zero based python values are converted to int32 """ return (self.serializer.read(name, self.savepoint) - 1).astype(int32) @@ -83,6 +85,9 @@ def _read_int32(self, name: str): """ return self.serializer.read(name, self.savepoint).astype(int32) + def _read_bool(self, name: str): + return self.serializer.read(name, self.savepoint).astype(bool) + def read_int(self, name: str): buffer = self.serializer.read(name, self.savepoint).astype(int) self.log.debug(f"{name} {buffer.shape}") @@ -146,48 +151,87 @@ def edge_end_index(self): def print_connectivity_info(self, name: str, ar: np.ndarray): self.log.debug(f" connectivity {name} {ar.shape}") - def refin_ctrl(self, dim: Dimension): - if dim == CellDim: - return self.serializer.read("c_refin_ctl", self.savepoint) - elif dim == EdgeDim: - return self.serializer.read("e_refin_ctl", self.savepoint) - elif dim == VertexDim: - return self.serializer.read("v_refin_ctl", self.savepoint) - else: - return None - def c2e(self): + return self._get_connectivity_array("c2e") - return self._get_connectiviy_array("c2e") - - def _get_connectiviy_array(self, name: str): + def _get_connectivity_array(self, name: str): connectivity = self.serializer.read(name, self.savepoint) - 1 self.log.debug(f" connectivity {name} : {connectivity.shape}") return connectivity def c2e2c(self): - return self._get_connectiviy_array("c2e2c") + return self._get_connectivity_array("c2e2c") def e2c(self): - return self._get_connectiviy_array("e2c") + return self._get_connectivity_array("e2c") def e2v(self): # array "e2v" is actually e2c2v - v_ = self._get_connectiviy_array("e2v")[:, 0:2] + v_ = self._get_connectivity_array("e2v")[:, 0:2] self.log.debug(f"real e2v {v_.shape}") return v_ def e2c2v(self): # array "e2v" is actually e2c2v - return self._get_connectiviy_array("e2v") + return self._get_connectivity_array("e2v") def v2e(self): - return self._get_connectiviy_array("v2e") + return self._get_connectivity_array("v2e") + + def refin_ctrl(self, dim: Dimension): + field_name = "refin_ctl" + return self._read_field_for_dim(field_name, self._read_int32, dim) + + def num(self, dim: Dimension): + match (dim): + case dimension.CellDim: + return self.serializer.read("num_cells", savepoint=self.savepoint).astype(int32)[0] + case dimension.EdgeDim: + return self.serializer.read("num_edges", savepoint=self.savepoint).astype(int32)[0] + case dimension.VertexDim: + return self.serializer.read("num_vert", savepoint=self.savepoint).astype(int32)[0] + case dimension.KDim: + return self.get_metadata("nlev") + case _: + raise NotImplementedError(f"only {CellDim, EdgeDim, VertexDim, KDim} are supported") + + + + + def _read_field_for_dim(self, field_name, read_func, dim): + match (dim): + case dimension.CellDim: + return read_func(f"c_{field_name}") + case dimension.EdgeDim: + return read_func(f"e_{field_name}") + case dimension.VertexDim: + return read_func(f"v_{field_name}") + case _: + raise NotImplementedError( + f"only {dimension.CellDim, dimension.EdgeDim, dimension.VertexDim} are handled") + + def owner_mask(self, dim: Dimension): + field_name = "owner_mask" + mask = self._read_field_for_dim(field_name, self._read_bool, dim) + return np.squeeze(mask) + + def global_index(self, dim: Dimension): + field_name = "glb_index" + return self._read_field_for_dim(field_name, self._read_int32_shift1, dim) + def decomp_domain(self, dim): + field_name = "decomp_domain" + return self._read_field_for_dim(field_name, self._read_int32, dim) + + def construct_decomposition_info(self): + index = self.global_index(CellDim) + num_cells = self.num(CellDim) + mask = self.owner_mask(CellDim)[0:num_cells] + + return DecompositionInfo().with_dimension(CellDim, index, mask) + def construct_icon_grid(self) -> IconGrid: - sp_meta = self.get_metadata( - "nproma", "nlev", "num_vert", "num_cells", "num_edges" - ) + cell_starts = self.cells_start_index() cell_ends = self.cells_end_index() vertex_starts = self.vertex_start_index() @@ -196,11 +240,11 @@ def construct_icon_grid(self) -> IconGrid: edge_ends = self.edge_end_index() config = MeshConfig( HorizontalMeshSize( - num_vertices=sp_meta["nproma"], # or rather "num_vert" - num_cells=sp_meta["nproma"], # or rather "num_cells" - num_edges=sp_meta["nproma"], # or rather "num_edges" + num_vertices=self.num(VertexDim), + num_cells=self.num(CellDim), + num_edges=self.num(EdgeDim), ), - VerticalMeshConfig(num_lev=sp_meta["nlev"]), + VerticalMeshConfig(num_lev=self.num(KDim)), ) c2e2c = self.c2e2c() c2e2c0 = np.column_stack((c2e2c, (np.asarray(range(c2e2c.shape[0]))))) @@ -405,8 +449,8 @@ def exner(self): class IconSerialDataProvider: - def __init__(self, fname_prefix, path=".", do_print=False): - self.rank = 0 + def __init__(self, fname_prefix, path=".", do_print=False, mpi_rank=0): + self.rank = mpi_rank self.serializer: ser.Serializer = None self.file_path: str = path self.fname = f"{fname_prefix}_rank{str(self.rank)}" @@ -427,9 +471,13 @@ def print_info(self): self.log.info(f"FIELDNAMES: {self.serializer.fieldnames()}") def from_savepoint_grid(self) -> IconGridSavePoint: - savepoint = self.serializer.savepoint["icon-grid"].id[1].as_savepoint() + savepoint = self._get_icon_grid_savepoint() return IconGridSavePoint(savepoint, self.serializer) + def _get_icon_grid_savepoint(self): + savepoint = self.serializer.savepoint["icon-grid"].id[1].as_savepoint() + return savepoint + def from_savepoint_diffusion_init( self, linit: bool, date: str ) -> IconDiffusionInitSavepoint: @@ -451,3 +499,5 @@ def from_savepoint_diffusion_exit( .as_savepoint() ) return IconDiffusionExitSavepoint(savepoint, self.serializer) + + From 900d3812ded2e5e25cd7997d14112bbe60b2e789 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 30 May 2023 09:02:50 +0200 Subject: [PATCH 105/263] add return types in icon_grid.py --- atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index a7a0b2528..57649c589 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -46,27 +46,27 @@ def __init__( self._horizontal = horizontal_config @property - def limited_area(self): + def limited_area(self)->bool: return self._limited_area @property - def num_k_levels(self): + def num_k_levels(self)->int: return self._vertical.num_lev @property - def n_shift_total(self): + def n_shift_total(self)->int: return self._n_shift_total @property - def num_vertices(self): + def num_vertices(self)->int: return self._horizontal.num_vertices @property - def num_edges(self): + def num_edges(self)->int: return self._horizontal.num_edges @property - def num_cells(self): + def num_cells(self)->int: return self._horizontal.num_cells From 18abe03431efa4481c3f181418c2d3f6b56cee81 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 30 May 2023 09:03:57 +0200 Subject: [PATCH 106/263] (fix/refactor) - use same error message for non sb serialization type, - get value for nlev --- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 2 +- atm_dyn_iconam/tests/test_io_utils.py | 2 +- testutils/src/icon4py/testutils/serialbox_utils.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index 647964bd9..85f7a3022 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -56,7 +56,7 @@ def read_icon_grid(path: Path, ser_type=SerializationType.SB) -> IconGrid: .construct_icon_grid() ) else: - raise NotImplementedError("Only ser_type='sb' is implemented so far.") + raise NotImplementedError(SERIALBOX_ONLY_MSG) def read_initial_state( diff --git a/atm_dyn_iconam/tests/test_io_utils.py b/atm_dyn_iconam/tests/test_io_utils.py index e98578628..6106b02a7 100644 --- a/atm_dyn_iconam/tests/test_io_utils.py +++ b/atm_dyn_iconam/tests/test_io_utils.py @@ -34,7 +34,7 @@ "read_fun", (read_geometry_fields, read_static_fields, read_icon_grid) ) def test_read_geometry_fields_not_implemented_type(read_fun): - with pytest.raises(NotImplementedError, match=r"Only ser_type='sb'"): + with pytest.raises(NotImplementedError): read_fun(path=test_data_path, ser_type=SerializationType.NC) diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 336867fad..fd9bf7a72 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -191,7 +191,7 @@ def num(self, dim: Dimension): case dimension.VertexDim: return self.serializer.read("num_vert", savepoint=self.savepoint).astype(int32)[0] case dimension.KDim: - return self.get_metadata("nlev") + return self.get_metadata("nlev")["nlev"] case _: raise NotImplementedError(f"only {CellDim, EdgeDim, VertexDim, KDim} are supported") From bcccb0cbe2fd94e92f7e1378d95657be812cdb51 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 30 May 2023 10:16:57 +0200 Subject: [PATCH 107/263] pre-commit fixes --- .../src/icon4py/decomposition/decomposed.py | 21 ++++++++-- .../src/icon4py/diffusion/diffusion.py | 2 - .../src/icon4py/diffusion/icon_grid.py | 14 +++---- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 30 ++++++++------ .../src/icon4py/driver/parallel_setup.py | 39 ++++++++++++------- atm_dyn_iconam/tests/conftest.py | 18 ++++++--- .../tests/mpi_tests/test_parallel_setup.py | 28 ++++++++++--- atm_dyn_iconam/tests/test_io_utils.py | 4 +- .../src/icon4py/testutils/serialbox_utils.py | 28 +++++++------ 9 files changed, 119 insertions(+), 65 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index 00db3eaeb..5bf300061 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -1,10 +1,24 @@ +# 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 mpi4py.MPI -class ProcessProperties(): - def __init__(self, name='', rank = 0): +class ProcessProperties: + def __init__(self, name="", rank=0): self._communicator_name: str = name self._rank: int = rank + @property def rank(self): return self._rank @@ -12,8 +26,7 @@ def rank(self): @property def comm_name(self): return self._communicator_name + @classmethod def from_mpi_comm(cls, comm: mpi4py.MPI.Comm): return ProcessProperties(comm.Get_name(), comm.Get_rank()) - - diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index f4282fe72..610402429 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -763,8 +763,6 @@ def _do_diffusion_step( HorizontalMarkerIndex.local(CellDim), ) - - edge_start_nudging_plus_one, edge_end_local = self.grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index 57649c589..937dcfc1d 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -46,32 +46,30 @@ def __init__( self._horizontal = horizontal_config @property - def limited_area(self)->bool: + def limited_area(self) -> bool: return self._limited_area @property - def num_k_levels(self)->int: + def num_k_levels(self) -> int: return self._vertical.num_lev @property - def n_shift_total(self)->int: + def n_shift_total(self) -> int: return self._n_shift_total @property - def num_vertices(self)->int: + def num_vertices(self) -> int: return self._horizontal.num_vertices @property - def num_edges(self)->int: + def num_edges(self) -> int: return self._horizontal.num_edges @property - def num_cells(self)->int: + def num_cells(self) -> int: return self._horizontal.num_cells - - class IconGrid: def __init__(self): self.config: MeshConfig = None diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index 85f7a3022..c2ed2d1a7 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -26,6 +26,7 @@ from icon4py.testutils import serialbox_utils from icon4py.testutils.serialbox_utils import IconSerialDataProvider + SERIALBOX_ONLY_MSG = "Only ser_type='sb (Serialbox)' is implemented so far." SIMULATION_START_DATE = "2021-06-20T12:00:10.000" @@ -113,24 +114,29 @@ def read_geometry_fields( # /home/magdalena/data/exclaim/dycore/mch_ch_r04b09_dsl/node2/mch_ch_r04b09_dsl/icon_grid def read_decomp_info( - path: Path, procs_props: ProcessProperties, ser_type = SerializationType.SB, -)->DecompositionInfo: + path: Path, + procs_props: ProcessProperties, + ser_type=SerializationType.SB, +) -> DecompositionInfo: if ser_type == SerializationType.SB: - sp = serialbox_utils.IconSerialDataProvider("icon_grid", str(path.absolute()), True, - procs_props.rank) - return \ - sp.from_savepoint_grid().construct_decomposition_info() + sp = serialbox_utils.IconSerialDataProvider( + "icon_grid", str(path.absolute()), True, procs_props.rank + ) + return sp.from_savepoint_grid().construct_decomposition_info() else: raise NotImplementedError(SERIALBOX_ONLY_MSG) + def read_grid( - path: Path, procs_props: ProcessProperties, ser_type = SerializationType.SB, -)->IconGrid: + path: Path, + procs_props: ProcessProperties, + ser_type=SerializationType.SB, +) -> IconGrid: if ser_type == SerializationType.SB: - sp = serialbox_utils.IconSerialDataProvider("icon_grid", str(path.absolute()), True, - procs_props.rank) - return \ - sp.from_savepoint_grid().construct_icon_grid() + sp = serialbox_utils.IconSerialDataProvider( + "icon_grid", str(path.absolute()), True, procs_props.rank + ) + return sp.from_savepoint_grid().construct_icon_grid() else: raise NotImplementedError(SERIALBOX_ONLY_MSG) diff --git a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py index c86b98705..5db9aaf5d 100644 --- a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py @@ -1,3 +1,16 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + from enum import Enum from typing import Union @@ -7,17 +20,19 @@ from gt4py.next.common import Dimension from mpi4py.MPI import Comm - from icon4py.decomposition.decomposed import ProcessProperties from icon4py.diffusion.utils import builder + + mpi4py.rc.initialize = False CommId = Union[int, Comm, None] + def get_processor_properties(comm_id: CommId = None): init_mpi() - def _get_current_comm_or_comm_world(comm_id: CommId)->Comm: + def _get_current_comm_or_comm_world(comm_id: CommId) -> Comm: if isinstance(comm_id, int): comm = Comm.f2py(comm_id) elif isinstance(comm_id, Comm): @@ -32,32 +47,30 @@ def _get_current_comm_or_comm_world(comm_id: CommId)->Comm: def init_mpi(): from mpi4py import MPI + if not MPI.Is_initialized(): MPI.Init() - - class DecompositionInfo: class EntryType(int, Enum): - ALL = 0, - OWNED = 1, + ALL = (0,) + OWNED = (1,) HALO = 2 + @builder - def with_dimension(self, dim:Dimension, global_index:np.ndarray, owner_mask:np.ndarray): + def with_dimension( + self, dim: Dimension, global_index: np.ndarray, owner_mask: np.ndarray + ): masked_global_index = ma.array(global_index, mask=owner_mask) self._global_index[dim] = masked_global_index def __init__(self): self._global_index = {} - - - def global_index(self, dim:Dimension, type:EntryType = EntryType.ALL): - match(type): + def global_index(self, dim: Dimension, type: EntryType = EntryType.ALL): + match (type): case DecompositionInfo.EntryType.ALL: return ma.getdata(self._global_index[dim], subok=False) case _: raise NotImplementedError() - - diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index e9cd9fd58..fe3003b42 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -19,18 +19,27 @@ import wget from icon4py.diffusion.diffusion import DiffusionConfig -from icon4py.driver.io_utils import read_decomp_info, SerializationType -from icon4py.driver.parallel_setup import get_processor_properties, DecompositionInfo -from icon4py.testutils.serialbox_utils import IconSerialDataProvider, IconGridSavePoint +from icon4py.driver.io_utils import SerializationType, read_decomp_info +from icon4py.driver.parallel_setup import ( + DecompositionInfo, + get_processor_properties, +) +from icon4py.testutils.serialbox_utils import ( + IconGridSavePoint, + IconSerialDataProvider, +) + data_uri = "https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download" data_path = Path(__file__).parent.joinpath("ser_icondata") extracted_path = data_path.joinpath("mch_ch_r04b09_dsl/ser_data") data_file = data_path.joinpath("mch_ch_r04b09_dsl_v2.tar.gz").name + @pytest.fixture def mpi(): from mpi4py import MPI + mpi4py.rc.initialize = False return MPI @@ -61,9 +70,6 @@ def data_provider(setup_icon_data) -> IconSerialDataProvider: return IconSerialDataProvider("icon_pydycore", str(extracted_path), True) - - - @pytest.fixture def linit(): """ diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index 4fcdd04c3..6c6afae67 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -1,3 +1,16 @@ +# 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 pathlib import mpi4py @@ -6,7 +19,11 @@ from icon4py.common.dimension import CellDim from icon4py.decomposition.decomposed import ProcessProperties from icon4py.driver.io_utils import SerializationType, read_decomp_info, read_grid -from icon4py.driver.parallel_setup import get_processor_properties, DecompositionInfo +from icon4py.driver.parallel_setup import ( + DecompositionInfo, + get_processor_properties, +) + """ running tests with mpi: @@ -19,8 +36,6 @@ """ - - @pytest.mark.mpi def test_processor_properties_from_comm_world(mpi): props = get_processor_properties() @@ -30,9 +45,12 @@ def test_processor_properties_from_comm_world(mpi): def test_decomposition_info(): path = pathlib.Path( - "/home/magdalena/data/exclaim/dycore/mch_ch_r04b09_dsl/node2/mch_ch_r04b09_dsl/icon_grid") + "/home/magdalena/data/exclaim/dycore/mch_ch_r04b09_dsl/node2/mch_ch_r04b09_dsl/icon_grid" + ) props = get_processor_properties() decomposition_info = read_decomp_info(path, props, SerializationType.SB) icon_grid = read_grid(path, props, SerializationType.SB) - global_index = decomposition_info.global_index(CellDim, DecompositionInfo.EntryType.ALL) + global_index = decomposition_info.global_index( + CellDim, DecompositionInfo.EntryType.ALL + ) assert global_index.shape[0] == icon_grid.num_cells() diff --git a/atm_dyn_iconam/tests/test_io_utils.py b/atm_dyn_iconam/tests/test_io_utils.py index 6106b02a7..0d3197020 100644 --- a/atm_dyn_iconam/tests/test_io_utils.py +++ b/atm_dyn_iconam/tests/test_io_utils.py @@ -25,6 +25,7 @@ ) from icon4py.driver.parallel_setup import get_processor_properties + test_data_path = pathlib.Path(__file__).parent.joinpath( "./ser_icondata/mch_ch_r04b09_dsl/ser_data" ) @@ -115,6 +116,3 @@ def assert_interpolation_state_fields(interpolation_state): assert interpolation_state.rbf_coeff_2 assert interpolation_state.rbf_coeff_1 assert interpolation_state.geofac_n2s_c - - - diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index fd9bf7a72..c4c66a02c 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -77,7 +77,7 @@ def _read_int32_shift1(self, name: str): def _read_int32(self, name: str): """ - Read a int field by name. + Read an int field by name. use this for end indices: because FORTRAN slices are inclusive [from:to] _and_ one based this accounts for being exclusive python exclusive bounds: [from:to) @@ -185,18 +185,23 @@ def refin_ctrl(self, dim: Dimension): def num(self, dim: Dimension): match (dim): case dimension.CellDim: - return self.serializer.read("num_cells", savepoint=self.savepoint).astype(int32)[0] + return self.serializer.read( + "num_cells", savepoint=self.savepoint + ).astype(int32)[0] case dimension.EdgeDim: - return self.serializer.read("num_edges", savepoint=self.savepoint).astype(int32)[0] + return self.serializer.read( + "num_edges", savepoint=self.savepoint + ).astype(int32)[0] case dimension.VertexDim: - return self.serializer.read("num_vert", savepoint=self.savepoint).astype(int32)[0] + return self.serializer.read( + "num_vert", savepoint=self.savepoint + ).astype(int32)[0] case dimension.KDim: return self.get_metadata("nlev")["nlev"] case _: - raise NotImplementedError(f"only {CellDim, EdgeDim, VertexDim, KDim} are supported") - - - + raise NotImplementedError( + f"only {CellDim, EdgeDim, VertexDim, KDim} are supported" + ) def _read_field_for_dim(self, field_name, read_func, dim): match (dim): @@ -208,7 +213,8 @@ def _read_field_for_dim(self, field_name, read_func, dim): return read_func(f"v_{field_name}") case _: raise NotImplementedError( - f"only {dimension.CellDim, dimension.EdgeDim, dimension.VertexDim} are handled") + f"only {dimension.CellDim, dimension.EdgeDim, dimension.VertexDim} are handled" + ) def owner_mask(self, dim: Dimension): field_name = "owner_mask" @@ -218,6 +224,7 @@ def owner_mask(self, dim: Dimension): def global_index(self, dim: Dimension): field_name = "glb_index" return self._read_field_for_dim(field_name, self._read_int32_shift1, dim) + def decomp_domain(self, dim): field_name = "decomp_domain" return self._read_field_for_dim(field_name, self._read_int32, dim) @@ -229,7 +236,6 @@ def construct_decomposition_info(self): return DecompositionInfo().with_dimension(CellDim, index, mask) - def construct_icon_grid(self) -> IconGrid: cell_starts = self.cells_start_index() @@ -499,5 +505,3 @@ def from_savepoint_diffusion_exit( .as_savepoint() ) return IconDiffusionExitSavepoint(savepoint, self.serializer) - - From 7926408ed660327bea6f3756fd2850fd1c999622 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 30 May 2023 14:05:44 +0200 Subject: [PATCH 108/263] skip tests that fail due changed data --- .../tests/mpi_tests/test_parallel_setup.py | 28 +++++++++++++------ atm_dyn_iconam/tests/test_diffusion.py | 6 ++-- atm_dyn_iconam/tests/test_io_utils.py | 2 +- 3 files changed, 24 insertions(+), 12 deletions(-) diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index 6c6afae67..5f08b7acc 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -16,7 +16,7 @@ import mpi4py import pytest -from icon4py.common.dimension import CellDim +from icon4py.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.decomposition.decomposed import ProcessProperties from icon4py.driver.io_utils import SerializationType, read_decomp_info, read_grid from icon4py.driver.parallel_setup import ( @@ -35,22 +35,34 @@ """ - +#TODO [magdalena] fix this @pytest.mark.mpi def test_processor_properties_from_comm_world(mpi): props = get_processor_properties() assert props.rank < mpi.COMM_WORLD.Get_size() assert props.comm_name == mpi.COMM_WORLD.Get_name() - -def test_decomposition_info(): +# TODO s [magdalena] extract fixture, more useful asserts..., fix run parallel (as is will not work on second node...) +@pytest.mark.parametrize(("dim, owned, total"), ((CellDim, 10448, 10611), (EdgeDim, 15820, 16065), (VertexDim, 5373, 5455))) +def test_decomposition_info_masked(dim, owned, total): path = pathlib.Path( "/home/magdalena/data/exclaim/dycore/mch_ch_r04b09_dsl/node2/mch_ch_r04b09_dsl/icon_grid" ) props = get_processor_properties() decomposition_info = read_decomp_info(path, props, SerializationType.SB) - icon_grid = read_grid(path, props, SerializationType.SB) - global_index = decomposition_info.global_index( - CellDim, DecompositionInfo.EntryType.ALL + owned_indices = decomposition_info.global_index( + dim, DecompositionInfo.EntryType.ALL ) - assert global_index.shape[0] == icon_grid.num_cells() + assert owned_indices.shape[0] == total + + owned_indices = decomposition_info.global_index( + dim, DecompositionInfo.EntryType.OWNED + ) + assert owned_indices.shape[0] == owned + + halo_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.HALO) + assert halo_indices.shape[0] == total - owned + + + + diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 1ee598828..6d243efea 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -218,7 +218,7 @@ def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): assert len(params.smagorinski_factor) == 4 assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor))) - +@pytest.mark.skip("fix: switch geofac_grg changed in new dataset") @pytest.mark.datatest def test_diffusion_init( diffusion_savepoint_init, @@ -353,7 +353,7 @@ def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) assert exptected_smag_offset == 0.0 - +@pytest.mark.skip("fix: switch geofac_grg changed in new dataset") @pytest.mark.datatest def test_verify_diffusion_init_against_first_regular_savepoint( diffusion_savepoint_init, @@ -399,7 +399,7 @@ def test_verify_diffusion_init_against_first_regular_savepoint( _verify_init_values_against_savepoint(savepoint, diffusion) - +@pytest.mark.skip("fix: switch geofac_grg changed in new dataset") @pytest.mark.datatest @pytest.mark.parametrize("step_date_init", ["2021-06-20T12:00:50.000"]) def test_verify_diffusion_init_against_other_regular_savepoint( diff --git a/atm_dyn_iconam/tests/test_io_utils.py b/atm_dyn_iconam/tests/test_io_utils.py index 0d3197020..8ef86be92 100644 --- a/atm_dyn_iconam/tests/test_io_utils.py +++ b/atm_dyn_iconam/tests/test_io_utils.py @@ -51,7 +51,7 @@ def assert_grid_size_and_connectivities(grid): assert grid.get_c2e2c_connectivity() assert grid.get_e2ecv_connectivity() - +@pytest.mark.skip("fix: switch geofac_grg changed in new dataset") @pytest.mark.datatest def test_read_icon_grid_for_type_sb(): grid = read_icon_grid(test_data_path, ser_type=SerializationType.SB) From 0ea7df7f1723a714b977d78327ae78a21132a754 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 30 May 2023 14:06:13 +0200 Subject: [PATCH 109/263] pre-commit fixes --- .../src/icon4py/driver/parallel_setup.py | 10 +++++++-- atm_dyn_iconam/tests/conftest.py | 11 +--------- .../tests/mpi_tests/test_parallel_setup.py | 21 ++++++++++--------- atm_dyn_iconam/tests/test_diffusion.py | 3 +++ atm_dyn_iconam/tests/test_io_utils.py | 4 ++-- .../src/icon4py/testutils/serialbox_utils.py | 15 +++++++++---- 6 files changed, 36 insertions(+), 28 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py index 5db9aaf5d..a94b9faed 100644 --- a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py @@ -68,9 +68,15 @@ def with_dimension( def __init__(self): self._global_index = {} - def global_index(self, dim: Dimension, type: EntryType = EntryType.ALL): - match (type): + def global_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL): + match (entry_type): case DecompositionInfo.EntryType.ALL: return ma.getdata(self._global_index[dim], subok=False) + case DecompositionInfo.EntryType.OWNED: + global_index = self._global_index[dim] + return ma.getdata(global_index[global_index.mask]) + case DecompositionInfo.EntryType.HALO: + global_index = self._global_index[dim] + return ma.getdata(global_index[~global_index.mask]) case _: raise NotImplementedError() diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index fe3003b42..a059bf924 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -10,7 +10,6 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -import pathlib import tarfile from pathlib import Path @@ -19,15 +18,7 @@ import wget from icon4py.diffusion.diffusion import DiffusionConfig -from icon4py.driver.io_utils import SerializationType, read_decomp_info -from icon4py.driver.parallel_setup import ( - DecompositionInfo, - get_processor_properties, -) -from icon4py.testutils.serialbox_utils import ( - IconGridSavePoint, - IconSerialDataProvider, -) +from icon4py.testutils.serialbox_utils import IconSerialDataProvider data_uri = "https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download" diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index 5f08b7acc..caa8a3e82 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -13,12 +13,10 @@ import pathlib -import mpi4py import pytest from icon4py.common.dimension import CellDim, EdgeDim, VertexDim -from icon4py.decomposition.decomposed import ProcessProperties -from icon4py.driver.io_utils import SerializationType, read_decomp_info, read_grid +from icon4py.driver.io_utils import SerializationType, read_decomp_info from icon4py.driver.parallel_setup import ( DecompositionInfo, get_processor_properties, @@ -35,15 +33,20 @@ """ -#TODO [magdalena] fix this + +# TODO [magdalena] fix this @pytest.mark.mpi def test_processor_properties_from_comm_world(mpi): props = get_processor_properties() assert props.rank < mpi.COMM_WORLD.Get_size() assert props.comm_name == mpi.COMM_WORLD.Get_name() + # TODO s [magdalena] extract fixture, more useful asserts..., fix run parallel (as is will not work on second node...) -@pytest.mark.parametrize(("dim, owned, total"), ((CellDim, 10448, 10611), (EdgeDim, 15820, 16065), (VertexDim, 5373, 5455))) +@pytest.mark.parametrize( + ("dim, owned, total"), + ((CellDim, 10448, 10611), (EdgeDim, 15820, 16065), (VertexDim, 5373, 5455)), +) def test_decomposition_info_masked(dim, owned, total): path = pathlib.Path( "/home/magdalena/data/exclaim/dycore/mch_ch_r04b09_dsl/node2/mch_ch_r04b09_dsl/icon_grid" @@ -60,9 +63,7 @@ def test_decomposition_info_masked(dim, owned, total): ) assert owned_indices.shape[0] == owned - halo_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.HALO) + halo_indices = decomposition_info.global_index( + dim, DecompositionInfo.EntryType.HALO + ) assert halo_indices.shape[0] == total - owned - - - - diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 6d243efea..219776ec7 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -218,6 +218,7 @@ def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): assert len(params.smagorinski_factor) == 4 assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor))) + @pytest.mark.skip("fix: switch geofac_grg changed in new dataset") @pytest.mark.datatest def test_diffusion_init( @@ -353,6 +354,7 @@ def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) assert exptected_smag_offset == 0.0 + @pytest.mark.skip("fix: switch geofac_grg changed in new dataset") @pytest.mark.datatest def test_verify_diffusion_init_against_first_regular_savepoint( @@ -399,6 +401,7 @@ def test_verify_diffusion_init_against_first_regular_savepoint( _verify_init_values_against_savepoint(savepoint, diffusion) + @pytest.mark.skip("fix: switch geofac_grg changed in new dataset") @pytest.mark.datatest @pytest.mark.parametrize("step_date_init", ["2021-06-20T12:00:50.000"]) diff --git a/atm_dyn_iconam/tests/test_io_utils.py b/atm_dyn_iconam/tests/test_io_utils.py index 8ef86be92..04ce8de57 100644 --- a/atm_dyn_iconam/tests/test_io_utils.py +++ b/atm_dyn_iconam/tests/test_io_utils.py @@ -23,7 +23,6 @@ read_icon_grid, read_static_fields, ) -from icon4py.driver.parallel_setup import get_processor_properties test_data_path = pathlib.Path(__file__).parent.joinpath( @@ -51,13 +50,14 @@ def assert_grid_size_and_connectivities(grid): assert grid.get_c2e2c_connectivity() assert grid.get_e2ecv_connectivity() -@pytest.mark.skip("fix: switch geofac_grg changed in new dataset") + @pytest.mark.datatest def test_read_icon_grid_for_type_sb(): grid = read_icon_grid(test_data_path, ser_type=SerializationType.SB) assert_grid_size_and_connectivities(grid) +@pytest.mark.skip("fix: switch geofac_grg changed in new dataset") @pytest.mark.datatest def test_read_static_fields_for_type_sb(): metric_state, interpolation_state = read_static_fields( diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index c4c66a02c..d368fc1de 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -230,11 +230,18 @@ def decomp_domain(self, dim): return self._read_field_for_dim(field_name, self._read_int32, dim) def construct_decomposition_info(self): - index = self.global_index(CellDim) - num_cells = self.num(CellDim) - mask = self.owner_mask(CellDim)[0:num_cells] + return ( + DecompositionInfo() + .with_dimension(*self._get_decomp_fields(CellDim)) + .with_dimension(*self._get_decomp_fields(EdgeDim)) + .with_dimension(*self._get_decomp_fields(VertexDim)) + ) - return DecompositionInfo().with_dimension(CellDim, index, mask) + def _get_decomp_fields(self, dim: Dimension): + index = self.global_index(dim) + number = self.num(dim) + mask = self.owner_mask(dim)[0:number] + return dim, index, mask def construct_icon_grid(self) -> IconGrid: From 4782b235ad8767dfffbdcb25039f7f62e2494265 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 31 May 2023 10:45:45 +0200 Subject: [PATCH 110/263] add decomposition info to driver --- atm_dyn_iconam/src/icon4py/driver/dycore_driver.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index 2ceada3f6..28657ed56 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -31,8 +31,9 @@ read_geometry_fields, read_icon_grid, read_initial_state, - read_static_fields, + read_static_fields, read_decomp_info, ) +from icon4py.driver.parallel_setup import get_processor_properties from icon4py.testutils.serialbox_utils import IconSerialDataProvider @@ -185,6 +186,10 @@ def initialize(n_time_steps, file_path: Path): prognostic_state: initial state fro prognostic and diagnostic variables diagnostic_state: """ + log.info("initialize parallel runtime") + parallel_props = get_processor_properties() + decomp_info = read_decomp_info("/home/magdalena/data/exclaim/dycore/mch_ch_r04b09_dsl/node2/mch_ch_r04b09_dsl/icon_grid", parallel_props) + experiment_name = "mch_ch_r04b09_dsl" log.info(f"reading configuration: experiment {experiment_name}") config = read_config(experiment_name, n_time_steps=n_time_steps) From 8443e8490da1c23065d4fb9db3f6745a664333f8 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 31 May 2023 11:10:16 +0200 Subject: [PATCH 111/263] fix local path --- atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index caa8a3e82..4a6ad8956 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -49,7 +49,7 @@ def test_processor_properties_from_comm_world(mpi): ) def test_decomposition_info_masked(dim, owned, total): path = pathlib.Path( - "/home/magdalena/data/exclaim/dycore/mch_ch_r04b09_dsl/node2/mch_ch_r04b09_dsl/icon_grid" + "/home/magdalena/data/exclaim/dycore/mch_ch_r04b09_dsl/mpitasks2/mch_ch_r04b09_dsl/ser_data" ) props = get_processor_properties() decomposition_info = read_decomp_info(path, props, SerializationType.SB) From b62e931ba443c78e1b23eb833939d060c222538e Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 31 May 2023 16:51:38 +0200 Subject: [PATCH 112/263] - adapt diffusion.py to the structure of recently serialized data. - update download link of seralized data to most recent dycore version --- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 10 +- atm_dyn_iconam/tests/conftest.py | 14 +- atm_dyn_iconam/tests/test_diffusion.py | 162 +++++------------- .../src/icon4py/testutils/serialbox_utils.py | 150 ++++++++-------- 4 files changed, 148 insertions(+), 188 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index 858f9c9e7..ea3df6686 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -123,11 +123,13 @@ def read_static_fields( """ if ser_type == SerializationType.SB: - sp = serialbox_utils.IconSerialDataProvider( + dataprovider = serialbox_utils.IconSerialDataProvider( "icon_pydycore", str(path.absolute()), False - ).from_savepoint_diffusion_init(linit=True, date=SIMULATION_START_DATE) - metric_state = sp.construct_metric_state() - interpolation_state = sp.construct_interpolation_state() + ) + interpolation_state = ( + dataprovider.from_interpolation_savepoint().construct_interpolation_state() + ) + metric_state = dataprovider.from_metrics_savepoint().construct_metric_state() return metric_state, interpolation_state else: raise NotImplementedError("Only ser_type='sb' is implemented so far.") diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index ba7d2b336..6c0f53d91 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -20,7 +20,7 @@ from icon4py.testutils.serialbox_utils import IconSerialDataProvider -data_uri = "https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download" +data_uri = "https://polybox.ethz.ch/index.php/s/LcAbscZqnsx4WCf/download" data_path = Path(__file__).parent.joinpath("ser_icondata") extracted_path = data_path.joinpath("mch_ch_r04b09_dsl/ser_data") data_file = data_path.joinpath("mch_ch_r04b09_dsl_v2.tar.gz").name @@ -107,6 +107,18 @@ def diffusion_savepoint_exit(data_provider, step_date_exit): return sp +@pytest.fixture +def interpolation_savepoint(data_provider): + """Load data from ICON interplation state savepoint.""" + return data_provider.from_interpolation_savepoint() + + +@pytest.fixture +def metrics_savepoint(data_provider): + """Load data from ICON mestric state savepoint.""" + return data_provider.from_metrics_savepoint() + + @pytest.fixture def icon_grid(grid_savepoint): """ diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 1ee598828..4b0a70874 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -18,8 +18,6 @@ from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.diffusion.horizontal import CellParams, EdgeParams from icon4py.diffusion.icon_grid import VerticalModelParams -from icon4py.diffusion.interpolation_state import InterpolationState -from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.utils import ( _en_smag_fac_for_zero_nshift, _setup_runtime_diff_multfac_vn, @@ -33,6 +31,9 @@ from icon4py.testutils.utils import random_field, zero_field +datarun_reduced_substeps = 2 + + def test_scale_k(): mesh = SimpleMesh() field = random_field(mesh, KDim) @@ -154,7 +155,6 @@ def test_set_zero_vertex_k(): assert np.allclose(0.0, f) -@pytest.mark.datatest def test_diffusion_coefficients_with_hdiff_efdt_ratio(r04b09_diffusion_config): config = r04b09_diffusion_config config.hdiff_efdt_ratio = 1.0 @@ -168,7 +168,6 @@ def test_diffusion_coefficients_with_hdiff_efdt_ratio(r04b09_diffusion_config): assert params.K4W == pytest.approx(1.0 / 72.0, abs=1e-12) -@pytest.mark.datatest def test_diffusion_coefficients_without_hdiff_efdt_ratio(r04b09_diffusion_config): config = r04b09_diffusion_config config.hdiff_efdt_ratio = 0.0 @@ -182,7 +181,6 @@ def test_diffusion_coefficients_without_hdiff_efdt_ratio(r04b09_diffusion_config assert params.K4W == 0.0 -@pytest.mark.datatest def test_smagorinski_factor_for_diffusion_type_4(r04b09_diffusion_config): config = r04b09_diffusion_config config.smagorinski_scaling_factor = 0.15 @@ -212,8 +210,7 @@ def test_smagorinski_heights_diffusion_type_5_are_consistent( def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): - config = r04b09_diffusion_config - params = DiffusionParams(config) + params = DiffusionParams(r04b09_diffusion_config) assert len(params.smagorinski_factor) == len(params.smagorinski_height) assert len(params.smagorinski_factor) == 4 assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor))) @@ -222,6 +219,8 @@ def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): @pytest.mark.datatest def test_diffusion_init( diffusion_savepoint_init, + interpolation_savepoint, + metrics_savepoint, grid_savepoint, icon_grid, r04b09_diffusion_config, @@ -229,6 +228,7 @@ def test_diffusion_init( damping_height, ): config = r04b09_diffusion_config + config.ndyn_substeps = datarun_reduced_substeps additional_parameters = DiffusionParams(config) vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) @@ -238,25 +238,8 @@ def test_diffusion_init( assert meta["linit"] is False assert meta["date"] == step_date_init - grg = diffusion_savepoint_init.geofac_grg() - interpolation_state = InterpolationState( - e_bln_c_s=diffusion_savepoint_init.e_bln_c_s(), - rbf_coeff_1=diffusion_savepoint_init.rbf_vec_coeff_v1(), - rbf_coeff_2=diffusion_savepoint_init.rbf_vec_coeff_v2(), - geofac_div=diffusion_savepoint_init.geofac_div(), - geofac_n2s=diffusion_savepoint_init.geofac_n2s(), - geofac_grg_x=grg[0], - geofac_grg_y=grg[1], - nudgecoeff_e=diffusion_savepoint_init.nudgecoeff_e(), - ) - metric_state = MetricState( - mask_hdiff=diffusion_savepoint_init.mask_diff(), - theta_ref_mc=diffusion_savepoint_init.theta_ref_mc(), - wgtfac_c=diffusion_savepoint_init.wgtfac_c(), - zd_intcoef=diffusion_savepoint_init.zd_intcoef(), - zd_vertidx=diffusion_savepoint_init.zd_vertoffset(), - zd_diffcoef=diffusion_savepoint_init.zd_diffcoef(), - ) + interpolation_state = interpolation_savepoint.construct_interpolation_state() + metric_state = metrics_savepoint.construct_metric_state() diffusion = Diffusion() diffusion.init( @@ -334,6 +317,7 @@ def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( ): savepoint = diffusion_savepoint_init config = r04b09_diffusion_config + config.ndyn_substeps = datarun_reduced_substeps params = DiffusionParams(config) expected_diff_multfac_vn = savepoint.diff_multfac_vn() @@ -357,36 +341,21 @@ def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( @pytest.mark.datatest def test_verify_diffusion_init_against_first_regular_savepoint( diffusion_savepoint_init, + interpolation_savepoint, + metrics_savepoint, grid_savepoint, r04b09_diffusion_config, icon_grid, damping_height, ): config = r04b09_diffusion_config + config.ndyn_substeps = datarun_reduced_substeps additional_parameters = DiffusionParams(config) - - savepoint = diffusion_savepoint_init vct_a = grid_savepoint.vct_a() - grg = savepoint.geofac_grg() - interpolation_state = InterpolationState( - e_bln_c_s=savepoint.e_bln_c_s(), - rbf_coeff_1=savepoint.rbf_vec_coeff_v1(), - rbf_coeff_2=savepoint.rbf_vec_coeff_v2(), - geofac_div=savepoint.geofac_div(), - geofac_n2s=savepoint.geofac_n2s(), - geofac_grg_x=grg[0], - geofac_grg_y=grg[1], - nudgecoeff_e=savepoint.nudgecoeff_e(), - ) - metric_state = MetricState( - mask_hdiff=savepoint.mask_diff(), - theta_ref_mc=savepoint.theta_ref_mc(), - wgtfac_c=savepoint.wgtfac_c(), - zd_intcoef=savepoint.zd_intcoef(), - zd_vertidx=savepoint.zd_vertoffset(), - zd_diffcoef=savepoint.zd_diffcoef(), - ) + interpolation_state = interpolation_savepoint.construct_interpolation_state() + metric_state = metrics_savepoint.construct_metric_state() + diffusion = Diffusion() diffusion.init( config=config, @@ -397,7 +366,7 @@ def test_verify_diffusion_init_against_first_regular_savepoint( interpolation_state=interpolation_state, ) - _verify_init_values_against_savepoint(savepoint, diffusion) + _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) @pytest.mark.datatest @@ -406,32 +375,18 @@ def test_verify_diffusion_init_against_other_regular_savepoint( r04b09_diffusion_config, grid_savepoint, icon_grid, + interpolation_savepoint, + metrics_savepoint, diffusion_savepoint_init, damping_height, ): config = r04b09_diffusion_config + config.ndyn_substeps = datarun_reduced_substeps additional_parameters = DiffusionParams(config) vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) - grg = diffusion_savepoint_init.geofac_grg() - interpolation_state = InterpolationState( - e_bln_c_s=diffusion_savepoint_init.e_bln_c_s(), - rbf_coeff_1=diffusion_savepoint_init.rbf_vec_coeff_v1(), - rbf_coeff_2=diffusion_savepoint_init.rbf_vec_coeff_v2(), - geofac_div=diffusion_savepoint_init.geofac_div(), - geofac_n2s=diffusion_savepoint_init.geofac_n2s(), - geofac_grg_x=grg[0], - geofac_grg_y=grg[1], - nudgecoeff_e=diffusion_savepoint_init.nudgecoeff_e(), - ) - metric_state = MetricState( - mask_hdiff=diffusion_savepoint_init.mask_diff(), - theta_ref_mc=diffusion_savepoint_init.theta_ref_mc(), - wgtfac_c=diffusion_savepoint_init.wgtfac_c(), - zd_intcoef=diffusion_savepoint_init.zd_intcoef(), - zd_vertidx=diffusion_savepoint_init.zd_vertoffset(), - zd_diffcoef=diffusion_savepoint_init.zd_diffcoef(), - ) + interpolation_state = interpolation_savepoint.construct_interpolation_state() + metric_state = metrics_savepoint.construct_metric_state() diffusion = Diffusion() diffusion.init( @@ -453,31 +408,32 @@ def test_run_diffusion_single_step( run_with_program, diffusion_savepoint_init, diffusion_savepoint_exit, + interpolation_savepoint, + metrics_savepoint, grid_savepoint, icon_grid, r04b09_diffusion_config, damping_height, ): - ( - dtime, - cell_geometry, - edge_geometry, - diagnostic_state, - interpolation_state, - metric_state, - prognostic_state, - ) = _read_fields(diffusion_savepoint_init, grid_savepoint) - + dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") + edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() + cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() + interpolation_state = interpolation_savepoint.construct_interpolation_state() + metric_state = metrics_savepoint.construct_metric_state() + diagnostic_state = diffusion_savepoint_init.construct_diagnostics() + prognostic_state = diffusion_savepoint_init.construct_prognostics() vct_a = grid_savepoint.vct_a() vertical_params = VerticalModelParams( vct_a=vct_a, rayleigh_damping_height=damping_height ) - additional_parameters = DiffusionParams(r04b09_diffusion_config) + config = r04b09_diffusion_config + config.ndyn_substeps = datarun_reduced_substeps + additional_parameters = DiffusionParams(config) diffusion = Diffusion(run_program=run_with_program) diffusion.init( grid=icon_grid, - config=r04b09_diffusion_config, + config=config, params=additional_parameters, vertical_params=vertical_params, metric_state=metric_state, @@ -512,31 +468,6 @@ def test_run_diffusion_single_step( ) -def _read_fields(diffusion_savepoint_init, grid_savepoint): - dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") - edge_geometry, cell_geometry = _read_geometry_fields(grid_savepoint) - - interpolation_state = diffusion_savepoint_init.construct_interpolation_state() - metric_state = diffusion_savepoint_init.construct_metric_state() - diagnostic_state = diffusion_savepoint_init.construct_diagnostics() - prognostic_state = diffusion_savepoint_init.construct_prognostics() - return ( - dtime, - cell_geometry, - edge_geometry, - diagnostic_state, - interpolation_state, - metric_state, - prognostic_state, - ) - - -def _read_geometry_fields(grid_savepoint): - edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() - cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() - return edge_geometry, cell_geometry - - @pytest.mark.skip("fix: diffusion_stencil_15") @pytest.mark.datatest def test_diffusion_five_steps( @@ -544,31 +475,32 @@ def test_diffusion_five_steps( r04b09_diffusion_config, icon_grid, grid_savepoint, + interpolation_savepoint, + metrics_savepoint, diffusion_savepoint_init, diffusion_savepoint_exit, linit=True, step_date_exit="2021-06-20T12:01:00.000", ): - ( - dtime, - cell_geometry, - edge_geometry, - diagnostic_state, - interpolation_state, - metric_state, - prognostic_state, - ) = _read_fields(diffusion_savepoint_init, grid_savepoint) + dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") + edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() + cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() + interpolation_state = interpolation_savepoint.construct_interpolation_state() + metric_state = metrics_savepoint.construct_metric_state() + diagnostic_state = diffusion_savepoint_init.construct_diagnostics() + prognostic_state = diffusion_savepoint_init.construct_prognostics() vertical_params = VerticalModelParams( vct_a=grid_savepoint.vct_a(), rayleigh_damping_height=damping_height ) additional_parameters = DiffusionParams(r04b09_diffusion_config) - + config = r04b09_diffusion_config + config.ndyn_substeps = datarun_reduced_substeps diffusion = Diffusion() diffusion.init( grid=icon_grid, - config=r04b09_diffusion_config, + config=config, params=additional_parameters, vertical_params=vertical_params, metric_state=metric_state, diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index ad312d679..a711e1543 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -249,30 +249,70 @@ def construct_cell_geometry(self) -> CellParams: return CellParams(area=self.cell_areas()) -class IconDiffusionInitSavepoint(IconSavepoint): - def hdef_ic(self): - return self._get_field("hdef_ic", CellDim, KDim) +class InterpolationSavepoint(IconSavepoint): + def geofac_grg(self): + grg = np.squeeze(self.serializer.read("geofac_grg", self.savepoint)) + return np_as_located_field(CellDim, C2E2CODim)( + grg[:, :, 0] + ), np_as_located_field(CellDim, C2E2CODim)(grg[:, :, 1]) - def div_ic(self): - return self._get_field("div_ic", CellDim, KDim) + def zd_intcoef(self): + return self._get_field("vcoef", CellDim, C2E2CDim, KDim) - def dwdx(self): - return self._get_field("dwdx", CellDim, KDim) + def e_bln_c_s(self): + return self._get_field("e_bln_c_s", CellDim, C2EDim) - def dwdy(self): - return self._get_field("dwdy", CellDim, KDim) + def geofac_div(self): + return self._get_field("geofac_div", CellDim, C2EDim) - def vn(self): - return self._get_field("vn", EdgeDim, KDim) + def geofac_n2s(self): + return self._get_field("geofac_n2s", CellDim, C2E2CODim) - def theta_v(self): - return self._get_field("theta_v", CellDim, KDim) + def rbf_vec_coeff_v1(self): + return self._get_field("rbf_vec_coeff_v1", VertexDim, V2EDim) - def w(self): - return self._get_field("w", CellDim, KDim) + def rbf_vec_coeff_v2(self): + return self._get_field("rbf_vec_coeff_v2", VertexDim, V2EDim) - def exner(self): - return self._get_field("exner", CellDim, KDim) + def nudgecoeff_e(self): + return self._get_field("nudgecoeff_e", EdgeDim) + + def construct_interpolation_state(self) -> InterpolationState: + grg = self.geofac_grg() + return InterpolationState( + e_bln_c_s=self.e_bln_c_s(), + rbf_coeff_1=self.rbf_vec_coeff_v1(), + rbf_coeff_2=self.rbf_vec_coeff_v2(), + geofac_div=self.geofac_div(), + geofac_n2s=self.geofac_n2s(), + geofac_grg_x=grg[0], + geofac_grg_y=grg[1], + nudgecoeff_e=self.nudgecoeff_e(), + ) + + +class MetricSavepoint(IconSavepoint): + def construct_metric_state(self) -> MetricState: + return MetricState( + mask_hdiff=self.mask_diff(), + theta_ref_mc=self.theta_ref_mc(), + wgtfac_c=self.wgtfac_c(), + zd_intcoef=self.zd_intcoef(), + zd_vertidx=self.zd_vertoffset(), + zd_diffcoef=self.zd_diffcoef(), + ) + + def zd_diffcoef(self): + return self._get_field("zd_diffcoef", CellDim, KDim) + + def zd_intcoef(self): + return self._get_field("vcoef", CellDim, C2E2CDim, KDim) + + def zd_vertidx(self): + return self._get_field("zd_vertidx", CellDim, C2E2CDim, dtype=int) + + def zd_vertoffset(self): + return self._get_field("zd_vertoffset", CellDim, C2E2CDim, KDim, dtype=int) def theta_ref_mc(self): return self._get_field("theta_ref_mc", CellDim, KDim) @@ -286,42 +326,31 @@ def wgtfac_e(self): def mask_diff(self): return self._get_field("mask_hdiff", CellDim, KDim, dtype=int) - def zd_diffcoef(self): - return self._get_field("zd_diffcoef", CellDim, KDim) - - def zd_intcoef(self): - return self._get_field("vcoef", CellDim, C2E2CDim, KDim) - - def e_bln_c_s(self): - return self._get_field("e_bln_c_s", CellDim, C2EDim) - def geofac_div(self): - return self._get_field("geofac_div", CellDim, C2EDim) +class IconDiffusionInitSavepoint(IconSavepoint): + def hdef_ic(self): + return self._get_field("hdef_ic", CellDim, KDim) - def geofac_n2s(self): - return self._get_field("geofac_n2s", CellDim, C2E2CODim) + def div_ic(self): + return self._get_field("div_ic", CellDim, KDim) - def geofac_grg(self): - grg = np.squeeze(self.serializer.read("geofac_grg", self.savepoint)) - return np_as_located_field(CellDim, C2E2CODim)( - grg[:, :, 0] - ), np_as_located_field(CellDim, C2E2CODim)(grg[:, :, 1]) + def dwdx(self): + return self._get_field("dwdx", CellDim, KDim) - def nudgecoeff_e(self): - return self._get_field("nudgecoeff_e", EdgeDim) + def dwdy(self): + return self._get_field("dwdy", CellDim, KDim) - def zd_vertidx(self): - # TODO fix this - return self._get_field("zd_vertidx", CellDim, C2E2CDim, dtype=int) + def vn(self): + return self._get_field("vn", EdgeDim, KDim) - def zd_vertoffset(self): - return self._get_field("zd_vertoffset", CellDim, C2E2CDim, KDim, dtype=int) + def theta_v(self): + return self._get_field("theta_v", CellDim, KDim) - def rbf_vec_coeff_v1(self): - return self._get_field("rbf_vec_coeff_v1", VertexDim, V2EDim) + def w(self): + return self._get_field("w", CellDim, KDim) - def rbf_vec_coeff_v2(self): - return self._get_field("rbf_vec_coeff_v2", VertexDim, V2EDim) + def exner(self): + return self._get_field("exner", CellDim, KDim) def diff_multfac_smag(self): return np.squeeze(self.serializer.read("diff_multfac_smag", self.savepoint)) @@ -350,29 +379,6 @@ def diff_multfac_w(self): def diff_multfac_vn(self): return self.serializer.read("diff_multfac_vn", self.savepoint) - def construct_interpolation_state(self) -> InterpolationState: - grg = self.geofac_grg() - return InterpolationState( - e_bln_c_s=self.e_bln_c_s(), - rbf_coeff_1=self.rbf_vec_coeff_v1(), - rbf_coeff_2=self.rbf_vec_coeff_v2(), - geofac_div=self.geofac_div(), - geofac_n2s=self.geofac_n2s(), - geofac_grg_x=grg[0], - geofac_grg_y=grg[1], - nudgecoeff_e=self.nudgecoeff_e(), - ) - - def construct_metric_state(self) -> MetricState: - return MetricState( - mask_hdiff=self.mask_diff(), - theta_ref_mc=self.theta_ref_mc(), - wgtfac_c=self.wgtfac_c(), - zd_intcoef=self.zd_intcoef(), - zd_vertidx=self.zd_vertoffset(), - zd_diffcoef=self.zd_diffcoef(), - ) - def construct_prognostics(self) -> PrognosticState: return PrognosticState( w=self.w(), @@ -441,6 +447,14 @@ def from_savepoint_diffusion_init( ) return IconDiffusionInitSavepoint(savepoint, self.serializer) + def from_interpolation_savepoint(self) -> InterpolationSavepoint: + savepoint = self.serializer.savepoint["interpolation_state"].as_savepoint() + return InterpolationSavepoint(savepoint, self.serializer) + + def from_metrics_savepoint(self) -> MetricSavepoint: + savepoint = self.serializer.savepoint["metric_state"].as_savepoint() + return MetricSavepoint(savepoint, self.serializer) + def from_savepoint_diffusion_exit( self, linit: bool, date: str ) -> IconDiffusionExitSavepoint: From 5517258cfd8c200fa41b54a70f8834697362cb3a Mon Sep 17 00:00:00 2001 From: Magdalena Date: Wed, 31 May 2023 17:05:16 +0200 Subject: [PATCH 113/263] remove `backend=..` specification from program definitions and use `.with_backend` at runtime (#209) --- ...ate_nabla2_and_smag_coefficients_for_vn.py | 3 +- .../fused_mo_nh_diffusion_stencil_02_03.py | 3 +- .../fused_mo_nh_diffusion_stencil_04_05_06.py | 3 +- ...sed_mo_nh_diffusion_stencil_07_08_09_10.py | 3 +- .../fused_mo_nh_diffusion_stencil_11_12.py | 3 +- .../fused_mo_nh_diffusion_stencil_13_14.py | 3 +- .../mo_intp_rbf_rbf_vec_interpol_vertex.py | 3 +- .../atm_dyn_iconam/update_theta_and_exner.py | 3 +- .../src/icon4py/diffusion/diffusion.py | 31 ++++++++++--------- .../icon4py/diffusion/diffusion_program.py | 2 +- atm_dyn_iconam/src/icon4py/diffusion/utils.py | 12 +++---- .../src/icon4py/driver/dycore_driver.py | 3 +- 12 files changed, 33 insertions(+), 39 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py index 9da38e186..74becbc26 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py @@ -13,7 +13,6 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, maximum, minimum, sqrt -from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.common.dimension import ( E2C2V, @@ -136,7 +135,7 @@ def _calculate_nabla2_and_smag_coefficients_for_vn( return kh_smag_e, kh_smag_ec, z_nabla2_e -@program(backend=gtfn_cpu.run_gtfn) +@program def calculate_nabla2_and_smag_coefficients_for_vn( diff_multfac_smag: Field[[KDim], float], tangent_orientation: Field[[EdgeDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py index dd75ec25a..3bbca2b79 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -13,7 +13,6 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulence import ( _calculate_diagnostics_for_turbulence, @@ -40,7 +39,7 @@ def _fused_mo_nh_diffusion_stencil_02_03( return div_ic, hdef_ic -@program(backend=gtfn_cpu.run_gtfn) +@program def fused_mo_nh_diffusion_stencil_02_03( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py index 15b75110c..897dca2e0 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py @@ -13,7 +13,6 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32, where -from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.apply_nabla2_and_nabla4_to_vn import ( _apply_nabla2_and_nabla4_to_vn, @@ -75,7 +74,7 @@ def _fused_mo_nh_diffusion_stencil_04_05_06( return vn -@program(backend=gtfn_cpu.run_gtfn) +@program def fused_mo_nh_diffusion_stencil_04_05_06( u_vert: Field[[VertexDim, KDim], float], v_vert: Field[[VertexDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py index e1318c068..b83b00dfd 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -13,7 +13,6 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where -from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.apply_nabla2_to_w import _apply_nabla2_to_w from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import ( @@ -81,7 +80,7 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( return w, dwdx, dwdy -@program(backend=gtfn_cpu.run_gtfn) +@program def fused_mo_nh_diffusion_stencil_07_08_09_10( area: Field[[CellDim], float], geofac_n2s: Field[[CellDim, C2E2CODim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py index 0ce4c8cdf..a9fb3dbd2 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py @@ -13,7 +13,6 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( _enhance_diffusion_coefficient_for_grid_point_cold_pools, @@ -40,7 +39,7 @@ def _fused_mo_nh_diffusion_stencil_11_12( return kh_smag_e -@program(backend=gtfn_cpu.run_gtfn) +@program def fused_mo_nh_diffusion_stencil_11_12( theta_v: Field[[CellDim, KDim], float], theta_ref_mc: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py index 4c9d39693..39b9c5c0d 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -13,7 +13,6 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.calculate_nabla2_for_z import _calculate_nabla2_for_z from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( @@ -34,7 +33,7 @@ def _fused_mo_nh_diffusion_stencil_13_14( return z_temp -@program(backend=gtfn_cpu.run_gtfn) +@program def fused_mo_nh_diffusion_stencil_13_14( kh_smag_e: Field[[EdgeDim, KDim], float], inv_dual_edge_length: Field[[EdgeDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py index e062c3f53..2d34018a6 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -13,7 +13,6 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.common.dimension import V2E, EdgeDim, KDim, V2EDim, VertexDim @@ -29,7 +28,7 @@ def _mo_intp_rbf_rbf_vec_interpol_vertex( return p_u_out, p_v_out -@program(backend=gtfn_cpu.run_gtfn) +@program def mo_intp_rbf_rbf_vec_interpol_vertex( p_e_in: Field[[EdgeDim, KDim], float], ptr_coeff_1: Field[[VertexDim, V2EDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py index 7c73b3ca3..c597ed3f2 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py @@ -13,7 +13,6 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.common.dimension import CellDim, KDim @@ -32,7 +31,7 @@ def _update_theta_and_exner( return theta_v, exner -@program(backend=gtfn_cpu.run_gtfn) +@program def update_theta_and_exner( z_temp: Field[[CellDim, KDim], float], area: Field[[CellDim], float], diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 813463602..62cff72c1 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -20,6 +20,7 @@ from gt4py.next.common import Dimension from gt4py.next.ffront.fbuiltins import Field, int32 from gt4py.next.iterator.embedded import np_as_located_field +from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn import icon4py.diffusion.diffusion_program as diff_prog from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( @@ -428,7 +429,7 @@ def init( 1.0 / 48.0, params.K4W * config.substep_as_float() ) - init_diffusion_local_fields_for_regular_timestep( + init_diffusion_local_fields_for_regular_timestep.with_backend(run_gtfn)( params.K4, config.substep_as_float(), *params.smagorinski_factor, @@ -504,7 +505,7 @@ def initial_step( diff_multfac_vn = zero_field(self.grid, KDim) smag_limit = zero_field(self.grid, KDim) - setup_fields_for_initial_step( + setup_fields_for_initial_step.with_backend(run_gtfn)( self.params.K4, self.config.hdiff_efdt_ratio, diff_multfac_vn, @@ -626,7 +627,7 @@ def time_step( HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, ) log.info("diffusion program: start") - diff_prog.diffusion_run( + diff_prog.diffusion_run.with_backend(run_gtfn)( diagnostic_hdef_ic=diagnostic_state.hdef_ic, diagnostic_div_ic=diagnostic_state.div_ic, diagnostic_dwdx=diagnostic_state.dwdx, @@ -810,11 +811,13 @@ def _do_diffusion_step( # 0b call timer start # # 0c. dtime dependent stuff: enh_smag_factor, - scale_k(self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={}) + scale_k.with_backend(run_gtfn)( + self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={} + ) # TODO: @magdalena is this needed?, if not remove - set_zero_v_k(self.u_vert, offset_provider={}) - set_zero_v_k(self.v_vert, offset_provider={}) + set_zero_v_k.with_backend(run_gtfn)(self.u_vert, offset_provider={}) + set_zero_v_k.with_backend(run_gtfn)(self.v_vert, offset_provider={}) log.debug("rbf interpolation: start") # # 1. CALL rbf_vec_interpol_vertex mo_intp_rbf_rbf_vec_interpol_vertex( @@ -834,7 +837,7 @@ def _do_diffusion_step( # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: start") - calculate_nabla2_and_smag_coefficients_for_vn( + calculate_nabla2_and_smag_coefficients_for_vn.with_backend(run_gtfn)( diff_multfac_smag=self.diff_multfac_smag, tangent_orientation=tangent_orientation, inv_primal_edge_length=inverse_primal_edge_lengths, @@ -862,7 +865,7 @@ def _do_diffusion_step( ) log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: end") log.debug("running fused stencil fused stencil 02_03: start") - fused_mo_nh_diffusion_stencil_02_03( + fused_mo_nh_diffusion_stencil_02_03.with_backend(run_gtfn)( kh_smag_ec=self.kh_smag_ec, vn=prognostic_state.vn, e_bln_c_s=self.interpolation_state.e_bln_c_s, @@ -886,7 +889,7 @@ def _do_diffusion_step( # # # 5. CALL rbf_vec_interpol_vertex_wp log.debug("rbf interpolation: start") - mo_intp_rbf_rbf_vec_interpol_vertex( + mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(run_gtfn)( p_e_in=self.z_nabla2_e, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, ptr_coeff_2=self.interpolation_state.rbf_coeff_2, @@ -906,7 +909,7 @@ def _do_diffusion_step( # log.debug("running fused stencil 04_05_06: start") - fused_mo_nh_diffusion_stencil_04_05_06( + fused_mo_nh_diffusion_stencil_04_05_06.with_backend(run_gtfn)( u_vert=self.u_vert, v_vert=self.v_vert, primal_normal_vert_v1=primal_normal_vert[0], @@ -937,7 +940,7 @@ def _do_diffusion_step( log.debug("running fused stencil 04_05_06: end") log.debug("running fused stencil 07_08_09_10: start") - fused_mo_nh_diffusion_stencil_07_08_09_10( + fused_mo_nh_diffusion_stencil_07_08_09_10.with_backend(run_gtfn)( area=cell_areas, geofac_n2s=self.interpolation_state.geofac_n2s, geofac_grg_x=self.interpolation_state.geofac_grg_x, @@ -973,7 +976,7 @@ def _do_diffusion_step( # # TODO @magdalena check: kh_smag_e is an out field, should not be calculated in init? # log.debug("running fused stencil 11_12: start") - fused_mo_nh_diffusion_stencil_11_12( + fused_mo_nh_diffusion_stencil_11_12.with_backend(run_gtfn)( theta_v=prognostic_state.theta_v, theta_ref_mc=self.metric_state.theta_ref_mc, thresh_tdiff=self.thresh_tdiff, @@ -989,7 +992,7 @@ def _do_diffusion_step( ) log.debug("running fused stencil 11_12: end") log.debug("running fused stencil 13_14: start") - fused_mo_nh_diffusion_stencil_13_14( + fused_mo_nh_diffusion_stencil_13_14.with_backend(run_gtfn)( kh_smag_e=self.kh_smag_e, inv_dual_edge_length=inverse_dual_edge_length, theta_v=prognostic_state.theta_v, @@ -1026,7 +1029,7 @@ def _do_diffusion_step( # ) log.debug("running fused stencil 15: end") log.debug("running fused stencil update_theta_and_exner: start") - update_theta_and_exner( + update_theta_and_exner.with_backend(run_gtfn)( z_temp=self.z_temp, area=cell_areas, theta_v=prognostic_state.theta_v, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py index 755dab9ba..d19dfd548 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py @@ -54,7 +54,7 @@ from icon4py.diffusion.utils import _scale_k, _set_zero_v_k -@program(backend=gtfn_cpu.run_gtfn) +@program def diffusion_run( diagnostic_hdef_ic: Field[[CellDim, KDim], float], diagnostic_div_ic: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/diffusion/utils.py b/atm_dyn_iconam/src/icon4py/diffusion/utils.py index 3f6a2b706..674a5db7a 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/utils.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/utils.py @@ -17,7 +17,6 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import broadcast, int32, maximum, minimum from gt4py.next.iterator.embedded import np_as_located_field -from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.common.dimension import CellDim, EdgeDim, KDim, Koff, VertexDim @@ -42,7 +41,7 @@ def _identity_e_k( return field -@program(backend=gtfn_cpu.run_gtfn) +@program def copy_diagnostic_and_prognostics( hdef_ic_new: Field[[CellDim, KDim], float], hdef_ic: Field[[CellDim, KDim], float], @@ -76,7 +75,7 @@ def _scale_k(field: Field[[KDim], float], factor: float) -> Field[[KDim], float] return field * factor -@program(backend=gtfn_cpu.run_gtfn) +@program def scale_k( field: Field[[KDim], float], factor: float, scaled_field: Field[[KDim], float] ): @@ -88,7 +87,7 @@ def _set_zero_v_k() -> Field[[VertexDim, KDim], float]: return broadcast(0.0, (VertexDim, KDim)) -@program(backend=gtfn_cpu.run_gtfn) +@program def set_zero_v_k(field: Field[[VertexDim, KDim], float]): _set_zero_v_k(out=field) @@ -107,7 +106,6 @@ def _setup_runtime_diff_multfac_vn( return broadcast(minimum(con, dyn), (KDim,)) -# @field_operator(backend=run_gtfn) @field_operator def _setup_initial_diff_multfac_vn( k4: float, hdiff_efdt_ratio: float @@ -124,7 +122,7 @@ def _setup_fields_for_initial_step( return diff_multfac_vn, smag_limit -@program(backend=gtfn_cpu.run_gtfn) +@program def setup_fields_for_initial_step( k4: float, hdiff_efdt_ratio: float, @@ -199,7 +197,7 @@ def _init_diffusion_local_fields_for_regular_timestemp( ) -@program(backend=gtfn_cpu.run_gtfn) +@program def init_diffusion_local_fields_for_regular_timestep( k4: float, dyn_substeps: float, diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index 57a8d378e..a6467faf1 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -18,6 +18,7 @@ import click import pytz from devtools import Timer +from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.diffusion import Diffusion, DiffusionParams @@ -69,7 +70,7 @@ def do_dynamics_substepping( ) new_p = sp.construct_prognostics() new_d = sp.construct_diagnostics() - copy_diagnostic_and_prognostics( + copy_diagnostic_and_prognostics.with_backend(run_gtfn)( new_d.hdef_ic, diagnostic_state.hdef_ic, new_d.div_ic, From bbae06cf36acf21e54e9b840e392e544d8cc10a4 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 1 Jun 2023 16:22:08 +0200 Subject: [PATCH 114/263] fix decomposition test for run on 2 procs --- .../src/icon4py/decomposition/decomposed.py | 6 ++++- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 2 +- .../tests/mpi_tests/test_parallel_setup.py | 22 ++++++++++++------- 3 files changed, 20 insertions(+), 10 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index 5bf300061..7324671fb 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -15,9 +15,10 @@ class ProcessProperties: - def __init__(self, name="", rank=0): + def __init__(self, name="", size = 0, rank=0): self._communicator_name: str = name self._rank: int = rank + self._comm_size = size @property def rank(self): @@ -26,6 +27,9 @@ def rank(self): @property def comm_name(self): return self._communicator_name + @property + def comm_size(self): + return self._comm_size @classmethod def from_mpi_comm(cls, comm: mpi4py.MPI.Comm): diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index 51e8e69cd..0bc855b71 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -120,7 +120,7 @@ def read_decomp_info( ) -> DecompositionInfo: if ser_type == SerializationType.SB: sp = serialbox_utils.IconSerialDataProvider( - "icon_grid", str(path.absolute()), True, procs_props.rank + "icon_pydycore", str(path.absolute()), True, procs_props.rank ) return sp.from_savepoint_grid().construct_decomposition_info() else: diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index 4a6ad8956..c2e12882d 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -32,9 +32,8 @@ """ +props = get_processor_properties() - -# TODO [magdalena] fix this @pytest.mark.mpi def test_processor_properties_from_comm_world(mpi): props = get_processor_properties() @@ -42,28 +41,35 @@ def test_processor_properties_from_comm_world(mpi): assert props.comm_name == mpi.COMM_WORLD.Get_name() -# TODO s [magdalena] extract fixture, more useful asserts..., fix run parallel (as is will not work on second node...) +# TODO s [magdalena] extract fixture, more useful asserts... +@pytest.mark.skipif(props.comm_size > 2, reason="input files available for 1 or 2 nodes") @pytest.mark.parametrize( ("dim, owned, total"), - ((CellDim, 10448, 10611), (EdgeDim, 15820, 16065), (VertexDim, 5373, 5455)), + ((CellDim, (10448, 10448), (10611, 10612)), (EdgeDim, (15820, 15738), (16065, 16067)), (VertexDim, (5373, 5290), (5455, 5456))) ) -def test_decomposition_info_masked(dim, owned, total): +def test_decomposition_info_masked(mpi, dim, owned, total): path = pathlib.Path( "/home/magdalena/data/exclaim/dycore/mch_ch_r04b09_dsl/mpitasks2/mch_ch_r04b09_dsl/ser_data" ) props = get_processor_properties() + + + my_rank = props.rank decomposition_info = read_decomp_info(path, props, SerializationType.SB) owned_indices = decomposition_info.global_index( dim, DecompositionInfo.EntryType.ALL ) - assert owned_indices.shape[0] == total + my_total = total[my_rank] + my_owned = owned[my_rank] + assert owned_indices.shape[0] == my_total owned_indices = decomposition_info.global_index( dim, DecompositionInfo.EntryType.OWNED ) - assert owned_indices.shape[0] == owned + assert owned_indices.shape[0] == my_owned halo_indices = decomposition_info.global_index( dim, DecompositionInfo.EntryType.HALO ) - assert halo_indices.shape[0] == total - owned + assert halo_indices.shape[0] == my_total - my_owned + From 8fe58eaba744d68775ae135cd1d0f1c776b25f69 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 2 Jun 2023 08:05:18 +0200 Subject: [PATCH 115/263] clean up tests (I) --- .../src/icon4py/decomposition/decomposed.py | 3 ++- .../src/icon4py/driver/dycore_driver.py | 8 +++++-- atm_dyn_iconam/tests/conftest.py | 7 ++++++ .../tests/mpi_tests/test_parallel_setup.py | 22 ++++++++++--------- 4 files changed, 27 insertions(+), 13 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index 7324671fb..29062eaeb 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -15,7 +15,7 @@ class ProcessProperties: - def __init__(self, name="", size = 0, rank=0): + def __init__(self, name="", size=0, rank=0): self._communicator_name: str = name self._rank: int = rank self._comm_size = size @@ -27,6 +27,7 @@ def rank(self): @property def comm_name(self): return self._communicator_name + @property def comm_size(self): return self._comm_size diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index 28657ed56..01b41c0b6 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -28,10 +28,11 @@ from icon4py.driver.io_utils import ( SIMULATION_START_DATE, configure_logging, + read_decomp_info, read_geometry_fields, read_icon_grid, read_initial_state, - read_static_fields, read_decomp_info, + read_static_fields, ) from icon4py.driver.parallel_setup import get_processor_properties from icon4py.testutils.serialbox_utils import IconSerialDataProvider @@ -188,7 +189,10 @@ def initialize(n_time_steps, file_path: Path): """ log.info("initialize parallel runtime") parallel_props = get_processor_properties() - decomp_info = read_decomp_info("/home/magdalena/data/exclaim/dycore/mch_ch_r04b09_dsl/node2/mch_ch_r04b09_dsl/icon_grid", parallel_props) + decomp_info = read_decomp_info( + "/home/magdalena/data/exclaim/dycore/mch_ch_r04b09_dsl/node2/mch_ch_r04b09_dsl/icon_grid", + parallel_props, + ) experiment_name = "mch_ch_r04b09_dsl" log.info(f"reading configuration: experiment {experiment_name}") diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 0947dfe42..d379c1a02 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -25,6 +25,13 @@ data_path = Path(__file__).parent.joinpath("ser_icondata") extracted_path = data_path.joinpath("mch_ch_r04b09_dsl/ser_data") data_file = data_path.joinpath("mch_ch_r04b09_dsl_v2.tar.gz").name +data_2_nodes = data_path.joinpath("mpitasks2/mch_ch_r04b09_dsl/ser_data") + + +@pytest.fixture +def datapath(num_nodes=2): + local_path = f"mpitasks{num_nodes}/mch_ch_r04b09_dsl/ser_data" + return data_path.joinpath(local_path) @pytest.fixture diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index c2e12882d..eea99d3e3 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -11,7 +11,6 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import pathlib import pytest @@ -34,6 +33,7 @@ """ props = get_processor_properties() + @pytest.mark.mpi def test_processor_properties_from_comm_world(mpi): props = get_processor_properties() @@ -42,20 +42,23 @@ def test_processor_properties_from_comm_world(mpi): # TODO s [magdalena] extract fixture, more useful asserts... -@pytest.mark.skipif(props.comm_size > 2, reason="input files available for 1 or 2 nodes") +@pytest.mark.skipif( + props.comm_size > 2, reason="input files available for 1 or 2 nodes" +) @pytest.mark.parametrize( ("dim, owned, total"), - ((CellDim, (10448, 10448), (10611, 10612)), (EdgeDim, (15820, 15738), (16065, 16067)), (VertexDim, (5373, 5290), (5455, 5456))) + ( + (CellDim, (10448, 10448), (10611, 10612)), + (EdgeDim, (15820, 15738), (16065, 16067)), + (VertexDim, (5373, 5290), (5455, 5456)), + ), ) -def test_decomposition_info_masked(mpi, dim, owned, total): - path = pathlib.Path( - "/home/magdalena/data/exclaim/dycore/mch_ch_r04b09_dsl/mpitasks2/mch_ch_r04b09_dsl/ser_data" - ) - props = get_processor_properties() +def test_decomposition_info_masked(mpi, datapath, dim, owned, total): + props = get_processor_properties() my_rank = props.rank - decomposition_info = read_decomp_info(path, props, SerializationType.SB) + decomposition_info = read_decomp_info(datapath, props, SerializationType.SB) owned_indices = decomposition_info.global_index( dim, DecompositionInfo.EntryType.ALL ) @@ -72,4 +75,3 @@ def test_decomposition_info_masked(mpi, dim, owned, total): dim, DecompositionInfo.EntryType.HALO ) assert halo_indices.shape[0] == my_total - my_owned - From 62579a54680769386c5a458adfa94e9373ebd951 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 2 Jun 2023 09:42:33 +0200 Subject: [PATCH 116/263] - fix usage of stencil_15 --- .../mo_nh_diffusion_stencil_15.py | 5 +++ .../src/icon4py/diffusion/diffusion.py | 40 ++++++++++--------- .../src/icon4py/diffusion/metric_state.py | 3 +- atm_dyn_iconam/tests/test_diffusion.py | 6 ++- .../src/icon4py/testutils/serialbox_utils.py | 21 +++++----- 5 files changed, 44 insertions(+), 31 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py index 9abed41ee..86d5ab3df 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py @@ -75,6 +75,10 @@ def mo_nh_diffusion_stencil_15( vcoef: Field[[CECDim, KDim], float], theta_v: Field[[CellDim, KDim], float], z_temp: Field[[CellDim, KDim], float], + horizontal_start: int, + horizontal_end:int, + vertical_start:int, + vertical_end:int, ): _mo_nh_diffusion_stencil_15( mask, @@ -86,4 +90,5 @@ def mo_nh_diffusion_stencil_15( theta_v, z_temp, out=z_temp, + domain = {CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end), } ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 0c9239e18..8969a9444 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -20,7 +20,7 @@ from gt4py.next.common import Dimension from gt4py.next.ffront.fbuiltins import Field, int32 from gt4py.next.iterator.embedded import np_as_located_field -from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn +from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn, run_gtfn_cached import icon4py.diffusion.diffusion_program as diff_prog from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( @@ -87,7 +87,9 @@ VectorTuple = namedtuple("VectorTuple", "x y") - +cached_backend = run_gtfn_cached +compiled_backend = run_gtfn +backend = compiled_backend # class DiffusionConfig: """ Contains necessary parameter to configure a diffusion run. @@ -379,8 +381,8 @@ def __init__(self, run_program=True): self.config: Optional[DiffusionConfig] = None self.params: Optional[DiffusionParams] = None self.vertical_params: Optional[VerticalModelParams] = None - self.interpolation_state = None - self.metric_state = None + self.interpolation_state:InterpolationState = None + self.metric_state:MetricState = None self.diff_multfac_w: Optional[float] = None self.diff_multfac_n2w: Field[[KDim], float] = None self.smag_offset: Optional[float] = None @@ -429,7 +431,7 @@ def init( 1.0 / 48.0, params.K4W * config.substep_as_float() ) - init_diffusion_local_fields_for_regular_timestep.with_backend(run_gtfn)( + init_diffusion_local_fields_for_regular_timestep.with_backend(backend)( params.K4, config.substep_as_float(), *params.smagorinski_factor, @@ -505,7 +507,7 @@ def initial_step( diff_multfac_vn = zero_field(self.grid, KDim) smag_limit = zero_field(self.grid, KDim) - setup_fields_for_initial_step.with_backend(run_gtfn)( + setup_fields_for_initial_step.with_backend(backend)( self.params.K4, self.config.hdiff_efdt_ratio, diff_multfac_vn, @@ -627,7 +629,7 @@ def run( HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, ) log.info("diffusion program: start") - diff_prog.diffusion_run.with_backend(run_gtfn)( + diff_prog.diffusion_run.with_backend(backend)( diagnostic_hdef_ic=diagnostic_state.hdef_ic, diagnostic_div_ic=diagnostic_state.div_ic, diagnostic_dwdx=diagnostic_state.dwdx, @@ -812,13 +814,13 @@ def _do_diffusion_step( # 0b call timer start # # 0c. dtime dependent stuff: enh_smag_factor, - scale_k.with_backend(run_gtfn)( + scale_k.with_backend(backend)( self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={} ) # TODO: @magdalena is this needed?, if not remove - set_zero_v_k.with_backend(run_gtfn)(self.u_vert, offset_provider={}) - set_zero_v_k.with_backend(run_gtfn)(self.v_vert, offset_provider={}) + set_zero_v_k.with_backend(backend)(self.u_vert, offset_provider={}) + set_zero_v_k.with_backend(backend)(self.v_vert, offset_provider={}) log.debug("rbf interpolation: start") # # 1. CALL rbf_vec_interpol_vertex mo_intp_rbf_rbf_vec_interpol_vertex( @@ -838,7 +840,7 @@ def _do_diffusion_step( # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: start") - calculate_nabla2_and_smag_coefficients_for_vn.with_backend(run_gtfn)( + calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend)( diff_multfac_smag=self.diff_multfac_smag, tangent_orientation=tangent_orientation, inv_primal_edge_length=inverse_primal_edge_lengths, @@ -866,7 +868,7 @@ def _do_diffusion_step( ) log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: end") log.debug("running fused stencil fused stencil 02_03: start") - fused_mo_nh_diffusion_stencil_02_03.with_backend(run_gtfn)( + fused_mo_nh_diffusion_stencil_02_03.with_backend(backend)( kh_smag_ec=self.kh_smag_ec, vn=prognostic_state.vn, e_bln_c_s=self.interpolation_state.e_bln_c_s, @@ -890,7 +892,7 @@ def _do_diffusion_step( # # # 5. CALL rbf_vec_interpol_vertex_wp log.debug("rbf interpolation: start") - mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(run_gtfn)( + mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( p_e_in=self.z_nabla2_e, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, ptr_coeff_2=self.interpolation_state.rbf_coeff_2, @@ -910,7 +912,7 @@ def _do_diffusion_step( # log.debug("running fused stencil 04_05_06: start") - fused_mo_nh_diffusion_stencil_04_05_06.with_backend(run_gtfn)( + fused_mo_nh_diffusion_stencil_04_05_06.with_backend(backend)( u_vert=self.u_vert, v_vert=self.v_vert, primal_normal_vert_v1=primal_normal_vert[0], @@ -941,7 +943,7 @@ def _do_diffusion_step( log.debug("running fused stencil 04_05_06: end") log.debug("running fused stencil 07_08_09_10: start") - fused_mo_nh_diffusion_stencil_07_08_09_10.with_backend(run_gtfn)( + fused_mo_nh_diffusion_stencil_07_08_09_10.with_backend(backend)( area=cell_areas, geofac_n2s=self.interpolation_state.geofac_n2s, geofac_grg_x=self.interpolation_state.geofac_grg_x, @@ -977,7 +979,7 @@ def _do_diffusion_step( # # TODO @magdalena check: kh_smag_e is an out field, should not be calculated in init? # log.debug("running fused stencil 11_12: start") - fused_mo_nh_diffusion_stencil_11_12.with_backend(run_gtfn)( + fused_mo_nh_diffusion_stencil_11_12.with_backend(backend)( theta_v=prognostic_state.theta_v, theta_ref_mc=self.metric_state.theta_ref_mc, thresh_tdiff=self.thresh_tdiff, @@ -993,7 +995,7 @@ def _do_diffusion_step( ) log.debug("running fused stencil 11_12: end") log.debug("running fused stencil 13_14: start") - fused_mo_nh_diffusion_stencil_13_14.with_backend(run_gtfn)( + fused_mo_nh_diffusion_stencil_13_14.with_backend(backend)( kh_smag_e=self.kh_smag_e, inv_dual_edge_length=inverse_dual_edge_length, theta_v=prognostic_state.theta_v, @@ -1010,7 +1012,7 @@ def _do_diffusion_step( ) log.debug("running fused stencil 13_14: end") log.debug("running fused stencil 15: start") - mo_nh_diffusion_stencil_15( + mo_nh_diffusion_stencil_15.with_backend(backend)( mask=self.metric_state.mask_hdiff, zd_vertoffset=self.metric_state.zd_vertidx, zd_diffcoef=self.metric_state.zd_diffcoef, @@ -1032,7 +1034,7 @@ def _do_diffusion_step( log.debug("running fused stencil 15: end") log.debug("running fused stencil update_theta_and_exner: start") - update_theta_and_exner.with_backend(run_gtfn)( + update_theta_and_exner.with_backend(backend)( z_temp=self.z_temp, area=cell_areas, theta_v=prognostic_state.theta_v, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py b/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py index 33d6672f9..bf4f55432 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py @@ -14,6 +14,7 @@ from dataclasses import dataclass from gt4py.next.common import Field +from numpy import int32 from icon4py.common.dimension import CECDim, CellDim, KDim @@ -25,6 +26,6 @@ class MetricState: [CellDim, KDim], float ] # weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_c) mask_hdiff: Field[[CellDim, KDim], bool] - zd_vertidx: Field[[CECDim, KDim], int] + zd_vertidx: Field[[CECDim, KDim], int32] zd_diffcoef: Field[[CellDim, KDim], float] zd_intcoef: Field[[CECDim, KDim], float] diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index bde5849ea..97b5af5b2 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -401,8 +401,10 @@ def test_verify_diffusion_init_against_other_regular_savepoint( _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) -@pytest.mark.skip("fix: diffusion_stencil_15") -@pytest.mark.parametrize("run_with_program", [True, False]) +#@pytest.mark.skip("fix: diffusion_stencil_15") +#@pytest.mark.parametrize("run_with_program", [True, False]) + +@pytest.mark.parametrize("run_with_program", [False]) @pytest.mark.datatest def test_run_diffusion_single_step( run_with_program, diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index a711e1543..9c3d25863 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -30,7 +30,7 @@ EdgeDim, KDim, V2EDim, - VertexDim, + VertexDim, CECDim, ) from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.diffusion import VectorTuple @@ -256,9 +256,6 @@ def geofac_grg(self): grg[:, :, 0] ), np_as_located_field(CellDim, C2E2CODim)(grg[:, :, 1]) - def zd_intcoef(self): - return self._get_field("vcoef", CellDim, C2E2CDim, KDim) - def e_bln_c_s(self): return self._get_field("e_bln_c_s", CellDim, C2EDim) @@ -306,13 +303,19 @@ def zd_diffcoef(self): return self._get_field("zd_diffcoef", CellDim, KDim) def zd_intcoef(self): - return self._get_field("vcoef", CellDim, C2E2CDim, KDim) + return self._from_cell_c2e2c_to_cec("vcoef") - def zd_vertidx(self): - return self._get_field("zd_vertidx", CellDim, C2E2CDim, dtype=int) + + def _from_cell_c2e2c_to_cec(self, field_name:str ): + ser_input = np.squeeze(self.serializer.read(field_name, self.savepoint)) + old_shape = ser_input.shape + return np_as_located_field(CECDim, KDim)( + ser_input.reshape(old_shape[0] * old_shape[1], old_shape[2]) + ) def zd_vertoffset(self): - return self._get_field("zd_vertoffset", CellDim, C2E2CDim, KDim, dtype=int) + return self._from_cell_c2e2c_to_cec("zd_vertoffset") + def theta_ref_mc(self): return self._get_field("theta_ref_mc", CellDim, KDim) @@ -324,7 +327,7 @@ def wgtfac_e(self): return self._get_field("wgtfac_e", EdgeDim, KDim) def mask_diff(self): - return self._get_field("mask_hdiff", CellDim, KDim, dtype=int) + return self._get_field("mask_hdiff", CellDim, KDim, dtype=bool) class IconDiffusionInitSavepoint(IconSavepoint): From 18a643b4465c6963642ce848053b3abd12bbbe75 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 2 Jun 2023 16:53:47 +0200 Subject: [PATCH 117/263] add missing fields to IconDiffusionExitSavepoint --- testutils/src/icon4py/testutils/serialbox_utils.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 9c3d25863..8660743b0 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -409,9 +409,19 @@ def theta_v(self): def w(self): return self._get_field("x_w", CellDim, KDim) + def dwdx(self): + return self._get_field("x_dwdx", CellDim, KDim) + + def dwdy(self): + return self._get_field("x_dwdy", CellDim, KDim) def exner(self): return self._get_field("x_exner", CellDim, KDim) + def div_ic(self): + return self._get_field("x_div_ic",CellDim, KDim ) + def hdef_ic(self): + return self._get_field("x_hdef_ic", CellDim, KDim) + class IconSerialDataProvider: def __init__(self, fname_prefix, path=".", do_print=False): From 024ecffe8aaa157c84305168e5c886ea732f86cb Mon Sep 17 00:00:00 2001 From: Magdalena Date: Tue, 6 Jun 2023 14:52:42 +0200 Subject: [PATCH 118/263] - adapt diffusion.py to the structure of recently serialized data. (#219) - update download link of seralized data to most recent dycore version --- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 10 +- atm_dyn_iconam/tests/conftest.py | 14 +- atm_dyn_iconam/tests/test_diffusion.py | 162 +++++------------- .../src/icon4py/testutils/serialbox_utils.py | 150 ++++++++-------- 4 files changed, 148 insertions(+), 188 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index 858f9c9e7..ea3df6686 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -123,11 +123,13 @@ def read_static_fields( """ if ser_type == SerializationType.SB: - sp = serialbox_utils.IconSerialDataProvider( + dataprovider = serialbox_utils.IconSerialDataProvider( "icon_pydycore", str(path.absolute()), False - ).from_savepoint_diffusion_init(linit=True, date=SIMULATION_START_DATE) - metric_state = sp.construct_metric_state() - interpolation_state = sp.construct_interpolation_state() + ) + interpolation_state = ( + dataprovider.from_interpolation_savepoint().construct_interpolation_state() + ) + metric_state = dataprovider.from_metrics_savepoint().construct_metric_state() return metric_state, interpolation_state else: raise NotImplementedError("Only ser_type='sb' is implemented so far.") diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index ba7d2b336..6c0f53d91 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -20,7 +20,7 @@ from icon4py.testutils.serialbox_utils import IconSerialDataProvider -data_uri = "https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download" +data_uri = "https://polybox.ethz.ch/index.php/s/LcAbscZqnsx4WCf/download" data_path = Path(__file__).parent.joinpath("ser_icondata") extracted_path = data_path.joinpath("mch_ch_r04b09_dsl/ser_data") data_file = data_path.joinpath("mch_ch_r04b09_dsl_v2.tar.gz").name @@ -107,6 +107,18 @@ def diffusion_savepoint_exit(data_provider, step_date_exit): return sp +@pytest.fixture +def interpolation_savepoint(data_provider): + """Load data from ICON interplation state savepoint.""" + return data_provider.from_interpolation_savepoint() + + +@pytest.fixture +def metrics_savepoint(data_provider): + """Load data from ICON mestric state savepoint.""" + return data_provider.from_metrics_savepoint() + + @pytest.fixture def icon_grid(grid_savepoint): """ diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 1ee598828..4b0a70874 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -18,8 +18,6 @@ from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.diffusion.horizontal import CellParams, EdgeParams from icon4py.diffusion.icon_grid import VerticalModelParams -from icon4py.diffusion.interpolation_state import InterpolationState -from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.utils import ( _en_smag_fac_for_zero_nshift, _setup_runtime_diff_multfac_vn, @@ -33,6 +31,9 @@ from icon4py.testutils.utils import random_field, zero_field +datarun_reduced_substeps = 2 + + def test_scale_k(): mesh = SimpleMesh() field = random_field(mesh, KDim) @@ -154,7 +155,6 @@ def test_set_zero_vertex_k(): assert np.allclose(0.0, f) -@pytest.mark.datatest def test_diffusion_coefficients_with_hdiff_efdt_ratio(r04b09_diffusion_config): config = r04b09_diffusion_config config.hdiff_efdt_ratio = 1.0 @@ -168,7 +168,6 @@ def test_diffusion_coefficients_with_hdiff_efdt_ratio(r04b09_diffusion_config): assert params.K4W == pytest.approx(1.0 / 72.0, abs=1e-12) -@pytest.mark.datatest def test_diffusion_coefficients_without_hdiff_efdt_ratio(r04b09_diffusion_config): config = r04b09_diffusion_config config.hdiff_efdt_ratio = 0.0 @@ -182,7 +181,6 @@ def test_diffusion_coefficients_without_hdiff_efdt_ratio(r04b09_diffusion_config assert params.K4W == 0.0 -@pytest.mark.datatest def test_smagorinski_factor_for_diffusion_type_4(r04b09_diffusion_config): config = r04b09_diffusion_config config.smagorinski_scaling_factor = 0.15 @@ -212,8 +210,7 @@ def test_smagorinski_heights_diffusion_type_5_are_consistent( def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): - config = r04b09_diffusion_config - params = DiffusionParams(config) + params = DiffusionParams(r04b09_diffusion_config) assert len(params.smagorinski_factor) == len(params.smagorinski_height) assert len(params.smagorinski_factor) == 4 assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor))) @@ -222,6 +219,8 @@ def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): @pytest.mark.datatest def test_diffusion_init( diffusion_savepoint_init, + interpolation_savepoint, + metrics_savepoint, grid_savepoint, icon_grid, r04b09_diffusion_config, @@ -229,6 +228,7 @@ def test_diffusion_init( damping_height, ): config = r04b09_diffusion_config + config.ndyn_substeps = datarun_reduced_substeps additional_parameters = DiffusionParams(config) vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) @@ -238,25 +238,8 @@ def test_diffusion_init( assert meta["linit"] is False assert meta["date"] == step_date_init - grg = diffusion_savepoint_init.geofac_grg() - interpolation_state = InterpolationState( - e_bln_c_s=diffusion_savepoint_init.e_bln_c_s(), - rbf_coeff_1=diffusion_savepoint_init.rbf_vec_coeff_v1(), - rbf_coeff_2=diffusion_savepoint_init.rbf_vec_coeff_v2(), - geofac_div=diffusion_savepoint_init.geofac_div(), - geofac_n2s=diffusion_savepoint_init.geofac_n2s(), - geofac_grg_x=grg[0], - geofac_grg_y=grg[1], - nudgecoeff_e=diffusion_savepoint_init.nudgecoeff_e(), - ) - metric_state = MetricState( - mask_hdiff=diffusion_savepoint_init.mask_diff(), - theta_ref_mc=diffusion_savepoint_init.theta_ref_mc(), - wgtfac_c=diffusion_savepoint_init.wgtfac_c(), - zd_intcoef=diffusion_savepoint_init.zd_intcoef(), - zd_vertidx=diffusion_savepoint_init.zd_vertoffset(), - zd_diffcoef=diffusion_savepoint_init.zd_diffcoef(), - ) + interpolation_state = interpolation_savepoint.construct_interpolation_state() + metric_state = metrics_savepoint.construct_metric_state() diffusion = Diffusion() diffusion.init( @@ -334,6 +317,7 @@ def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( ): savepoint = diffusion_savepoint_init config = r04b09_diffusion_config + config.ndyn_substeps = datarun_reduced_substeps params = DiffusionParams(config) expected_diff_multfac_vn = savepoint.diff_multfac_vn() @@ -357,36 +341,21 @@ def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( @pytest.mark.datatest def test_verify_diffusion_init_against_first_regular_savepoint( diffusion_savepoint_init, + interpolation_savepoint, + metrics_savepoint, grid_savepoint, r04b09_diffusion_config, icon_grid, damping_height, ): config = r04b09_diffusion_config + config.ndyn_substeps = datarun_reduced_substeps additional_parameters = DiffusionParams(config) - - savepoint = diffusion_savepoint_init vct_a = grid_savepoint.vct_a() - grg = savepoint.geofac_grg() - interpolation_state = InterpolationState( - e_bln_c_s=savepoint.e_bln_c_s(), - rbf_coeff_1=savepoint.rbf_vec_coeff_v1(), - rbf_coeff_2=savepoint.rbf_vec_coeff_v2(), - geofac_div=savepoint.geofac_div(), - geofac_n2s=savepoint.geofac_n2s(), - geofac_grg_x=grg[0], - geofac_grg_y=grg[1], - nudgecoeff_e=savepoint.nudgecoeff_e(), - ) - metric_state = MetricState( - mask_hdiff=savepoint.mask_diff(), - theta_ref_mc=savepoint.theta_ref_mc(), - wgtfac_c=savepoint.wgtfac_c(), - zd_intcoef=savepoint.zd_intcoef(), - zd_vertidx=savepoint.zd_vertoffset(), - zd_diffcoef=savepoint.zd_diffcoef(), - ) + interpolation_state = interpolation_savepoint.construct_interpolation_state() + metric_state = metrics_savepoint.construct_metric_state() + diffusion = Diffusion() diffusion.init( config=config, @@ -397,7 +366,7 @@ def test_verify_diffusion_init_against_first_regular_savepoint( interpolation_state=interpolation_state, ) - _verify_init_values_against_savepoint(savepoint, diffusion) + _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) @pytest.mark.datatest @@ -406,32 +375,18 @@ def test_verify_diffusion_init_against_other_regular_savepoint( r04b09_diffusion_config, grid_savepoint, icon_grid, + interpolation_savepoint, + metrics_savepoint, diffusion_savepoint_init, damping_height, ): config = r04b09_diffusion_config + config.ndyn_substeps = datarun_reduced_substeps additional_parameters = DiffusionParams(config) vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) - grg = diffusion_savepoint_init.geofac_grg() - interpolation_state = InterpolationState( - e_bln_c_s=diffusion_savepoint_init.e_bln_c_s(), - rbf_coeff_1=diffusion_savepoint_init.rbf_vec_coeff_v1(), - rbf_coeff_2=diffusion_savepoint_init.rbf_vec_coeff_v2(), - geofac_div=diffusion_savepoint_init.geofac_div(), - geofac_n2s=diffusion_savepoint_init.geofac_n2s(), - geofac_grg_x=grg[0], - geofac_grg_y=grg[1], - nudgecoeff_e=diffusion_savepoint_init.nudgecoeff_e(), - ) - metric_state = MetricState( - mask_hdiff=diffusion_savepoint_init.mask_diff(), - theta_ref_mc=diffusion_savepoint_init.theta_ref_mc(), - wgtfac_c=diffusion_savepoint_init.wgtfac_c(), - zd_intcoef=diffusion_savepoint_init.zd_intcoef(), - zd_vertidx=diffusion_savepoint_init.zd_vertoffset(), - zd_diffcoef=diffusion_savepoint_init.zd_diffcoef(), - ) + interpolation_state = interpolation_savepoint.construct_interpolation_state() + metric_state = metrics_savepoint.construct_metric_state() diffusion = Diffusion() diffusion.init( @@ -453,31 +408,32 @@ def test_run_diffusion_single_step( run_with_program, diffusion_savepoint_init, diffusion_savepoint_exit, + interpolation_savepoint, + metrics_savepoint, grid_savepoint, icon_grid, r04b09_diffusion_config, damping_height, ): - ( - dtime, - cell_geometry, - edge_geometry, - diagnostic_state, - interpolation_state, - metric_state, - prognostic_state, - ) = _read_fields(diffusion_savepoint_init, grid_savepoint) - + dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") + edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() + cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() + interpolation_state = interpolation_savepoint.construct_interpolation_state() + metric_state = metrics_savepoint.construct_metric_state() + diagnostic_state = diffusion_savepoint_init.construct_diagnostics() + prognostic_state = diffusion_savepoint_init.construct_prognostics() vct_a = grid_savepoint.vct_a() vertical_params = VerticalModelParams( vct_a=vct_a, rayleigh_damping_height=damping_height ) - additional_parameters = DiffusionParams(r04b09_diffusion_config) + config = r04b09_diffusion_config + config.ndyn_substeps = datarun_reduced_substeps + additional_parameters = DiffusionParams(config) diffusion = Diffusion(run_program=run_with_program) diffusion.init( grid=icon_grid, - config=r04b09_diffusion_config, + config=config, params=additional_parameters, vertical_params=vertical_params, metric_state=metric_state, @@ -512,31 +468,6 @@ def test_run_diffusion_single_step( ) -def _read_fields(diffusion_savepoint_init, grid_savepoint): - dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") - edge_geometry, cell_geometry = _read_geometry_fields(grid_savepoint) - - interpolation_state = diffusion_savepoint_init.construct_interpolation_state() - metric_state = diffusion_savepoint_init.construct_metric_state() - diagnostic_state = diffusion_savepoint_init.construct_diagnostics() - prognostic_state = diffusion_savepoint_init.construct_prognostics() - return ( - dtime, - cell_geometry, - edge_geometry, - diagnostic_state, - interpolation_state, - metric_state, - prognostic_state, - ) - - -def _read_geometry_fields(grid_savepoint): - edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() - cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() - return edge_geometry, cell_geometry - - @pytest.mark.skip("fix: diffusion_stencil_15") @pytest.mark.datatest def test_diffusion_five_steps( @@ -544,31 +475,32 @@ def test_diffusion_five_steps( r04b09_diffusion_config, icon_grid, grid_savepoint, + interpolation_savepoint, + metrics_savepoint, diffusion_savepoint_init, diffusion_savepoint_exit, linit=True, step_date_exit="2021-06-20T12:01:00.000", ): - ( - dtime, - cell_geometry, - edge_geometry, - diagnostic_state, - interpolation_state, - metric_state, - prognostic_state, - ) = _read_fields(diffusion_savepoint_init, grid_savepoint) + dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") + edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() + cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() + interpolation_state = interpolation_savepoint.construct_interpolation_state() + metric_state = metrics_savepoint.construct_metric_state() + diagnostic_state = diffusion_savepoint_init.construct_diagnostics() + prognostic_state = diffusion_savepoint_init.construct_prognostics() vertical_params = VerticalModelParams( vct_a=grid_savepoint.vct_a(), rayleigh_damping_height=damping_height ) additional_parameters = DiffusionParams(r04b09_diffusion_config) - + config = r04b09_diffusion_config + config.ndyn_substeps = datarun_reduced_substeps diffusion = Diffusion() diffusion.init( grid=icon_grid, - config=r04b09_diffusion_config, + config=config, params=additional_parameters, vertical_params=vertical_params, metric_state=metric_state, diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index ad312d679..a711e1543 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -249,30 +249,70 @@ def construct_cell_geometry(self) -> CellParams: return CellParams(area=self.cell_areas()) -class IconDiffusionInitSavepoint(IconSavepoint): - def hdef_ic(self): - return self._get_field("hdef_ic", CellDim, KDim) +class InterpolationSavepoint(IconSavepoint): + def geofac_grg(self): + grg = np.squeeze(self.serializer.read("geofac_grg", self.savepoint)) + return np_as_located_field(CellDim, C2E2CODim)( + grg[:, :, 0] + ), np_as_located_field(CellDim, C2E2CODim)(grg[:, :, 1]) - def div_ic(self): - return self._get_field("div_ic", CellDim, KDim) + def zd_intcoef(self): + return self._get_field("vcoef", CellDim, C2E2CDim, KDim) - def dwdx(self): - return self._get_field("dwdx", CellDim, KDim) + def e_bln_c_s(self): + return self._get_field("e_bln_c_s", CellDim, C2EDim) - def dwdy(self): - return self._get_field("dwdy", CellDim, KDim) + def geofac_div(self): + return self._get_field("geofac_div", CellDim, C2EDim) - def vn(self): - return self._get_field("vn", EdgeDim, KDim) + def geofac_n2s(self): + return self._get_field("geofac_n2s", CellDim, C2E2CODim) - def theta_v(self): - return self._get_field("theta_v", CellDim, KDim) + def rbf_vec_coeff_v1(self): + return self._get_field("rbf_vec_coeff_v1", VertexDim, V2EDim) - def w(self): - return self._get_field("w", CellDim, KDim) + def rbf_vec_coeff_v2(self): + return self._get_field("rbf_vec_coeff_v2", VertexDim, V2EDim) - def exner(self): - return self._get_field("exner", CellDim, KDim) + def nudgecoeff_e(self): + return self._get_field("nudgecoeff_e", EdgeDim) + + def construct_interpolation_state(self) -> InterpolationState: + grg = self.geofac_grg() + return InterpolationState( + e_bln_c_s=self.e_bln_c_s(), + rbf_coeff_1=self.rbf_vec_coeff_v1(), + rbf_coeff_2=self.rbf_vec_coeff_v2(), + geofac_div=self.geofac_div(), + geofac_n2s=self.geofac_n2s(), + geofac_grg_x=grg[0], + geofac_grg_y=grg[1], + nudgecoeff_e=self.nudgecoeff_e(), + ) + + +class MetricSavepoint(IconSavepoint): + def construct_metric_state(self) -> MetricState: + return MetricState( + mask_hdiff=self.mask_diff(), + theta_ref_mc=self.theta_ref_mc(), + wgtfac_c=self.wgtfac_c(), + zd_intcoef=self.zd_intcoef(), + zd_vertidx=self.zd_vertoffset(), + zd_diffcoef=self.zd_diffcoef(), + ) + + def zd_diffcoef(self): + return self._get_field("zd_diffcoef", CellDim, KDim) + + def zd_intcoef(self): + return self._get_field("vcoef", CellDim, C2E2CDim, KDim) + + def zd_vertidx(self): + return self._get_field("zd_vertidx", CellDim, C2E2CDim, dtype=int) + + def zd_vertoffset(self): + return self._get_field("zd_vertoffset", CellDim, C2E2CDim, KDim, dtype=int) def theta_ref_mc(self): return self._get_field("theta_ref_mc", CellDim, KDim) @@ -286,42 +326,31 @@ def wgtfac_e(self): def mask_diff(self): return self._get_field("mask_hdiff", CellDim, KDim, dtype=int) - def zd_diffcoef(self): - return self._get_field("zd_diffcoef", CellDim, KDim) - - def zd_intcoef(self): - return self._get_field("vcoef", CellDim, C2E2CDim, KDim) - - def e_bln_c_s(self): - return self._get_field("e_bln_c_s", CellDim, C2EDim) - def geofac_div(self): - return self._get_field("geofac_div", CellDim, C2EDim) +class IconDiffusionInitSavepoint(IconSavepoint): + def hdef_ic(self): + return self._get_field("hdef_ic", CellDim, KDim) - def geofac_n2s(self): - return self._get_field("geofac_n2s", CellDim, C2E2CODim) + def div_ic(self): + return self._get_field("div_ic", CellDim, KDim) - def geofac_grg(self): - grg = np.squeeze(self.serializer.read("geofac_grg", self.savepoint)) - return np_as_located_field(CellDim, C2E2CODim)( - grg[:, :, 0] - ), np_as_located_field(CellDim, C2E2CODim)(grg[:, :, 1]) + def dwdx(self): + return self._get_field("dwdx", CellDim, KDim) - def nudgecoeff_e(self): - return self._get_field("nudgecoeff_e", EdgeDim) + def dwdy(self): + return self._get_field("dwdy", CellDim, KDim) - def zd_vertidx(self): - # TODO fix this - return self._get_field("zd_vertidx", CellDim, C2E2CDim, dtype=int) + def vn(self): + return self._get_field("vn", EdgeDim, KDim) - def zd_vertoffset(self): - return self._get_field("zd_vertoffset", CellDim, C2E2CDim, KDim, dtype=int) + def theta_v(self): + return self._get_field("theta_v", CellDim, KDim) - def rbf_vec_coeff_v1(self): - return self._get_field("rbf_vec_coeff_v1", VertexDim, V2EDim) + def w(self): + return self._get_field("w", CellDim, KDim) - def rbf_vec_coeff_v2(self): - return self._get_field("rbf_vec_coeff_v2", VertexDim, V2EDim) + def exner(self): + return self._get_field("exner", CellDim, KDim) def diff_multfac_smag(self): return np.squeeze(self.serializer.read("diff_multfac_smag", self.savepoint)) @@ -350,29 +379,6 @@ def diff_multfac_w(self): def diff_multfac_vn(self): return self.serializer.read("diff_multfac_vn", self.savepoint) - def construct_interpolation_state(self) -> InterpolationState: - grg = self.geofac_grg() - return InterpolationState( - e_bln_c_s=self.e_bln_c_s(), - rbf_coeff_1=self.rbf_vec_coeff_v1(), - rbf_coeff_2=self.rbf_vec_coeff_v2(), - geofac_div=self.geofac_div(), - geofac_n2s=self.geofac_n2s(), - geofac_grg_x=grg[0], - geofac_grg_y=grg[1], - nudgecoeff_e=self.nudgecoeff_e(), - ) - - def construct_metric_state(self) -> MetricState: - return MetricState( - mask_hdiff=self.mask_diff(), - theta_ref_mc=self.theta_ref_mc(), - wgtfac_c=self.wgtfac_c(), - zd_intcoef=self.zd_intcoef(), - zd_vertidx=self.zd_vertoffset(), - zd_diffcoef=self.zd_diffcoef(), - ) - def construct_prognostics(self) -> PrognosticState: return PrognosticState( w=self.w(), @@ -441,6 +447,14 @@ def from_savepoint_diffusion_init( ) return IconDiffusionInitSavepoint(savepoint, self.serializer) + def from_interpolation_savepoint(self) -> InterpolationSavepoint: + savepoint = self.serializer.savepoint["interpolation_state"].as_savepoint() + return InterpolationSavepoint(savepoint, self.serializer) + + def from_metrics_savepoint(self) -> MetricSavepoint: + savepoint = self.serializer.savepoint["metric_state"].as_savepoint() + return MetricSavepoint(savepoint, self.serializer) + def from_savepoint_diffusion_exit( self, linit: bool, date: str ) -> IconDiffusionExitSavepoint: From 2d5febc26593932e5c3ac0d7f7b103d31b5f9de7 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 6 Jun 2023 23:08:56 +0200 Subject: [PATCH 119/263] verifying up to theta_v and exner --- .../atm_dyn_iconam/apply_nabla2_to_w.py | 6 +- ...pply_nabla2_to_w_in_upper_damping_layer.py | 6 +- ...ate_horizontal_gradients_for_turbulence.py | 6 +- .../atm_dyn_iconam/calculate_nabla2_for_w.py | 6 +- .../src/icon4py/diffusion/diffusion.py | 290 +++++++++++------- .../icon4py/diffusion/diffusion_program.py | 2 +- .../icon4py/diffusion/interpolation_state.py | 4 +- ...ate_horizontal_gradients_for_turbulence.py | 2 +- atm_dyn_iconam/tests/test_diffusion.py | 38 ++- atm_dyn_iconam/tests/test_icon_grid.py | 2 + atm_dyn_iconam/tests/test_vertical.py | 7 +- .../src/icon4py/testutils/serialbox_utils.py | 7 +- 12 files changed, 236 insertions(+), 140 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py index e77fe2872..2d2e0399a 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py @@ -38,5 +38,9 @@ def apply_nabla2_to_w( geofac_n2s: Field[[CellDim, C2E2CODim], float], w: Field[[CellDim, KDim], float], diff_multfac_w: float, + horizontal_start:int, + horizontal_end:int, + vertical_start:int, + vertical_end:int ): - _apply_nabla2_to_w(area, z_nabla2_c, geofac_n2s, w, diff_multfac_w, out=w) + _apply_nabla2_to_w(area, z_nabla2_c, geofac_n2s, w, diff_multfac_w, out=w, domain={CellDim:(horizontal_start, horizontal_end), KDim:(vertical_start, vertical_end)}) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py index 07561e618..9efc07f37 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py @@ -34,7 +34,11 @@ def apply_nabla2_to_w_in_upper_damping_layer( diff_multfac_n2w: Field[[KDim], float], cell_area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], + horizontal_start: int, + horizontal_end: int, + vertical_start: int, + vertical_end: int ): _apply_nabla2_to_w_in_upper_damping_layer( - w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w + w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w, domain={CellDim:(horizontal_start, horizontal_end), KDim:(vertical_start, vertical_end)} ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py index 5ba95c84a..5fa26dcd9 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py @@ -35,7 +35,11 @@ def calculate_horizontal_gradients_for_turbulence( geofac_grg_y: Field[[CellDim, C2E2CODim], float], dwdx: Field[[CellDim, KDim], float], dwdy: Field[[CellDim, KDim], float], +horizontal_start: int, + horizontal_end: int, + vertical_start: int, + vertical_end: int ): _calculate_horizontal_gradients_for_turbulence( - w, geofac_grg_x, geofac_grg_y, out=(dwdx, dwdy) + w, geofac_grg_x, geofac_grg_y, out=(dwdx, dwdy), domain={CellDim:(horizontal_start, horizontal_end), KDim:(vertical_start, vertical_end)} ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py index fefa04f4c..ea6ca06c9 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py @@ -30,5 +30,9 @@ def calculate_nabla2_for_w( w: Field[[CellDim, KDim], float], geofac_n2s: Field[[CellDim, C2E2CODim], float], z_nabla2_c: Field[[CellDim, KDim], float], + horizontal_start: int, + horizontal_end: int, + vertical_start: int, + vertical_end: int ): - _calculate_nabla2_for_w(w, geofac_n2s, out=z_nabla2_c) + _calculate_nabla2_for_w(w, geofac_n2s, out=z_nabla2_c, domain={CellDim:(horizontal_start, horizontal_end), KDim:(vertical_start, vertical_end)}) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 8969a9444..de67555ef 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -23,9 +23,15 @@ from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn, run_gtfn_cached import icon4py.diffusion.diffusion_program as diff_prog +from icon4py.atm_dyn_iconam.apply_nabla2_to_w import apply_nabla2_to_w +from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import \ + apply_nabla2_to_w_in_upper_damping_layer +from icon4py.atm_dyn_iconam.calculate_horizontal_gradients_for_turbulence import \ + calculate_horizontal_gradients_for_turbulence from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( calculate_nabla2_and_smag_coefficients_for_vn, ) +from icon4py.atm_dyn_iconam.calculate_nabla2_for_w import calculate_nabla2_for_w from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import ( fused_mo_nh_diffusion_stencil_02_03, ) @@ -81,7 +87,6 @@ zero_field, ) - # flake8: noqa log = logging.getLogger(__name__) @@ -89,7 +94,9 @@ cached_backend = run_gtfn_cached compiled_backend = run_gtfn -backend = compiled_backend # +backend = compiled_backend # + + class DiffusionConfig: """ Contains necessary parameter to configure a diffusion run. @@ -381,8 +388,8 @@ def __init__(self, run_program=True): self.config: Optional[DiffusionConfig] = None self.params: Optional[DiffusionParams] = None self.vertical_params: Optional[VerticalModelParams] = None - self.interpolation_state:InterpolationState = None - self.metric_state:MetricState = None + self.interpolation_state: InterpolationState = None + self.metric_state: MetricState = None self.diff_multfac_w: Optional[float] = None self.diff_multfac_n2w: Field[[KDim], float] = None self.smag_offset: Optional[float] = None @@ -767,10 +774,10 @@ def _do_diffusion_step( HorizontalMarkerIndex.local(CellDim), ) - cell_start_nudging, _ = self.grid.get_indices_from_to( + cell_start_nudging, cell_end_local_minus1 = self.grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.local(CellDim), + HorizontalMarkerIndex.local(CellDim) - 1, ) edge_start_nudging_plus_one, edge_end_local = self.grid.get_indices_from_to( @@ -778,6 +785,11 @@ def _do_diffusion_step( HorizontalMarkerIndex.nudging(EdgeDim) + 1, HorizontalMarkerIndex.local(EdgeDim), ) + edge_start_interior, edge_end_nudging = self.grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.interior(EdgeDim), + HorizontalMarkerIndex.nudging(EdgeDim), + ) edge_start_lb_plus4, _ = self.grid.get_indices_from_to( EdgeDim, @@ -823,14 +835,14 @@ def _do_diffusion_step( set_zero_v_k.with_backend(backend)(self.v_vert, offset_provider={}) log.debug("rbf interpolation: start") # # 1. CALL rbf_vec_interpol_vertex - mo_intp_rbf_rbf_vec_interpol_vertex( + mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( p_e_in=prognostic_state.vn, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, ptr_coeff_2=self.interpolation_state.rbf_coeff_2, p_u_out=self.u_vert, p_v_out=self.v_vert, - horizontal_start=vertex_start_local_boundary_plus3, - horizontal_end=vertex_end_local_minus1, + horizontal_start=vertex_start_local_boundary_plus1, + horizontal_end=vertex_end_local, vertical_start=0, vertical_end=klevels, offset_provider={"V2E": self.grid.get_v2e_connectivity()}, @@ -879,7 +891,7 @@ def _do_diffusion_step( hdef_ic=diagnostic_state.hdef_ic, horizontal_start=cell_start_nudging, horizontal_end=cell_end_local, - vertical_start=0, + vertical_start=1, vertical_end=klevels, offset_provider={ "C2E": self.grid.get_c2e_connectivity(), @@ -898,19 +910,19 @@ def _do_diffusion_step( ptr_coeff_2=self.interpolation_state.rbf_coeff_2, p_u_out=self.u_vert, p_v_out=self.v_vert, - horizontal_start=vertex_start_local_boundary_plus3, + horizontal_start=vertex_start_local_boundary_plus1, horizontal_end=vertex_end_local, vertical_start=0, vertical_end=klevels, offset_provider={"V2E": self.grid.get_v2e_connectivity()}, ) log.debug("rbf interpolation: end") - # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult + # # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult + # # + # # # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 + # # # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 + # # # - # # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 - # # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 - # - log.debug("running fused stencil 04_05_06: start") fused_mo_nh_diffusion_stencil_04_05_06.with_backend(backend)( u_vert=self.u_vert, @@ -928,8 +940,8 @@ def _do_diffusion_step( horz_idx=self.horizontal_edge_index, nudgezone_diff=self.nudgezone_diff, fac_bdydiff_v=self.fac_bdydiff_v, - start_2nd_nudge_line_idx_e=int32(edge_start_nudging_minus1), - horizontal_start=edge_start_nudging_plus_one, + start_2nd_nudge_line_idx_e=int32(edge_start_nudging_plus_one), + horizontal_start=edge_start_lb_plus4, horizontal_end=edge_end_local, vertical_start=0, vertical_end=klevels, @@ -938,113 +950,165 @@ def _do_diffusion_step( "E2ECV": self.grid.get_e2ecv_connectivity(), }, ) + log.debug("running fused stencil 04_05_06: end") # # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, # # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 - log.debug("running fused stencil 04_05_06: end") log.debug("running fused stencil 07_08_09_10: start") - fused_mo_nh_diffusion_stencil_07_08_09_10.with_backend(backend)( - area=cell_areas, - geofac_n2s=self.interpolation_state.geofac_n2s, + w_old = prognostic_state.w + calculate_horizontal_gradients_for_turbulence.with_backend(backend)( + w=prognostic_state.w, geofac_grg_x=self.interpolation_state.geofac_grg_x, geofac_grg_y=self.interpolation_state.geofac_grg_y, - w_old=prognostic_state.w, - w=prognostic_state.w, dwdx=diagnostic_state.dwdx, dwdy=diagnostic_state.dwdy, - diff_multfac_w=self.diff_multfac_w, - diff_multfac_n2w=self.diff_multfac_n2w, - vert_idx=self.vertical_index, - horz_idx=self.horizontal_cell_index, - nrdmax=self.vertical_params.index_of_damping_layer, - interior_idx=int32( - cell_start_interior - ), # h end index for stencil_09 and stencil_10 - halo_idx=int32( - cell_end_local - ), # h end index for stencil_09 and stencil_10, - horizontal_start=cell_start_nudging, # h start index for stencil_07 and stencil_08 - horizontal_end=cell_end_local_plus1, # h end index for stencil_07 and stencil_08 - vertical_start=0, - vertical_end=klevels, - offset_provider={ - "C2E2CO": self.grid.get_c2e2co_connectivity(), - }, - ) - log.debug("running fused stencil 07_08_09_10: start") - # # 8. HALO EXCHANGE: CALL sync_patch_array - # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, - # # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 - # - # # TODO @magdalena check: kh_smag_e is an out field, should not be calculated in init? - # - log.debug("running fused stencil 11_12: start") - fused_mo_nh_diffusion_stencil_11_12.with_backend(backend)( - theta_v=prognostic_state.theta_v, - theta_ref_mc=self.metric_state.theta_ref_mc, - thresh_tdiff=self.thresh_tdiff, - kh_smag_e=self.kh_smag_e, - horizontal_start=edge_start_nudging_plus_one, - horizontal_end=edge_end_local, - vertical_start=k_start_end_minus2, + vertical_start=1, vertical_end=klevels, - offset_provider={ - "E2C": self.grid.get_e2c_connectivity(), - "C2E2C": self.grid.get_c2e2c_connectivity(), - }, - ) - log.debug("running fused stencil 11_12: end") - log.debug("running fused stencil 13_14: start") - fused_mo_nh_diffusion_stencil_13_14.with_backend(backend)( - kh_smag_e=self.kh_smag_e, - inv_dual_edge_length=inverse_dual_edge_length, - theta_v=prognostic_state.theta_v, - geofac_div=self.interpolation_state.geofac_div, - z_temp=self.z_temp, horizontal_start=cell_start_nudging, - horizontal_end=cell_end_local, - vertical_start=0, - vertical_end=klevels, - offset_provider={ - "C2E": self.grid.get_c2e_connectivity(), - "E2C": self.grid.get_e2c_connectivity(), - }, - ) - log.debug("running fused stencil 13_14: end") - log.debug("running fused stencil 15: start") - mo_nh_diffusion_stencil_15.with_backend(backend)( - mask=self.metric_state.mask_hdiff, - zd_vertoffset=self.metric_state.zd_vertidx, - zd_diffcoef=self.metric_state.zd_diffcoef, - geofac_n2s_c=self.interpolation_state.geofac_n2s_c, - geofac_n2s_nbh=self.interpolation_state.geofac_n2s_nbh, - vcoef=self.metric_state.zd_intcoef, - theta_v=prognostic_state.theta_v, - z_temp=self.z_temp, - horizontal_start=cell_start_nudging, - horizontal_end=cell_end_local, - vertical_start=0, - vertical_end=klevels, - offset_provider={ - "C2CEC": self.grid.get_c2cec_connectivity(), - "C2E2C": self.grid.get_c2e2c_connectivity(), - "Koff": KDim, - }, + horizontal_end=cell_end_local_minus1, + offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()} ) - log.debug("running fused stencil 15: end") - log.debug("running fused stencil update_theta_and_exner: start") - update_theta_and_exner.with_backend(backend)( - z_temp=self.z_temp, - area=cell_areas, - theta_v=prognostic_state.theta_v, - exner=prognostic_state.exner_pressure, - rd_o_cvd=self.rd_o_cvd, - horizontal_start=cell_start_nudging, - horizontal_end=cell_end_local, - vertical_start=0, - vertical_end=klevels, - offset_provider={}, + z_nabla2_c = zero_field(self.grid, CellDim, KDim, dtype=float) + + calculate_nabla2_for_w.with_backend(backend)(w=prognostic_state.w, + geofac_n2s=self.interpolation_state.geofac_n2s, + z_nabla2_c=z_nabla2_c, + vertical_start=0, + vertical_end=klevels, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local_minus1, + offset_provider={ + "C2E2CO": self.grid.get_c2e2co_connectivity()} + ) + apply_nabla2_to_w.with_backend(backend)(area=cell_areas, + z_nabla2_c=z_nabla2_c, + w=prognostic_state.w, + diff_multfac_w=self.diff_multfac_w, + geofac_n2s=self.interpolation_state.geofac_n2s, + vertical_start=0, + vertical_end=klevels, + horizontal_start=cell_start_interior, + horizontal_end=cell_end_local, + offset_provider={ + "C2E2CO": self.grid.get_c2e2co_connectivity()} + ) + + apply_nabla2_to_w_in_upper_damping_layer.with_backend(backend)( + w=prognostic_state.w, + diff_multfac_n2w=self.diff_multfac_n2w, + cell_area=cell_areas, + z_nabla2_c=z_nabla2_c, + vertical_start=1, + vertical_end=int(self.vertical_params.index_of_damping_layer + 1), + horizontal_start=int(cell_start_interior), + horizontal_end=int(cell_end_local), + offset_provider = {} + ) + # fused_mo_nh_diffusion_stencil_07_08_09_10.with_backend(backend)( + # area=cell_areas, + # geofac_n2s=self.interpolation_state.geofac_n2s, + # geofac_grg_x=self.interpolation_state.geofac_grg_x, + # geofac_grg_y=self.interpolation_state.geofac_grg_y, + # w_old=w_old, + # w=prognostic_state.w, + # dwdx=diagnostic_state.dwdx, + # dwdy=diagnostic_state.dwdy, + # diff_multfac_w=self.diff_multfac_w, + # diff_multfac_n2w=self.diff_multfac_n2w, + # vert_idx=self.vertical_index, + # horz_idx=self.horizontal_cell_index, + # nrdmax=self.vertical_params.index_of_damping_layer, + # interior_idx=int32( + # cell_start_interior + # ), # h end index for stencil_09 and stencil_10 + # halo_idx=int32( + # cell_end_local + # ), # h end index for stencil_09 and stencil_10, + # horizontal_start=cell_start_nudging, # h start index for stencil_07 and stencil_08 + # horizontal_end=cell_end_local_minus1, # h end index for stencil_07 and stencil_08 + # vertical_start=0, + # vertical_end=klevels, + # offset_provider={ + # "C2E2CO": self.grid.get_c2e2co_connectivity(), + # }, + # ) + log.debug("running fused stencil 07_08_09_10: end") + # # # 8. HALO EXCHANGE: CALL sync_patch_array + # # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, + # # # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 + # # + # # # TODO @magdalena check: kh_smag_e is an out field, should not be calculated in init? + # # + # log.debug("running fused stencil 11_12: start") + # fused_mo_nh_diffusion_stencil_11_12.with_backend(backend)( + # theta_v=prognostic_state.theta_v, + # theta_ref_mc=self.metric_state.theta_ref_mc, + # thresh_tdiff=self.thresh_tdiff, + # kh_smag_e=self.kh_smag_e, + # horizontal_start=edge_start_nudging_plus_one, + # horizontal_end=edge_end_local, + # vertical_start=k_start_end_minus2, + # vertical_end=klevels, + # offset_provider={ + # "E2C": self.grid.get_e2c_connectivity(), + # "C2E2C": self.grid.get_c2e2c_connectivity(), + # }, + # ) + # log.debug("running fused stencil 11_12: end") + # log.debug("running fused stencil 13_14: start") + # fused_mo_nh_diffusion_stencil_13_14.with_backend(backend)( + # kh_smag_e=self.kh_smag_e, + # inv_dual_edge_length=inverse_dual_edge_length, + # theta_v=prognostic_state.theta_v, + # geofac_div=self.interpolation_state.geofac_div, + # z_temp=self.z_temp, + # horizontal_start=cell_start_nudging, + # horizontal_end=cell_end_local, + # vertical_start=0, + # vertical_end=klevels, + # offset_provider={ + # "C2E": self.grid.get_c2e_connectivity(), + # "E2C": self.grid.get_e2c_connectivity(), + # }, + # ) + # log.debug("running fused stencil 13_14: end") + # log.debug("running fused stencil 15: start") + # mo_nh_diffusion_stencil_15.with_backend(backend)( + # mask=self.metric_state.mask_hdiff, + # zd_vertoffset=self.metric_state.zd_vertidx, + # zd_diffcoef=self.metric_state.zd_diffcoef, + # geofac_n2s_c=self.interpolation_state.geofac_n2s_c, + # geofac_n2s_nbh=self.interpolation_state.geofac_n2s_nbh, + # vcoef=self.metric_state.zd_intcoef, + # theta_v=prognostic_state.theta_v, + # z_temp=self.z_temp, + # horizontal_start=cell_start_nudging, + # horizontal_end=cell_end_local, + # vertical_start=0, + # vertical_end=klevels, + # offset_provider={ + # "C2CEC": self.grid.get_c2cec_connectivity(), + # "C2E2C": self.grid.get_c2e2c_connectivity(), + # "Koff": KDim, + # }, + # ) + # + # log.debug("running fused stencil 15: end") + # log.debug("running fused stencil update_theta_and_exner: start") + # update_theta_and_exner.with_backend(backend)( + # z_temp=self.z_temp, + # area=cell_areas, + # theta_v=prognostic_state.theta_v, + # exner=prognostic_state.exner_pressure, + # rd_o_cvd=self.rd_o_cvd, + # horizontal_start=cell_start_nudging, + # horizontal_end=cell_end_local, + # vertical_start=0, + # vertical_end=klevels, + # offset_provider={}, + # ) + log.debug("running fused stencil update_theta_and_exner: end") # 10. HALO EXCHANGE sync_patch_array diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py index 98ac5058e..3e220787a 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py @@ -188,7 +188,7 @@ def diffusion_run( ), domain={ CellDim: (cell_startindex_nudging, cell_endindex_local), - KDim: (0, nlev), + KDim: (1, nlev), }, ) # diff --git a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py index 4a864cb72..c9abbd707 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py @@ -55,8 +55,8 @@ class InterpolationState: ] # factor for nabla2-scalar (nproma,cell_type+1,nblks_c) geofac_grg_x: Field[ [CellDim, C2E2CODim], float - ] # factor for green gauss gradient (nproma,4,nblks_c,2) - geofac_grg_y: Field[[CellDim, C2E2CODim], float] + ] + geofac_grg_y: Field[[CellDim, C2E2CODim], float] # factors for green gauss gradient (nproma,4,nblks_c,2) nudgecoeff_e: Field[[EdgeDim], float] # Nudgeing coeffients for edges @property diff --git a/atm_dyn_iconam/tests/test_calculate_horizontal_gradients_for_turbulence.py b/atm_dyn_iconam/tests/test_calculate_horizontal_gradients_for_turbulence.py index 1232df961..bc6597d9b 100644 --- a/atm_dyn_iconam/tests/test_calculate_horizontal_gradients_for_turbulence.py +++ b/atm_dyn_iconam/tests/test_calculate_horizontal_gradients_for_turbulence.py @@ -53,9 +53,9 @@ def test_calculate_horizontal_gradients_for_turbulence(): geofac_grg_y, dwdx, dwdy, + 0, mesh.n_cells, 0, mesh.k_level, offset_provider={ "C2E2CO": mesh.get_c2e2cO_offset_provider(), - "C2E2CODim": C2E2CODim, }, ) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 97b5af5b2..2e0a79354 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -454,20 +454,30 @@ def test_run_diffusion_single_step( edge_areas=edge_geometry.edge_areas, cell_areas=cell_geometry.area, ) - - icon_result_exner = diffusion_savepoint_exit.exner() - icon_result_vn = diffusion_savepoint_exit.vn() - icon_result_w = diffusion_savepoint_exit.w() - icon_result_theta_w = diffusion_savepoint_exit.theta_v() - - assert np.allclose(icon_result_w, np.asarray(prognostic_state.w)) - assert np.allclose(np.asarray(icon_result_vn), np.asarray(prognostic_state.vn)) - assert np.allclose( - np.asarray(icon_result_theta_w), np.asarray(prognostic_state.theta_v) - ) - assert np.allclose( - np.asarray(icon_result_exner), np.asarray(prognostic_state.exner_pressure) - ) + assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v + icon_result_div_ic = diffusion_savepoint_exit.div_ic() + assert np.allclose(np.asarray(icon_result_div_ic), np.asarray(diagnostic_state.div_ic)) + icon_result_hdef_ic = diffusion_savepoint_exit.hdef_ic() + assert np.allclose(np.asarray(icon_result_hdef_ic), np.asarray(diagnostic_state.hdef_ic)) + + ref_w = np.asarray(diffusion_savepoint_exit.w()) + val_w = np.asarray(prognostic_state.w) + ref_dwdx = np.asarray(diffusion_savepoint_exit.dwdx()) + val_dwdx = np.asarray(diagnostic_state.dwdx) + ref_dwdy = np.asarray(diffusion_savepoint_exit.dwdy()) + val_dwdy = np.asarray(diagnostic_state.dwdy) + ref_vn = np.asarray(diffusion_savepoint_exit.vn()) + val_vn = np.asarray(prognostic_state.vn) + + assert np.allclose(ref_vn, val_vn) + assert np.allclose(ref_dwdx, val_dwdx) + assert np.allclose(ref_dwdy, val_dwdy) + assert np.allclose(ref_w, np.asarray(val_w)) + ref_exner = np.asarray(diffusion_savepoint_exit.exner()) + ref_theta_v = np.asarray(diffusion_savepoint_exit.theta_v()) + val_theta_v = np.asarray(prognostic_state.theta_v) + val_exner = np.asarray(prognostic_state.exner_pressure) + # @pytest.mark.skip("fix: diffusion_stencil_15") diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index c3de4fbb1..6d81ef680 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -238,3 +238,5 @@ def test_horizontal_vertex_indices(icon_grid): HorizontalMarkerIndex.interior(VertexDim), HorizontalMarkerIndex.interior(VertexDim), ) == (2071, 10663) + + diff --git a/atm_dyn_iconam/tests/test_vertical.py b/atm_dyn_iconam/tests/test_vertical.py index ccf0b7dbf..7cd3335dc 100644 --- a/atm_dyn_iconam/tests/test_vertical.py +++ b/atm_dyn_iconam/tests/test_vertical.py @@ -36,11 +36,12 @@ def test_nrdmax_calculation(max_h, damping, delta): @pytest.mark.datatest def test_nrdmax_calculation_from_icon_input(icon_grid, grid_savepoint): a = grid_savepoint.vct_a() + nrdmax = grid_savepoint.nrdmax() damping_height = 12500 vertical_params = VerticalModelParams( rayleigh_damping_height=damping_height, vct_a=a ) - assert 9 == vertical_params.index_of_damping_layer + assert nrdmax == vertical_params.index_of_damping_layer a_array = np.asarray(a) - assert a_array[9] > damping_height - assert a_array[10] < damping_height + assert a_array[nrdmax] > damping_height + assert a_array[nrdmax + 1] < damping_height diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 8660743b0..8ea350448 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -157,7 +157,6 @@ def refin_ctrl(self, dim: Dimension): return None def c2e(self): - return self._get_connectiviy_array("c2e") def _get_connectiviy_array(self, name: str): @@ -184,6 +183,9 @@ def e2c2v(self): def v2e(self): return self._get_connectiviy_array("v2e") + def nrdmax(self): + return self._get_connectiviy_array("nrdmax") + def construct_icon_grid(self) -> IconGrid: sp_meta = self.get_metadata( "nproma", "nlev", "num_vert", "num_cells", "num_edges" @@ -203,7 +205,8 @@ def construct_icon_grid(self) -> IconGrid: VerticalMeshConfig(num_lev=sp_meta["nlev"]), ) c2e2c = self.c2e2c() - c2e2c0 = np.column_stack((c2e2c, (np.asarray(range(c2e2c.shape[0]))))) + origin = np.asarray(range(c2e2c.shape[0])) + c2e2c0 = np.column_stack((origin, c2e2c)) grid = ( IconGrid() .with_config(config) From e994ab31cfb44b14e417a50a4b7ef7b5b1e35bce Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 7 Jun 2023 09:46:45 +0200 Subject: [PATCH 120/263] add validation of theta_v and exner --- .../src/icon4py/diffusion/diffusion.py | 146 +++++++++--------- atm_dyn_iconam/tests/test_diffusion.py | 5 +- 2 files changed, 76 insertions(+), 75 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index de67555ef..192e977ee 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -1035,80 +1035,80 @@ def _do_diffusion_step( # }, # ) log.debug("running fused stencil 07_08_09_10: end") - # # # 8. HALO EXCHANGE: CALL sync_patch_array - # # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, - # # # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 - # # - # # # TODO @magdalena check: kh_smag_e is an out field, should not be calculated in init? - # # - # log.debug("running fused stencil 11_12: start") - # fused_mo_nh_diffusion_stencil_11_12.with_backend(backend)( - # theta_v=prognostic_state.theta_v, - # theta_ref_mc=self.metric_state.theta_ref_mc, - # thresh_tdiff=self.thresh_tdiff, - # kh_smag_e=self.kh_smag_e, - # horizontal_start=edge_start_nudging_plus_one, - # horizontal_end=edge_end_local, - # vertical_start=k_start_end_minus2, - # vertical_end=klevels, - # offset_provider={ - # "E2C": self.grid.get_e2c_connectivity(), - # "C2E2C": self.grid.get_c2e2c_connectivity(), - # }, - # ) - # log.debug("running fused stencil 11_12: end") - # log.debug("running fused stencil 13_14: start") - # fused_mo_nh_diffusion_stencil_13_14.with_backend(backend)( - # kh_smag_e=self.kh_smag_e, - # inv_dual_edge_length=inverse_dual_edge_length, - # theta_v=prognostic_state.theta_v, - # geofac_div=self.interpolation_state.geofac_div, - # z_temp=self.z_temp, - # horizontal_start=cell_start_nudging, - # horizontal_end=cell_end_local, - # vertical_start=0, - # vertical_end=klevels, - # offset_provider={ - # "C2E": self.grid.get_c2e_connectivity(), - # "E2C": self.grid.get_e2c_connectivity(), - # }, - # ) - # log.debug("running fused stencil 13_14: end") - # log.debug("running fused stencil 15: start") - # mo_nh_diffusion_stencil_15.with_backend(backend)( - # mask=self.metric_state.mask_hdiff, - # zd_vertoffset=self.metric_state.zd_vertidx, - # zd_diffcoef=self.metric_state.zd_diffcoef, - # geofac_n2s_c=self.interpolation_state.geofac_n2s_c, - # geofac_n2s_nbh=self.interpolation_state.geofac_n2s_nbh, - # vcoef=self.metric_state.zd_intcoef, - # theta_v=prognostic_state.theta_v, - # z_temp=self.z_temp, - # horizontal_start=cell_start_nudging, - # horizontal_end=cell_end_local, - # vertical_start=0, - # vertical_end=klevels, - # offset_provider={ - # "C2CEC": self.grid.get_c2cec_connectivity(), - # "C2E2C": self.grid.get_c2e2c_connectivity(), - # "Koff": KDim, - # }, - # ) + # # 8. HALO EXCHANGE: CALL sync_patch_array + # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, + # # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 # - # log.debug("running fused stencil 15: end") - # log.debug("running fused stencil update_theta_and_exner: start") - # update_theta_and_exner.with_backend(backend)( - # z_temp=self.z_temp, - # area=cell_areas, - # theta_v=prognostic_state.theta_v, - # exner=prognostic_state.exner_pressure, - # rd_o_cvd=self.rd_o_cvd, - # horizontal_start=cell_start_nudging, - # horizontal_end=cell_end_local, - # vertical_start=0, - # vertical_end=klevels, - # offset_provider={}, - # ) + # # TODO @magdalena check: kh_smag_e is an out field, should not be calculated in init? + # + log.debug("running fused stencil 11_12: start") + fused_mo_nh_diffusion_stencil_11_12.with_backend(backend)( + theta_v=prognostic_state.theta_v, + theta_ref_mc=self.metric_state.theta_ref_mc, + thresh_tdiff=self.thresh_tdiff, + kh_smag_e=self.kh_smag_e, + horizontal_start=edge_start_nudging_plus_one, + horizontal_end=edge_end_local, + vertical_start=k_start_end_minus2, + vertical_end=klevels, + offset_provider={ + "E2C": self.grid.get_e2c_connectivity(), + "C2E2C": self.grid.get_c2e2c_connectivity(), + }, + ) + log.debug("running fused stencil 11_12: end") + log.debug("running fused stencil 13_14: start") + fused_mo_nh_diffusion_stencil_13_14.with_backend(backend)( + kh_smag_e=self.kh_smag_e, + inv_dual_edge_length=inverse_dual_edge_length, + theta_v=prognostic_state.theta_v, + geofac_div=self.interpolation_state.geofac_div, + z_temp=self.z_temp, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={ + "C2E": self.grid.get_c2e_connectivity(), + "E2C": self.grid.get_e2c_connectivity(), + }, + ) + log.debug("running fused stencil 13_14: end") + log.debug("running fused stencil 15: start") + mo_nh_diffusion_stencil_15.with_backend(backend)( + mask=self.metric_state.mask_hdiff, + zd_vertoffset=self.metric_state.zd_vertidx, + zd_diffcoef=self.metric_state.zd_diffcoef, + geofac_n2s_c=self.interpolation_state.geofac_n2s_c, + geofac_n2s_nbh=self.interpolation_state.geofac_n2s_nbh, + vcoef=self.metric_state.zd_intcoef, + theta_v=prognostic_state.theta_v, + z_temp=self.z_temp, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={ + "C2CEC": self.grid.get_c2cec_connectivity(), + "C2E2C": self.grid.get_c2e2c_connectivity(), + "Koff": KDim, + }, + ) + + log.debug("running fused stencil 15: end") + log.debug("running fused stencil update_theta_and_exner: start") + update_theta_and_exner.with_backend(backend)( + z_temp=self.z_temp, + area=cell_areas, + theta_v=prognostic_state.theta_v, + exner=prognostic_state.exner_pressure, + rd_o_cvd=self.rd_o_cvd, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={}, + ) log.debug("running fused stencil update_theta_and_exner: end") # 10. HALO EXCHANGE sync_patch_array diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 2e0a79354..839af3778 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -472,12 +472,13 @@ def test_run_diffusion_single_step( assert np.allclose(ref_vn, val_vn) assert np.allclose(ref_dwdx, val_dwdx) assert np.allclose(ref_dwdy, val_dwdy) - assert np.allclose(ref_w, np.asarray(val_w)) + assert np.allclose(ref_w, val_w) ref_exner = np.asarray(diffusion_savepoint_exit.exner()) ref_theta_v = np.asarray(diffusion_savepoint_exit.theta_v()) val_theta_v = np.asarray(prognostic_state.theta_v) val_exner = np.asarray(prognostic_state.exner_pressure) - # + assert np.allclose(ref_theta_v, val_theta_v) + assert np.allclose(ref_exner, val_exner) @pytest.mark.skip("fix: diffusion_stencil_15") From d56de306b223497597955a825da948dc589e5f3c Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 7 Jun 2023 17:02:41 +0200 Subject: [PATCH 121/263] rename metrics state zd_vertidx fields --- atm_dyn_iconam/src/icon4py/diffusion/diffusion.py | 14 ++++++++++---- .../src/icon4py/diffusion/diffusion_program.py | 6 +++--- .../src/icon4py/diffusion/metric_state.py | 2 +- atm_dyn_iconam/tests/test_io_utils.py | 2 +- testutils/src/icon4py/testutils/serialbox_utils.py | 6 +++++- 5 files changed, 20 insertions(+), 10 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 192e977ee..00033bf43 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -648,7 +648,7 @@ def run( metric_theta_ref_mc=self.metric_state.theta_ref_mc, metric_wgtfac_c=self.metric_state.wgtfac_c, metric_mask_hdiff=self.metric_state.mask_hdiff, # - metric_zd_vertidx=self.metric_state.zd_vertidx, # + metric_zd_vertoffset=self.metric_state.zd_vertoffset, # metric_zd_diffcoef=self.metric_state.zd_diffcoef, # metric_zd_intcoef=self.metric_state.zd_intcoef, # interpolation_e_bln_c_s=self.interpolation_state.e_bln_c_s, @@ -791,6 +791,12 @@ def _do_diffusion_step( HorizontalMarkerIndex.nudging(EdgeDim), ) + edge_start_nudging, edge_end_halo = self.grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.nudging(EdgeDim), + HorizontalMarkerIndex.halo(EdgeDim), + ) + edge_start_lb_plus4, _ = self.grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, @@ -1047,8 +1053,8 @@ def _do_diffusion_step( theta_ref_mc=self.metric_state.theta_ref_mc, thresh_tdiff=self.thresh_tdiff, kh_smag_e=self.kh_smag_e, - horizontal_start=edge_start_nudging_plus_one, - horizontal_end=edge_end_local, + horizontal_start=edge_start_nudging, + horizontal_end=edge_end_halo, vertical_start=k_start_end_minus2, vertical_end=klevels, offset_provider={ @@ -1077,7 +1083,7 @@ def _do_diffusion_step( log.debug("running fused stencil 15: start") mo_nh_diffusion_stencil_15.with_backend(backend)( mask=self.metric_state.mask_hdiff, - zd_vertoffset=self.metric_state.zd_vertidx, + zd_vertoffset=self.metric_state.zd_vertoffset, zd_diffcoef=self.metric_state.zd_diffcoef, geofac_n2s_c=self.interpolation_state.geofac_n2s_c, geofac_n2s_nbh=self.interpolation_state.geofac_n2s_nbh, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py index 3e220787a..efdcfacb8 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py @@ -16,7 +16,7 @@ from gt4py.next.common import Field from gt4py.next.ffront.decorator import program from gt4py.next.ffront.fbuiltins import int32 -from gt4py.next.program_processors.runners import gtfn_cpu + from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( _calculate_nabla2_and_smag_coefficients_for_vn, @@ -71,7 +71,7 @@ def diffusion_run( metric_theta_ref_mc: Field[[CellDim, KDim], float], metric_wgtfac_c: Field[[CellDim, KDim], float], metric_mask_hdiff: Field[[CellDim, KDim], bool], - metric_zd_vertidx: Field[[CECDim, KDim], int32], + metric_zd_vertoffset: Field[[CECDim, KDim], int32], metric_zd_diffcoef: Field[[CellDim, KDim], float], metric_zd_intcoef: Field[[CECDim, KDim], float], interpolation_e_bln_c_s: Field[[CellDim, C2EDim], float], @@ -295,7 +295,7 @@ def diffusion_run( # MO_NH_DIFFUSION_STENCIL_15: as_offset index fields! _mo_nh_diffusion_stencil_15( metric_mask_hdiff, - metric_zd_vertidx, + metric_zd_vertoffset, metric_zd_diffcoef, interpolation_geofac_n2s_c, interpolation_geofac_n2s_nbh, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py b/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py index bf4f55432..4fcdb0830 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py @@ -26,6 +26,6 @@ class MetricState: [CellDim, KDim], float ] # weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_c) mask_hdiff: Field[[CellDim, KDim], bool] - zd_vertidx: Field[[CECDim, KDim], int32] + zd_vertoffset: Field[[CECDim, KDim], int32] zd_diffcoef: Field[[CellDim, KDim], float] zd_intcoef: Field[[CECDim, KDim], float] diff --git a/atm_dyn_iconam/tests/test_io_utils.py b/atm_dyn_iconam/tests/test_io_utils.py index a03d810db..c31ce40d8 100644 --- a/atm_dyn_iconam/tests/test_io_utils.py +++ b/atm_dyn_iconam/tests/test_io_utils.py @@ -101,7 +101,7 @@ def assert_metric_state_fields(metric_state): assert metric_state.zd_diffcoef assert metric_state.theta_ref_mc assert metric_state.mask_hdiff - assert metric_state.zd_vertidx + assert metric_state.zd_vertoffset def assert_interpolation_state_fields(interpolation_state): diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 8ea350448..1aea3dba2 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -298,7 +298,7 @@ def construct_metric_state(self) -> MetricState: theta_ref_mc=self.theta_ref_mc(), wgtfac_c=self.wgtfac_c(), zd_intcoef=self.zd_intcoef(), - zd_vertidx=self.zd_vertoffset(), + zd_vertoffset=self.zd_vertoffset(), zd_diffcoef=self.zd_diffcoef(), ) @@ -420,8 +420,12 @@ def dwdy(self): def exner(self): return self._get_field("x_exner", CellDim, KDim) + def z_temp(self): + return self._get_field("x_z_temp", CellDim, KDim) + def div_ic(self): return self._get_field("x_div_ic",CellDim, KDim ) + def hdef_ic(self): return self._get_field("x_hdef_ic", CellDim, KDim) From f5df4614973a366a3eeb4d6789092ea7f51d26f7 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 7 Jun 2023 17:03:25 +0200 Subject: [PATCH 122/263] fix unit tests of stencils that now take offset args --- atm_dyn_iconam/tests/test_apply_nabla2_to_w.py | 1 + .../tests/test_apply_nabla2_to_w_in_upper_damping_layer.py | 1 + atm_dyn_iconam/tests/test_calculate_nabla2_for_w.py | 1 + 3 files changed, 3 insertions(+) diff --git a/atm_dyn_iconam/tests/test_apply_nabla2_to_w.py b/atm_dyn_iconam/tests/test_apply_nabla2_to_w.py index 90eedf2e7..0831226ba 100644 --- a/atm_dyn_iconam/tests/test_apply_nabla2_to_w.py +++ b/atm_dyn_iconam/tests/test_apply_nabla2_to_w.py @@ -59,6 +59,7 @@ def test_apply_nabla2_to_w(): geofac_n2s, w, diff_multfac_w, + 0, mesh.n_cells, 0, mesh.k_level, offset_provider={ "C2E2CO": mesh.get_c2e2cO_offset_provider(), "C2E2CODim": C2E2CODim, diff --git a/atm_dyn_iconam/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py b/atm_dyn_iconam/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py index 8b7596d6c..5170d7008 100644 --- a/atm_dyn_iconam/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py +++ b/atm_dyn_iconam/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py @@ -51,6 +51,7 @@ def test_apply_nabla2_to_w_in_upper_damping_layer(): diff_multfac_n2w, cell_area, z_nabla2_c, + 0, mesh.n_cells, 0, mesh.k_level, offset_provider={}, ) assert np.allclose(w, ref) diff --git a/atm_dyn_iconam/tests/test_calculate_nabla2_for_w.py b/atm_dyn_iconam/tests/test_calculate_nabla2_for_w.py index 5b40fac8c..2bb626f9a 100644 --- a/atm_dyn_iconam/tests/test_calculate_nabla2_for_w.py +++ b/atm_dyn_iconam/tests/test_calculate_nabla2_for_w.py @@ -41,6 +41,7 @@ def test_calculate_nabla2_for_w(): w, geofac_n2s, z_nabla2_c, + 0, mesh.n_cells, 0, mesh.k_level, offset_provider={ "C2E2CO": mesh.get_c2e2cO_offset_provider(), "C2E2CODim": C2E2CODim, From 040d986abcfa7767f302f1e130814f37bba0cb54 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 7 Jun 2023 17:04:13 +0200 Subject: [PATCH 123/263] refactoring in horizontal.py --- .../src/icon4py/diffusion/horizontal.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py index 82c3a206e..2c01c7bf5 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py @@ -42,19 +42,22 @@ class HorizontalMarkerIndex: _LOCAL_BOUNDARY_EDGES: Final[int] = 1 + _ICON_INDEX_OFFSET_EDGES _INTERIOR_EDGES: Final[int] = _ICON_INDEX_OFFSET_EDGES _NUDGING_EDGES: Final[int] = _GRF_BOUNDARY_WIDTH_EDGES + _ICON_INDEX_OFFSET_EDGES - _HALO_EDGES: Final[int] = _MIN_RL_EDGE_INT + _ICON_INDEX_OFFSET_EDGES + _HALO_EDGES: Final[int] = _MIN_RL_EDGE_INT -1 + _ICON_INDEX_OFFSET_EDGES + _LOCAL_EDGES: Final[int] = _MIN_RL_EDGE_INT + _ICON_INDEX_OFFSET_EDGES _END_EDGES: Final[int] = 0 _LOCAL_BOUNDARY_CELLS: Final[int] = 1 + _ICON_INDEX_OFFSET_CELLS _INTERIOR_CELLS: Final[int] = _ICON_INDEX_OFFSET_CELLS _NUDGING_CELLS: Final[int] = _GRF_BOUNDARY_WIDTH_CELL + 1 + _ICON_INDEX_OFFSET_CELLS - _HALO_CELLS: Final[int] = _MIN_RL_CELL_INT + _ICON_INDEX_OFFSET_CELLS + _HALO_CELLS: Final[int] = _MIN_RL_CELL_INT -1 + _ICON_INDEX_OFFSET_CELLS + _LOCAL_CELLS: Final[int] = _MIN_RL_CELL_INT + _ICON_INDEX_OFFSET_CELLS _END_CELLS: Final[int] = 0 _LOCAL_BOUNDARY_VERTICES = 1 + _ICON_INDEX_OFFSET_VERTEX _INTERIOR_VERTICES: Final[int] = _ICON_INDEX_OFFSET_VERTEX _NUDGING_VERTICES: Final[int] = 0 - _HALO_VERTICES: Final[int] = _MIN_RL_VERTEX_INT + _ICON_INDEX_OFFSET_VERTEX + _HALO_VERTICES: Final[int] = _MIN_RL_VERTEX_INT -1 + _ICON_INDEX_OFFSET_VERTEX + _LOCAL_VERTICES: Final[int] = _MIN_RL_VERTEX_INT+ _ICON_INDEX_OFFSET_VERTEX _END_VERTICES: Final[int] = 0 @classmethod @@ -69,6 +72,16 @@ def lateral_boundary(cls, dim: Dimension) -> int: @classmethod def local(cls, dim: Dimension) -> int: + match (dim): + case (dimension.CellDim): + return cls._LOCAL_CELLS + case (dimension.EdgeDim): + return cls._LOCAL_EDGES + case (dimension.VertexDim): + return cls._LOCAL_VERTICES + + @classmethod + def halo(cls, dim: Dimension) -> int: match (dim): case (dimension.CellDim): return cls._HALO_CELLS @@ -76,7 +89,6 @@ def local(cls, dim: Dimension) -> int: return cls._HALO_EDGES case (dimension.VertexDim): return cls._HALO_VERTICES - @classmethod def nudging(cls, dim: Dimension) -> int: match (dim): From eaa838230528e985836d7889547cd59d4128e632 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 7 Jun 2023 17:04:37 +0200 Subject: [PATCH 124/263] verify theta_v and exner only for non steep points --- atm_dyn_iconam/tests/test_diffusion.py | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 839af3778..0bb789d9d 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -401,9 +401,7 @@ def test_verify_diffusion_init_against_other_regular_savepoint( _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) -#@pytest.mark.skip("fix: diffusion_stencil_15") #@pytest.mark.parametrize("run_with_program", [True, False]) - @pytest.mark.parametrize("run_with_program", [False]) @pytest.mark.datatest def test_run_diffusion_single_step( @@ -455,10 +453,13 @@ def test_run_diffusion_single_step( cell_areas=cell_geometry.area, ) assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v - icon_result_div_ic = diffusion_savepoint_exit.div_ic() - assert np.allclose(np.asarray(icon_result_div_ic), np.asarray(diagnostic_state.div_ic)) - icon_result_hdef_ic = diffusion_savepoint_exit.hdef_ic() - assert np.allclose(np.asarray(icon_result_hdef_ic), np.asarray(diagnostic_state.hdef_ic)) + ref_div_ic = np.asarray(diffusion_savepoint_exit.div_ic()) + val_div_ic = np.asarray(diagnostic_state.div_ic) + ref_hdef_ic = np.asarray(diffusion_savepoint_exit.hdef_ic()) + val_hdef_ic = np.asarray(diagnostic_state.hdef_ic) + + assert np.allclose(ref_div_ic, val_div_ic) + assert np.allclose(ref_hdef_ic, val_hdef_ic) ref_w = np.asarray(diffusion_savepoint_exit.w()) val_w = np.asarray(prognostic_state.w) @@ -477,8 +478,12 @@ def test_run_diffusion_single_step( ref_theta_v = np.asarray(diffusion_savepoint_exit.theta_v()) val_theta_v = np.asarray(prognostic_state.theta_v) val_exner = np.asarray(prognostic_state.exner_pressure) - assert np.allclose(ref_theta_v, val_theta_v) - assert np.allclose(ref_exner, val_exner) + val_z_temp = np.asarray(diffusion.z_temp) + ref_z_temp = np.asarray(diffusion_savepoint_exit.z_temp()) + #assert np.allclose(ref_z_temp, val_z_temp) + flat_points = ~np.asarray(diffusion.metric_state.mask_hdiff) + assert np.allclose(ref_theta_v[flat_points], val_theta_v[flat_points]) + assert np.allclose(ref_exner[flat_points], val_exner[flat_points]) @pytest.mark.skip("fix: diffusion_stencil_15") From 18973282048cdb5620ea6b012693d843e4829794 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 8 Jun 2023 13:05:02 +0200 Subject: [PATCH 125/263] - verification of theta_v and exner: separate steep and flat points - (trial) allow for offset when reading (index) field - --- atm_dyn_iconam/tests/test_diffusion.py | 10 ++++++---- testutils/src/icon4py/testutils/serialbox_utils.py | 6 +++--- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 0bb789d9d..69672a5b7 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -480,10 +480,12 @@ def test_run_diffusion_single_step( val_exner = np.asarray(prognostic_state.exner_pressure) val_z_temp = np.asarray(diffusion.z_temp) ref_z_temp = np.asarray(diffusion_savepoint_exit.z_temp()) - #assert np.allclose(ref_z_temp, val_z_temp) - flat_points = ~np.asarray(diffusion.metric_state.mask_hdiff) - assert np.allclose(ref_theta_v[flat_points], val_theta_v[flat_points]) - assert np.allclose(ref_exner[flat_points], val_exner[flat_points]) + steep_points = np.asarray(diffusion.metric_state.mask_hdiff) + assert np.allclose(ref_z_temp[~steep_points], val_z_temp[~steep_points], atol=1e-7) + assert np.allclose(ref_theta_v[~steep_points], val_theta_v[~steep_points]) + assert np.allclose(ref_exner[~steep_points], val_exner[~steep_points]) + assert np.allclose(ref_theta_v[steep_points], val_theta_v[steep_points]) + assert np.allclose(ref_exner[steep_points], val_exner[steep_points]) @pytest.mark.skip("fix: diffusion_stencil_15") diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 1aea3dba2..03ff67fab 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -309,15 +309,15 @@ def zd_intcoef(self): return self._from_cell_c2e2c_to_cec("vcoef") - def _from_cell_c2e2c_to_cec(self, field_name:str ): - ser_input = np.squeeze(self.serializer.read(field_name, self.savepoint)) + def _from_cell_c2e2c_to_cec(self, field_name:str , offset:int = 0): + ser_input = np.squeeze(self.serializer.read(field_name, self.savepoint)) + offset old_shape = ser_input.shape return np_as_located_field(CECDim, KDim)( ser_input.reshape(old_shape[0] * old_shape[1], old_shape[2]) ) def zd_vertoffset(self): - return self._from_cell_c2e2c_to_cec("zd_vertoffset") + return self._from_cell_c2e2c_to_cec("zd_vertoffset", 0) def theta_ref_mc(self): From ee4e3cdac4b46e8f135d7bf45ebdd44a5bc003b4 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 8 Jun 2023 13:11:49 +0200 Subject: [PATCH 126/263] pre-commit hook --- .../atm_dyn_iconam/apply_nabla2_to_w.py | 21 ++++-- ...pply_nabla2_to_w_in_upper_damping_layer.py | 12 +++- ...ate_horizontal_gradients_for_turbulence.py | 13 +++- .../atm_dyn_iconam/calculate_nabla2_for_w.py | 12 +++- .../mo_nh_diffusion_stencil_15.py | 11 ++-- .../src/icon4py/diffusion/diffusion.py | 65 ++++++++++--------- .../icon4py/diffusion/diffusion_program.py | 1 - .../src/icon4py/diffusion/horizontal.py | 9 +-- .../icon4py/diffusion/interpolation_state.py | 6 +- .../tests/test_apply_nabla2_to_w.py | 5 +- ...pply_nabla2_to_w_in_upper_damping_layer.py | 5 +- ...ate_horizontal_gradients_for_turbulence.py | 5 +- .../tests/test_calculate_nabla2_for_w.py | 5 +- atm_dyn_iconam/tests/test_diffusion.py | 2 +- atm_dyn_iconam/tests/test_icon_grid.py | 2 - .../src/icon4py/testutils/serialbox_utils.py | 14 ++-- 16 files changed, 121 insertions(+), 67 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py index 2d2e0399a..93bb21450 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py @@ -38,9 +38,20 @@ def apply_nabla2_to_w( geofac_n2s: Field[[CellDim, C2E2CODim], float], w: Field[[CellDim, KDim], float], diff_multfac_w: float, - horizontal_start:int, - horizontal_end:int, - vertical_start:int, - vertical_end:int + horizontal_start: int, + horizontal_end: int, + vertical_start: int, + vertical_end: int, ): - _apply_nabla2_to_w(area, z_nabla2_c, geofac_n2s, w, diff_multfac_w, out=w, domain={CellDim:(horizontal_start, horizontal_end), KDim:(vertical_start, vertical_end)}) + _apply_nabla2_to_w( + area, + z_nabla2_c, + geofac_n2s, + w, + diff_multfac_w, + out=w, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py index 9efc07f37..0483c7d7b 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py @@ -37,8 +37,16 @@ def apply_nabla2_to_w_in_upper_damping_layer( horizontal_start: int, horizontal_end: int, vertical_start: int, - vertical_end: int + vertical_end: int, ): _apply_nabla2_to_w_in_upper_damping_layer( - w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w, domain={CellDim:(horizontal_start, horizontal_end), KDim:(vertical_start, vertical_end)} + w, + diff_multfac_n2w, + cell_area, + z_nabla2_c, + out=w, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py index 5fa26dcd9..ce8a44758 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py @@ -35,11 +35,18 @@ def calculate_horizontal_gradients_for_turbulence( geofac_grg_y: Field[[CellDim, C2E2CODim], float], dwdx: Field[[CellDim, KDim], float], dwdy: Field[[CellDim, KDim], float], -horizontal_start: int, + horizontal_start: int, horizontal_end: int, vertical_start: int, - vertical_end: int + vertical_end: int, ): _calculate_horizontal_gradients_for_turbulence( - w, geofac_grg_x, geofac_grg_y, out=(dwdx, dwdy), domain={CellDim:(horizontal_start, horizontal_end), KDim:(vertical_start, vertical_end)} + w, + geofac_grg_x, + geofac_grg_y, + out=(dwdx, dwdy), + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py index ea6ca06c9..2a9dc2c41 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py @@ -33,6 +33,14 @@ def calculate_nabla2_for_w( horizontal_start: int, horizontal_end: int, vertical_start: int, - vertical_end: int + vertical_end: int, ): - _calculate_nabla2_for_w(w, geofac_n2s, out=z_nabla2_c, domain={CellDim:(horizontal_start, horizontal_end), KDim:(vertical_start, vertical_end)}) + _calculate_nabla2_for_w( + w, + geofac_n2s, + out=z_nabla2_c, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py index 86d5ab3df..0d12a990d 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py @@ -76,9 +76,9 @@ def mo_nh_diffusion_stencil_15( theta_v: Field[[CellDim, KDim], float], z_temp: Field[[CellDim, KDim], float], horizontal_start: int, - horizontal_end:int, - vertical_start:int, - vertical_end:int, + horizontal_end: int, + vertical_start: int, + vertical_end: int, ): _mo_nh_diffusion_stencil_15( mask, @@ -90,5 +90,8 @@ def mo_nh_diffusion_stencil_15( theta_v, z_temp, out=z_temp, - domain = {CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end), } + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 00033bf43..4f07a7fb3 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -20,14 +20,19 @@ from gt4py.next.common import Dimension from gt4py.next.ffront.fbuiltins import Field, int32 from gt4py.next.iterator.embedded import np_as_located_field -from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn, run_gtfn_cached +from gt4py.next.program_processors.runners.gtfn_cpu import ( + run_gtfn, + run_gtfn_cached, +) import icon4py.diffusion.diffusion_program as diff_prog from icon4py.atm_dyn_iconam.apply_nabla2_to_w import apply_nabla2_to_w -from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import \ - apply_nabla2_to_w_in_upper_damping_layer -from icon4py.atm_dyn_iconam.calculate_horizontal_gradients_for_turbulence import \ - calculate_horizontal_gradients_for_turbulence +from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import ( + apply_nabla2_to_w_in_upper_damping_layer, +) +from icon4py.atm_dyn_iconam.calculate_horizontal_gradients_for_turbulence import ( + calculate_horizontal_gradients_for_turbulence, +) from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( calculate_nabla2_and_smag_coefficients_for_vn, ) @@ -87,6 +92,7 @@ zero_field, ) + # flake8: noqa log = logging.getLogger(__name__) @@ -972,33 +978,33 @@ def _do_diffusion_step( vertical_end=klevels, horizontal_start=cell_start_nudging, horizontal_end=cell_end_local_minus1, - offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()} + offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, ) z_nabla2_c = zero_field(self.grid, CellDim, KDim, dtype=float) - calculate_nabla2_for_w.with_backend(backend)(w=prognostic_state.w, - geofac_n2s=self.interpolation_state.geofac_n2s, - z_nabla2_c=z_nabla2_c, - vertical_start=0, - vertical_end=klevels, - horizontal_start=cell_start_nudging, - horizontal_end=cell_end_local_minus1, - offset_provider={ - "C2E2CO": self.grid.get_c2e2co_connectivity()} - ) - apply_nabla2_to_w.with_backend(backend)(area=cell_areas, - z_nabla2_c=z_nabla2_c, - w=prognostic_state.w, - diff_multfac_w=self.diff_multfac_w, - geofac_n2s=self.interpolation_state.geofac_n2s, - vertical_start=0, - vertical_end=klevels, - horizontal_start=cell_start_interior, - horizontal_end=cell_end_local, - offset_provider={ - "C2E2CO": self.grid.get_c2e2co_connectivity()} - ) + calculate_nabla2_for_w.with_backend(backend)( + w=prognostic_state.w, + geofac_n2s=self.interpolation_state.geofac_n2s, + z_nabla2_c=z_nabla2_c, + vertical_start=0, + vertical_end=klevels, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local_minus1, + offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, + ) + apply_nabla2_to_w.with_backend(backend)( + area=cell_areas, + z_nabla2_c=z_nabla2_c, + w=prognostic_state.w, + diff_multfac_w=self.diff_multfac_w, + geofac_n2s=self.interpolation_state.geofac_n2s, + vertical_start=0, + vertical_end=klevels, + horizontal_start=cell_start_interior, + horizontal_end=cell_end_local, + offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, + ) apply_nabla2_to_w_in_upper_damping_layer.with_backend(backend)( w=prognostic_state.w, @@ -1009,8 +1015,7 @@ def _do_diffusion_step( vertical_end=int(self.vertical_params.index_of_damping_layer + 1), horizontal_start=int(cell_start_interior), horizontal_end=int(cell_end_local), - offset_provider = {} - + offset_provider={}, ) # fused_mo_nh_diffusion_stencil_07_08_09_10.with_backend(backend)( # area=cell_areas, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py index efdcfacb8..53a91d209 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py @@ -17,7 +17,6 @@ from gt4py.next.ffront.decorator import program from gt4py.next.ffront.fbuiltins import int32 - from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( _calculate_nabla2_and_smag_coefficients_for_vn, ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py index 2c01c7bf5..4095f4636 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py @@ -42,22 +42,22 @@ class HorizontalMarkerIndex: _LOCAL_BOUNDARY_EDGES: Final[int] = 1 + _ICON_INDEX_OFFSET_EDGES _INTERIOR_EDGES: Final[int] = _ICON_INDEX_OFFSET_EDGES _NUDGING_EDGES: Final[int] = _GRF_BOUNDARY_WIDTH_EDGES + _ICON_INDEX_OFFSET_EDGES - _HALO_EDGES: Final[int] = _MIN_RL_EDGE_INT -1 + _ICON_INDEX_OFFSET_EDGES + _HALO_EDGES: Final[int] = _MIN_RL_EDGE_INT - 1 + _ICON_INDEX_OFFSET_EDGES _LOCAL_EDGES: Final[int] = _MIN_RL_EDGE_INT + _ICON_INDEX_OFFSET_EDGES _END_EDGES: Final[int] = 0 _LOCAL_BOUNDARY_CELLS: Final[int] = 1 + _ICON_INDEX_OFFSET_CELLS _INTERIOR_CELLS: Final[int] = _ICON_INDEX_OFFSET_CELLS _NUDGING_CELLS: Final[int] = _GRF_BOUNDARY_WIDTH_CELL + 1 + _ICON_INDEX_OFFSET_CELLS - _HALO_CELLS: Final[int] = _MIN_RL_CELL_INT -1 + _ICON_INDEX_OFFSET_CELLS + _HALO_CELLS: Final[int] = _MIN_RL_CELL_INT - 1 + _ICON_INDEX_OFFSET_CELLS _LOCAL_CELLS: Final[int] = _MIN_RL_CELL_INT + _ICON_INDEX_OFFSET_CELLS _END_CELLS: Final[int] = 0 _LOCAL_BOUNDARY_VERTICES = 1 + _ICON_INDEX_OFFSET_VERTEX _INTERIOR_VERTICES: Final[int] = _ICON_INDEX_OFFSET_VERTEX _NUDGING_VERTICES: Final[int] = 0 - _HALO_VERTICES: Final[int] = _MIN_RL_VERTEX_INT -1 + _ICON_INDEX_OFFSET_VERTEX - _LOCAL_VERTICES: Final[int] = _MIN_RL_VERTEX_INT+ _ICON_INDEX_OFFSET_VERTEX + _HALO_VERTICES: Final[int] = _MIN_RL_VERTEX_INT - 1 + _ICON_INDEX_OFFSET_VERTEX + _LOCAL_VERTICES: Final[int] = _MIN_RL_VERTEX_INT + _ICON_INDEX_OFFSET_VERTEX _END_VERTICES: Final[int] = 0 @classmethod @@ -89,6 +89,7 @@ def halo(cls, dim: Dimension) -> int: return cls._HALO_EDGES case (dimension.VertexDim): return cls._HALO_VERTICES + @classmethod def nudging(cls, dim: Dimension) -> int: match (dim): diff --git a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py index c9abbd707..c79308646 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py @@ -53,10 +53,10 @@ class InterpolationState: geofac_n2s: Field[ [CellDim, C2E2CODim], float ] # factor for nabla2-scalar (nproma,cell_type+1,nblks_c) - geofac_grg_x: Field[ + geofac_grg_x: Field[[CellDim, C2E2CODim], float] + geofac_grg_y: Field[ [CellDim, C2E2CODim], float - ] - geofac_grg_y: Field[[CellDim, C2E2CODim], float] # factors for green gauss gradient (nproma,4,nblks_c,2) + ] # factors for green gauss gradient (nproma,4,nblks_c,2) nudgecoeff_e: Field[[EdgeDim], float] # Nudgeing coeffients for edges @property diff --git a/atm_dyn_iconam/tests/test_apply_nabla2_to_w.py b/atm_dyn_iconam/tests/test_apply_nabla2_to_w.py index 0831226ba..e7452794d 100644 --- a/atm_dyn_iconam/tests/test_apply_nabla2_to_w.py +++ b/atm_dyn_iconam/tests/test_apply_nabla2_to_w.py @@ -59,7 +59,10 @@ def test_apply_nabla2_to_w(): geofac_n2s, w, diff_multfac_w, - 0, mesh.n_cells, 0, mesh.k_level, + 0, + mesh.n_cells, + 0, + mesh.k_level, offset_provider={ "C2E2CO": mesh.get_c2e2cO_offset_provider(), "C2E2CODim": C2E2CODim, diff --git a/atm_dyn_iconam/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py b/atm_dyn_iconam/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py index 5170d7008..e58548e9b 100644 --- a/atm_dyn_iconam/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py +++ b/atm_dyn_iconam/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py @@ -51,7 +51,10 @@ def test_apply_nabla2_to_w_in_upper_damping_layer(): diff_multfac_n2w, cell_area, z_nabla2_c, - 0, mesh.n_cells, 0, mesh.k_level, + 0, + mesh.n_cells, + 0, + mesh.k_level, offset_provider={}, ) assert np.allclose(w, ref) diff --git a/atm_dyn_iconam/tests/test_calculate_horizontal_gradients_for_turbulence.py b/atm_dyn_iconam/tests/test_calculate_horizontal_gradients_for_turbulence.py index bc6597d9b..627fede5a 100644 --- a/atm_dyn_iconam/tests/test_calculate_horizontal_gradients_for_turbulence.py +++ b/atm_dyn_iconam/tests/test_calculate_horizontal_gradients_for_turbulence.py @@ -53,7 +53,10 @@ def test_calculate_horizontal_gradients_for_turbulence(): geofac_grg_y, dwdx, dwdy, - 0, mesh.n_cells, 0, mesh.k_level, + 0, + mesh.n_cells, + 0, + mesh.k_level, offset_provider={ "C2E2CO": mesh.get_c2e2cO_offset_provider(), }, diff --git a/atm_dyn_iconam/tests/test_calculate_nabla2_for_w.py b/atm_dyn_iconam/tests/test_calculate_nabla2_for_w.py index 2bb626f9a..93fd9bed9 100644 --- a/atm_dyn_iconam/tests/test_calculate_nabla2_for_w.py +++ b/atm_dyn_iconam/tests/test_calculate_nabla2_for_w.py @@ -41,7 +41,10 @@ def test_calculate_nabla2_for_w(): w, geofac_n2s, z_nabla2_c, - 0, mesh.n_cells, 0, mesh.k_level, + 0, + mesh.n_cells, + 0, + mesh.k_level, offset_provider={ "C2E2CO": mesh.get_c2e2cO_offset_provider(), "C2E2CODim": C2E2CODim, diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 69672a5b7..e70969b55 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -401,7 +401,7 @@ def test_verify_diffusion_init_against_other_regular_savepoint( _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) -#@pytest.mark.parametrize("run_with_program", [True, False]) +# @pytest.mark.parametrize("run_with_program", [True, False]) @pytest.mark.parametrize("run_with_program", [False]) @pytest.mark.datatest def test_run_diffusion_single_step( diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index 6d81ef680..c3de4fbb1 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -238,5 +238,3 @@ def test_horizontal_vertex_indices(icon_grid): HorizontalMarkerIndex.interior(VertexDim), HorizontalMarkerIndex.interior(VertexDim), ) == (2071, 10663) - - diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index 03ff67fab..c7703757b 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -22,6 +22,7 @@ C2E2CDim, C2E2CODim, C2EDim, + CECDim, CellDim, E2C2VDim, E2CDim, @@ -30,7 +31,7 @@ EdgeDim, KDim, V2EDim, - VertexDim, CECDim, + VertexDim, ) from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.diffusion import VectorTuple @@ -308,9 +309,10 @@ def zd_diffcoef(self): def zd_intcoef(self): return self._from_cell_c2e2c_to_cec("vcoef") - - def _from_cell_c2e2c_to_cec(self, field_name:str , offset:int = 0): - ser_input = np.squeeze(self.serializer.read(field_name, self.savepoint)) + offset + def _from_cell_c2e2c_to_cec(self, field_name: str, offset: int = 0): + ser_input = ( + np.squeeze(self.serializer.read(field_name, self.savepoint)) + offset + ) old_shape = ser_input.shape return np_as_located_field(CECDim, KDim)( ser_input.reshape(old_shape[0] * old_shape[1], old_shape[2]) @@ -319,7 +321,6 @@ def _from_cell_c2e2c_to_cec(self, field_name:str , offset:int = 0): def zd_vertoffset(self): return self._from_cell_c2e2c_to_cec("zd_vertoffset", 0) - def theta_ref_mc(self): return self._get_field("theta_ref_mc", CellDim, KDim) @@ -417,6 +418,7 @@ def dwdx(self): def dwdy(self): return self._get_field("x_dwdy", CellDim, KDim) + def exner(self): return self._get_field("x_exner", CellDim, KDim) @@ -424,7 +426,7 @@ def z_temp(self): return self._get_field("x_z_temp", CellDim, KDim) def div_ic(self): - return self._get_field("x_div_ic",CellDim, KDim ) + return self._get_field("x_div_ic", CellDim, KDim) def hdef_ic(self): return self._get_field("x_hdef_ic", CellDim, KDim) From 2885fe58391d88425d571582c2e04311fda71aff Mon Sep 17 00:00:00 2001 From: Magdalena Date: Thu, 8 Jun 2023 14:53:40 +0200 Subject: [PATCH 127/263] (fix) icon has origin at first place in C2E2CO (#223) --- testutils/src/icon4py/testutils/serialbox_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/testutils/src/icon4py/testutils/serialbox_utils.py b/testutils/src/icon4py/testutils/serialbox_utils.py index a711e1543..7cd15ec30 100644 --- a/testutils/src/icon4py/testutils/serialbox_utils.py +++ b/testutils/src/icon4py/testutils/serialbox_utils.py @@ -157,7 +157,6 @@ def refin_ctrl(self, dim: Dimension): return None def c2e(self): - return self._get_connectiviy_array("c2e") def _get_connectiviy_array(self, name: str): @@ -203,7 +202,7 @@ def construct_icon_grid(self) -> IconGrid: VerticalMeshConfig(num_lev=sp_meta["nlev"]), ) c2e2c = self.c2e2c() - c2e2c0 = np.column_stack((c2e2c, (np.asarray(range(c2e2c.shape[0]))))) + c2e2c0 = np.column_stack(((np.asarray(range(c2e2c.shape[0]))), c2e2c)) grid = ( IconGrid() .with_config(config) From 576d381dcbf62113d3477974cbdb2ea20a20b7e8 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 8 Jun 2023 16:56:31 +0200 Subject: [PATCH 128/263] remove program, add test that verifies the init step --- .../src/icon4py/diffusion/diffusion.py | 222 ++---------- .../icon4py/diffusion/diffusion_program.py | 323 ------------------ atm_dyn_iconam/tests/conftest.py | 8 +- atm_dyn_iconam/tests/test_diffusion.py | 101 +++--- 4 files changed, 85 insertions(+), 569 deletions(-) delete mode 100644 atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 4f07a7fb3..96ca9cce3 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -25,7 +25,6 @@ run_gtfn_cached, ) -import icon4py.diffusion.diffusion_program as diff_prog from icon4py.atm_dyn_iconam.apply_nabla2_to_w import apply_nabla2_to_w from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import ( apply_nabla2_to_w_in_upper_damping_layer, @@ -389,7 +388,6 @@ def __init__(self, run_program=True): self.thresh_tdiff: float = ( -5.0 ) # threshold temperature deviation from neighboring grid points hat activates extra diffusion against runaway cooling - self._run_program = run_program self.grid: Optional[IconGrid] = None self.config: Optional[DiffusionConfig] = None self.params: Optional[DiffusionParams] = None @@ -564,168 +562,22 @@ def run( runs a diffusion step for the parameter linit=False, within regular time loop. """ - if not self._run_program: - self._do_diffusion_step( - diagnostic_state=diagnostic_state, - prognostic_state=prognostic_state, - dtime=dtime, - tangent_orientation=tangent_orientation, - inverse_primal_edge_lengths=inverse_primal_edge_lengths, - inverse_dual_edge_length=inverse_dual_edge_length, - inverse_vertex_vertex_lengths=inverse_vert_vert_lengths, - primal_normal_vert=primal_normal_vert, - dual_normal_vert=dual_normal_vert, - edge_areas=edge_areas, - cell_areas=cell_areas, - diff_multfac_vn=self.diff_multfac_vn, - smag_limit=self.smag_limit, - smag_offset=self.smag_offset, - ) - else: - log.info("running diffusion_program") - ( - cell_startindex_nudging, - cell_endindex_local, - ) = self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.local(CellDim), - ) - - ( - cell_startindex_interior, - cell_endindex_local_plus1, - ) = self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.interior(CellDim), - HorizontalMarkerIndex.local(CellDim) - 1, - ) - - ( - edge_startindex_nudging_plus1, - edge_endindex_local, - ) = self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim) + 1, - HorizontalMarkerIndex.local(EdgeDim), - ) - - ( - edge_startindex_nudging_minus1, - edge_endindex_local_minus2, - ) = self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim) - 1, - HorizontalMarkerIndex.local(EdgeDim) - 2, - ) - - ( - vertex_startindex_lb_plus3, - vertex_endindex_local, - ) = self.grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, - HorizontalMarkerIndex.local(VertexDim), - ) - - ( - vertex_startindex_lb_plus1, - vertex_endindex_local_minus1, - ) = self.grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, - HorizontalMarkerIndex.local(VertexDim) - 1, - ) - edge_start_lb_plus4, _ = self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, - ) - log.info("diffusion program: start") - diff_prog.diffusion_run.with_backend(backend)( - diagnostic_hdef_ic=diagnostic_state.hdef_ic, - diagnostic_div_ic=diagnostic_state.div_ic, - diagnostic_dwdx=diagnostic_state.dwdx, - diagnostic_dwdy=diagnostic_state.dwdy, - prognostic_w=prognostic_state.w, - prognostic_vn=prognostic_state.vn, - prognostic_exner_pressure=prognostic_state.exner_pressure, - prognostic_theta_v=prognostic_state.theta_v, - metric_theta_ref_mc=self.metric_state.theta_ref_mc, - metric_wgtfac_c=self.metric_state.wgtfac_c, - metric_mask_hdiff=self.metric_state.mask_hdiff, # - metric_zd_vertoffset=self.metric_state.zd_vertoffset, # - metric_zd_diffcoef=self.metric_state.zd_diffcoef, # - metric_zd_intcoef=self.metric_state.zd_intcoef, # - interpolation_e_bln_c_s=self.interpolation_state.e_bln_c_s, - interpolation_rbf_coeff_1=self.interpolation_state.rbf_coeff_1, - interpolation_rbf_coeff_2=self.interpolation_state.rbf_coeff_2, - interpolation_geofac_div=self.interpolation_state.geofac_div, - interpolation_geofac_grg_x=self.interpolation_state.geofac_grg_x, - interpolation_geofac_grg_y=self.interpolation_state.geofac_grg_y, - interpolation_nudgecoeff_e=self.interpolation_state.nudgecoeff_e, - interpolation_geofac_n2s=self.interpolation_state.geofac_n2s, - interpolation_geofac_n2s_c=self.interpolation_state.geofac_n2s_c, - interpolation_geofac_n2s_nbh=self.interpolation_state.geofac_n2s_nbh, # - tangent_orientation=tangent_orientation, - inverse_primal_edge_lengths=inverse_primal_edge_lengths, - inverse_dual_edge_lengths=inverse_dual_edge_length, - inverse_vert_vert_lengths=inverse_vert_vert_lengths, - primal_normal_vert_1=primal_normal_vert[0], - primal_normal_vert_2=primal_normal_vert[1], - dual_normal_vert_1=dual_normal_vert[0], - dual_normal_vert_2=dual_normal_vert[1], - edge_areas=edge_areas, - cell_areas=cell_areas, - diff_multfac_vn=self.diff_multfac_vn, - dtime=dtime, - rd_o_cvd=self.rd_o_cvd, - local_thresh_tdiff=self.thresh_tdiff, - local_smag_limit=self.smag_limit, - local_u_vert=self.u_vert, - local_v_vert=self.v_vert, - local_enh_smag_fac=self.enh_smag_fac, - local_kh_smag_e=self.kh_smag_e, - local_kh_smag_ec=self.kh_smag_ec, - local_z_nabla2_e=self.z_nabla2_e, - local_z_temp=self.z_temp, - local_diff_multfac_smag=self.diff_multfac_smag, - local_diff_multfac_n2w=self.diff_multfac_n2w, - local_smag_offset=self.smag_offset, - local_nudgezone_diff=self.nudgezone_diff, - local_fac_bdydiff_v=self.fac_bdydiff_v, - local_diff_multfac_w=self.diff_multfac_w, - local_vertical_index=self.vertical_index, - local_horizontal_cell_index=self.horizontal_cell_index, - local_horizontal_edge_index=self.horizontal_edge_index, - cell_startindex_interior=int32(cell_startindex_interior), - cell_startindex_nudging=cell_startindex_nudging, - cell_endindex_local_plus1=cell_endindex_local_plus1, - cell_endindex_local=cell_endindex_local, - cell_halo_idx=int32(cell_endindex_local), - edge_startindex_nudging_plus1=edge_startindex_nudging_plus1, - edge_startindex_nudging_minus1=int32(edge_startindex_nudging_minus1), - edge_endindex_local=edge_endindex_local, - edge_endindex_local_minus2=edge_endindex_local_minus2, - vertex_startindex_lb_plus3=vertex_startindex_lb_plus3, - vertex_startindex_lb_plus1=vertex_startindex_lb_plus1, - vertex_endindex_local=vertex_endindex_local, - vertex_endindex_local_minus1=vertex_endindex_local_minus1, - index_of_damping_height=self.vertical_params.index_of_damping_layer, - nlev=self.grid.n_lev(), - boundary_diffusion_start_index_edges=edge_start_lb_plus4, - offset_provider={ - "V2E": self.grid.get_v2e_connectivity(), - "E2C2V": self.grid.get_e2c2v_connectivity(), - "E2ECV": self.grid.get_e2ecv_connectivity(), - "C2E": self.grid.get_c2e_connectivity(), - "E2C": self.grid.get_e2c_connectivity(), - "C2E2C": self.grid.get_c2e2c_connectivity(), - "C2CEC": self.grid.get_c2cec_connectivity(), - "C2E2CO": self.grid.get_c2e2co_connectivity(), - "Koff": KDim, - }, - ) + self._do_diffusion_step( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + tangent_orientation=tangent_orientation, + inverse_primal_edge_lengths=inverse_primal_edge_lengths, + inverse_dual_edge_length=inverse_dual_edge_length, + inverse_vertex_vertex_lengths=inverse_vert_vert_lengths, + primal_normal_vert=primal_normal_vert, + dual_normal_vert=dual_normal_vert, + edge_areas=edge_areas, + cell_areas=cell_areas, + diff_multfac_vn=self.diff_multfac_vn, + smag_limit=self.smag_limit, + smag_offset=self.smag_offset, + ) log.info("diffusion program: end") def _do_diffusion_step( @@ -766,24 +618,16 @@ def _do_diffusion_step( """ klevels = self.grid.n_lev() - k_start_end_minus2 = klevels - 2 - - cell_start_nudging_minus1, cell_end_local_plus1 = self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim) - 1, - HorizontalMarkerIndex.local(CellDim) - 1, - ) - cell_start_interior, cell_end_local = self.grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.interior(CellDim), HorizontalMarkerIndex.local(CellDim), ) - cell_start_nudging, cell_end_local_minus1 = self.grid.get_indices_from_to( + cell_start_nudging, cell_end_halo = self.grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.local(CellDim) - 1, + HorizontalMarkerIndex.halo(CellDim), ) edge_start_nudging_plus_one, edge_end_local = self.grid.get_indices_from_to( @@ -791,11 +635,6 @@ def _do_diffusion_step( HorizontalMarkerIndex.nudging(EdgeDim) + 1, HorizontalMarkerIndex.local(EdgeDim), ) - edge_start_interior, edge_end_nudging = self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.interior(EdgeDim), - HorizontalMarkerIndex.nudging(EdgeDim), - ) edge_start_nudging, edge_end_halo = self.grid.get_indices_from_to( EdgeDim, @@ -827,7 +666,7 @@ def _do_diffusion_step( HorizontalMarkerIndex.local(VertexDim), ) ( - vertex_start_local_boundary_plus1, + vertex_start_lb_plus1, vertex_end_local_minus1, ) = self.grid.get_indices_from_to( VertexDim, @@ -853,7 +692,7 @@ def _do_diffusion_step( ptr_coeff_2=self.interpolation_state.rbf_coeff_2, p_u_out=self.u_vert, p_v_out=self.v_vert, - horizontal_start=vertex_start_local_boundary_plus1, + horizontal_start=vertex_start_lb_plus1, horizontal_end=vertex_end_local, vertical_start=0, vertical_end=klevels, @@ -922,7 +761,7 @@ def _do_diffusion_step( ptr_coeff_2=self.interpolation_state.rbf_coeff_2, p_u_out=self.u_vert, p_v_out=self.v_vert, - horizontal_start=vertex_start_local_boundary_plus1, + horizontal_start=vertex_start_lb_plus1, horizontal_end=vertex_end_local, vertical_start=0, vertical_end=klevels, @@ -966,8 +805,7 @@ def _do_diffusion_step( # # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, # # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 - log.debug("running fused stencil 07_08_09_10: start") - w_old = prognostic_state.w + log.debug("running stencils 07_08_09_10: start") calculate_horizontal_gradients_for_turbulence.with_backend(backend)( w=prognostic_state.w, geofac_grg_x=self.interpolation_state.geofac_grg_x, @@ -977,7 +815,7 @@ def _do_diffusion_step( vertical_start=1, vertical_end=klevels, horizontal_start=cell_start_nudging, - horizontal_end=cell_end_local_minus1, + horizontal_end=cell_end_halo, offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, ) @@ -990,7 +828,7 @@ def _do_diffusion_step( vertical_start=0, vertical_end=klevels, horizontal_start=cell_start_nudging, - horizontal_end=cell_end_local_minus1, + horizontal_end=cell_end_halo, offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, ) apply_nabla2_to_w.with_backend(backend)( @@ -1006,6 +844,7 @@ def _do_diffusion_step( offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, ) + # TODO @magdalena: fix this python offset problem (+1): int(self.vertical_params.index_of_damping_layer + 1) apply_nabla2_to_w_in_upper_damping_layer.with_backend(backend)( w=prognostic_state.w, diff_multfac_n2w=self.diff_multfac_n2w, @@ -1017,6 +856,7 @@ def _do_diffusion_step( horizontal_end=int(cell_end_local), offset_provider={}, ) + # w_old = prognostic_state.w # fused_mo_nh_diffusion_stencil_07_08_09_10.with_backend(backend)( # area=cell_areas, # geofac_n2s=self.interpolation_state.geofac_n2s, @@ -1030,15 +870,15 @@ def _do_diffusion_step( # diff_multfac_n2w=self.diff_multfac_n2w, # vert_idx=self.vertical_index, # horz_idx=self.horizontal_cell_index, - # nrdmax=self.vertical_params.index_of_damping_layer, + # nrdmax=int32(self.vertical_params.index_of_damping_layer +1) , # interior_idx=int32( - # cell_start_interior - # ), # h end index for stencil_09 and stencil_10 + # cell_start_interior -1 + # ), # h end index for stencil_09 and stencil_10 # TODO: in ICON: start_interior_idx_c -1 ?? # halo_idx=int32( # cell_end_local # ), # h end index for stencil_09 and stencil_10, # horizontal_start=cell_start_nudging, # h start index for stencil_07 and stencil_08 - # horizontal_end=cell_end_local_minus1, # h end index for stencil_07 and stencil_08 + # horizontal_end=cell_end_halo, # h end index for stencil_07 and stencil_08 # vertical_start=0, # vertical_end=klevels, # offset_provider={ @@ -1060,7 +900,7 @@ def _do_diffusion_step( kh_smag_e=self.kh_smag_e, horizontal_start=edge_start_nudging, horizontal_end=edge_end_halo, - vertical_start=k_start_end_minus2, + vertical_start=(klevels - 2), vertical_end=klevels, offset_provider={ "E2C": self.grid.get_e2c_connectivity(), diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py deleted file mode 100644 index 53a91d209..000000000 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py +++ /dev/null @@ -1,323 +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 - -# flake8: noqa - -from gt4py.next.common import Field -from gt4py.next.ffront.decorator import program -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( - _calculate_nabla2_and_smag_coefficients_for_vn, -) -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import ( - _fused_mo_nh_diffusion_stencil_02_03, -) -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import ( - _fused_mo_nh_diffusion_stencil_04_05_06, -) -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_07_08_09_10 import ( - _fused_mo_nh_diffusion_stencil_07_08_09_10, -) -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_11_12 import ( - _fused_mo_nh_diffusion_stencil_11_12, -) -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_13_14 import ( - _fused_mo_nh_diffusion_stencil_13_14, -) -from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( - _mo_intp_rbf_rbf_vec_interpol_vertex, -) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_15 import ( - _mo_nh_diffusion_stencil_15, -) -from icon4py.atm_dyn_iconam.update_theta_and_exner import _update_theta_and_exner -from icon4py.common.dimension import ( - C2E2CDim, - C2E2CODim, - C2EDim, - CECDim, - CellDim, - ECVDim, - EdgeDim, - KDim, - V2EDim, - VertexDim, -) -from icon4py.diffusion.utils import _scale_k, _set_zero_v_k - - -@program -def diffusion_run( - diagnostic_hdef_ic: Field[[CellDim, KDim], float], - diagnostic_div_ic: Field[[CellDim, KDim], float], - diagnostic_dwdx: Field[[CellDim, KDim], float], - diagnostic_dwdy: Field[[CellDim, KDim], float], - prognostic_w: Field[[CellDim, KDim], float], - prognostic_vn: Field[[EdgeDim, KDim], float], - prognostic_exner_pressure: Field[[CellDim, KDim], float], - prognostic_theta_v: Field[[CellDim, KDim], float], - metric_theta_ref_mc: Field[[CellDim, KDim], float], - metric_wgtfac_c: Field[[CellDim, KDim], float], - metric_mask_hdiff: Field[[CellDim, KDim], bool], - metric_zd_vertoffset: Field[[CECDim, KDim], int32], - metric_zd_diffcoef: Field[[CellDim, KDim], float], - metric_zd_intcoef: Field[[CECDim, KDim], float], - interpolation_e_bln_c_s: Field[[CellDim, C2EDim], float], - interpolation_rbf_coeff_1: Field[[VertexDim, V2EDim], float], - interpolation_rbf_coeff_2: Field[[VertexDim, V2EDim], float], - interpolation_geofac_div: Field[[CellDim, C2EDim], float], - interpolation_geofac_grg_x: Field[[CellDim, C2E2CODim], float], - interpolation_geofac_grg_y: Field[[CellDim, C2E2CODim], float], - interpolation_nudgecoeff_e: Field[[EdgeDim], float], - interpolation_geofac_n2s: Field[[CellDim, C2E2CODim], float], - interpolation_geofac_n2s_c: Field[[CellDim], float], - interpolation_geofac_n2s_nbh: Field[[CECDim], float], - tangent_orientation: Field[[EdgeDim], float], - inverse_primal_edge_lengths: Field[[EdgeDim], float], - inverse_dual_edge_lengths: Field[[EdgeDim], float], - inverse_vert_vert_lengths: Field[[EdgeDim], float], - primal_normal_vert_1: Field[[ECVDim], float], - primal_normal_vert_2: Field[[ECVDim], float], - dual_normal_vert_1: Field[[ECVDim], float], - dual_normal_vert_2: Field[[ECVDim], float], - edge_areas: Field[[EdgeDim], float], - cell_areas: Field[[CellDim], float], - diff_multfac_vn: Field[[KDim], float], - dtime: float, - rd_o_cvd: float, - local_thresh_tdiff: float, - local_smag_limit: Field[[KDim], float], - local_u_vert: Field[[VertexDim, KDim], float], - local_v_vert: Field[[VertexDim, KDim], float], - local_enh_smag_fac: Field[[KDim], float], - local_kh_smag_e: Field[[EdgeDim, KDim], float], - local_kh_smag_ec: Field[[EdgeDim, KDim], float], - local_z_nabla2_e: Field[[EdgeDim, KDim], float], - local_z_temp: Field[[CellDim, KDim], float], - local_diff_multfac_smag: Field[[KDim], float], - local_diff_multfac_n2w: Field[[KDim], float], - local_smag_offset: float, - local_nudgezone_diff: float, - local_fac_bdydiff_v: float, - local_diff_multfac_w: float, - local_vertical_index: Field[[KDim], int32], - local_horizontal_cell_index: Field[[CellDim], int32], - local_horizontal_edge_index: Field[[EdgeDim], int32], - cell_startindex_interior: int32, - cell_halo_idx: int32, - cell_startindex_nudging: int, - cell_endindex_local_plus1: int, - cell_endindex_local: int, - edge_startindex_nudging_plus1: int, - edge_startindex_nudging_minus1: int32, - edge_endindex_local: int, - edge_endindex_local_minus2: int, - vertex_startindex_lb_plus3: int, - vertex_startindex_lb_plus1: int, - vertex_endindex_local: int, - vertex_endindex_local_minus1: int, - index_of_damping_height: int32, - nlev: int, - boundary_diffusion_start_index_edges: int, -): - _scale_k(local_enh_smag_fac, dtime, out=local_diff_multfac_smag) - - # TODO: @magdalena is this needed?, if not remove - _set_zero_v_k(out=local_u_vert) - _set_zero_v_k(out=local_v_vert) - - # # 1. CALL rbf_vec_interpol_vertex - _mo_intp_rbf_rbf_vec_interpol_vertex( - prognostic_vn, - interpolation_rbf_coeff_1, - interpolation_rbf_coeff_2, - out=(local_u_vert, local_v_vert), - domain={ - VertexDim: (vertex_startindex_lb_plus1, vertex_endindex_local_minus1), - KDim: (0, nlev), - }, - ) - - # 2. HALO EXCHANGE -- CALL sync_patch_array_mult - # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 - - _calculate_nabla2_and_smag_coefficients_for_vn( - local_diff_multfac_smag, - tangent_orientation, - inverse_primal_edge_lengths, - inverse_vert_vert_lengths, - local_u_vert, - local_v_vert, - primal_normal_vert_1, - primal_normal_vert_2, - dual_normal_vert_1, - dual_normal_vert_2, - prognostic_vn, - local_smag_limit, - local_smag_offset, - out=(local_kh_smag_e, local_kh_smag_ec, local_z_nabla2_e), - domain={ - # TODO @magdalena wrong start index?? - EdgeDim: (boundary_diffusion_start_index_edges, edge_endindex_local_minus2), - KDim: (0, nlev), - }, - ) - - _fused_mo_nh_diffusion_stencil_02_03( - local_kh_smag_ec, - prognostic_vn, - interpolation_e_bln_c_s, - interpolation_geofac_div, - local_diff_multfac_smag, - metric_wgtfac_c, - out=( - diagnostic_div_ic, - diagnostic_hdef_ic, - ), - domain={ - CellDim: (cell_startindex_nudging, cell_endindex_local), - KDim: (1, nlev), - }, - ) - # - # # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH - # - # # 5. CALL rbf_vec_interpol_vertex_wp - _mo_intp_rbf_rbf_vec_interpol_vertex( - local_z_nabla2_e, - interpolation_rbf_coeff_1, - interpolation_rbf_coeff_2, - out=(local_u_vert, local_v_vert), - domain={ - VertexDim: (vertex_startindex_lb_plus3, vertex_endindex_local), - KDim: (0, nlev), - }, - ) - # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult - # - # # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 - # # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 - # - _fused_mo_nh_diffusion_stencil_04_05_06( - local_u_vert, - local_v_vert, - primal_normal_vert_1, - primal_normal_vert_2, - local_z_nabla2_e, - inverse_vert_vert_lengths, - inverse_primal_edge_lengths, - edge_areas, - local_kh_smag_e, - diff_multfac_vn, - interpolation_nudgecoeff_e, - prognostic_vn, - local_horizontal_edge_index, - local_nudgezone_diff, - local_fac_bdydiff_v, - edge_startindex_nudging_minus1, - out=prognostic_vn, - domain={ - EdgeDim: (edge_startindex_nudging_plus1, edge_endindex_local), - KDim: (0, nlev), - }, - ) - # # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, - # # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 - - _fused_mo_nh_diffusion_stencil_07_08_09_10( - cell_areas, - interpolation_geofac_n2s, - interpolation_geofac_grg_x, - interpolation_geofac_grg_y, - prognostic_w, - prognostic_w, - diagnostic_dwdx, - diagnostic_dwdy, - local_diff_multfac_w, - local_diff_multfac_n2w, - local_vertical_index, - local_horizontal_cell_index, - index_of_damping_height, - cell_startindex_interior, - cell_halo_idx, - out=( - prognostic_w, - diagnostic_dwdx, - diagnostic_dwdy, - ), - domain={ - CellDim: (cell_startindex_nudging, cell_endindex_local_plus1), - KDim: (0, nlev), - }, - ) - # # 8. HALO EXCHANGE: CALL sync_patch_array - # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, - # # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 - # - # # TODO @magdalena check: kh_smag_e is an out field, should not be calculated in init? - # - - _fused_mo_nh_diffusion_stencil_11_12( - prognostic_theta_v, - metric_theta_ref_mc, - local_thresh_tdiff, - local_kh_smag_e, - out=local_kh_smag_e, - domain={ - EdgeDim: (edge_startindex_nudging_plus1, edge_endindex_local), - KDim: (nlev - 2, nlev), - }, - ) - _fused_mo_nh_diffusion_stencil_13_14( - local_kh_smag_e, - inverse_dual_edge_lengths, - prognostic_theta_v, - interpolation_geofac_div, - out=local_z_temp, - domain={ - CellDim: (cell_startindex_nudging, cell_endindex_local), - KDim: (0, nlev), - }, - ) - - # MO_NH_DIFFUSION_STENCIL_15: as_offset index fields! - _mo_nh_diffusion_stencil_15( - metric_mask_hdiff, - metric_zd_vertoffset, - metric_zd_diffcoef, - interpolation_geofac_n2s_c, - interpolation_geofac_n2s_nbh, - metric_zd_intcoef, - prognostic_theta_v, - local_z_temp, - domain={ - CellDim: (cell_startindex_nudging, cell_endindex_local), - KDim: (0, nlev), - }, - out=local_z_temp, - ) - - _update_theta_and_exner( - local_z_temp, - cell_areas, - prognostic_theta_v, - prognostic_exner_pressure, - rd_o_cvd, - out=(prognostic_theta_v, prognostic_exner_pressure), - domain={ - CellDim: (cell_startindex_nudging, cell_endindex_local), - KDim: (0, nlev), - }, - ) - # 10. HALO EXCHANGE sync_patch_array diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 6c0f53d91..1b2b90bf4 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -87,8 +87,8 @@ def diffusion_savepoint_init(data_provider, linit, step_date_init): """ Load data from ICON savepoint at start of diffusion module. - date of the timestamp to be selected can be set seperately by overriding the 'step_data' - fixture, passing 'step_data=' + date of the timestamp to be selected can be set seperately by overriding the 'step_date_init' + fixture, passing 'step_date_init=' linit flag can be set by overriding the 'linit' fixture """ @@ -96,14 +96,14 @@ def diffusion_savepoint_init(data_provider, linit, step_date_init): @pytest.fixture -def diffusion_savepoint_exit(data_provider, step_date_exit): +def diffusion_savepoint_exit(data_provider,linit, step_date_exit): """ Load data from ICON savepoint at exist of diffusion module. date of the timestamp to be selected can be set seperately by overriding the 'step_data' fixture, passing 'step_data=' """ - sp = data_provider.from_savepoint_diffusion_exit(linit=False, date=step_date_exit) + sp = data_provider.from_savepoint_diffusion_exit(linit=linit, date=step_date_exit) return sp diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index e70969b55..3e6da9ac8 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -401,11 +401,10 @@ def test_verify_diffusion_init_against_other_regular_savepoint( _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) -# @pytest.mark.parametrize("run_with_program", [True, False]) -@pytest.mark.parametrize("run_with_program", [False]) + @pytest.mark.datatest +@pytest.mark.parametrize("step_date_init, step_date_exit", [("2021-06-20T12:00:10.000", "2021-06-20T12:00:10.000"), ("2021-06-20T12:00:20.000", "2021-06-20T12:00:20.000"), ("2021-06-20T12:01:00.000", "2021-06-20T12:01:00.000")] ) def test_run_diffusion_single_step( - run_with_program, diffusion_savepoint_init, diffusion_savepoint_exit, interpolation_savepoint, @@ -430,7 +429,7 @@ def test_run_diffusion_single_step( config.ndyn_substeps = datarun_reduced_substeps additional_parameters = DiffusionParams(config) - diffusion = Diffusion(run_program=run_with_program) + diffusion = Diffusion() diffusion.init( grid=icon_grid, config=config, @@ -478,29 +477,24 @@ def test_run_diffusion_single_step( ref_theta_v = np.asarray(diffusion_savepoint_exit.theta_v()) val_theta_v = np.asarray(prognostic_state.theta_v) val_exner = np.asarray(prognostic_state.exner_pressure) - val_z_temp = np.asarray(diffusion.z_temp) - ref_z_temp = np.asarray(diffusion_savepoint_exit.z_temp()) steep_points = np.asarray(diffusion.metric_state.mask_hdiff) - assert np.allclose(ref_z_temp[~steep_points], val_z_temp[~steep_points], atol=1e-7) assert np.allclose(ref_theta_v[~steep_points], val_theta_v[~steep_points]) assert np.allclose(ref_exner[~steep_points], val_exner[~steep_points]) - assert np.allclose(ref_theta_v[steep_points], val_theta_v[steep_points]) - assert np.allclose(ref_exner[steep_points], val_exner[steep_points]) + #assert np.allclose(ref_theta_v[steep_points], val_theta_v[steep_points]) + #assert np.allclose(ref_exner[steep_points], val_exner[steep_points]) -@pytest.mark.skip("fix: diffusion_stencil_15") @pytest.mark.datatest -def test_diffusion_five_steps( - damping_height, - r04b09_diffusion_config, - icon_grid, - grid_savepoint, - interpolation_savepoint, - metrics_savepoint, +@pytest.mark.parametrize("linit", [True]) +def test_run_diffusion_initial_step( diffusion_savepoint_init, diffusion_savepoint_exit, - linit=True, - step_date_exit="2021-06-20T12:01:00.000", + interpolation_savepoint, + metrics_savepoint, + grid_savepoint, + icon_grid, + r04b09_diffusion_config, + damping_height, ): dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() @@ -509,14 +503,14 @@ def test_diffusion_five_steps( metric_state = metrics_savepoint.construct_metric_state() diagnostic_state = diffusion_savepoint_init.construct_diagnostics() prognostic_state = diffusion_savepoint_init.construct_prognostics() - + vct_a = grid_savepoint.vct_a() vertical_params = VerticalModelParams( - vct_a=grid_savepoint.vct_a(), rayleigh_damping_height=damping_height + vct_a=vct_a, rayleigh_damping_height=damping_height ) - - additional_parameters = DiffusionParams(r04b09_diffusion_config) config = r04b09_diffusion_config config.ndyn_substeps = datarun_reduced_substeps + additional_parameters = DiffusionParams(config) + diffusion = Diffusion() diffusion.init( grid=icon_grid, @@ -526,7 +520,7 @@ def test_diffusion_five_steps( metric_state=metric_state, interpolation_state=interpolation_state, ) - diffusion.initial_run( + diffusion.initial_step( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime, @@ -539,30 +533,35 @@ def test_diffusion_five_steps( edge_areas=edge_geometry.edge_areas, cell_areas=cell_geometry.area, ) - for _ in range(4): - diffusion.run( - diagnostic_state=diagnostic_state, - prognostic_state=prognostic_state, - dtime=dtime, - tangent_orientation=edge_geometry.tangent_orientation, - inverse_primal_edge_lengths=edge_geometry.inverse_primal_edge_lengths, - inverse_dual_edge_length=edge_geometry.inverse_dual_edge_lengths, - inverse_vert_vert_lengths=edge_geometry.inverse_vertex_vertex_lengths, - primal_normal_vert=edge_geometry.primal_normal_vert, - dual_normal_vert=edge_geometry.dual_normal_vert, - edge_areas=edge_geometry.edge_areas, - cell_areas=cell_geometry.area, - ) - - icon_result_exner = diffusion_savepoint_exit.exner() - icon_result_vn = diffusion_savepoint_exit.vn() - icon_result_w = diffusion_savepoint_exit.w() - icon_result_theta_w = diffusion_savepoint_exit.theta_v() - assert np.allclose(icon_result_w, np.asarray(prognostic_state.w)) - assert np.allclose(np.asarray(icon_result_vn), np.asarray(prognostic_state.vn)) - assert np.allclose( - np.asarray(icon_result_theta_w), np.asarray(prognostic_state.theta_v) - ) - assert np.allclose( - np.asarray(icon_result_exner), np.asarray(prognostic_state.exner_pressure) - ) + assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v + ref_div_ic = np.asarray(diffusion_savepoint_exit.div_ic()) + val_div_ic = np.asarray(diagnostic_state.div_ic) + ref_hdef_ic = np.asarray(diffusion_savepoint_exit.hdef_ic()) + val_hdef_ic = np.asarray(diagnostic_state.hdef_ic) + + assert np.allclose(ref_div_ic, val_div_ic) + assert np.allclose(ref_hdef_ic, val_hdef_ic) + + ref_w = np.asarray(diffusion_savepoint_exit.w()) + val_w = np.asarray(prognostic_state.w) + ref_dwdx = np.asarray(diffusion_savepoint_exit.dwdx()) + val_dwdx = np.asarray(diagnostic_state.dwdx) + ref_dwdy = np.asarray(diffusion_savepoint_exit.dwdy()) + val_dwdy = np.asarray(diagnostic_state.dwdy) + ref_vn = np.asarray(diffusion_savepoint_exit.vn()) + val_vn = np.asarray(prognostic_state.vn) + + assert np.allclose(ref_vn, val_vn) + assert np.allclose(ref_dwdx, val_dwdx) + assert np.allclose(ref_dwdy, val_dwdy) + assert np.allclose(ref_w, val_w) + ref_exner = np.asarray(diffusion_savepoint_exit.exner()) + ref_theta_v = np.asarray(diffusion_savepoint_exit.theta_v()) + val_theta_v = np.asarray(prognostic_state.theta_v) + val_exner = np.asarray(prognostic_state.exner_pressure) + steep_points = np.asarray(diffusion.metric_state.mask_hdiff) + assert np.allclose(ref_theta_v[~steep_points], val_theta_v[~steep_points]) + assert np.allclose(ref_exner[~steep_points], val_exner[~steep_points]) + assert np.allclose(ref_theta_v[steep_points], val_theta_v[steep_points]) + assert np.allclose(ref_exner[steep_points], val_exner[steep_points]) + From fb0eca7876f6dd8efeb49fd1fe90372e4a0d59b0 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 9 Jun 2023 09:05:54 +0200 Subject: [PATCH 129/263] add numpy test to check transformation of geofac_n2s_nbh (1d sparse field) --- .../src/icon4py/diffusion/interpolation_state.py | 6 +++--- atm_dyn_iconam/tests/test_diffusion.py | 4 ++-- atm_dyn_iconam/tests/test_interpolation_state.py | 14 ++++++++++++++ 3 files changed, 19 insertions(+), 5 deletions(-) create mode 100644 atm_dyn_iconam/tests/test_interpolation_state.py diff --git a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py index c79308646..cad91a1d4 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py @@ -65,8 +65,8 @@ def geofac_n2s_c(self) -> Field[[CellDim], float]: @property def geofac_n2s_nbh(self) -> Field[[CECDim], float]: - np_array = np.asarray(self.geofac_n2s)[:, 1:] - old_shape = np_array.shape + geofac_nbh_ar = np.asarray(self.geofac_n2s)[:, 1:] + old_shape = geofac_nbh_ar.shape return np_as_located_field(CECDim)( - np_array.reshape(old_shape[0] * old_shape[1]) + geofac_nbh_ar.reshape(old_shape[0] * old_shape[1],) ) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 3e6da9ac8..d8529b642 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -562,6 +562,6 @@ def test_run_diffusion_initial_step( steep_points = np.asarray(diffusion.metric_state.mask_hdiff) assert np.allclose(ref_theta_v[~steep_points], val_theta_v[~steep_points]) assert np.allclose(ref_exner[~steep_points], val_exner[~steep_points]) - assert np.allclose(ref_theta_v[steep_points], val_theta_v[steep_points]) - assert np.allclose(ref_exner[steep_points], val_exner[steep_points]) + #assert np.allclose(ref_theta_v[steep_points], val_theta_v[steep_points]) + #assert np.allclose(ref_exner[steep_points], val_exner[steep_points]) diff --git a/atm_dyn_iconam/tests/test_interpolation_state.py b/atm_dyn_iconam/tests/test_interpolation_state.py new file mode 100644 index 000000000..dde88e1ad --- /dev/null +++ b/atm_dyn_iconam/tests/test_interpolation_state.py @@ -0,0 +1,14 @@ +import numpy as np +import pytest + + +@pytest.mark.datatest +def test_cecdim(interpolation_savepoint, icon_grid): + interpolation_fields = interpolation_savepoint.construct_interpolation_state() + geofac_n2s = np.asarray(interpolation_fields.geofac_n2s) + geofac_n2s_nbh = np.asarray(interpolation_fields.geofac_n2s_nbh) + assert np.count_nonzero(geofac_n2s_nbh) > 0 + c2cec= icon_grid.get_c2cec_connectivity().table + ported = geofac_n2s_nbh[c2cec] + assert ported.shape == geofac_n2s[:, 1:].shape + assert np.allclose(ported, geofac_n2s[:, 1:]) From 3909ef3ee4c2f5e544a902bb13eefc9d817f8958 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 9 Jun 2023 16:09:13 +0200 Subject: [PATCH 130/263] switch to renamed fused stencils and remove unused files --- .../fused_mo_nh_diffusion_stencil_02_03.py | 69 ---------- ...sed_mo_nh_diffusion_stencil_07_08_09_10.py | 126 ------------------ .../fused_mo_nh_diffusion_stencil_11_12.py | 63 --------- .../fused_mo_nh_diffusion_stencil_13_14.py | 58 -------- .../src/icon4py/diffusion/diffusion.py | 12 +- 5 files changed, 7 insertions(+), 321 deletions(-) delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py deleted file mode 100644 index 3bbca2b79..000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ /dev/null @@ -1,69 +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 - -from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field - -from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulence import ( - _calculate_diagnostics_for_turbulence, -) -from icon4py.atm_dyn_iconam.temporary_fields_for_turbulence_diagnostics import ( - _temporary_fields_for_turbulence_diagnostics, -) -from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim - - -@field_operator -def _fused_mo_nh_diffusion_stencil_02_03( - kh_smag_ec: Field[[EdgeDim, KDim], float], - vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], - diff_multfac_smag: Field[[KDim], float], - wgtfac_c: Field[[CellDim, KDim], float], -) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - kh_c, div = _temporary_fields_for_turbulence_diagnostics( - kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag - ) - div_ic, hdef_ic = _calculate_diagnostics_for_turbulence(div, kh_c, wgtfac_c) - return div_ic, hdef_ic - - -@program -def fused_mo_nh_diffusion_stencil_02_03( - kh_smag_ec: Field[[EdgeDim, KDim], float], - vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], - diff_multfac_smag: Field[[KDim], float], - wgtfac_c: Field[[CellDim, KDim], float], - div_ic: Field[[CellDim, KDim], float], - hdef_ic: Field[[CellDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, -): - _fused_mo_nh_diffusion_stencil_02_03( - kh_smag_ec, - vn, - e_bln_c_s, - geofac_div, - diff_multfac_smag, - wgtfac_c, - out=(div_ic, hdef_ic), - domain={ - CellDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, - ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py deleted file mode 100644 index b83b00dfd..000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ /dev/null @@ -1,126 +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 - -from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where - -from icon4py.atm_dyn_iconam.apply_nabla2_to_w import _apply_nabla2_to_w -from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import ( - _apply_nabla2_to_w_in_upper_damping_layer, -) -from icon4py.atm_dyn_iconam.calculate_horizontal_gradients_for_turbulence import ( - _calculate_horizontal_gradients_for_turbulence, -) -from icon4py.atm_dyn_iconam.calculate_nabla2_for_w import _calculate_nabla2_for_w -from icon4py.common.dimension import C2E2CODim, CellDim, KDim - - -@field_operator -def _fused_mo_nh_diffusion_stencil_07_08_09_10( - area: Field[[CellDim], float], - geofac_n2s: Field[[CellDim, C2E2CODim], float], - geofac_grg_x: Field[[CellDim, C2E2CODim], float], - geofac_grg_y: Field[[CellDim, C2E2CODim], float], - w_old: Field[[CellDim, KDim], float], - w: Field[[CellDim, KDim], float], - dwdx: Field[[CellDim, KDim], float], - dwdy: Field[[CellDim, KDim], float], - diff_multfac_w: float, - diff_multfac_n2w: Field[[KDim], float], - vert_idx: Field[[KDim], int32], - horz_idx: Field[[CellDim], int32], - nrdmax: int32, - interior_idx: int32, - halo_idx: int32, -) -> tuple[ - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], -]: - - vert_idx = broadcast(vert_idx, (CellDim, KDim)) - - dwdx, dwdy = where( - vert_idx > int32(0), - _calculate_horizontal_gradients_for_turbulence( - w_old, geofac_grg_x, geofac_grg_y - ), - (dwdx, dwdy), - ) - - z_nabla2_c = _calculate_nabla2_for_w(w_old, geofac_n2s) - - w = where( - (horz_idx >= interior_idx) & (horz_idx < halo_idx), - _apply_nabla2_to_w(area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w), - w_old, - ) - - w = where( - (vert_idx > int32(0)) - & (vert_idx < nrdmax) - & (horz_idx >= interior_idx) - & (horz_idx < halo_idx), - _apply_nabla2_to_w_in_upper_damping_layer( - w, diff_multfac_n2w, area, z_nabla2_c - ), - w, - ) - - return w, dwdx, dwdy - - -@program -def fused_mo_nh_diffusion_stencil_07_08_09_10( - area: Field[[CellDim], float], - geofac_n2s: Field[[CellDim, C2E2CODim], float], - geofac_grg_x: Field[[CellDim, C2E2CODim], float], - geofac_grg_y: Field[[CellDim, C2E2CODim], float], - w_old: Field[[CellDim, KDim], float], - w: Field[[CellDim, KDim], float], - dwdx: Field[[CellDim, KDim], float], - dwdy: Field[[CellDim, KDim], float], - diff_multfac_w: float, - diff_multfac_n2w: Field[[KDim], float], - vert_idx: Field[[KDim], int32], - horz_idx: Field[[CellDim], int32], - nrdmax: int32, - interior_idx: int32, - halo_idx: int32, - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, -): - _fused_mo_nh_diffusion_stencil_07_08_09_10( - area, - geofac_n2s, - geofac_grg_x, - geofac_grg_y, - w_old, - w, - dwdx, - dwdy, - diff_multfac_w, - diff_multfac_n2w, - vert_idx, - horz_idx, - nrdmax, - interior_idx, - halo_idx, - out=(w, dwdx, dwdy), - domain={ - CellDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, - ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py deleted file mode 100644 index a9fb3dbd2..000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py +++ /dev/null @@ -1,63 +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 - -from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field - -from icon4py.atm_dyn_iconam.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( - _enhance_diffusion_coefficient_for_grid_point_cold_pools, -) -from icon4py.atm_dyn_iconam.temporary_field_for_grid_point_cold_pools_enhancement import ( - _temporary_field_for_grid_point_cold_pools_enhancement, -) -from icon4py.common.dimension import CellDim, EdgeDim, KDim - - -@field_operator -def _fused_mo_nh_diffusion_stencil_11_12( - theta_v: Field[[CellDim, KDim], float], - theta_ref_mc: Field[[CellDim, KDim], float], - thresh_tdiff: float, - kh_smag_e: Field[[EdgeDim, KDim], float], -) -> Field[[EdgeDim, KDim], float]: - enh_diffu_3d = _temporary_field_for_grid_point_cold_pools_enhancement( - theta_v, theta_ref_mc, thresh_tdiff - ) - kh_smag_e = _enhance_diffusion_coefficient_for_grid_point_cold_pools( - kh_smag_e, enh_diffu_3d - ) - return kh_smag_e - - -@program -def fused_mo_nh_diffusion_stencil_11_12( - theta_v: Field[[CellDim, KDim], float], - theta_ref_mc: Field[[CellDim, KDim], float], - thresh_tdiff: float, - kh_smag_e: Field[[EdgeDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, -): - _fused_mo_nh_diffusion_stencil_11_12( - theta_v, - theta_ref_mc, - thresh_tdiff, - kh_smag_e, - out=kh_smag_e, - domain={ - EdgeDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, - ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py deleted file mode 100644 index d13d7c8e1..000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ /dev/null @@ -1,58 +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 - -from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field - -from icon4py.atm_dyn_iconam.calculate_nabla2_for_z import _calculate_nabla2_for_z -from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( - _calculate_nabla2_of_theta, -) -from icon4py.common.dimension import CEDim, CellDim, EdgeDim, KDim - - -@field_operator -def _fused_mo_nh_diffusion_stencil_13_14( - kh_smag_e: Field[[EdgeDim, KDim], float], - inv_dual_edge_length: Field[[EdgeDim], float], - theta_v: Field[[CellDim, KDim], float], - geofac_div: Field[[CEDim], float], -) -> Field[[CellDim, KDim], float]: - z_nabla2_e = _calculate_nabla2_for_z(kh_smag_e, inv_dual_edge_length, theta_v) - z_temp = _calculate_nabla2_of_theta(z_nabla2_e, geofac_div) - return z_temp - - -@program -def fused_mo_nh_diffusion_stencil_13_14( - kh_smag_e: Field[[EdgeDim, KDim], float], - inv_dual_edge_length: Field[[EdgeDim], float], - theta_v: Field[[CellDim, KDim], float], - geofac_div: Field[[CEDim], float], - z_temp: Field[[CellDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, -): - _fused_mo_nh_diffusion_stencil_13_14( - kh_smag_e, - inv_dual_edge_length, - theta_v, - geofac_div, - out=z_temp, - domain={ - CellDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, - ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index ec5a30a4f..d81695070 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -25,10 +25,14 @@ run_gtfn_cached, ) +from icon4py.atm_dyn_iconam.calculate_diagnostic_quantities_for_turbulence import calculate_diagnostic_quantities_for_turbulence from icon4py.atm_dyn_iconam.apply_nabla2_to_w import apply_nabla2_to_w from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import ( apply_nabla2_to_w_in_upper_damping_layer, ) +from icon4py.atm_dyn_iconam.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( + apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance, +) from icon4py.atm_dyn_iconam.calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools import ( calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools, ) @@ -42,9 +46,7 @@ calculate_nabla2_for_theta, ) from icon4py.atm_dyn_iconam.calculate_nabla2_for_w import calculate_nabla2_for_w -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_02_03 import ( - fused_mo_nh_diffusion_stencil_02_03, -) + from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import ( fused_mo_nh_diffusion_stencil_04_05_06, ) @@ -716,7 +718,7 @@ def _do_diffusion_step( ) log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: end") log.debug("running fused stencil fused stencil 02_03: start") - fused_mo_nh_diffusion_stencil_02_03.with_backend(backend)( + calculate_diagnostic_quantities_for_turbulence.with_backend(backend)( kh_smag_ec=self.kh_smag_ec, vn=prognostic_state.vn, e_bln_c_s=self.interpolation_state.e_bln_c_s, @@ -842,7 +844,7 @@ def _do_diffusion_step( offset_provider={}, ) # w_old = prognostic_state.w - # fused_mo_nh_diffusion_stencil_07_08_09_10.with_backend(backend)( + # apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.with_backend(backend)( # area=cell_areas, # geofac_n2s=self.interpolation_state.geofac_n2s, # geofac_grg_x=self.interpolation_state.geofac_grg_x, From ef67246d3bffa359ba47d8f751a7879c62404d21 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 9 Jun 2023 16:46:39 +0200 Subject: [PATCH 131/263] unify usage of geofac_div within diffusion.py: switch to 1d sparse field where ever used. --- ...late_diagnostic_quantities_for_turbulence.py | 10 +++++----- ...mporary_fields_for_turbulence_diagnostics.py | 14 +++++++------- .../src/icon4py/diffusion/diffusion.py | 3 ++- .../icon4py/diffusion/interpolation_state.py | 13 +++---------- atm_dyn_iconam/tests/test_io_utils.py | 2 +- ...mporary_fields_for_turbulence_diagnostics.py | 17 +++++++++-------- .../tests/test_utils/serialbox_utils.py | 9 ++++----- 7 files changed, 31 insertions(+), 37 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py index 12d42bf80..84530b8d4 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py @@ -20,15 +20,15 @@ from icon4py.atm_dyn_iconam.temporary_fields_for_turbulence_diagnostics import ( _temporary_fields_for_turbulence_diagnostics, ) -from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim +from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim, CEDim @field_operator def _calculate_diagnostic_quantities_for_turbulence( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], + e_bln_c_s: Field[[CEDim], float], + geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], wgtfac_c: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: @@ -43,8 +43,8 @@ def _calculate_diagnostic_quantities_for_turbulence( def calculate_diagnostic_quantities_for_turbulence( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], + e_bln_c_s: Field[[CEDim], float], + geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], wgtfac_c: Field[[CellDim, KDim], float], div_ic: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_fields_for_turbulence_diagnostics.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_fields_for_turbulence_diagnostics.py index c0b536319..708671269 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_fields_for_turbulence_diagnostics.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_fields_for_turbulence_diagnostics.py @@ -14,19 +14,19 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim +from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim, CEDim, C2CE @field_operator def _temporary_fields_for_turbulence_diagnostics( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], + e_bln_c_s: Field[[CEDim], float], + geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - kh_c = neighbor_sum(kh_smag_ec(C2E) * e_bln_c_s, axis=C2EDim) / diff_multfac_smag - div = neighbor_sum(vn(C2E) * geofac_div, axis=C2EDim) + kh_c = neighbor_sum(kh_smag_ec(C2E) * e_bln_c_s(C2CE), axis=C2EDim) / diff_multfac_smag + div = neighbor_sum(vn(C2E) * geofac_div(C2CE), axis=C2EDim) return kh_c, div @@ -34,8 +34,8 @@ def _temporary_fields_for_turbulence_diagnostics( def temporary_fields_for_turbulence_diagnostics( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], + e_bln_c_s: Field[[CEDim], float], + geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], kh_c: Field[[CellDim, KDim], float], div: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index d81695070..4611c1079 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -722,7 +722,7 @@ def _do_diffusion_step( kh_smag_ec=self.kh_smag_ec, vn=prognostic_state.vn, e_bln_c_s=self.interpolation_state.e_bln_c_s, - geofac_div=self.interpolation_state._geofac_div, + geofac_div=self.interpolation_state.geofac_div, diff_multfac_smag=self.diff_multfac_smag, wgtfac_c=self.metric_state.wgtfac_c, div_ic=diagnostic_state.div_ic, @@ -733,6 +733,7 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={ "C2E": self.grid.get_c2e_connectivity(), + "C2CE": self.grid.get_c2ce_connectivity(), "Koff": KDim, }, ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py index 3eb818f7e..22b8a2198 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py @@ -38,7 +38,7 @@ class InterpolationState: """ e_bln_c_s: Field[ - [CellDim, C2EDim], float + [CEDim], float ] # coefficent for bilinear interpolation from edge to cell () rbf_coeff_1: Field[ [VertexDim, V2EDim], float @@ -47,8 +47,8 @@ class InterpolationState: [VertexDim, V2EDim], float ] # rbf_vec_coeff_v_2(nproma, rbf_vec_dim_v, nblks_v) - _geofac_div: Field[ - [CellDim, C2EDim], float + geofac_div: Field[ + [CEDim], float ] # factor for divergence (nproma,cell_type,nblks_c) geofac_n2s: Field[ @@ -60,13 +60,6 @@ class InterpolationState: ] # factors for green gauss gradient (nproma,4,nblks_c,2) nudgecoeff_e: Field[[EdgeDim], float] # Nudgeing coeffients for edges - @property - def geofac_div(self): - old_shape = np.asarray(self._geofac_div).shape - return np_as_located_field(CEDim)( - np.asarray(self._geofac_div).reshape((old_shape[0] * old_shape[1],)) - ) - @property def geofac_n2s_c(self) -> Field[[CellDim], float]: return np_as_located_field(CellDim)(np.asarray(self.geofac_n2s)[:, 0]) diff --git a/atm_dyn_iconam/tests/test_io_utils.py b/atm_dyn_iconam/tests/test_io_utils.py index 9cc61a7b2..c31ce40d8 100644 --- a/atm_dyn_iconam/tests/test_io_utils.py +++ b/atm_dyn_iconam/tests/test_io_utils.py @@ -109,7 +109,7 @@ def assert_interpolation_state_fields(interpolation_state): assert interpolation_state.e_bln_c_s assert interpolation_state.nudgecoeff_e assert interpolation_state.geofac_n2s_nbh - assert interpolation_state._geofac_div + assert interpolation_state.geofac_div assert interpolation_state.geofac_grg_y assert interpolation_state.geofac_grg_x assert interpolation_state.rbf_coeff_2 diff --git a/atm_dyn_iconam/tests/test_temporary_fields_for_turbulence_diagnostics.py b/atm_dyn_iconam/tests/test_temporary_fields_for_turbulence_diagnostics.py index 58bf65b17..59c97209e 100644 --- a/atm_dyn_iconam/tests/test_temporary_fields_for_turbulence_diagnostics.py +++ b/atm_dyn_iconam/tests/test_temporary_fields_for_turbulence_diagnostics.py @@ -12,17 +12,18 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +from gt4py.next import StridedNeighborOffsetProvider from icon4py.atm_dyn_iconam.temporary_fields_for_turbulence_diagnostics import ( temporary_fields_for_turbulence_diagnostics, ) -from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim +from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim, CEDim -from .test_utils.helpers import random_field, zero_field +from .test_utils.helpers import random_field, zero_field, as_1D_sparse_field from .test_utils.simple_mesh import SimpleMesh -def mo_nh_diffusion_stencil_02_numpy( +def temporary_fields_for_turbulence_diagnostics_numpy( c2e: np.array, kh_smag_ec: np.array, vn: np.array, @@ -43,7 +44,7 @@ def mo_nh_diffusion_stencil_02_numpy( return div, kh_c -def test_mo_nh_diffusion_stencil_02(): +def test_temporary_fields_for_turbulence_diagnostics_numpy(): mesh = SimpleMesh() vn = random_field(mesh, EdgeDim, KDim) @@ -55,7 +56,7 @@ def test_mo_nh_diffusion_stencil_02(): kh_c = zero_field(mesh, CellDim, KDim) div = zero_field(mesh, CellDim, KDim) - div_ref, kh_c_ref = mo_nh_diffusion_stencil_02_numpy( + div_ref, kh_c_ref = temporary_fields_for_turbulence_diagnostics_numpy( mesh.c2e, np.asarray(kh_smag_ec), np.asarray(vn), @@ -67,14 +68,14 @@ def test_mo_nh_diffusion_stencil_02(): temporary_fields_for_turbulence_diagnostics( kh_smag_ec, vn, - e_bln_c_s, - geofac_div, + as_1D_sparse_field(e_bln_c_s, CEDim), + as_1D_sparse_field(geofac_div, CEDim), diff_multfac_smag, kh_c, div, offset_provider={ "C2E": mesh.get_c2e_offset_provider(), - "C2EDim": C2EDim, + "C2CE": StridedNeighborOffsetProvider(CellDim, CEDim, mesh.n_c2e), }, ) assert np.allclose(kh_c, kh_c_ref) diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index 00469930c..a7e898e44 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -18,7 +18,7 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import np_as_located_field -from atm_dyn_iconam.tests.test_utils.helpers import as_1D_sparse_field +from .helpers import as_1D_sparse_field from icon4py.common.dimension import ( C2E2CDim, C2E2CODim, @@ -32,7 +32,7 @@ EdgeDim, KDim, V2EDim, - VertexDim, + VertexDim, CEDim, ) from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.diffusion import VectorTuple @@ -89,7 +89,6 @@ def read_int(self, name: str): self.log.debug(f"{name} {buffer.shape}") return buffer - class IconGridSavePoint(IconSavepoint): def vct_a(self): return self._get_field("vct_a", KDim) @@ -283,10 +282,10 @@ def nudgecoeff_e(self): def construct_interpolation_state(self) -> InterpolationState: grg = self.geofac_grg() return InterpolationState( - e_bln_c_s=self.e_bln_c_s(), + e_bln_c_s=as_1D_sparse_field(self.e_bln_c_s(), CEDim), rbf_coeff_1=self.rbf_vec_coeff_v1(), rbf_coeff_2=self.rbf_vec_coeff_v2(), - _geofac_div=self.geofac_div(), + geofac_div=as_1D_sparse_field(self.geofac_div(), CEDim), geofac_n2s=self.geofac_n2s(), geofac_grg_x=grg[0], geofac_grg_y=grg[1], From 5366e9460843c68c2f2fa55f9a64e333048bdfba Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 13 Jun 2023 10:58:09 +0200 Subject: [PATCH 132/263] pre-commit fixes --- ...lculate_diagnostic_quantities_for_turbulence.py | 2 +- .../temporary_fields_for_turbulence_diagnostics.py | 14 ++++++++++++-- atm_dyn_iconam/src/icon4py/diffusion/diffusion.py | 14 ++++++++------ .../src/icon4py/diffusion/interpolation_state.py | 1 - atm_dyn_iconam/src/icon4py/driver/io_utils.py | 1 - ..._temporary_fields_for_turbulence_diagnostics.py | 4 ++-- atm_dyn_iconam/tests/test_utils/serialbox_utils.py | 6 ++++-- 7 files changed, 27 insertions(+), 15 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py index 84530b8d4..e2974fffa 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py @@ -20,7 +20,7 @@ from icon4py.atm_dyn_iconam.temporary_fields_for_turbulence_diagnostics import ( _temporary_fields_for_turbulence_diagnostics, ) -from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim, CEDim +from icon4py.common.dimension import CEDim, CellDim, EdgeDim, KDim @field_operator diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_fields_for_turbulence_diagnostics.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_fields_for_turbulence_diagnostics.py index 708671269..0ca6c9352 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_fields_for_turbulence_diagnostics.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_fields_for_turbulence_diagnostics.py @@ -14,7 +14,15 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim, CEDim, C2CE +from icon4py.common.dimension import ( + C2CE, + C2E, + C2EDim, + CEDim, + CellDim, + EdgeDim, + KDim, +) @field_operator @@ -25,7 +33,9 @@ def _temporary_fields_for_turbulence_diagnostics( geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - kh_c = neighbor_sum(kh_smag_ec(C2E) * e_bln_c_s(C2CE), axis=C2EDim) / diff_multfac_smag + kh_c = ( + neighbor_sum(kh_smag_ec(C2E) * e_bln_c_s(C2CE), axis=C2EDim) / diff_multfac_smag + ) div = neighbor_sum(vn(C2E) * geofac_div(C2CE), axis=C2EDim) return kh_c, div diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 1cc85fb5c..1cb50e69d 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -25,14 +25,15 @@ run_gtfn_cached, ) -from icon4py.atm_dyn_iconam.calculate_nabla2_for_theta import calculate_nabla2_for_theta -from icon4py.atm_dyn_iconam.calculate_diagnostic_quantities_for_turbulence import calculate_diagnostic_quantities_for_turbulence +from icon4py.atm_dyn_iconam.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( + apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance, +) from icon4py.atm_dyn_iconam.apply_nabla2_to_w import apply_nabla2_to_w from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import ( apply_nabla2_to_w_in_upper_damping_layer, ) -from icon4py.atm_dyn_iconam.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( - apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance, +from icon4py.atm_dyn_iconam.calculate_diagnostic_quantities_for_turbulence import ( + calculate_diagnostic_quantities_for_turbulence, ) from icon4py.atm_dyn_iconam.calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools import ( calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools, @@ -43,9 +44,10 @@ from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( calculate_nabla2_and_smag_coefficients_for_vn, ) - +from icon4py.atm_dyn_iconam.calculate_nabla2_for_theta import ( + calculate_nabla2_for_theta, +) from icon4py.atm_dyn_iconam.calculate_nabla2_for_w import calculate_nabla2_for_w - from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import ( fused_mo_nh_diffusion_stencil_04_05_06, ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py index 22b8a2198..98586b90c 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py @@ -19,7 +19,6 @@ from icon4py.common.dimension import ( C2E2CODim, - C2EDim, CECDim, CEDim, CellDim, diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index 78b414038..96a796fb2 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -27,7 +27,6 @@ from icon4py.driver.parallel_setup import DecompositionInfo - SERIALBOX_ONLY_MSG = "Only ser_type='sb (Serialbox)' is implemented so far." SIMULATION_START_DATE = "2021-06-20T12:00:10.000" diff --git a/atm_dyn_iconam/tests/test_temporary_fields_for_turbulence_diagnostics.py b/atm_dyn_iconam/tests/test_temporary_fields_for_turbulence_diagnostics.py index 59c97209e..94558279a 100644 --- a/atm_dyn_iconam/tests/test_temporary_fields_for_turbulence_diagnostics.py +++ b/atm_dyn_iconam/tests/test_temporary_fields_for_turbulence_diagnostics.py @@ -17,9 +17,9 @@ from icon4py.atm_dyn_iconam.temporary_fields_for_turbulence_diagnostics import ( temporary_fields_for_turbulence_diagnostics, ) -from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim, CEDim +from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from .test_utils.helpers import random_field, zero_field, as_1D_sparse_field +from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field from .test_utils.simple_mesh import SimpleMesh diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index 4ed88f3c9..705cb56a2 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -18,13 +18,13 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import np_as_located_field -from .helpers import as_1D_sparse_field from icon4py.common import dimension from icon4py.common.dimension import ( C2E2CDim, C2E2CODim, C2EDim, CECDim, + CEDim, CellDim, E2C2VDim, E2CDim, @@ -33,7 +33,7 @@ EdgeDim, KDim, V2EDim, - VertexDim, CEDim, + VertexDim, ) from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.diffusion import VectorTuple @@ -48,6 +48,7 @@ from icon4py.diffusion.prognostic_state import PrognosticState from icon4py.driver.parallel_setup import DecompositionInfo +from .helpers import as_1D_sparse_field class IconSavepoint: @@ -95,6 +96,7 @@ def read_int(self, name: str): self.log.debug(f"{name} {buffer.shape}") return buffer + class IconGridSavePoint(IconSavepoint): def vct_a(self): return self._get_field("vct_a", KDim) From d4393db543f9573bcdfaa190182b7ec43fff6cb7 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 13 Jun 2023 11:06:04 +0200 Subject: [PATCH 133/263] fix (local) data path --- atm_dyn_iconam/src/icon4py/driver/dycore_driver.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index 6ac9d0670..bebf76b3a 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -36,7 +36,7 @@ read_initial_state, read_static_fields, ) - +from icon4py.driver.parallel_setup import get_processor_properties log = logging.getLogger(__name__) @@ -188,15 +188,18 @@ def initialize(n_time_steps, file_path: Path): diagnostic_state: """ log.info("initialize parallel runtime") + experiment_name = "mch_ch_r04b09_dsl" + log.info(f"reading configuration: experiment {experiment_name}") + config = read_config(experiment_name, n_time_steps=n_time_steps) + parallel_props = get_processor_properties() decomp_info = read_decomp_info( - "/home/magdalena/data/exclaim/dycore/mch_ch_r04b09_dsl/node2/mch_ch_r04b09_dsl/icon_grid", + f"/home/magdalena/data/exclaim/dycore/{experiment_name}/mpitasks2/{experiment_name}/ser_data", parallel_props, ) - experiment_name = "mch_ch_r04b09_dsl" - log.info(f"reading configuration: experiment {experiment_name}") - config = read_config(experiment_name, n_time_steps=n_time_steps) + + log.info("initializing the grid") icon_grid = read_icon_grid(file_path) log.info("reading input fields") From 7a0b0f3248730829169119eee889d53542449d1c Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 14 Jun 2023 17:53:51 +0200 Subject: [PATCH 134/263] setup exchange properties add sync calls in diffusion (WIP) --- .../src/icon4py/decomposition/decomposed.py | 15 +- .../src/icon4py/diffusion/diffusion.py | 331 +++++++++--------- .../src/icon4py/driver/dycore_driver.py | 5 +- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 16 +- .../src/icon4py/driver/parallel_setup.py | 88 ++++- atm_dyn_iconam/tests/conftest.py | 6 +- .../tests/mpi_tests/test_parallel_setup.py | 149 +++++++- .../tests/test_utils/serialbox_utils.py | 2 +- 8 files changed, 424 insertions(+), 188 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index 29062eaeb..95fa174b6 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -15,10 +15,11 @@ class ProcessProperties: - def __init__(self, name="", size=0, rank=0): - self._communicator_name: str = name - self._rank: int = rank - self._comm_size = size + def __init__(self, comm: mpi4py.MPI.Comm): + self._communicator_name: str = comm.Get_name() + self._rank: int = comm.Get_rank() + self._comm_size = comm.Get_size() + self._comm = comm @property def rank(self): @@ -32,6 +33,10 @@ def comm_name(self): def comm_size(self): return self._comm_size + @property + def comm(self): + return self._comm + @classmethod def from_mpi_comm(cls, comm: mpi4py.MPI.Comm): - return ProcessProperties(comm.Get_name(), comm.Get_rank()) + return ProcessProperties(comm) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 1cb50e69d..1cc5f1fe5 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -78,6 +78,7 @@ setup_fields_for_initial_step, zero_field, ) +from icon4py.driver.parallel_setup import Exchange # flake8: noqa @@ -369,8 +370,8 @@ def _diffusion_type_5_smagorinski_factor(config: DiffusionConfig): class Diffusion: """Class that configures diffusion and does one diffusion step.""" - def __init__(self, run_program=True): - + def __init__(self, exchange: Optional[Exchange] = None): + self._exchange = exchange self._initialized = False self.rd_o_cvd: float = GAS_CONSTANT_DRY_AIR / (CPD - GAS_CONSTANT_DRY_AIR) self.thresh_tdiff: float = ( @@ -566,7 +567,7 @@ def run( smag_limit=self.smag_limit, smag_offset=self.smag_offset, ) - log.info("diffusion program: end") + log.info("diffusion step: done") def _do_diffusion_step( self, @@ -687,164 +688,167 @@ def _do_diffusion_step( offset_provider={"V2E": self.grid.get_v2e_connectivity()}, ) log.debug("rbf interpolation: end") - # 2. HALO EXCHANGE -- CALL sync_patch_array_mult + # 2. HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert + self._sync_fields(VertexDim, self.u_vert, self.v_vert) + # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 - log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: start") - calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend)( - diff_multfac_smag=self.diff_multfac_smag, - tangent_orientation=tangent_orientation, - inv_primal_edge_length=inverse_primal_edge_lengths, - inv_vert_vert_length=inverse_vertex_vertex_lengths, - u_vert=self.u_vert, - v_vert=self.v_vert, - primal_normal_vert_x=primal_normal_vert[0], - primal_normal_vert_y=primal_normal_vert[1], - dual_normal_vert_x=dual_normal_vert[0], - dual_normal_vert_y=dual_normal_vert[1], - vn=prognostic_state.vn, - smag_limit=smag_limit, - kh_smag_e=self.kh_smag_e, - kh_smag_ec=self.kh_smag_ec, - z_nabla2_e=self.z_nabla2_e, - smag_offset=smag_offset, - horizontal_start=edge_start_lb_plus4, - horizontal_end=edge_end_local_minus2, - vertical_start=0, - vertical_end=klevels, - offset_provider={ - "E2C2V": self.grid.get_e2c2v_connectivity(), - "E2ECV": self.grid.get_e2ecv_connectivity(), - }, - ) - log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: end") - log.debug("running fused stencil fused stencil 02_03: start") - calculate_diagnostic_quantities_for_turbulence.with_backend(backend)( - kh_smag_ec=self.kh_smag_ec, - vn=prognostic_state.vn, - e_bln_c_s=self.interpolation_state.e_bln_c_s, - geofac_div=self.interpolation_state.geofac_div, - diff_multfac_smag=self.diff_multfac_smag, - wgtfac_c=self.metric_state.wgtfac_c, - div_ic=diagnostic_state.div_ic, - hdef_ic=diagnostic_state.hdef_ic, - horizontal_start=cell_start_nudging, - horizontal_end=cell_end_local, - vertical_start=1, - vertical_end=klevels, - offset_provider={ - "C2E": self.grid.get_c2e_connectivity(), - "C2CE": self.grid.get_c2ce_connectivity(), - "Koff": KDim, - }, - ) - log.debug("running fused stencil fused stencil 02_03: end") + # log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: start") + # calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend)( + # diff_multfac_smag=self.diff_multfac_smag, + # tangent_orientation=tangent_orientation, + # inv_primal_edge_length=inverse_primal_edge_lengths, + # inv_vert_vert_length=inverse_vertex_vertex_lengths, + # u_vert=self.u_vert, + # v_vert=self.v_vert, + # primal_normal_vert_x=primal_normal_vert[0], + # primal_normal_vert_y=primal_normal_vert[1], + # dual_normal_vert_x=dual_normal_vert[0], + # dual_normal_vert_y=dual_normal_vert[1], + # vn=prognostic_state.vn, + # smag_limit=smag_limit, + # kh_smag_e=self.kh_smag_e, + # kh_smag_ec=self.kh_smag_ec, + # z_nabla2_e=self.z_nabla2_e, + # smag_offset=smag_offset, + # horizontal_start=edge_start_lb_plus4, + # horizontal_end=edge_end_local_minus2, + # vertical_start=0, + # vertical_end=klevels, + # offset_provider={ + # "E2C2V": self.grid.get_e2c2v_connectivity(), + # "E2ECV": self.grid.get_e2ecv_connectivity(), + # }, + # ) + # log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: end") + # log.debug("running fused stencil fused stencil 02_03: start") + # calculate_diagnostic_quantities_for_turbulence.with_backend(backend)( + # kh_smag_ec=self.kh_smag_ec, + # vn=prognostic_state.vn, + # e_bln_c_s=self.interpolation_state.e_bln_c_s, + # geofac_div=self.interpolation_state.geofac_div, + # diff_multfac_smag=self.diff_multfac_smag, + # wgtfac_c=self.metric_state.wgtfac_c, + # div_ic=diagnostic_state.div_ic, + # hdef_ic=diagnostic_state.hdef_ic, + # horizontal_start=cell_start_nudging, + # horizontal_end=cell_end_local, + # vertical_start=1, + # vertical_end=klevels, + # offset_provider={ + # "C2E": self.grid.get_c2e_connectivity(), + # "C2CE": self.grid.get_c2ce_connectivity(), + # "Koff": KDim, + # }, + # ) + # log.debug("running fused stencil fused stencil 02_03: end") + # # + # # # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH + # #self._sync_fields(EdgeDim, self.z_nabla2_e) # - # # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH + # # # 5. CALL rbf_vec_interpol_vertex_wp + # log.debug("rbf interpolation: start") + # mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( + # p_e_in=self.z_nabla2_e, + # ptr_coeff_1=self.interpolation_state.rbf_coeff_1, + # ptr_coeff_2=self.interpolation_state.rbf_coeff_2, + # p_u_out=self.u_vert, + # p_v_out=self.v_vert, + # horizontal_start=vertex_start_lb_plus1, + # horizontal_end=vertex_end_local, + # vertical_start=0, + # vertical_end=klevels, + # offset_provider={"V2E": self.grid.get_v2e_connectivity()}, + # ) + # log.debug("rbf interpolation: end") + # # # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult + # # # + # self._sync_fields(VertexDim, self.u_vert, self.v_vert) + # # # # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 + # # # # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 + # log.debug("running fused stencil 04_05_06: start") + # fused_mo_nh_diffusion_stencil_04_05_06.with_backend(backend)( + # u_vert=self.u_vert, + # v_vert=self.v_vert, + # primal_normal_vert_v1=primal_normal_vert[0], + # primal_normal_vert_v2=primal_normal_vert[1], + # z_nabla2_e=self.z_nabla2_e, + # inv_vert_vert_length=inverse_vertex_vertex_lengths, + # inv_primal_edge_length=inverse_primal_edge_lengths, + # area_edge=edge_areas, + # kh_smag_e=self.kh_smag_e, + # diff_multfac_vn=diff_multfac_vn, + # nudgecoeff_e=self.interpolation_state.nudgecoeff_e, + # vn=prognostic_state.vn, + # horz_idx=self.horizontal_edge_index, + # nudgezone_diff=self.nudgezone_diff, + # fac_bdydiff_v=self.fac_bdydiff_v, + # start_2nd_nudge_line_idx_e=int32(edge_start_nudging_plus_one), + # horizontal_start=edge_start_lb_plus4, + # horizontal_end=edge_end_local, + # vertical_start=0, + # vertical_end=klevels, + # offset_provider={ + # "E2C2V": self.grid.get_e2c2v_connectivity(), + # "E2ECV": self.grid.get_e2ecv_connectivity(), + # }, + # ) + # log.debug("running fused stencil 04_05_06: end") + # # # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, + # # # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 # - # # 5. CALL rbf_vec_interpol_vertex_wp - log.debug("rbf interpolation: start") - mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( - p_e_in=self.z_nabla2_e, - ptr_coeff_1=self.interpolation_state.rbf_coeff_1, - ptr_coeff_2=self.interpolation_state.rbf_coeff_2, - p_u_out=self.u_vert, - p_v_out=self.v_vert, - horizontal_start=vertex_start_lb_plus1, - horizontal_end=vertex_end_local, - vertical_start=0, - vertical_end=klevels, - offset_provider={"V2E": self.grid.get_v2e_connectivity()}, - ) - log.debug("rbf interpolation: end") - # # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult - # # - # # # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 - # # # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 - # # + # log.debug("running stencils 07_08_09_10: start") + # calculate_horizontal_gradients_for_turbulence.with_backend(backend)( + # w=prognostic_state.w, + # geofac_grg_x=self.interpolation_state.geofac_grg_x, + # geofac_grg_y=self.interpolation_state.geofac_grg_y, + # dwdx=diagnostic_state.dwdx, + # dwdy=diagnostic_state.dwdy, + # vertical_start=1, + # vertical_end=klevels, + # horizontal_start=cell_start_nudging, + # horizontal_end=cell_end_halo, + # offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, + # ) + # + # z_nabla2_c = zero_field(self.grid, CellDim, KDim, dtype=float) + # + # calculate_nabla2_for_w.with_backend(backend)( + # w=prognostic_state.w, + # geofac_n2s=self.interpolation_state.geofac_n2s, + # z_nabla2_c=z_nabla2_c, + # vertical_start=0, + # vertical_end=klevels, + # horizontal_start=cell_start_nudging, + # horizontal_end=cell_end_halo, + # offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, + # ) + # apply_nabla2_to_w.with_backend(backend)( + # area=cell_areas, + # z_nabla2_c=z_nabla2_c, + # w=prognostic_state.w, + # diff_multfac_w=self.diff_multfac_w, + # geofac_n2s=self.interpolation_state.geofac_n2s, + # vertical_start=0, + # vertical_end=klevels, + # horizontal_start=cell_start_interior, + # horizontal_end=cell_end_local, + # offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, + # ) + # + # # TODO @magdalena: fix this python offset problem (+1): int(self.vertical_params.index_of_damping_layer + 1) + # apply_nabla2_to_w_in_upper_damping_layer.with_backend(backend)( + # w=prognostic_state.w, + # diff_multfac_n2w=self.diff_multfac_n2w, + # cell_area=cell_areas, + # z_nabla2_c=z_nabla2_c, + # vertical_start=1, + # vertical_end=int(self.vertical_params.index_of_damping_layer + 1), + # horizontal_start=int(cell_start_interior), + # horizontal_end=int(cell_end_local), + # offset_provider={}, + # ) # - log.debug("running fused stencil 04_05_06: start") - fused_mo_nh_diffusion_stencil_04_05_06.with_backend(backend)( - u_vert=self.u_vert, - v_vert=self.v_vert, - primal_normal_vert_v1=primal_normal_vert[0], - primal_normal_vert_v2=primal_normal_vert[1], - z_nabla2_e=self.z_nabla2_e, - inv_vert_vert_length=inverse_vertex_vertex_lengths, - inv_primal_edge_length=inverse_primal_edge_lengths, - area_edge=edge_areas, - kh_smag_e=self.kh_smag_e, - diff_multfac_vn=diff_multfac_vn, - nudgecoeff_e=self.interpolation_state.nudgecoeff_e, - vn=prognostic_state.vn, - horz_idx=self.horizontal_edge_index, - nudgezone_diff=self.nudgezone_diff, - fac_bdydiff_v=self.fac_bdydiff_v, - start_2nd_nudge_line_idx_e=int32(edge_start_nudging_plus_one), - horizontal_start=edge_start_lb_plus4, - horizontal_end=edge_end_local, - vertical_start=0, - vertical_end=klevels, - offset_provider={ - "E2C2V": self.grid.get_e2c2v_connectivity(), - "E2ECV": self.grid.get_e2ecv_connectivity(), - }, - ) - log.debug("running fused stencil 04_05_06: end") - # # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, - # # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 - - log.debug("running stencils 07_08_09_10: start") - calculate_horizontal_gradients_for_turbulence.with_backend(backend)( - w=prognostic_state.w, - geofac_grg_x=self.interpolation_state.geofac_grg_x, - geofac_grg_y=self.interpolation_state.geofac_grg_y, - dwdx=diagnostic_state.dwdx, - dwdy=diagnostic_state.dwdy, - vertical_start=1, - vertical_end=klevels, - horizontal_start=cell_start_nudging, - horizontal_end=cell_end_halo, - offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, - ) - - z_nabla2_c = zero_field(self.grid, CellDim, KDim, dtype=float) - - calculate_nabla2_for_w.with_backend(backend)( - w=prognostic_state.w, - geofac_n2s=self.interpolation_state.geofac_n2s, - z_nabla2_c=z_nabla2_c, - vertical_start=0, - vertical_end=klevels, - horizontal_start=cell_start_nudging, - horizontal_end=cell_end_halo, - offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, - ) - apply_nabla2_to_w.with_backend(backend)( - area=cell_areas, - z_nabla2_c=z_nabla2_c, - w=prognostic_state.w, - diff_multfac_w=self.diff_multfac_w, - geofac_n2s=self.interpolation_state.geofac_n2s, - vertical_start=0, - vertical_end=klevels, - horizontal_start=cell_start_interior, - horizontal_end=cell_end_local, - offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, - ) - - # TODO @magdalena: fix this python offset problem (+1): int(self.vertical_params.index_of_damping_layer + 1) - apply_nabla2_to_w_in_upper_damping_layer.with_backend(backend)( - w=prognostic_state.w, - diff_multfac_n2w=self.diff_multfac_n2w, - cell_area=cell_areas, - z_nabla2_c=z_nabla2_c, - vertical_start=1, - vertical_end=int(self.vertical_params.index_of_damping_layer + 1), - horizontal_start=int(cell_start_interior), - horizontal_end=int(cell_end_local), - offset_provider={}, - ) # w_old = prognostic_state.w # apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.with_backend(backend)( # area=cell_areas, @@ -876,6 +880,8 @@ def _do_diffusion_step( # ) log.debug("running fused stencil 07_08_09_10: end") # # 8. HALO EXCHANGE: CALL sync_patch_array + ## self._sync_fields(EdgeDim, prognostic_state.vn) + # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, # # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 # @@ -955,4 +961,15 @@ def _do_diffusion_step( offset_provider={}, ) log.debug("running fused stencil update_theta_and_exner: end") - # 10. HALO EXCHANGE sync_patch_array + # 10. HALO EXCHANGE sync_patch_array (Cell fields) + # TODO @magdalena why not trigger the exchange of w earlier? + #TODO if condition: IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN + self._sync_fields(CellDim, + prognostic_state.theta_v, + prognostic_state.exner_pressure, + prognostic_state.w, + ) + + def _sync_fields(self, dim: Dimension, *fields): + if self._exchange: + self._exchange.exchange(dim, *fields) diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index bebf76b3a..f5a1229d1 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -38,6 +38,7 @@ ) from icon4py.driver.parallel_setup import get_processor_properties + log = logging.getLogger(__name__) @@ -198,8 +199,6 @@ def initialize(n_time_steps, file_path: Path): parallel_props, ) - - log.info("initializing the grid") icon_grid = read_icon_grid(file_path) log.info("reading input fields") @@ -208,7 +207,7 @@ def initialize(n_time_steps, file_path: Path): log.info("initializing diffusion") diffusion_params = DiffusionParams(config.diffusion_config) - diffusion = Diffusion(run_program=False) + diffusion = Diffusion() diffusion.init( icon_grid, config.diffusion_config, diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index 96a796fb2..7068322b5 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -38,7 +38,7 @@ class SerializationType(str, Enum): NC = "netcdf" -def read_icon_grid(path: Path, ser_type=SerializationType.SB) -> IconGrid: +def read_icon_grid(path: Path, rank=0, ser_type=SerializationType.SB) -> IconGrid: """ Return IconGrid parsed from a given input type. @@ -51,7 +51,7 @@ def read_icon_grid(path: Path, ser_type=SerializationType.SB) -> IconGrid: if ser_type == SerializationType.SB: return ( serialbox_utils.IconSerialDataProvider( - "icon_pydycore", str(path.absolute()), False + "icon_pydycore", str(path.absolute()), False, mpi_rank=rank ) .from_savepoint_grid() .construct_icon_grid() @@ -61,7 +61,7 @@ def read_icon_grid(path: Path, ser_type=SerializationType.SB) -> IconGrid: def read_initial_state( - gridfile_path: Path, + gridfile_path: Path, rank=0 ) -> tuple[IconSerialDataProvider, DiagnosticState, PrognosticState]: """ Read prognostic and diagnostic state from serialized data. @@ -75,7 +75,7 @@ def read_initial_state( """ data_provider = serialbox_utils.IconSerialDataProvider( - "icon_pydycore", str(gridfile_path), False + "icon_pydycore", str(gridfile_path), False, mpi_rank=rank ) init_savepoint = data_provider.from_savepoint_diffusion_init( linit=True, date=SIMULATION_START_DATE @@ -86,7 +86,7 @@ def read_initial_state( def read_geometry_fields( - path: Path, ser_type=SerializationType.SB + path: Path, ser_type=SerializationType.SB, rank=0 ) -> tuple[EdgeParams, CellParams, VerticalModelParams]: """ Read fields containing grid properties. @@ -100,7 +100,7 @@ def read_geometry_fields( """ if ser_type == SerializationType.SB: sp = serialbox_utils.IconSerialDataProvider( - "icon_pydycore", str(path.absolute()), False + "icon_pydycore", str(path.absolute()), False, mpi_rank=rank ).from_savepoint_grid() edge_geometry = sp.construct_edge_geometry() cell_geometry = sp.construct_cell_geometry() @@ -142,7 +142,7 @@ def read_grid( def read_static_fields( - path: Path, ser_type=SerializationType.SB + path: Path, ser_type=SerializationType.SB, rank=0 ) -> tuple[MetricState, InterpolationState]: """ Read fields for metric and interpolation state. @@ -157,7 +157,7 @@ def read_static_fields( """ if ser_type == SerializationType.SB: dataprovider = serialbox_utils.IconSerialDataProvider( - "icon_pydycore", str(path.absolute()), False + "icon_pydycore", str(path.absolute()), False, mpi_rank=rank ) interpolation_state = ( dataprovider.from_interpolation_savepoint().construct_interpolation_state() diff --git a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py index a94b9faed..c3ac299a0 100644 --- a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py @@ -10,16 +10,18 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - +import logging from enum import Enum from typing import Union +import ghex.unstructured as ghex import mpi4py import numpy as np import numpy.ma as ma from gt4py.next.common import Dimension from mpi4py.MPI import Comm +from icon4py.common.dimension import CellDim, VertexDim, EdgeDim from icon4py.decomposition.decomposed import ProcessProperties from icon4py.diffusion.utils import builder @@ -27,6 +29,7 @@ mpi4py.rc.initialize = False CommId = Union[int, Comm, None] +log = logging.getLogger(__name__) def get_processor_properties(comm_id: CommId = None): @@ -65,8 +68,31 @@ def with_dimension( masked_global_index = ma.array(global_index, mask=owner_mask) self._global_index[dim] = masked_global_index - def __init__(self): + def __init__(self, klevels: int): self._global_index = {} + self._klevels = klevels + + @property + def klevels(self): + return self._klevels + + def local_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL): + match (entry_type): + case DecompositionInfo.EntryType.ALL: + return self._to_local_index(dim) + case DecompositionInfo.EntryType.HALO: + index = self._to_local_index(dim) + mask = self._global_index[dim].mask + return index[~mask] + case DecompositionInfo.EntryType.OWNED: + index = self._to_local_index(dim) + mask = self._global_index[dim].mask + return index[mask] + + def _to_local_index(self, dim): + data = ma.getdata(self._global_index[dim], subok=False) + assert data.ndim == 1 + return np.arange(data.shape[0]) def global_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL): match (entry_type): @@ -80,3 +106,61 @@ def global_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL): return ma.getdata(global_index[~global_index.mask]) case _: raise NotImplementedError() + + +class Exchange: + def __init__(self, context, domain_decomposition: DecompositionInfo): + self._context = context + self._decomposition_info = domain_decomposition + self._domain_descriptors = { + CellDim: self._create_domain_descriptor(CellDim), + VertexDim: self._create_domain_descriptor(VertexDim), + EdgeDim: self._create_domain_descriptor(EdgeDim) + } + log.info(f"exchange patterns initialized {self._domain_descriptors}") + + self._patterns = { + CellDim: self._create_pattern(CellDim), + VertexDim: self._create_pattern(VertexDim), + EdgeDim: self._create_pattern(EdgeDim) + } + log.info(f"exchange patterns initialized {self._patterns}") + + def get_size(self): + return self._context.size() + def _create_domain_descriptor(self, dim: Dimension): + all_global = self._decomposition_info.global_index( + dim, DecompositionInfo.EntryType.ALL + ) + local_halo = self._decomposition_info.local_index( + dim, DecompositionInfo.EntryType.HALO + ) + domain_descr = ghex.domain_descriptor( + self._context.rank(), + all_global.tolist(), + local_halo.tolist(), + self._decomposition_info.klevels, + ) + return domain_descr + + def _create_pattern(self, dim): + halo_generator = ghex.halo_generator_with_gids( + self._decomposition_info.global_index(dim, DecompositionInfo.EntryType.HALO) + ) + pattern = ghex.make_pattern( + self._context, halo_generator, [self._domain_descriptors[dim]] + ) + return pattern + + def exchange(self, dim: Dimension, *fields): + assert dim in [CellDim, EdgeDim, VertexDim] + log.info(f"exchanging fields for dim={dim}:") + for f in fields: + log.info(f"{f.shape}") + domain_descriptor = self._domain_descriptors[dim] + pattern = self._patterns[dim] + communicator = ghex.make_co(self._context, pattern) + pattern_of_fields = [ + pattern(ghex.field_descriptor(domain_descriptor, np.asarray(f))) for f in fields + ] + communicator.exchange(pattern_of_fields) diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 0c9679f92..a597b9141 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -65,8 +65,10 @@ def setup_icon_data(): @pytest.fixture -def data_provider(setup_icon_data) -> IconSerialDataProvider: - return IconSerialDataProvider("icon_pydycore", str(extracted_path), True) +def data_provider( + setup_icon_data, path=extracted_path, rank=0 +) -> IconSerialDataProvider: + return IconSerialDataProvider("icon_pydycore", str(path), True) @pytest.fixture diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index eea99d3e3..2c6aebd4a 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -10,14 +10,26 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import logging +from pathlib import Path - +import ghex +import numpy as np import pytest +from atm_dyn_iconam.tests.test_utils.serialbox_utils import IconSerialDataProvider from icon4py.common.dimension import CellDim, EdgeDim, VertexDim -from icon4py.driver.io_utils import SerializationType, read_decomp_info +from icon4py.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.driver.io_utils import ( + SerializationType, + read_decomp_info, + read_geometry_fields, + read_icon_grid, + read_static_fields, +) from icon4py.driver.parallel_setup import ( DecompositionInfo, + Exchange, get_processor_properties, ) @@ -25,7 +37,7 @@ """ running tests with mpi: -mpirun -np 2 python -m pytest --with-mpi tests/test_parallel_setup.py +mpirun -np 2 python -m pytest --with-mpi tests/mpi_tests/test_parallel_setup.py mpirun -np 2 pytest -v --with-mpi tests/mpi_tests/ @@ -42,6 +54,11 @@ def test_processor_properties_from_comm_world(mpi): # TODO s [magdalena] extract fixture, more useful asserts... + + + + + @pytest.mark.skipif( props.comm_size > 2, reason="input files available for 1 or 2 nodes" ) @@ -54,17 +71,12 @@ def test_processor_properties_from_comm_world(mpi): ), ) def test_decomposition_info_masked(mpi, datapath, dim, owned, total): - - props = get_processor_properties() - my_rank = props.rank decomposition_info = read_decomp_info(datapath, props, SerializationType.SB) - owned_indices = decomposition_info.global_index( - dim, DecompositionInfo.EntryType.ALL - ) + all_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.ALL) my_total = total[my_rank] my_owned = owned[my_rank] - assert owned_indices.shape[0] == my_total + assert all_indices.shape[0] == my_total owned_indices = decomposition_info.global_index( dim, DecompositionInfo.EntryType.OWNED @@ -75,3 +87,120 @@ def test_decomposition_info_masked(mpi, datapath, dim, owned, total): dim, DecompositionInfo.EntryType.HALO ) assert halo_indices.shape[0] == my_total - my_owned + _assert_index_partitioning(all_indices, halo_indices, owned_indices) + +@pytest.mark.skipif( + props.comm_size > 2, reason="input files available for 1 or 2 nodes" +) +@pytest.mark.parametrize( + ("dim, owned, total"), + ( + (CellDim, (10448, 10448), (10611, 10612)), + (EdgeDim, (15820, 15738), (16065, 16067)), + (VertexDim, (5373, 5290), (5455, 5456)), + ), +) +def test_decomposition_info_local_index(mpi, datapath, dim, owned, total): + props = get_processor_properties() + + my_rank = props.rank + decomposition_info = read_decomp_info(datapath, props, SerializationType.SB) + all_indices = decomposition_info.local_index(dim, DecompositionInfo.EntryType.ALL) + my_total = total[my_rank] + my_owned = owned[my_rank] + + assert all_indices.shape[0] == my_total + assert np.array_equal(all_indices, np.arange(0, my_total)) + halo_indices = decomposition_info.local_index(dim, DecompositionInfo.EntryType.HALO) + assert halo_indices.shape[0] == my_total - my_owned + assert halo_indices.shape[0] < all_indices.shape[0] + assert np.alltrue(halo_indices <= np.max(all_indices)) + + owned_indices = decomposition_info.local_index( + dim, DecompositionInfo.EntryType.OWNED + ) + assert owned_indices.shape[0] == my_owned + assert owned_indices.shape[0] <= all_indices.shape[0] + assert np.alltrue(owned_indices <= np.max(all_indices)) + _assert_index_partitioning(all_indices, halo_indices, owned_indices) + + + +def _assert_index_partitioning(all_indices, halo_indices, owned_indices): + owned_list = owned_indices.tolist() + halos_list = halo_indices.tolist() + all_list = all_indices.tolist() + assert set(owned_list) & set(halos_list) == set() + assert set(owned_list) & set(all_list) == set(owned_list) + assert set(halos_list) & set(all_list) == set(halos_list) + assert set(halos_list) | set(owned_list) == set(all_list) + + + +@pytest.mark.mpi +def test_decomposition_info_matches_gridsize(datapath): + decomposition_info = read_decomp_info(datapath, props, SerializationType.SB) + icon_grid = read_icon_grid(datapath, props.rank) + assert \ + decomposition_info.global_index(dim=CellDim, entry_type=DecompositionInfo.EntryType.ALL).shape[ + 0] == icon_grid.num_cells() + assert decomposition_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape[ + 0] == icon_grid.num_vertices() + assert decomposition_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape[ + 0] == icon_grid.num_edges() + +#@pytest.mark.mpi +@pytest.mark.skip +def test_parallel_diffusion(r04b09_diffusion_config, parallel_props, step_date_init, caplog): + caplog.set_level(logging.INFO) + experiment_name = "mch_ch_r04b09_dsl" + path = Path( + f"/home/magdalena/data/exclaim/dycore/{experiment_name}/mpitasks2/{experiment_name}/ser_data" + ) + + icon_grid = read_icon_grid(path, rank=parallel_props.rank) + decomp_info = read_decomp_info( + path, + parallel_props, + ) + context = ghex.context(ghex.mpi_comm(parallel_props.comm), True) + #assert context.size() == 2 + + diffusion_params = DiffusionParams(r04b09_diffusion_config) + diffusion_initial_data = IconSerialDataProvider( + "icon_pydycore", str(path), True, mpi_rank=parallel_props.rank + ).from_savepoint_diffusion_init(linit=True, date=step_date_init) + (edge_geometry, cell_geometry, vertical_geometry) = read_geometry_fields( + path, rank=parallel_props.rank + ) + (metric_state, interpolation_state) = read_static_fields( + path, rank=parallel_props.rank + ) + + dtime = diffusion_initial_data.get_metadata("dtime").get("dtime") + exchange = Exchange(context, decomp_info) + print(f"exchange setup: for {exchange.get_size()} nodes") + + diffusion = Diffusion(exchange) + diffusion.init( + icon_grid, + r04b09_diffusion_config, + diffusion_params, + vertical_geometry, + metric_state, + interpolation_state, + ) + print("diffusion initialized") + diffusion.run( + diagnostic_state=diffusion_initial_data.construct_diagnostics(), + prognostic_state=diffusion_initial_data.construct_prognostics(), + dtime=dtime, + tangent_orientation=edge_geometry.tangent_orientation, + inverse_primal_edge_lengths=edge_geometry.inverse_primal_edge_lengths, + inverse_dual_edge_length=edge_geometry.inverse_dual_edge_lengths, + inverse_vert_vert_lengths=edge_geometry.inverse_vertex_vertex_lengths, + primal_normal_vert=edge_geometry.primal_normal_vert, + dual_normal_vert=edge_geometry.dual_normal_vert, + edge_areas=edge_geometry.edge_areas, + cell_areas=cell_geometry.area, + ) diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index 705cb56a2..ec9c1e2a4 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -234,7 +234,7 @@ def decomp_domain(self, dim): def construct_decomposition_info(self): return ( - DecompositionInfo() + DecompositionInfo(klevels=self.num(KDim)) .with_dimension(*self._get_decomp_fields(CellDim)) .with_dimension(*self._get_decomp_fields(EdgeDim)) .with_dimension(*self._get_decomp_fields(VertexDim)) From 91895634b08c22bbf5911389e7d2272726df3d3c Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 15 Jun 2023 12:53:02 +0200 Subject: [PATCH 135/263] field descriptors for k and k+1 fields --- .../src/icon4py/diffusion/diffusion.py | 2 + .../src/icon4py/driver/parallel_setup.py | 52 +++++++++++++++---- .../tests/mpi_tests/test_parallel_setup.py | 16 +++--- 3 files changed, 52 insertions(+), 18 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 1cc5f1fe5..1ae921d62 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -23,6 +23,7 @@ from gt4py.next.program_processors.runners.gtfn_cpu import ( run_gtfn, run_gtfn_cached, +run_gtfn_imperative, ) from icon4py.atm_dyn_iconam.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( @@ -88,6 +89,7 @@ cached_backend = run_gtfn_cached compiled_backend = run_gtfn +imperative_backend = run_gtfn_imperative backend = compiled_backend # diff --git a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py index c3ac299a0..cd3205689 100644 --- a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py @@ -21,7 +21,7 @@ from gt4py.next.common import Dimension from mpi4py.MPI import Comm -from icon4py.common.dimension import CellDim, VertexDim, EdgeDim +from icon4py.common.dimension import CellDim, VertexDim, EdgeDim, KDim, KHalfDim from icon4py.decomposition.decomposed import ProcessProperties from icon4py.diffusion.utils import builder @@ -115,7 +115,8 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): self._domain_descriptors = { CellDim: self._create_domain_descriptor(CellDim), VertexDim: self._create_domain_descriptor(VertexDim), - EdgeDim: self._create_domain_descriptor(EdgeDim) + EdgeDim: self._create_domain_descriptor(EdgeDim), + } log.info(f"exchange patterns initialized {self._domain_descriptors}") @@ -125,9 +126,25 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): EdgeDim: self._create_pattern(EdgeDim) } log.info(f"exchange patterns initialized {self._patterns}") + self._field_size = { + CellDim: self._decomposition_info.global_index( + CellDim, DecompositionInfo.EntryType.ALL + ).shape[0], + EdgeDim: self._decomposition_info.global_index( + EdgeDim, DecompositionInfo.EntryType.ALL + ).shape[0], + VertexDim: self._decomposition_info.global_index( + VertexDim, DecompositionInfo.EntryType.ALL + ).shape[0], + KDim: domain_decomposition.klevels, + KHalfDim: domain_decomposition.klevels + 1 + } + def get_size(self): return self._context.size() + + # TODO [magdalena] is the tolist() necessary? def _create_domain_descriptor(self, dim: Dimension): all_global = self._decomposition_info.global_index( dim, DecompositionInfo.EntryType.ALL @@ -135,32 +152,47 @@ def _create_domain_descriptor(self, dim: Dimension): local_halo = self._decomposition_info.local_index( dim, DecompositionInfo.EntryType.HALO ) - domain_descr = ghex.domain_descriptor( + domain_desc_full_levels = ghex.domain_descriptor( self._context.rank(), all_global.tolist(), local_halo.tolist(), - self._decomposition_info.klevels, + self._decomposition_info.klevels ) - return domain_descr + domain_descr_half_levels = ghex.domain_descriptor( + self._context.rank(), + all_global.tolist(), + local_halo.tolist(), + self._decomposition_info.klevels + 1 + ) + return (domain_desc_full_levels, domain_descr_half_levels) def _create_pattern(self, dim): halo_generator = ghex.halo_generator_with_gids( self._decomposition_info.global_index(dim, DecompositionInfo.EntryType.HALO) ) pattern = ghex.make_pattern( - self._context, halo_generator, [self._domain_descriptors[dim]] + self._context, halo_generator, self._domain_descriptors[dim] ) return pattern def exchange(self, dim: Dimension, *fields): + + def _get_descriptor_type(f): + return 0 if f.shape[1] == self._decomposition_info.klevels else 1 + assert dim in [CellDim, EdgeDim, VertexDim] log.info(f"exchanging fields for dim={dim}:") - for f in fields: - log.info(f"{f.shape}") - domain_descriptor = self._domain_descriptors[dim] + horizontal_size = self._field_size[dim] pattern = self._patterns[dim] communicator = ghex.make_co(self._context, pattern) + + fields = [np.asarray(f)[:horizontal_size, :] for f in fields] + for f in fields: + log.info(f"{f.shape}") pattern_of_fields = [ - pattern(ghex.field_descriptor(domain_descriptor, np.asarray(f))) for f in fields + pattern(ghex.field_descriptor( + self._domain_descriptors[dim][_get_descriptor_type(f)], f)) for f in fields ] communicator.exchange(pattern_of_fields) + + diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index 2c6aebd4a..08e1dc8ab 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -150,31 +150,31 @@ def test_decomposition_info_matches_gridsize(datapath): 0] == icon_grid.num_edges() #@pytest.mark.mpi -@pytest.mark.skip -def test_parallel_diffusion(r04b09_diffusion_config, parallel_props, step_date_init, caplog): +#@pytest.mark.skip +def test_parallel_diffusion(r04b09_diffusion_config, step_date_init, caplog): caplog.set_level(logging.INFO) experiment_name = "mch_ch_r04b09_dsl" path = Path( f"/home/magdalena/data/exclaim/dycore/{experiment_name}/mpitasks2/{experiment_name}/ser_data" ) - icon_grid = read_icon_grid(path, rank=parallel_props.rank) + icon_grid = read_icon_grid(path, rank=props.rank) decomp_info = read_decomp_info( path, - parallel_props, + props, ) - context = ghex.context(ghex.mpi_comm(parallel_props.comm), True) + context = ghex.context(ghex.mpi_comm(props.comm), True) #assert context.size() == 2 diffusion_params = DiffusionParams(r04b09_diffusion_config) diffusion_initial_data = IconSerialDataProvider( - "icon_pydycore", str(path), True, mpi_rank=parallel_props.rank + "icon_pydycore", str(path), True, mpi_rank=props.rank ).from_savepoint_diffusion_init(linit=True, date=step_date_init) (edge_geometry, cell_geometry, vertical_geometry) = read_geometry_fields( - path, rank=parallel_props.rank + path, rank=props.rank ) (metric_state, interpolation_state) = read_static_fields( - path, rank=parallel_props.rank + path, rank=props.rank ) dtime = diffusion_initial_data.get_metadata("dtime").get("dtime") From 11aab95c6097be387d34b2571e7cde79baea3388 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 15 Jun 2023 17:40:34 +0200 Subject: [PATCH 136/263] can do exchanges except for cells... --- .../src/icon4py/diffusion/diffusion.py | 345 +++++++++--------- .../src/icon4py/driver/parallel_setup.py | 114 +++--- .../tests/mpi_tests/test_parallel_setup.py | 72 ++-- 3 files changed, 282 insertions(+), 249 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 1ae921d62..159cda4db 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -64,7 +64,7 @@ DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO, GAS_CONSTANT_DRY_AIR, ) -from icon4py.common.dimension import CellDim, ECVDim, EdgeDim, KDim, VertexDim +from icon4py.common.dimension import CellDim, ECVDim, EdgeDim, KDim, VertexDim, KHalfDim from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.horizontal import HorizontalMarkerIndex from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams @@ -691,166 +691,170 @@ def _do_diffusion_step( ) log.debug("rbf interpolation: end") # 2. HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert - self._sync_fields(VertexDim, self.u_vert, self.v_vert) + res = self._sync_fields((VertexDim, KDim), self.u_vert, self.v_vert) + self._wait(res) # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 - # log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: start") - # calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend)( - # diff_multfac_smag=self.diff_multfac_smag, - # tangent_orientation=tangent_orientation, - # inv_primal_edge_length=inverse_primal_edge_lengths, - # inv_vert_vert_length=inverse_vertex_vertex_lengths, - # u_vert=self.u_vert, - # v_vert=self.v_vert, - # primal_normal_vert_x=primal_normal_vert[0], - # primal_normal_vert_y=primal_normal_vert[1], - # dual_normal_vert_x=dual_normal_vert[0], - # dual_normal_vert_y=dual_normal_vert[1], - # vn=prognostic_state.vn, - # smag_limit=smag_limit, - # kh_smag_e=self.kh_smag_e, - # kh_smag_ec=self.kh_smag_ec, - # z_nabla2_e=self.z_nabla2_e, - # smag_offset=smag_offset, - # horizontal_start=edge_start_lb_plus4, - # horizontal_end=edge_end_local_minus2, - # vertical_start=0, - # vertical_end=klevels, - # offset_provider={ - # "E2C2V": self.grid.get_e2c2v_connectivity(), - # "E2ECV": self.grid.get_e2ecv_connectivity(), - # }, - # ) - # log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: end") - # log.debug("running fused stencil fused stencil 02_03: start") - # calculate_diagnostic_quantities_for_turbulence.with_backend(backend)( - # kh_smag_ec=self.kh_smag_ec, - # vn=prognostic_state.vn, - # e_bln_c_s=self.interpolation_state.e_bln_c_s, - # geofac_div=self.interpolation_state.geofac_div, - # diff_multfac_smag=self.diff_multfac_smag, - # wgtfac_c=self.metric_state.wgtfac_c, - # div_ic=diagnostic_state.div_ic, - # hdef_ic=diagnostic_state.hdef_ic, - # horizontal_start=cell_start_nudging, - # horizontal_end=cell_end_local, - # vertical_start=1, - # vertical_end=klevels, - # offset_provider={ - # "C2E": self.grid.get_c2e_connectivity(), - # "C2CE": self.grid.get_c2ce_connectivity(), - # "Koff": KDim, - # }, - # ) - # log.debug("running fused stencil fused stencil 02_03: end") - # # - # # # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH - # #self._sync_fields(EdgeDim, self.z_nabla2_e) - # - # # # 5. CALL rbf_vec_interpol_vertex_wp - # log.debug("rbf interpolation: start") - # mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( - # p_e_in=self.z_nabla2_e, - # ptr_coeff_1=self.interpolation_state.rbf_coeff_1, - # ptr_coeff_2=self.interpolation_state.rbf_coeff_2, - # p_u_out=self.u_vert, - # p_v_out=self.v_vert, - # horizontal_start=vertex_start_lb_plus1, - # horizontal_end=vertex_end_local, - # vertical_start=0, - # vertical_end=klevels, - # offset_provider={"V2E": self.grid.get_v2e_connectivity()}, - # ) - # log.debug("rbf interpolation: end") - # # # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult - # # # - # self._sync_fields(VertexDim, self.u_vert, self.v_vert) - # # # # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 - # # # # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 - # log.debug("running fused stencil 04_05_06: start") - # fused_mo_nh_diffusion_stencil_04_05_06.with_backend(backend)( - # u_vert=self.u_vert, - # v_vert=self.v_vert, - # primal_normal_vert_v1=primal_normal_vert[0], - # primal_normal_vert_v2=primal_normal_vert[1], - # z_nabla2_e=self.z_nabla2_e, - # inv_vert_vert_length=inverse_vertex_vertex_lengths, - # inv_primal_edge_length=inverse_primal_edge_lengths, - # area_edge=edge_areas, - # kh_smag_e=self.kh_smag_e, - # diff_multfac_vn=diff_multfac_vn, - # nudgecoeff_e=self.interpolation_state.nudgecoeff_e, - # vn=prognostic_state.vn, - # horz_idx=self.horizontal_edge_index, - # nudgezone_diff=self.nudgezone_diff, - # fac_bdydiff_v=self.fac_bdydiff_v, - # start_2nd_nudge_line_idx_e=int32(edge_start_nudging_plus_one), - # horizontal_start=edge_start_lb_plus4, - # horizontal_end=edge_end_local, - # vertical_start=0, - # vertical_end=klevels, - # offset_provider={ - # "E2C2V": self.grid.get_e2c2v_connectivity(), - # "E2ECV": self.grid.get_e2ecv_connectivity(), - # }, - # ) - # log.debug("running fused stencil 04_05_06: end") - # # # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, - # # # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 + log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: start") + calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend)( + diff_multfac_smag=self.diff_multfac_smag, + tangent_orientation=tangent_orientation, + inv_primal_edge_length=inverse_primal_edge_lengths, + inv_vert_vert_length=inverse_vertex_vertex_lengths, + u_vert=self.u_vert, + v_vert=self.v_vert, + primal_normal_vert_x=primal_normal_vert[0], + primal_normal_vert_y=primal_normal_vert[1], + dual_normal_vert_x=dual_normal_vert[0], + dual_normal_vert_y=dual_normal_vert[1], + vn=prognostic_state.vn, + smag_limit=smag_limit, + kh_smag_e=self.kh_smag_e, + kh_smag_ec=self.kh_smag_ec, + z_nabla2_e=self.z_nabla2_e, + smag_offset=smag_offset, + horizontal_start=edge_start_lb_plus4, + horizontal_end=edge_end_local_minus2, + vertical_start=0, + vertical_end=klevels, + offset_provider={ + "E2C2V": self.grid.get_e2c2v_connectivity(), + "E2ECV": self.grid.get_e2ecv_connectivity(), + }, + ) + log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: end") + log.debug("running fused stencil fused stencil 02_03: start") + calculate_diagnostic_quantities_for_turbulence.with_backend(backend)( + kh_smag_ec=self.kh_smag_ec, + vn=prognostic_state.vn, + e_bln_c_s=self.interpolation_state.e_bln_c_s, + geofac_div=self.interpolation_state.geofac_div, + diff_multfac_smag=self.diff_multfac_smag, + wgtfac_c=self.metric_state.wgtfac_c, + div_ic=diagnostic_state.div_ic, + hdef_ic=diagnostic_state.hdef_ic, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local, + vertical_start=1, + vertical_end=klevels, + offset_provider={ + "C2E": self.grid.get_c2e_connectivity(), + "C2CE": self.grid.get_c2ce_connectivity(), + "Koff": KDim, + }, + ) + log.debug("running fused stencil fused stencil 02_03: end") # + # # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH + res = self._sync_fields((EdgeDim, KDim), self.z_nabla2_e) + self._wait(res) + + # # 5. CALL rbf_vec_interpol_vertex_wp + log.debug("rbf interpolation: start") + mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( + p_e_in=self.z_nabla2_e, + ptr_coeff_1=self.interpolation_state.rbf_coeff_1, + ptr_coeff_2=self.interpolation_state.rbf_coeff_2, + p_u_out=self.u_vert, + p_v_out=self.v_vert, + horizontal_start=vertex_start_lb_plus1, + horizontal_end=vertex_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={"V2E": self.grid.get_v2e_connectivity()}, + ) + log.debug("rbf interpolation: end") + # # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult + # # + res = self._sync_fields((VertexDim, KDim), self.u_vert, self.v_vert) + self._wait(res) + + # # # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 + # # # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 + log.debug("running fused stencil 04_05_06: start") + fused_mo_nh_diffusion_stencil_04_05_06.with_backend(backend)( + u_vert=self.u_vert, + v_vert=self.v_vert, + primal_normal_vert_v1=primal_normal_vert[0], + primal_normal_vert_v2=primal_normal_vert[1], + z_nabla2_e=self.z_nabla2_e, + inv_vert_vert_length=inverse_vertex_vertex_lengths, + inv_primal_edge_length=inverse_primal_edge_lengths, + area_edge=edge_areas, + kh_smag_e=self.kh_smag_e, + diff_multfac_vn=diff_multfac_vn, + nudgecoeff_e=self.interpolation_state.nudgecoeff_e, + vn=prognostic_state.vn, + horz_idx=self.horizontal_edge_index, + nudgezone_diff=self.nudgezone_diff, + fac_bdydiff_v=self.fac_bdydiff_v, + start_2nd_nudge_line_idx_e=int32(edge_start_nudging_plus_one), + horizontal_start=edge_start_lb_plus4, + horizontal_end=edge_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={ + "E2C2V": self.grid.get_e2c2v_connectivity(), + "E2ECV": self.grid.get_e2ecv_connectivity(), + }, + ) + log.debug("running fused stencil 04_05_06: end") + # # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, + # # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 + # log.debug("running stencils 07_08_09_10: start") - # calculate_horizontal_gradients_for_turbulence.with_backend(backend)( - # w=prognostic_state.w, - # geofac_grg_x=self.interpolation_state.geofac_grg_x, - # geofac_grg_y=self.interpolation_state.geofac_grg_y, - # dwdx=diagnostic_state.dwdx, - # dwdy=diagnostic_state.dwdy, - # vertical_start=1, - # vertical_end=klevels, - # horizontal_start=cell_start_nudging, - # horizontal_end=cell_end_halo, - # offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, - # ) - # - # z_nabla2_c = zero_field(self.grid, CellDim, KDim, dtype=float) - # - # calculate_nabla2_for_w.with_backend(backend)( - # w=prognostic_state.w, - # geofac_n2s=self.interpolation_state.geofac_n2s, - # z_nabla2_c=z_nabla2_c, - # vertical_start=0, - # vertical_end=klevels, - # horizontal_start=cell_start_nudging, - # horizontal_end=cell_end_halo, - # offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, - # ) - # apply_nabla2_to_w.with_backend(backend)( - # area=cell_areas, - # z_nabla2_c=z_nabla2_c, - # w=prognostic_state.w, - # diff_multfac_w=self.diff_multfac_w, - # geofac_n2s=self.interpolation_state.geofac_n2s, - # vertical_start=0, - # vertical_end=klevels, - # horizontal_start=cell_start_interior, - # horizontal_end=cell_end_local, - # offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, - # ) - # - # # TODO @magdalena: fix this python offset problem (+1): int(self.vertical_params.index_of_damping_layer + 1) - # apply_nabla2_to_w_in_upper_damping_layer.with_backend(backend)( - # w=prognostic_state.w, - # diff_multfac_n2w=self.diff_multfac_n2w, - # cell_area=cell_areas, - # z_nabla2_c=z_nabla2_c, - # vertical_start=1, - # vertical_end=int(self.vertical_params.index_of_damping_layer + 1), - # horizontal_start=int(cell_start_interior), - # horizontal_end=int(cell_end_local), - # offset_provider={}, - # ) - # + calculate_horizontal_gradients_for_turbulence.with_backend(backend)( + w=prognostic_state.w, + geofac_grg_x=self.interpolation_state.geofac_grg_x, + geofac_grg_y=self.interpolation_state.geofac_grg_y, + dwdx=diagnostic_state.dwdx, + dwdy=diagnostic_state.dwdy, + vertical_start=1, + vertical_end=klevels, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_halo, + offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, + ) + + z_nabla2_c = zero_field(self.grid, CellDim, KDim, dtype=float) + + calculate_nabla2_for_w.with_backend(backend)( + w=prognostic_state.w, + geofac_n2s=self.interpolation_state.geofac_n2s, + z_nabla2_c=z_nabla2_c, + vertical_start=0, + vertical_end=klevels, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_halo, + offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, + ) + apply_nabla2_to_w.with_backend(backend)( + area=cell_areas, + z_nabla2_c=z_nabla2_c, + w=prognostic_state.w, + diff_multfac_w=self.diff_multfac_w, + geofac_n2s=self.interpolation_state.geofac_n2s, + vertical_start=0, + vertical_end=klevels, + horizontal_start=cell_start_interior, + horizontal_end=cell_end_local, + offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, + ) + + # TODO @magdalena: fix this python offset problem (+1): int(self.vertical_params.index_of_damping_layer + 1) + apply_nabla2_to_w_in_upper_damping_layer.with_backend(backend)( + w=prognostic_state.w, + diff_multfac_n2w=self.diff_multfac_n2w, + cell_area=cell_areas, + z_nabla2_c=z_nabla2_c, + vertical_start=1, + vertical_end=int(self.vertical_params.index_of_damping_layer + 1), + horizontal_start=int(cell_start_interior), + horizontal_end=int(cell_end_local), + offset_provider={}, + ) + # w_old = prognostic_state.w # apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.with_backend(backend)( # area=cell_areas, @@ -882,7 +886,9 @@ def _do_diffusion_step( # ) log.debug("running fused stencil 07_08_09_10: end") # # 8. HALO EXCHANGE: CALL sync_patch_array - ## self._sync_fields(EdgeDim, prognostic_state.vn) + comm_res = self._sync_fields((EdgeDim,KDim), prognostic_state.vn) + self._wait(comm_res) + # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, # # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 @@ -965,13 +971,22 @@ def _do_diffusion_step( log.debug("running fused stencil update_theta_and_exner: end") # 10. HALO EXCHANGE sync_patch_array (Cell fields) # TODO @magdalena why not trigger the exchange of w earlier? - #TODO if condition: IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN - self._sync_fields(CellDim, - prognostic_state.theta_v, - prognostic_state.exner_pressure, - prognostic_state.w, - ) + # TODO if condition: IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN + # TODO magdalena: why does this crash? + # res = self._sync_fields((CellDim,KDim), + # prognostic_state.theta_v, + # prognostic_state.exner_pressure, + # ) + res = self._sync_fields((CellDim, KHalfDim), prognostic_state.w) + self._wait(res) - def _sync_fields(self, dim: Dimension, *fields): + def _sync_fields(self, dim: tuple[Dimension, Dimension], *field): if self._exchange: - self._exchange.exchange(dim, *fields) + return self._exchange.exchange(dim, *field) + + + + + def _wait(self, comm_handle): + if comm_handle: + comm_handle.wait() diff --git a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py index cd3205689..37b566469 100644 --- a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py @@ -52,9 +52,16 @@ def init_mpi(): from mpi4py import MPI if not MPI.Is_initialized(): + log.info("initializing MPI") MPI.Init() +def finalize_mpi(): + from mpi4py import MPI + if not MPI.Is_finalized(): + log.info("finalizing MPI") + MPI.Finalize() + class DecompositionInfo: class EntryType(int, Enum): ALL = (0,) @@ -113,86 +120,87 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): self._context = context self._decomposition_info = domain_decomposition self._domain_descriptors = { - CellDim: self._create_domain_descriptor(CellDim), - VertexDim: self._create_domain_descriptor(VertexDim), - EdgeDim: self._create_domain_descriptor(EdgeDim), + (CellDim, KDim): self._create_domain_descriptor(CellDim, domain_decomposition.klevels), + (CellDim, KHalfDim): self._create_domain_descriptor(CellDim, domain_decomposition.klevels + 1), + (VertexDim, KDim): self._create_domain_descriptor(VertexDim, domain_decomposition.klevels), + (VertexDim, KHalfDim): self._create_domain_descriptor(VertexDim, + domain_decomposition.klevels + 1), + (EdgeDim, KDim): self._create_domain_descriptor(EdgeDim, domain_decomposition.klevels), + (EdgeDim, KHalfDim): self._create_domain_descriptor(EdgeDim, + domain_decomposition.klevels + 1), } - log.info(f"exchange patterns initialized {self._domain_descriptors}") - - self._patterns = { - CellDim: self._create_pattern(CellDim), - VertexDim: self._create_pattern(VertexDim), - EdgeDim: self._create_pattern(EdgeDim) - } - log.info(f"exchange patterns initialized {self._patterns}") + print(f"rank={self._context.rank()}/{self._context.size()} :domain descriptors initialized") self._field_size = { CellDim: self._decomposition_info.global_index( - CellDim, DecompositionInfo.EntryType.ALL - ).shape[0], - EdgeDim: self._decomposition_info.global_index( - EdgeDim, DecompositionInfo.EntryType.ALL - ).shape[0], - VertexDim: self._decomposition_info.global_index( - VertexDim, DecompositionInfo.EntryType.ALL - ).shape[0], - KDim: domain_decomposition.klevels, - KHalfDim: domain_decomposition.klevels + 1 + CellDim, DecompositionInfo.EntryType.ALL + ).shape[0], + EdgeDim: self._decomposition_info.global_index( + EdgeDim, DecompositionInfo.EntryType.ALL + ).shape[0], + VertexDim: self._decomposition_info.global_index( + VertexDim, DecompositionInfo.EntryType.ALL + ).shape[0], + KDim: domain_decomposition.klevels, + KHalfDim: domain_decomposition.klevels + 1 } + print(f"rank={self._context.rank()}/{self._context.size()} : field sizes = {self._field_size}") + + self._patterns = { + (CellDim, KDim): self._create_pattern(CellDim, KDim), + (CellDim, KHalfDim): self._create_pattern(CellDim, KHalfDim), + (VertexDim, KDim): self._create_pattern(VertexDim, KDim), + (VertexDim, KHalfDim): self._create_pattern(VertexDim, KHalfDim), + (EdgeDim, KDim): self._create_pattern(EdgeDim, KDim), + (EdgeDim, KHalfDim): self._create_pattern(EdgeDim, KHalfDim), + } + print(f"rank={self._context.rank()}/{self._context.size()} : patterns initialized ") + print(f"rank={self._context.rank()}/{self._context.size()} : exchange initialized") + def _domain_descriptor_info(self, descr): + return f" id={descr.domain_id()} levels = {descr.levels()}, size={descr.size()}, inner_size={descr.inner_size()}" def get_size(self): return self._context.size() # TODO [magdalena] is the tolist() necessary? - def _create_domain_descriptor(self, dim: Dimension): + def _create_domain_descriptor(self, dim: Dimension, levels:int): all_global = self._decomposition_info.global_index( dim, DecompositionInfo.EntryType.ALL ) local_halo = self._decomposition_info.local_index( dim, DecompositionInfo.EntryType.HALO ) - domain_desc_full_levels = ghex.domain_descriptor( + domain_desc = ghex.domain_descriptor( self._context.rank(), all_global.tolist(), local_halo.tolist(), - self._decomposition_info.klevels + levels ) - domain_descr_half_levels = ghex.domain_descriptor( - self._context.rank(), - all_global.tolist(), - local_halo.tolist(), - self._decomposition_info.klevels + 1 - ) - return (domain_desc_full_levels, domain_descr_half_levels) + print(f"rank={self._context.rank()}/{self._context.size()}: domain descriptor for dim {dim} with properties {self._domain_descriptor_info(domain_desc)}") + + return domain_desc - def _create_pattern(self, dim): + def _create_pattern(self, horizontal_dim:Dimension, vertical_dim:Dimension): halo_generator = ghex.halo_generator_with_gids( - self._decomposition_info.global_index(dim, DecompositionInfo.EntryType.HALO) + self._decomposition_info.global_index(horizontal_dim, DecompositionInfo.EntryType.HALO) ) + print(f"rank={self._context.rank()}/{self._context.size()}: halo generator for dim={horizontal_dim} created") + dimensions = (horizontal_dim, vertical_dim) pattern = ghex.make_pattern( - self._context, halo_generator, self._domain_descriptors[dim] + self._context, halo_generator,[self._domain_descriptors[dimensions]] ) + print(f"rank={self._context.rank()}/{self._context.size()}: pattern for dim={dimensions} and {self._domain_descriptors[dimensions]} created") return pattern - def exchange(self, dim: Dimension, *fields): - - def _get_descriptor_type(f): - return 0 if f.shape[1] == self._decomposition_info.klevels else 1 - - assert dim in [CellDim, EdgeDim, VertexDim] - log.info(f"exchanging fields for dim={dim}:") - horizontal_size = self._field_size[dim] - pattern = self._patterns[dim] - communicator = ghex.make_co(self._context, pattern) - + def exchange(self, dims: tuple[Dimension, Dimension], *fields): + assert dims[0] in [CellDim, EdgeDim, VertexDim] + horizontal_size = self._field_size[dims[0]] + pattern = self._patterns[dims] + assert pattern is not None fields = [np.asarray(f)[:horizontal_size, :] for f in fields] - for f in fields: - log.info(f"{f.shape}") - pattern_of_fields = [ - pattern(ghex.field_descriptor( - self._domain_descriptors[dim][_get_descriptor_type(f)], f)) for f in fields - ] - communicator.exchange(pattern_of_fields) - + print(f"rank = {self._context.rank()}/{self._context.size()}: communicating fields of dim = {dims} (shape = {fields[0].shape})") + communicator = ghex.make_co(self._context, pattern) + patterns_of_field = [pattern(ghex.field_descriptor(self._domain_descriptors[dims], f)) for f in fields] + return communicator.exchange(patterns_of_field) diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index 08e1dc8ab..5ffdf4d6d 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -30,37 +30,23 @@ from icon4py.driver.parallel_setup import ( DecompositionInfo, Exchange, - get_processor_properties, + get_processor_properties, finalize_mpi, ) """ running tests with mpi: -mpirun -np 2 python -m pytest --with-mpi tests/mpi_tests/test_parallel_setup.py +mpirun -np 2 python -m pytest -v --with-mpi tests/mpi_tests/test_parallel_setup.py mpirun -np 2 pytest -v --with-mpi tests/mpi_tests/ """ -props = get_processor_properties() - - -@pytest.mark.mpi -def test_processor_properties_from_comm_world(mpi): - props = get_processor_properties() - assert props.rank < mpi.COMM_WORLD.Get_size() - assert props.comm_name == mpi.COMM_WORLD.Get_name() - - -# TODO s [magdalena] extract fixture, more useful asserts... - - - - +props = get_processor_properties() @pytest.mark.skipif( - props.comm_size > 2, reason="input files available for 1 or 2 nodes" + props.comm_size > 2, reason="input files only available for 1 or 2 nodes" ) @pytest.mark.parametrize( ("dim, owned, total"), @@ -70,7 +56,8 @@ def test_processor_properties_from_comm_world(mpi): (VertexDim, (5373, 5290), (5455, 5456)), ), ) -def test_decomposition_info_masked(mpi, datapath, dim, owned, total): +def test_decomposition_info_masked(mpi, datapath, dim, owned, total, caplog): + props = get_processor_properties() my_rank = props.rank decomposition_info = read_decomp_info(datapath, props, SerializationType.SB) all_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.ALL) @@ -89,8 +76,9 @@ def test_decomposition_info_masked(mpi, datapath, dim, owned, total): assert halo_indices.shape[0] == my_total - my_owned _assert_index_partitioning(all_indices, halo_indices, owned_indices) + @pytest.mark.skipif( - props.comm_size > 2, reason="input files available for 1 or 2 nodes" + props.comm_size > 2, reason="input files only available for 1 or 2 nodes" ) @pytest.mark.parametrize( ("dim, owned, total"), @@ -100,7 +88,7 @@ def test_decomposition_info_masked(mpi, datapath, dim, owned, total): (VertexDim, (5373, 5290), (5455, 5456)), ), ) -def test_decomposition_info_local_index(mpi, datapath, dim, owned, total): +def test_decomposition_info_local_index(mpi, datapath, dim, owned, total, caplog): props = get_processor_properties() my_rank = props.rank @@ -135,11 +123,19 @@ def _assert_index_partitioning(all_indices, halo_indices, owned_indices): assert set(halos_list) & set(all_list) == set(halos_list) assert set(halos_list) | set(owned_list) == set(all_list) +@pytest.mark.mpi +def test_processor_properties_from_comm_world(mpi, caplog): + caplog.set_level(logging.DEBUG) + props = get_processor_properties() + + assert props.rank < mpi.COMM_WORLD.Get_size() + assert props.comm_name == mpi.COMM_WORLD.Get_name() @pytest.mark.mpi -def test_decomposition_info_matches_gridsize(datapath): - decomposition_info = read_decomp_info(datapath, props, SerializationType.SB) +def test_decomposition_info_matches_gridsize(datapath, caplog): + props = get_processor_properties() + decomposition_info = read_decomp_info(datapath, props, SerializationType.SB,) icon_grid = read_icon_grid(datapath, props.rank) assert \ decomposition_info.global_index(dim=CellDim, entry_type=DecompositionInfo.EntryType.ALL).shape[ @@ -149,22 +145,32 @@ def test_decomposition_info_matches_gridsize(datapath): assert decomposition_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape[ 0] == icon_grid.num_edges() -#@pytest.mark.mpi -#@pytest.mark.skip -def test_parallel_diffusion(r04b09_diffusion_config, step_date_init, caplog): - caplog.set_level(logging.INFO) +@pytest.mark.mpi +def test_parallel_diffusion(r04b09_diffusion_config, step_date_init): + # caplog.set_level(logging.DEBUG) + props = get_processor_properties() + experiment_name = "mch_ch_r04b09_dsl" path = Path( f"/home/magdalena/data/exclaim/dycore/{experiment_name}/mpitasks2/{experiment_name}/ser_data" ) + print(f"rank={props.rank}/{props.comm_size}: inializing diffusion for experiment {experiment_name}") + - icon_grid = read_icon_grid(path, rank=props.rank) decomp_info = read_decomp_info( path, props, ) + print(f"rank={props.rank}/{props.comm_size}: decomposition info : klevels = {decomp_info.klevels}, " + f"local cells = {decomp_info.global_index(CellDim, DecompositionInfo.EntryType.ALL).shape} " + f"local edges = {decomp_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape} local vertices = {decomp_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape}") context = ghex.context(ghex.mpi_comm(props.comm), True) - #assert context.size() == 2 + print(f"rank={props.rank}/{props.comm_size}: GHEX context setup: from {props.comm_name} with {props.comm_size} nodes") + assert context.size() == 2 + + + icon_grid = read_icon_grid(path, rank=props.rank) + print(f"rank={props.rank}: using local grid with {icon_grid.num_cells()} Cells, {icon_grid.num_edges()} Edges, {icon_grid.num_vertices()} Vertices") diffusion_params = DiffusionParams(r04b09_diffusion_config) diffusion_initial_data = IconSerialDataProvider( @@ -178,10 +184,11 @@ def test_parallel_diffusion(r04b09_diffusion_config, step_date_init, caplog): ) dtime = diffusion_initial_data.get_metadata("dtime").get("dtime") + print(f"rank={props.rank}/{props.comm_size}: setup: using {props.comm_name} with {props.comm_size} nodes") exchange = Exchange(context, decomp_info) - print(f"exchange setup: for {exchange.get_size()} nodes") diffusion = Diffusion(exchange) + diffusion.init( icon_grid, r04b09_diffusion_config, @@ -190,7 +197,7 @@ def test_parallel_diffusion(r04b09_diffusion_config, step_date_init, caplog): metric_state, interpolation_state, ) - print("diffusion initialized") + print(f"rank={props.rank}/{props.comm_size}: diffusion initialized ") diffusion.run( diagnostic_state=diffusion_initial_data.construct_diagnostics(), prognostic_state=diffusion_initial_data.construct_prognostics(), @@ -204,3 +211,6 @@ def test_parallel_diffusion(r04b09_diffusion_config, step_date_init, caplog): edge_areas=edge_geometry.edge_areas, cell_areas=cell_geometry.area, ) + print(f"rank={props.rank}/{props.comm_size}: diffusion run ") + + From 2aa01b0e76ec7c98fff7c65dbdf347c7dba7c9b0 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 16 Jun 2023 10:39:36 +0200 Subject: [PATCH 137/263] try to understand out, when it crashes... --- atm_dyn_iconam/src/icon4py/diffusion/diffusion.py | 13 +++++++------ atm_dyn_iconam/src/icon4py/driver/parallel_setup.py | 1 + 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 159cda4db..9770055fa 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -973,12 +973,13 @@ def _do_diffusion_step( # TODO @magdalena why not trigger the exchange of w earlier? # TODO if condition: IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN # TODO magdalena: why does this crash? - # res = self._sync_fields((CellDim,KDim), - # prognostic_state.theta_v, - # prognostic_state.exner_pressure, - # ) - res = self._sync_fields((CellDim, KHalfDim), prognostic_state.w) - self._wait(res) + res_1 = self._sync_fields((CellDim,KDim), + prognostic_state.theta_v, + prognostic_state.exner_pressure, + ) + self._wait(res_1) + res_w = self._sync_fields((CellDim, KHalfDim), prognostic_state.w) + #self._wait(res_w) def _sync_fields(self, dim: tuple[Dimension, Dimension], *field): if self._exchange: diff --git a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py index 37b566469..83d97778a 100644 --- a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py @@ -27,6 +27,7 @@ mpi4py.rc.initialize = False +#mpi4py.rc.finalize = False CommId = Union[int, Comm, None] log = logging.getLogger(__name__) From 668b44364daca6b581a97720957308b7cf911eaa Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 16 Jun 2023 12:08:34 +0200 Subject: [PATCH 138/263] pre construct ghex communicators for each pattern --- atm_dyn_iconam/src/icon4py/driver/parallel_setup.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py index 83d97778a..1daaf7892 100644 --- a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py @@ -155,7 +155,9 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): (EdgeDim, KDim): self._create_pattern(EdgeDim, KDim), (EdgeDim, KHalfDim): self._create_pattern(EdgeDim, KHalfDim), } - print(f"rank={self._context.rank()}/{self._context.size()} : patterns initialized ") + + self._comms = {k: ghex.make_co(context, v) for k, v in self._patterns.items()} + print(f"rank={self._context.rank()}/{self._context.size()} : patterns and communicators initialized ") print(f"rank={self._context.rank()}/{self._context.size()} : exchange initialized") @@ -201,7 +203,7 @@ def exchange(self, dims: tuple[Dimension, Dimension], *fields): assert pattern is not None fields = [np.asarray(f)[:horizontal_size, :] for f in fields] print(f"rank = {self._context.rank()}/{self._context.size()}: communicating fields of dim = {dims} (shape = {fields[0].shape})") - communicator = ghex.make_co(self._context, pattern) + communicator = self._comms[dims] patterns_of_field = [pattern(ghex.field_descriptor(self._domain_descriptors[dims], f)) for f in fields] return communicator.exchange(patterns_of_field) From 6cbd6b0437d0fdde583725feecae68e0532bb702 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 16 Jun 2023 14:11:20 +0200 Subject: [PATCH 139/263] use int32 instead of unqualified int in domain args --- ...ute_horizontal_gradients_for_turbulance.py | 8 ++--- .../apply_nabla2_and_nabla4_to_vn.py | 10 +++--- ...te_diagnostic_quantities_for_turbulence.py | 10 +++--- ..._coefficients_for_grid_point_cold_pools.py | 10 +++--- ...ate_nabla2_and_smag_coefficients_for_vn.py | 10 +++--- .../calculate_nabla2_for_theta.py | 10 +++--- .../fused_mo_nh_diffusion_stencil_04_05_06.py | 8 ++--- .../fused_mo_nh_diffusion_stencil_13_14.py | 2 +- .../mo_intp_rbf_rbf_vec_interpol_vertex.py | 10 +++--- .../atm_dyn_iconam/update_theta_and_exner.py | 10 +++--- .../src/icon4py/diffusion/diffusion.py | 24 +++++++------- .../icon4py/diffusion/diffusion_program.py | 31 +++++++++---------- .../src/icon4py/diffusion/icon_grid.py | 6 ++-- 13 files changed, 75 insertions(+), 74 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py index 3a11f8008..49091e909 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py @@ -96,10 +96,10 @@ def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( nrdmax: int32, interior_idx: int32, halo_idx: int32, - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( area, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py index be7a0bbf8..5cbe2b123 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, maximum +from gt4py.next.ffront.fbuiltins import Field, int32, maximum from icon4py.common.dimension import EdgeDim, KDim @@ -39,10 +39,10 @@ def mo_nh_diffusion_stencil_05_global_mode( z_nabla4_e2: Field[[EdgeDim, KDim], float], diff_multfac_vn: Field[[KDim], float], vn: Field[[EdgeDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _mo_nh_diffusion_stencil_05_global_mode( area_edge, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py index 12d42bf80..9db5202cc 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulence import ( _calculate_diagnostics_for_turbulence, @@ -49,10 +49,10 @@ def calculate_diagnostic_quantities_for_turbulence( wgtfac_c: Field[[CellDim, KDim], float], div_ic: Field[[CellDim, KDim], float], hdef_ic: Field[[CellDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _calculate_diagnostic_quantities_for_turbulence( kh_smag_ec, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py index db2c74490..43dccd7e9 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.atm_dyn_iconam.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( _enhance_diffusion_coefficient_for_grid_point_cold_pools, @@ -45,10 +45,10 @@ def calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( theta_ref_mc: Field[[CellDim, KDim], float], thresh_tdiff: float, kh_smag_e: Field[[EdgeDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( theta_v, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py index 74becbc26..b1a281a45 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, maximum, minimum, sqrt +from gt4py.next.ffront.fbuiltins import Field, int32, maximum, minimum, sqrt from icon4py.common.dimension import ( E2C2V, @@ -153,10 +153,10 @@ def calculate_nabla2_and_smag_coefficients_for_vn( kh_smag_ec: Field[[EdgeDim, KDim], float], z_nabla2_e: Field[[EdgeDim, KDim], float], smag_offset: float, - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _calculate_nabla2_and_smag_coefficients_for_vn( diff_multfac_smag, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py index 3285918ed..4827cee88 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.atm_dyn_iconam.calculate_nabla2_for_z import _calculate_nabla2_for_z from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( @@ -40,10 +40,10 @@ def calculate_nabla2_for_theta( theta_v: Field[[CellDim, KDim], float], geofac_div: Field[[CEDim], float], z_temp: Field[[CellDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _calculate_nabla2_for_theta( kh_smag_e, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py index 897dca2e0..93bcc055e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py @@ -92,10 +92,10 @@ def fused_mo_nh_diffusion_stencil_04_05_06( nudgezone_diff: float, fac_bdydiff_v: float, start_2nd_nudge_line_idx_e: int32, - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _fused_mo_nh_diffusion_stencil_04_05_06( u_vert, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py index 45aa27a2c..d13d7c8e1 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -18,7 +18,7 @@ from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( _calculate_nabla2_of_theta, ) -from icon4py.common.dimension import CellDim, EdgeDim, KDim, CEDim +from icon4py.common.dimension import CEDim, CellDim, EdgeDim, KDim @field_operator diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py index 2d34018a6..0dbcf5c71 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, neighbor_sum +from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum from icon4py.common.dimension import V2E, EdgeDim, KDim, V2EDim, VertexDim @@ -35,10 +35,10 @@ def mo_intp_rbf_rbf_vec_interpol_vertex( ptr_coeff_2: Field[[VertexDim, V2EDim], float], p_u_out: Field[[VertexDim, KDim], float], p_v_out: Field[[VertexDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _mo_intp_rbf_rbf_vec_interpol_vertex( p_e_in, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py index c597ed3f2..53811c8aa 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.common.dimension import CellDim, KDim @@ -38,10 +38,10 @@ def update_theta_and_exner( theta_v: Field[[CellDim, KDim], float], exner: Field[[CellDim, KDim], float], rd_o_cvd: float, - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _update_theta_and_exner( z_temp, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index af58f4156..d99d39f83 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -23,22 +23,24 @@ from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn import icon4py.diffusion.diffusion_program as diff_prog -from icon4py.atm_dyn_iconam import calculate_diagnostic_quantities_for_turbulence -from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( - calculate_nabla2_and_smag_coefficients_for_vn, -) -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import ( - fused_mo_nh_diffusion_stencil_04_05_06, -) -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_07_08_09_10 import ( +from icon4py.atm_dyn_iconam.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance, ) -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_11_12 import ( +from icon4py.atm_dyn_iconam.calculate_diagnostic_quantities_for_turbulence import ( + calculate_diagnostic_quantities_for_turbulence, +) +from icon4py.atm_dyn_iconam.calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools import ( calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools, ) -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_13_14 import ( +from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( + calculate_nabla2_and_smag_coefficients_for_vn, +) +from icon4py.atm_dyn_iconam.calculate_nabla2_for_theta import ( calculate_nabla2_for_theta, ) +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import ( + fused_mo_nh_diffusion_stencil_04_05_06, +) from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) @@ -815,7 +817,7 @@ def _do_diffusion_step( set_zero_v_k.with_backend(run_gtfn)(self.v_vert, offset_provider={}) log.debug("rbf interpolation: start") # # 1. CALL rbf_vec_interpol_vertex - mo_intp_rbf_rbf_vec_interpol_vertex( + mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(run_gtfn)( p_e_in=prognostic_state.vn, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, ptr_coeff_2=self.interpolation_state.rbf_coeff_2, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py index 5085e524e..a4f74c4c7 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py @@ -16,7 +16,6 @@ from gt4py.next.common import Field from gt4py.next.ffront.decorator import program from gt4py.next.ffront.fbuiltins import int32 -from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance, @@ -24,15 +23,15 @@ from icon4py.atm_dyn_iconam.calculate_diagnostic_quantities_for_turbulence import ( _calculate_diagnostic_quantities_for_turbulence, ) +from icon4py.atm_dyn_iconam.calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools import ( + _calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools, +) from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( _calculate_nabla2_and_smag_coefficients_for_vn, ) from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import ( _fused_mo_nh_diffusion_stencil_04_05_06, ) -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_11_12 import ( - _calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools, -) from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( _mo_intp_rbf_rbf_vec_interpol_vertex, ) @@ -110,20 +109,20 @@ def diffusion_run( local_horizontal_edge_index: Field[[EdgeDim], int32], cell_startindex_interior: int32, cell_halo_idx: int32, - cell_startindex_nudging: int, - cell_endindex_local_plus1: int, - cell_endindex_local: int, - edge_startindex_nudging_plus1: int, + cell_startindex_nudging: int32, + cell_endindex_local_plus1: int32, + cell_endindex_local: int32, + edge_startindex_nudging_plus1: int32, edge_startindex_nudging_minus1: int32, - edge_endindex_local: int, - edge_endindex_local_minus2: int, - vertex_startindex_lb_plus3: int, - vertex_startindex_lb_plus1: int, - vertex_endindex_local: int, - vertex_endindex_local_minus1: int, + edge_endindex_local: int32, + edge_endindex_local_minus2: int32, + vertex_startindex_lb_plus3: int32, + vertex_startindex_lb_plus1: int32, + vertex_endindex_local: int32, + vertex_endindex_local_minus1: int32, index_of_damping_height: int32, - nlev: int, - boundary_diffusion_start_index_edges: int, + nlev: int32, + boundary_diffusion_start_index_edges: int32, ): _scale_k(local_enh_smag_fac, dtime, out=local_diff_multfac_smag) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index 07f6625e5..2cc3dc56f 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -100,8 +100,8 @@ def with_config(self, config: MeshConfig): def with_start_end_indices( self, dim: Dimension, start_indices: np.ndarray, end_indices: np.ndarray ): - self.start_indices[dim] = start_indices.astype(int) - self.end_indices[dim] = end_indices.astype(int) + self.start_indices[dim] = start_indices.astype(int32) + self.end_indices[dim] = end_indices.astype(int32) @builder def with_connectivities(self, connectivity: Dict[Dimension, np.ndarray]): @@ -128,7 +128,7 @@ def num_edges(self): def get_indices_from_to( self, dim: Dimension, start_marker: int, end_marker: int - ) -> Tuple[int, int]: + ) -> Tuple[int32, int32]: """ Use to specify domains of a field for field_operator. From 505355cc3aa649614f30b3e48654bc988feee156 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 16 Jun 2023 14:27:46 +0200 Subject: [PATCH 140/263] fixes after merge of fix_greenline_remove_int_type --- .../atm_dyn_iconam/apply_nabla2_to_w.py | 10 +-- ...pply_nabla2_to_w_in_upper_damping_layer.py | 10 +-- ...ate_horizontal_gradients_for_turbulence.py | 10 +-- .../atm_dyn_iconam/calculate_nabla2_for_w.py | 10 +-- ...fusion_nabla_of_theta_over_steep_points.py | 8 +- .../src/icon4py/diffusion/diffusion.py | 29 +++--- .../src/icon4py/driver/parallel_setup.py | 90 ++++++++++++------- .../tests/mpi_tests/test_parallel_setup.py | 73 +++++++++------ 8 files changed, 147 insertions(+), 93 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py index 93bb21450..437c10003 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, neighbor_sum +from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim @@ -38,10 +38,10 @@ def apply_nabla2_to_w( geofac_n2s: Field[[CellDim, C2E2CODim], float], w: Field[[CellDim, KDim], float], diff_multfac_w: float, - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _apply_nabla2_to_w( area, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py index 0483c7d7b..345f7caf3 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.common.dimension import CellDim, KDim @@ -34,10 +34,10 @@ def apply_nabla2_to_w_in_upper_damping_layer( diff_multfac_n2w: Field[[KDim], float], cell_area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _apply_nabla2_to_w_in_upper_damping_layer( w, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py index ce8a44758..5fe2c465d 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, neighbor_sum +from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim @@ -35,10 +35,10 @@ def calculate_horizontal_gradients_for_turbulence( geofac_grg_y: Field[[CellDim, C2E2CODim], float], dwdx: Field[[CellDim, KDim], float], dwdy: Field[[CellDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _calculate_horizontal_gradients_for_turbulence( w, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py index 2a9dc2c41..5347a8e4e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, neighbor_sum +from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim @@ -30,10 +30,10 @@ def calculate_nabla2_for_w( w: Field[[CellDim, KDim], float], geofac_n2s: Field[[CellDim, C2E2CODim], float], z_nabla2_c: Field[[CellDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _calculate_nabla2_for_w( w, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index f79aa8f07..663fa0efe 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -75,10 +75,10 @@ def truly_horizontal_diffusion_nabla_of_theta_over_steep_points( vcoef: Field[[CECDim, KDim], float], theta_v: Field[[CellDim, KDim], float], z_temp: Field[[CellDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _truly_horizontal_diffusion_nabla_of_theta_over_steep_points( mask, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index a4260e234..57f83f5cb 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -23,7 +23,7 @@ from gt4py.next.program_processors.runners.gtfn_cpu import ( run_gtfn, run_gtfn_cached, -run_gtfn_imperative, + run_gtfn_imperative, ) from icon4py.atm_dyn_iconam.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( @@ -64,7 +64,14 @@ DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO, GAS_CONSTANT_DRY_AIR, ) -from icon4py.common.dimension import CellDim, ECVDim, EdgeDim, KDim, VertexDim, KHalfDim +from icon4py.common.dimension import ( + CellDim, + ECVDim, + EdgeDim, + KDim, + KHalfDim, + VertexDim, +) from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.horizontal import HorizontalMarkerIndex from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams @@ -379,7 +386,6 @@ def __init__(self, exchange: Optional[Exchange] = None): self.thresh_tdiff: float = ( -5.0 ) # threshold temperature deviation from neighboring grid points hat activates extra diffusion against runaway cooling - self._run_program = run_program self.grid: Optional[IconGrid] = None self.config: Optional[DiffusionConfig] = None self.params: Optional[DiffusionParams] = None @@ -887,10 +893,9 @@ def _do_diffusion_step( # ) log.debug("running fused stencil 07_08_09_10: end") # # 8. HALO EXCHANGE: CALL sync_patch_array - comm_res = self._sync_fields((EdgeDim,KDim), prognostic_state.vn) + comm_res = self._sync_fields((EdgeDim, KDim), prognostic_state.vn) self._wait(comm_res) - # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, # # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 # @@ -974,21 +979,19 @@ def _do_diffusion_step( # TODO @magdalena why not trigger the exchange of w earlier? # TODO if condition: IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN # TODO magdalena: why does this crash? - res_1 = self._sync_fields((CellDim,KDim), - prognostic_state.theta_v, - prognostic_state.exner_pressure, - ) + res_1 = self._sync_fields( + (CellDim, KDim), + prognostic_state.theta_v, + prognostic_state.exner_pressure, + ) self._wait(res_1) res_w = self._sync_fields((CellDim, KHalfDim), prognostic_state.w) - #self._wait(res_w) + # self._wait(res_w) def _sync_fields(self, dim: tuple[Dimension, Dimension], *field): if self._exchange: return self._exchange.exchange(dim, *field) - - - def _wait(self, comm_handle): if comm_handle: comm_handle.wait() diff --git a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py index 1daaf7892..85e1e56b1 100644 --- a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py @@ -21,13 +21,13 @@ from gt4py.next.common import Dimension from mpi4py.MPI import Comm -from icon4py.common.dimension import CellDim, VertexDim, EdgeDim, KDim, KHalfDim +from icon4py.common.dimension import CellDim, EdgeDim, KDim, KHalfDim, VertexDim from icon4py.decomposition.decomposed import ProcessProperties from icon4py.diffusion.utils import builder mpi4py.rc.initialize = False -#mpi4py.rc.finalize = False +# mpi4py.rc.finalize = False CommId = Union[int, Comm, None] log = logging.getLogger(__name__) @@ -59,10 +59,12 @@ def init_mpi(): def finalize_mpi(): from mpi4py import MPI + if not MPI.Is_finalized(): log.info("finalizing MPI") MPI.Finalize() + class DecompositionInfo: class EntryType(int, Enum): ALL = (0,) @@ -121,17 +123,28 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): self._context = context self._decomposition_info = domain_decomposition self._domain_descriptors = { - (CellDim, KDim): self._create_domain_descriptor(CellDim, domain_decomposition.klevels), - (CellDim, KHalfDim): self._create_domain_descriptor(CellDim, domain_decomposition.klevels + 1), - (VertexDim, KDim): self._create_domain_descriptor(VertexDim, domain_decomposition.klevels), - (VertexDim, KHalfDim): self._create_domain_descriptor(VertexDim, - domain_decomposition.klevels + 1), - (EdgeDim, KDim): self._create_domain_descriptor(EdgeDim, domain_decomposition.klevels), - (EdgeDim, KHalfDim): self._create_domain_descriptor(EdgeDim, - domain_decomposition.klevels + 1), - + (CellDim, KDim): self._create_domain_descriptor( + CellDim, domain_decomposition.klevels + ), + (CellDim, KHalfDim): self._create_domain_descriptor( + CellDim, domain_decomposition.klevels + 1 + ), + (VertexDim, KDim): self._create_domain_descriptor( + VertexDim, domain_decomposition.klevels + ), + (VertexDim, KHalfDim): self._create_domain_descriptor( + VertexDim, domain_decomposition.klevels + 1 + ), + (EdgeDim, KDim): self._create_domain_descriptor( + EdgeDim, domain_decomposition.klevels + ), + (EdgeDim, KHalfDim): self._create_domain_descriptor( + EdgeDim, domain_decomposition.klevels + 1 + ), } - print(f"rank={self._context.rank()}/{self._context.size()} :domain descriptors initialized") + print( + f"rank={self._context.rank()}/{self._context.size()} :domain descriptors initialized" + ) self._field_size = { CellDim: self._decomposition_info.global_index( CellDim, DecompositionInfo.EntryType.ALL @@ -143,9 +156,11 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): VertexDim, DecompositionInfo.EntryType.ALL ).shape[0], KDim: domain_decomposition.klevels, - KHalfDim: domain_decomposition.klevels + 1 + KHalfDim: domain_decomposition.klevels + 1, } - print(f"rank={self._context.rank()}/{self._context.size()} : field sizes = {self._field_size}") + print( + f"rank={self._context.rank()}/{self._context.size()} : field sizes = {self._field_size}" + ) self._patterns = { (CellDim, KDim): self._create_pattern(CellDim, KDim), @@ -157,17 +172,21 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): } self._comms = {k: ghex.make_co(context, v) for k, v in self._patterns.items()} - print(f"rank={self._context.rank()}/{self._context.size()} : patterns and communicators initialized ") - print(f"rank={self._context.rank()}/{self._context.size()} : exchange initialized") - + print( + f"rank={self._context.rank()}/{self._context.size()} : patterns and communicators initialized " + ) + print( + f"rank={self._context.rank()}/{self._context.size()} : exchange initialized" + ) def _domain_descriptor_info(self, descr): return f" id={descr.domain_id()} levels = {descr.levels()}, size={descr.size()}, inner_size={descr.inner_size()}" + def get_size(self): return self._context.size() # TODO [magdalena] is the tolist() necessary? - def _create_domain_descriptor(self, dim: Dimension, levels:int): + def _create_domain_descriptor(self, dim: Dimension, levels: int): all_global = self._decomposition_info.global_index( dim, DecompositionInfo.EntryType.ALL ) @@ -175,25 +194,30 @@ def _create_domain_descriptor(self, dim: Dimension, levels:int): dim, DecompositionInfo.EntryType.HALO ) domain_desc = ghex.domain_descriptor( - self._context.rank(), - all_global.tolist(), - local_halo.tolist(), - levels + self._context.rank(), all_global.tolist(), local_halo.tolist(), levels + ) + print( + f"rank={self._context.rank()}/{self._context.size()}: domain descriptor for dim {dim} with properties {self._domain_descriptor_info(domain_desc)}" ) - print(f"rank={self._context.rank()}/{self._context.size()}: domain descriptor for dim {dim} with properties {self._domain_descriptor_info(domain_desc)}") return domain_desc - def _create_pattern(self, horizontal_dim:Dimension, vertical_dim:Dimension): + def _create_pattern(self, horizontal_dim: Dimension, vertical_dim: Dimension): halo_generator = ghex.halo_generator_with_gids( - self._decomposition_info.global_index(horizontal_dim, DecompositionInfo.EntryType.HALO) + self._decomposition_info.global_index( + horizontal_dim, DecompositionInfo.EntryType.HALO + ) + ) + print( + f"rank={self._context.rank()}/{self._context.size()}: halo generator for dim={horizontal_dim} created" ) - print(f"rank={self._context.rank()}/{self._context.size()}: halo generator for dim={horizontal_dim} created") dimensions = (horizontal_dim, vertical_dim) pattern = ghex.make_pattern( - self._context, halo_generator,[self._domain_descriptors[dimensions]] + self._context, halo_generator, [self._domain_descriptors[dimensions]] + ) + print( + f"rank={self._context.rank()}/{self._context.size()}: pattern for dim={dimensions} and {self._domain_descriptors[dimensions]} created" ) - print(f"rank={self._context.rank()}/{self._context.size()}: pattern for dim={dimensions} and {self._domain_descriptors[dimensions]} created") return pattern def exchange(self, dims: tuple[Dimension, Dimension], *fields): @@ -202,8 +226,12 @@ def exchange(self, dims: tuple[Dimension, Dimension], *fields): pattern = self._patterns[dims] assert pattern is not None fields = [np.asarray(f)[:horizontal_size, :] for f in fields] - print(f"rank = {self._context.rank()}/{self._context.size()}: communicating fields of dim = {dims} (shape = {fields[0].shape})") + print( + f"rank = {self._context.rank()}/{self._context.size()}: communicating fields of dim = {dims} (shape = {fields[0].shape})" + ) communicator = self._comms[dims] - patterns_of_field = [pattern(ghex.field_descriptor(self._domain_descriptors[dims], f)) for f in fields] + patterns_of_field = [ + pattern(ghex.field_descriptor(self._domain_descriptors[dims], f)) + for f in fields + ] return communicator.exchange(patterns_of_field) - diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index 5ffdf4d6d..f98be6c5b 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -30,7 +30,8 @@ from icon4py.driver.parallel_setup import ( DecompositionInfo, Exchange, - get_processor_properties, finalize_mpi, + finalize_mpi, + get_processor_properties, ) @@ -45,6 +46,8 @@ """ props = get_processor_properties() + + @pytest.mark.skipif( props.comm_size > 2, reason="input files only available for 1 or 2 nodes" ) @@ -113,7 +116,6 @@ def test_decomposition_info_local_index(mpi, datapath, dim, owned, total, caplog _assert_index_partitioning(all_indices, halo_indices, owned_indices) - def _assert_index_partitioning(all_indices, halo_indices, owned_indices): owned_list = owned_indices.tolist() halos_list = halo_indices.tolist() @@ -123,6 +125,7 @@ def _assert_index_partitioning(all_indices, halo_indices, owned_indices): assert set(halos_list) & set(all_list) == set(halos_list) assert set(halos_list) | set(owned_list) == set(all_list) + @pytest.mark.mpi def test_processor_properties_from_comm_world(mpi, caplog): caplog.set_level(logging.DEBUG) @@ -135,42 +138,64 @@ def test_processor_properties_from_comm_world(mpi, caplog): @pytest.mark.mpi def test_decomposition_info_matches_gridsize(datapath, caplog): props = get_processor_properties() - decomposition_info = read_decomp_info(datapath, props, SerializationType.SB,) + decomposition_info = read_decomp_info( + datapath, + props, + SerializationType.SB, + ) icon_grid = read_icon_grid(datapath, props.rank) - assert \ - decomposition_info.global_index(dim=CellDim, entry_type=DecompositionInfo.EntryType.ALL).shape[ - 0] == icon_grid.num_cells() - assert decomposition_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape[ - 0] == icon_grid.num_vertices() - assert decomposition_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape[ - 0] == icon_grid.num_edges() + assert ( + decomposition_info.global_index( + dim=CellDim, entry_type=DecompositionInfo.EntryType.ALL + ).shape[0] + == icon_grid.num_cells() + ) + assert ( + decomposition_info.global_index( + VertexDim, DecompositionInfo.EntryType.ALL + ).shape[0] + == icon_grid.num_vertices() + ) + assert ( + decomposition_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape[ + 0 + ] + == icon_grid.num_edges() + ) + @pytest.mark.mpi def test_parallel_diffusion(r04b09_diffusion_config, step_date_init): - # caplog.set_level(logging.DEBUG) + # caplog.set_level(logging.DEBUG) props = get_processor_properties() experiment_name = "mch_ch_r04b09_dsl" path = Path( f"/home/magdalena/data/exclaim/dycore/{experiment_name}/mpitasks2/{experiment_name}/ser_data" ) - print(f"rank={props.rank}/{props.comm_size}: inializing diffusion for experiment {experiment_name}") - + print( + f"rank={props.rank}/{props.comm_size}: inializing diffusion for experiment {experiment_name}" + ) decomp_info = read_decomp_info( path, props, ) - print(f"rank={props.rank}/{props.comm_size}: decomposition info : klevels = {decomp_info.klevels}, " - f"local cells = {decomp_info.global_index(CellDim, DecompositionInfo.EntryType.ALL).shape} " - f"local edges = {decomp_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape} local vertices = {decomp_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape}") + print( + f"rank={props.rank}/{props.comm_size}: decomposition info : klevels = {decomp_info.klevels}, " + f"local cells = {decomp_info.global_index(CellDim, DecompositionInfo.EntryType.ALL).shape} " + f"local edges = {decomp_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape} local vertices = {decomp_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape}" + ) context = ghex.context(ghex.mpi_comm(props.comm), True) - print(f"rank={props.rank}/{props.comm_size}: GHEX context setup: from {props.comm_name} with {props.comm_size} nodes") + print( + f"rank={props.rank}/{props.comm_size}: GHEX context setup: from {props.comm_name} with {props.comm_size} nodes" + ) assert context.size() == 2 - icon_grid = read_icon_grid(path, rank=props.rank) - print(f"rank={props.rank}: using local grid with {icon_grid.num_cells()} Cells, {icon_grid.num_edges()} Edges, {icon_grid.num_vertices()} Vertices") + print( + f"rank={props.rank}: using local grid with {icon_grid.num_cells()} Cells, {icon_grid.num_edges()} Edges, {icon_grid.num_vertices()} Vertices" + ) diffusion_params = DiffusionParams(r04b09_diffusion_config) diffusion_initial_data = IconSerialDataProvider( @@ -179,12 +204,12 @@ def test_parallel_diffusion(r04b09_diffusion_config, step_date_init): (edge_geometry, cell_geometry, vertical_geometry) = read_geometry_fields( path, rank=props.rank ) - (metric_state, interpolation_state) = read_static_fields( - path, rank=props.rank - ) + (metric_state, interpolation_state) = read_static_fields(path, rank=props.rank) dtime = diffusion_initial_data.get_metadata("dtime").get("dtime") - print(f"rank={props.rank}/{props.comm_size}: setup: using {props.comm_name} with {props.comm_size} nodes") + print( + f"rank={props.rank}/{props.comm_size}: setup: using {props.comm_name} with {props.comm_size} nodes" + ) exchange = Exchange(context, decomp_info) diffusion = Diffusion(exchange) @@ -212,5 +237,3 @@ def test_parallel_diffusion(r04b09_diffusion_config, step_date_init): cell_areas=cell_geometry.area, ) print(f"rank={props.rank}/{props.comm_size}: diffusion run ") - - From e473cff8fafddf0c0f035f7d34605a14b3d7049e Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 16 Jun 2023 14:46:52 +0200 Subject: [PATCH 141/263] add wait after exchange call (follow up later) --- atm_dyn_iconam/src/icon4py/diffusion/diffusion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 57f83f5cb..447414981 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -986,7 +986,7 @@ def _do_diffusion_step( ) self._wait(res_1) res_w = self._sync_fields((CellDim, KHalfDim), prognostic_state.w) - # self._wait(res_w) + self._wait(res_w) def _sync_fields(self, dim: tuple[Dimension, Dimension], *field): if self._exchange: From 7b853d52c1f0fdd90c98366077e7bb5bfe9e98ce Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 16 Jun 2023 15:15:19 +0200 Subject: [PATCH 142/263] read exit data for verification --- .../tests/mpi_tests/test_parallel_setup.py | 53 ++++++++++++++++--- 1 file changed, 47 insertions(+), 6 deletions(-) diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index f98be6c5b..e66ac002d 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -168,10 +168,12 @@ def test_decomposition_info_matches_gridsize(datapath, caplog): def test_parallel_diffusion(r04b09_diffusion_config, step_date_init): # caplog.set_level(logging.DEBUG) props = get_processor_properties() + num_nodes = props.comm_size experiment_name = "mch_ch_r04b09_dsl" + path = Path( - f"/home/magdalena/data/exclaim/dycore/{experiment_name}/mpitasks2/{experiment_name}/ser_data" + f"/home/magdalena/data/exclaim/dycore/{experiment_name}/mpitasks{num_nodes}/{experiment_name}/ser_data" ) print( f"rank={props.rank}/{props.comm_size}: inializing diffusion for experiment {experiment_name}" @@ -190,17 +192,19 @@ def test_parallel_diffusion(r04b09_diffusion_config, step_date_init): print( f"rank={props.rank}/{props.comm_size}: GHEX context setup: from {props.comm_name} with {props.comm_size} nodes" ) - assert context.size() == 2 + #assert context.size() == 2 icon_grid = read_icon_grid(path, rank=props.rank) print( f"rank={props.rank}: using local grid with {icon_grid.num_cells()} Cells, {icon_grid.num_edges()} Edges, {icon_grid.num_vertices()} Vertices" ) - + initial_run = False + r04b09_diffusion_config.ndyn_substeps = 2 diffusion_params = DiffusionParams(r04b09_diffusion_config) + diffusion_initial_data = IconSerialDataProvider( "icon_pydycore", str(path), True, mpi_rank=props.rank - ).from_savepoint_diffusion_init(linit=True, date=step_date_init) + ).from_savepoint_diffusion_init(linit=initial_run, date=step_date_init) (edge_geometry, cell_geometry, vertical_geometry) = read_geometry_fields( path, rank=props.rank ) @@ -223,9 +227,11 @@ def test_parallel_diffusion(r04b09_diffusion_config, step_date_init): interpolation_state, ) print(f"rank={props.rank}/{props.comm_size}: diffusion initialized ") + diagnostic_state = diffusion_initial_data.construct_diagnostics() + prognostic_state = diffusion_initial_data.construct_prognostics() diffusion.run( - diagnostic_state=diffusion_initial_data.construct_diagnostics(), - prognostic_state=diffusion_initial_data.construct_prognostics(), + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, dtime=dtime, tangent_orientation=edge_geometry.tangent_orientation, inverse_primal_edge_lengths=edge_geometry.inverse_primal_edge_lengths, @@ -237,3 +243,38 @@ def test_parallel_diffusion(r04b09_diffusion_config, step_date_init): cell_areas=cell_geometry.area, ) print(f"rank={props.rank}/{props.comm_size}: diffusion run ") + + diffusion_savepoint_exit = IconSerialDataProvider( + "icon_pydycore", str(path), True, mpi_rank=props.rank + ).from_savepoint_diffusion_init(linit=initial_run, date=step_date_init) + verify_fields(diffusion_savepoint_exit, diagnostic_state, prognostic_state, + diffusion.metric_state.mask_hdiff) + + +def verify_fields(diffusion_savepoint_exit, diagnostic_state, prognostic_state, steep_points): + ref_div_ic = np.asarray(diffusion_savepoint_exit.div_ic()) + val_div_ic = np.asarray(diagnostic_state.div_ic) + ref_hdef_ic = np.asarray(diffusion_savepoint_exit.hdef_ic()) + val_hdef_ic = np.asarray(diagnostic_state.hdef_ic) + assert np.allclose(ref_div_ic, val_div_ic) + assert np.allclose(ref_hdef_ic, val_hdef_ic) + ref_w = np.asarray(diffusion_savepoint_exit.w()) + val_w = np.asarray(prognostic_state.w) + ref_dwdx = np.asarray(diffusion_savepoint_exit.dwdx()) + val_dwdx = np.asarray(prognostic_state.dwdx) + ref_dwdy = np.asarray(diffusion_savepoint_exit.dwdy()) + val_dwdy = np.asarray(prognostic_state.dwdy) + ref_vn = np.asarray(diffusion_savepoint_exit.vn()) + val_vn = np.asarray(prognostic_state.vn) + assert np.allclose(ref_vn, val_vn) + assert np.allclose(ref_dwdx, val_dwdx) + assert np.allclose(ref_dwdy, val_dwdy) + assert np.allclose(ref_w, val_w) + ref_exner = np.asarray(diffusion_savepoint_exit.exner()) + ref_theta_v = np.asarray(diffusion_savepoint_exit.theta_v()) + val_theta_v = np.asarray(prognostic_state.theta_v) + val_exner = np.asarray(prognostic_state.exner_pressure) + assert np.allclose(ref_theta_v[~steep_points], val_theta_v[~steep_points]) + assert np.allclose(ref_exner[~steep_points], val_exner[~steep_points]) + # assert np.allclose(ref_theta_v[steep_points], val_theta_v[steep_points]) + # assert np.allclose(ref_exner[steep_points], val_exner[steep_points]) From 396936eefc66abc573eba96e4c04d9ec7baee5cf Mon Sep 17 00:00:00 2001 From: Magdalena Date: Tue, 20 Jun 2023 08:23:16 +0200 Subject: [PATCH 143/263] use int32 instead of unqualified int in domain args (#231) --- ...ute_horizontal_gradients_for_turbulance.py | 8 ++--- .../apply_nabla2_and_nabla4_to_vn.py | 10 +++--- ...te_diagnostic_quantities_for_turbulence.py | 10 +++--- ..._coefficients_for_grid_point_cold_pools.py | 10 +++--- ...ate_nabla2_and_smag_coefficients_for_vn.py | 10 +++--- .../calculate_nabla2_for_theta.py | 10 +++--- .../fused_mo_nh_diffusion_stencil_04_05_06.py | 8 ++--- .../fused_mo_nh_diffusion_stencil_13_14.py | 2 +- .../mo_intp_rbf_rbf_vec_interpol_vertex.py | 10 +++--- .../atm_dyn_iconam/update_theta_and_exner.py | 10 +++--- .../src/icon4py/diffusion/diffusion.py | 24 +++++++------- .../icon4py/diffusion/diffusion_program.py | 31 +++++++++---------- .../src/icon4py/diffusion/icon_grid.py | 6 ++-- 13 files changed, 75 insertions(+), 74 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py index 3a11f8008..49091e909 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py @@ -96,10 +96,10 @@ def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( nrdmax: int32, interior_idx: int32, halo_idx: int32, - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( area, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py index be7a0bbf8..5cbe2b123 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, maximum +from gt4py.next.ffront.fbuiltins import Field, int32, maximum from icon4py.common.dimension import EdgeDim, KDim @@ -39,10 +39,10 @@ def mo_nh_diffusion_stencil_05_global_mode( z_nabla4_e2: Field[[EdgeDim, KDim], float], diff_multfac_vn: Field[[KDim], float], vn: Field[[EdgeDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _mo_nh_diffusion_stencil_05_global_mode( area_edge, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py index 12d42bf80..9db5202cc 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulence import ( _calculate_diagnostics_for_turbulence, @@ -49,10 +49,10 @@ def calculate_diagnostic_quantities_for_turbulence( wgtfac_c: Field[[CellDim, KDim], float], div_ic: Field[[CellDim, KDim], float], hdef_ic: Field[[CellDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _calculate_diagnostic_quantities_for_turbulence( kh_smag_ec, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py index db2c74490..43dccd7e9 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.atm_dyn_iconam.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( _enhance_diffusion_coefficient_for_grid_point_cold_pools, @@ -45,10 +45,10 @@ def calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( theta_ref_mc: Field[[CellDim, KDim], float], thresh_tdiff: float, kh_smag_e: Field[[EdgeDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( theta_v, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py index 74becbc26..b1a281a45 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_and_smag_coefficients_for_vn.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, maximum, minimum, sqrt +from gt4py.next.ffront.fbuiltins import Field, int32, maximum, minimum, sqrt from icon4py.common.dimension import ( E2C2V, @@ -153,10 +153,10 @@ def calculate_nabla2_and_smag_coefficients_for_vn( kh_smag_ec: Field[[EdgeDim, KDim], float], z_nabla2_e: Field[[EdgeDim, KDim], float], smag_offset: float, - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _calculate_nabla2_and_smag_coefficients_for_vn( diff_multfac_smag, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py index 3285918ed..4827cee88 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.atm_dyn_iconam.calculate_nabla2_for_z import _calculate_nabla2_for_z from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( @@ -40,10 +40,10 @@ def calculate_nabla2_for_theta( theta_v: Field[[CellDim, KDim], float], geofac_div: Field[[CEDim], float], z_temp: Field[[CellDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _calculate_nabla2_for_theta( kh_smag_e, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py index 897dca2e0..93bcc055e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py @@ -92,10 +92,10 @@ def fused_mo_nh_diffusion_stencil_04_05_06( nudgezone_diff: float, fac_bdydiff_v: float, start_2nd_nudge_line_idx_e: int32, - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _fused_mo_nh_diffusion_stencil_04_05_06( u_vert, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py index 45aa27a2c..d13d7c8e1 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -18,7 +18,7 @@ from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( _calculate_nabla2_of_theta, ) -from icon4py.common.dimension import CellDim, EdgeDim, KDim, CEDim +from icon4py.common.dimension import CEDim, CellDim, EdgeDim, KDim @field_operator diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py index 2d34018a6..0dbcf5c71 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, neighbor_sum +from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum from icon4py.common.dimension import V2E, EdgeDim, KDim, V2EDim, VertexDim @@ -35,10 +35,10 @@ def mo_intp_rbf_rbf_vec_interpol_vertex( ptr_coeff_2: Field[[VertexDim, V2EDim], float], p_u_out: Field[[VertexDim, KDim], float], p_v_out: Field[[VertexDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _mo_intp_rbf_rbf_vec_interpol_vertex( p_e_in, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py index c597ed3f2..53811c8aa 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/update_theta_and_exner.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.common.dimension import CellDim, KDim @@ -38,10 +38,10 @@ def update_theta_and_exner( theta_v: Field[[CellDim, KDim], float], exner: Field[[CellDim, KDim], float], rd_o_cvd: float, - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _update_theta_and_exner( z_temp, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index af58f4156..d99d39f83 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -23,22 +23,24 @@ from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn import icon4py.diffusion.diffusion_program as diff_prog -from icon4py.atm_dyn_iconam import calculate_diagnostic_quantities_for_turbulence -from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( - calculate_nabla2_and_smag_coefficients_for_vn, -) -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import ( - fused_mo_nh_diffusion_stencil_04_05_06, -) -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_07_08_09_10 import ( +from icon4py.atm_dyn_iconam.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance, ) -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_11_12 import ( +from icon4py.atm_dyn_iconam.calculate_diagnostic_quantities_for_turbulence import ( + calculate_diagnostic_quantities_for_turbulence, +) +from icon4py.atm_dyn_iconam.calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools import ( calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools, ) -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_13_14 import ( +from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( + calculate_nabla2_and_smag_coefficients_for_vn, +) +from icon4py.atm_dyn_iconam.calculate_nabla2_for_theta import ( calculate_nabla2_for_theta, ) +from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import ( + fused_mo_nh_diffusion_stencil_04_05_06, +) from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) @@ -815,7 +817,7 @@ def _do_diffusion_step( set_zero_v_k.with_backend(run_gtfn)(self.v_vert, offset_provider={}) log.debug("rbf interpolation: start") # # 1. CALL rbf_vec_interpol_vertex - mo_intp_rbf_rbf_vec_interpol_vertex( + mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(run_gtfn)( p_e_in=prognostic_state.vn, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, ptr_coeff_2=self.interpolation_state.rbf_coeff_2, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py index 5085e524e..a4f74c4c7 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py @@ -16,7 +16,6 @@ from gt4py.next.common import Field from gt4py.next.ffront.decorator import program from gt4py.next.ffront.fbuiltins import int32 -from gt4py.next.program_processors.runners import gtfn_cpu from icon4py.atm_dyn_iconam.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance, @@ -24,15 +23,15 @@ from icon4py.atm_dyn_iconam.calculate_diagnostic_quantities_for_turbulence import ( _calculate_diagnostic_quantities_for_turbulence, ) +from icon4py.atm_dyn_iconam.calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools import ( + _calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools, +) from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( _calculate_nabla2_and_smag_coefficients_for_vn, ) from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import ( _fused_mo_nh_diffusion_stencil_04_05_06, ) -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_11_12 import ( - _calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools, -) from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( _mo_intp_rbf_rbf_vec_interpol_vertex, ) @@ -110,20 +109,20 @@ def diffusion_run( local_horizontal_edge_index: Field[[EdgeDim], int32], cell_startindex_interior: int32, cell_halo_idx: int32, - cell_startindex_nudging: int, - cell_endindex_local_plus1: int, - cell_endindex_local: int, - edge_startindex_nudging_plus1: int, + cell_startindex_nudging: int32, + cell_endindex_local_plus1: int32, + cell_endindex_local: int32, + edge_startindex_nudging_plus1: int32, edge_startindex_nudging_minus1: int32, - edge_endindex_local: int, - edge_endindex_local_minus2: int, - vertex_startindex_lb_plus3: int, - vertex_startindex_lb_plus1: int, - vertex_endindex_local: int, - vertex_endindex_local_minus1: int, + edge_endindex_local: int32, + edge_endindex_local_minus2: int32, + vertex_startindex_lb_plus3: int32, + vertex_startindex_lb_plus1: int32, + vertex_endindex_local: int32, + vertex_endindex_local_minus1: int32, index_of_damping_height: int32, - nlev: int, - boundary_diffusion_start_index_edges: int, + nlev: int32, + boundary_diffusion_start_index_edges: int32, ): _scale_k(local_enh_smag_fac, dtime, out=local_diff_multfac_smag) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index 07f6625e5..2cc3dc56f 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -100,8 +100,8 @@ def with_config(self, config: MeshConfig): def with_start_end_indices( self, dim: Dimension, start_indices: np.ndarray, end_indices: np.ndarray ): - self.start_indices[dim] = start_indices.astype(int) - self.end_indices[dim] = end_indices.astype(int) + self.start_indices[dim] = start_indices.astype(int32) + self.end_indices[dim] = end_indices.astype(int32) @builder def with_connectivities(self, connectivity: Dict[Dimension, np.ndarray]): @@ -128,7 +128,7 @@ def num_edges(self): def get_indices_from_to( self, dim: Dimension, start_marker: int, end_marker: int - ) -> Tuple[int, int]: + ) -> Tuple[int32, int32]: """ Use to specify domains of a field for field_operator. From 86e31a1cefd46a6b17e34e474185915bcdb5e5f5 Mon Sep 17 00:00:00 2001 From: Magdalena Date: Tue, 20 Jun 2023 08:24:10 +0200 Subject: [PATCH 144/263] - introduce _HALO_XX == min_rlxx_int -1 and _LOCAL_XX == min_rlxx_int in the constants in horizontal.py (#225) - add halo() marker function that actually returns the marker for the halo lines (called halo+1 in dusk) --- .../src/icon4py/diffusion/horizontal.py | 28 +- atm_dyn_iconam/tests/test_icon_grid.py | 467 ++++++++++-------- 2 files changed, 291 insertions(+), 204 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py index 82c3a206e..aa59ec458 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py @@ -42,23 +42,31 @@ class HorizontalMarkerIndex: _LOCAL_BOUNDARY_EDGES: Final[int] = 1 + _ICON_INDEX_OFFSET_EDGES _INTERIOR_EDGES: Final[int] = _ICON_INDEX_OFFSET_EDGES _NUDGING_EDGES: Final[int] = _GRF_BOUNDARY_WIDTH_EDGES + _ICON_INDEX_OFFSET_EDGES - _HALO_EDGES: Final[int] = _MIN_RL_EDGE_INT + _ICON_INDEX_OFFSET_EDGES + _HALO_EDGES: Final[int] = _MIN_RL_EDGE_INT - 1 + _ICON_INDEX_OFFSET_EDGES + _LOCAL_EDGES: Final[int] = _MIN_RL_EDGE_INT + _ICON_INDEX_OFFSET_EDGES _END_EDGES: Final[int] = 0 _LOCAL_BOUNDARY_CELLS: Final[int] = 1 + _ICON_INDEX_OFFSET_CELLS _INTERIOR_CELLS: Final[int] = _ICON_INDEX_OFFSET_CELLS _NUDGING_CELLS: Final[int] = _GRF_BOUNDARY_WIDTH_CELL + 1 + _ICON_INDEX_OFFSET_CELLS - _HALO_CELLS: Final[int] = _MIN_RL_CELL_INT + _ICON_INDEX_OFFSET_CELLS + _HALO_CELLS: Final[int] = _MIN_RL_CELL_INT - 1 + _ICON_INDEX_OFFSET_CELLS + _LOCAL_CELLS: Final[int] = _MIN_RL_CELL_INT + _ICON_INDEX_OFFSET_CELLS _END_CELLS: Final[int] = 0 _LOCAL_BOUNDARY_VERTICES = 1 + _ICON_INDEX_OFFSET_VERTEX _INTERIOR_VERTICES: Final[int] = _ICON_INDEX_OFFSET_VERTEX _NUDGING_VERTICES: Final[int] = 0 - _HALO_VERTICES: Final[int] = _MIN_RL_VERTEX_INT + _ICON_INDEX_OFFSET_VERTEX + _HALO_VERTICES: Final[int] = _MIN_RL_VERTEX_INT - 1 + _ICON_INDEX_OFFSET_VERTEX + _LOCAL_VERTICES: Final[int] = _MIN_RL_VERTEX_INT + _ICON_INDEX_OFFSET_VERTEX _END_VERTICES: Final[int] = 0 @classmethod def lateral_boundary(cls, dim: Dimension) -> int: + """Indicate lateral boundary. + + These points correspond to the sorted points in ICON, the marker can be incremented in order + to accesss higher boundary lines + """ match (dim): case (dimension.CellDim): return cls._LOCAL_BOUNDARY_CELLS @@ -69,6 +77,18 @@ def lateral_boundary(cls, dim: Dimension) -> int: @classmethod def local(cls, dim: Dimension) -> int: + """Indicate points that are owned by the processing unit, i.e. no halo points.""" + match (dim): + case (dimension.CellDim): + return cls._LOCAL_CELLS + case (dimension.EdgeDim): + return cls._LOCAL_EDGES + case (dimension.VertexDim): + return cls._LOCAL_VERTICES + + @classmethod + def halo(cls, dim: Dimension) -> int: + """Indicate the halo points.""" match (dim): case (dimension.CellDim): return cls._HALO_CELLS @@ -79,6 +99,7 @@ def local(cls, dim: Dimension) -> int: @classmethod def nudging(cls, dim: Dimension) -> int: + """Indicate the nudging zone.""" match (dim): case (dimension.CellDim): return cls._NUDGING_CELLS @@ -89,6 +110,7 @@ def nudging(cls, dim: Dimension) -> int: @classmethod def interior(cls, dim: Dimension) -> int: + """Indicate interior i.e. unordered prognostic cells in ICON.""" match (dim): case (dimension.CellDim): return cls._INTERIOR_CELLS diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index c3de4fbb1..3c04aeb69 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -17,224 +17,289 @@ @pytest.mark.datatest -def test_horizontal_grid_cell_indices(icon_grid): +@pytest.mark.parametrize( + "start_marker, end_marker, expected_bounds", + [ + ( + HorizontalMarkerIndex.lateral_boundary(CellDim), + HorizontalMarkerIndex.lateral_boundary(CellDim), + (0, 850), + ), + ( + HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, + (850, 1688), + ), + ( + HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, + (1688, 2511), + ), + ( + HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, + (2511, 3316), + ), + ( + HorizontalMarkerIndex.interior(CellDim), + HorizontalMarkerIndex.interior(CellDim), + (4104, 20896), + ), + ( + HorizontalMarkerIndex.interior(CellDim) + 1, + HorizontalMarkerIndex.interior(CellDim) + 1, + (0, 850), + ), + ( + HorizontalMarkerIndex.nudging(CellDim), + HorizontalMarkerIndex.nudging(CellDim), + ( + 3316, + 4104, + ), + ), + ( + HorizontalMarkerIndex.end(CellDim), + HorizontalMarkerIndex.end(CellDim), + ( + 20896, + 20896, + ), + ), + ( + HorizontalMarkerIndex.halo(CellDim), + HorizontalMarkerIndex.halo(CellDim), + ( + 20896, + 20896, + ), + ), + ( + HorizontalMarkerIndex.local(CellDim), + HorizontalMarkerIndex.local(CellDim), + (-1, 20896), + ), + ], +) +def test_horizontal_cell_markers(icon_grid, start_marker, end_marker, expected_bounds): + assert ( + icon_grid.get_indices_from_to( + CellDim, + start_marker, + end_marker, + ) + == expected_bounds + ) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "start_marker, end_marker, expected_bounds", + [ + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim), + HorizontalMarkerIndex.lateral_boundary(EdgeDim), + (0, 428), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, + (428, 1278), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, + (1278, 1700), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, + (1700, 2538), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, + (2538, 2954), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, + (2954, 3777), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, + (3777, 4184), + ), + ( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, + (4184, 4989), + ), + ( + HorizontalMarkerIndex.interior(EdgeDim), + HorizontalMarkerIndex.interior(EdgeDim), + (6176, 31558), + ), + ( + HorizontalMarkerIndex.nudging(EdgeDim), + HorizontalMarkerIndex.nudging(EdgeDim), + ( + 4989, + 5387, + ), + ), + ( + HorizontalMarkerIndex.nudging(EdgeDim) + 1, + HorizontalMarkerIndex.nudging(EdgeDim) + 1, + (5387, 6176), + ), + ( + HorizontalMarkerIndex.end(EdgeDim), + HorizontalMarkerIndex.end(EdgeDim), + ( + 31558, + 31558, + ), + ), + ( + HorizontalMarkerIndex.halo(EdgeDim), + HorizontalMarkerIndex.halo(EdgeDim), + ( + 31558, + 31558, + ), + ), + ( + HorizontalMarkerIndex.local(EdgeDim), + HorizontalMarkerIndex.local(EdgeDim), + (-1, 31558), + ), + ], +) +def test_horizontal_edge_markers(icon_grid, start_marker, end_marker, expected_bounds): + assert ( + icon_grid.get_indices_from_to( + EdgeDim, + start_marker, + end_marker, + ) + == expected_bounds + ) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "start_marker, end_marker, expected_bounds", + [ + ( + HorizontalMarkerIndex.lateral_boundary(VertexDim), + HorizontalMarkerIndex.lateral_boundary(VertexDim), + (0, 428), + ), + ( + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, + (428, 850), + ), + ( + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, + (850, 1266), + ), + ( + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, + HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, + (1266, 1673), + ), + ( + HorizontalMarkerIndex.interior(VertexDim), + HorizontalMarkerIndex.interior(VertexDim), + (2071, 10663), + ), + ( + HorizontalMarkerIndex.interior(VertexDim) + 1, + HorizontalMarkerIndex.interior(VertexDim) + 1, + (0, 428), + ), + ( + HorizontalMarkerIndex.end(CellDim), + HorizontalMarkerIndex.end(CellDim), + ( + 10663, + 10663, + ), + ), + ( + HorizontalMarkerIndex.halo(VertexDim), + HorizontalMarkerIndex.halo(VertexDim), + ( + 10663, + 10663, + ), + ), + ( + HorizontalMarkerIndex.local(VertexDim), + HorizontalMarkerIndex.local(VertexDim), + (-1, 10663), + ), + ], +) +def test_horizontal_vertex_markers( + icon_grid, start_marker, end_marker, expected_bounds +): + assert ( + icon_grid.get_indices_from_to( + VertexDim, + start_marker, + end_marker, + ) + == expected_bounds + ) + + +@pytest.mark.datatest +def test_cross_check_marker_equivalences(icon_grid): + """ + Check actual equivalences of calculated markers. + + TODO [magdalena] This should go away once we refactor these markers in a good way, such that no + calculation need to be done with them anymore. + """ assert icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.local(CellDim) - 1, HorizontalMarkerIndex.local(CellDim) - 1, - ) == ( - 20896, - 20896, - ) # halo + 1 - - assert icon_grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.local(CellDim), - HorizontalMarkerIndex.local(CellDim), - ) == ( - -1, - 20896, - ) # halo in icon is (1,20896) - assert icon_grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.interior(CellDim), - HorizontalMarkerIndex.interior(CellDim), - ) == ( - 4104, - 20896, - ) # interior - assert icon_grid.get_indices_from_to( + ) == icon_grid.get_indices_from_to( CellDim, - HorizontalMarkerIndex.interior(CellDim) + 1, - HorizontalMarkerIndex.interior(CellDim) + 1, - ) == ( - 0, - 850, - ) # lb+1 - assert icon_grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, - HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, - ) == (850, 1688) - assert icon_grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, - HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, - ) == ( - 1688, - 2511, - ) # lb+2 + HorizontalMarkerIndex.halo(CellDim), + HorizontalMarkerIndex.halo(CellDim), + ) assert icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.nudging(CellDim) - 1, HorizontalMarkerIndex.nudging(CellDim) - 1, - ) == ( - 2511, - 3316, - ) # lb+3 - assert icon_grid.get_indices_from_to( + ) == icon_grid.get_indices_from_to( CellDim, - HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.nudging(CellDim), - ) == ( - 3316, - 4104, - ) # nudging - - -@pytest.mark.datatest -def test_horizontal_edge_indices(icon_grid): - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.interior(EdgeDim), - HorizontalMarkerIndex.interior(EdgeDim), - ) == (6176, 31558) - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.local(EdgeDim) - 2, - HorizontalMarkerIndex.local(EdgeDim) - 2, - ) == (31558, 31558) + HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, + ) assert icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 1, HorizontalMarkerIndex.local(EdgeDim) - 1, - ) == (31558, 31558) - assert icon_grid.get_indices_from_to( + ) == icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.local(EdgeDim), - HorizontalMarkerIndex.local(EdgeDim), - ) == ( - -1, - 31558, - ) # halo in icon is (1, 31558) + HorizontalMarkerIndex.halo(EdgeDim), + HorizontalMarkerIndex.halo(EdgeDim), + ) + assert icon_grid.get_indices_from_to( EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim) + 1, - HorizontalMarkerIndex.nudging(EdgeDim) + 1, - ) == ( - 5387, - 6176, - ) # nudging +1 - assert icon_grid.get_indices_from_to( + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 8, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 8, + ) == icon_grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), HorizontalMarkerIndex.nudging(EdgeDim), - ) == ( - 4989, - 5387, - ) # nudging - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, - ) == ( - 4184, - 4989, - ) # lb +7 - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, - ) == ( - 3777, - 4184, - ) # lb +6 - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, - ) == ( - 2954, - 3777, - ) # lb +5 - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, - ) == ( - 2538, - 2954, - ) # lb +4 - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, - ) == ( - 1700, - 2538, - ) # lb +3 - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, - ) == ( - 1278, - 1700, - ) # lb +2 - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, - ) == ( - 428, - 1278, - ) # lb +1 - assert icon_grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim), - HorizontalMarkerIndex.lateral_boundary(EdgeDim), - ) == ( - 0, - 428, - ) # lb +0 - - -@pytest.mark.datatest -def test_horizontal_vertex_indices(icon_grid): - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.end(VertexDim), - HorizontalMarkerIndex.end(VertexDim), - ) == (10663, 10663) - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.local(VertexDim), - HorizontalMarkerIndex.local(VertexDim), - ) == (-1, 10663) - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.local(VertexDim) - 1, - HorizontalMarkerIndex.local(VertexDim) - 1, - ) == (10663, 10663) - - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim), - HorizontalMarkerIndex.lateral_boundary(VertexDim), - ) == (0, 428) - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, - ) == (428, 850) - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, - ) == (850, 1266) - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, - ) == (1266, 1673) - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, - ) == (1673, 2071) - - assert icon_grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.interior(VertexDim), - HorizontalMarkerIndex.interior(VertexDim), - ) == (2071, 10663) + ) From e6bafb8e21276adb0cfbdb2e628600de965aa1db Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 20 Jun 2023 09:32:17 +0200 Subject: [PATCH 145/263] adding comments from discussion with Fabian --- atm_dyn_iconam/src/icon4py/driver/parallel_setup.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py index 85e1e56b1..9a6cce873 100644 --- a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py @@ -170,7 +170,8 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): (EdgeDim, KDim): self._create_pattern(EdgeDim, KDim), (EdgeDim, KHalfDim): self._create_pattern(EdgeDim, KHalfDim), } - + # TODO [magdalena] reuse communication object as much as possible. The pattern argument is needed for template args and should go away here + # if the pattern arg weren't there, we would need only as many communiciation objects as there are overlapping exchanges not as many as there are different patterns self._comms = {k: ghex.make_co(context, v) for k, v in self._patterns.items()} print( f"rank={self._context.rank()}/{self._context.size()} : patterns and communicators initialized " @@ -193,6 +194,9 @@ def _create_domain_descriptor(self, dim: Dimension, levels: int): local_halo = self._decomposition_info.local_index( dim, DecompositionInfo.EntryType.HALO ) + # TOOD [magdalena] first arg is the domain ID which builds up an MPI Tag, doesn't need to be the MPI rank. + # on the contrary it is safer if those are different for all domain descriptors (otherwise the system deadlocks if 2 parallel exchanges are done + # with the same domain-id domain_desc = ghex.domain_descriptor( self._context.rank(), all_global.tolist(), local_halo.tolist(), levels ) From 4228ac6fdbc878d95eae00bcfb9a85a1967a4e82 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 20 Jun 2023 13:36:23 +0200 Subject: [PATCH 146/263] udpate of ghex: - remove levels from field_descriptor - remove pattern from communcation object construction --- .../src/icon4py/diffusion/diffusion.py | 34 +++++----- .../src/icon4py/driver/parallel_setup.py | 64 +++++++------------ .../tests/mpi_tests/test_parallel_setup.py | 15 +++-- 3 files changed, 49 insertions(+), 64 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 447414981..026e8c636 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -698,8 +698,8 @@ def _do_diffusion_step( ) log.debug("rbf interpolation: end") # 2. HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert - res = self._sync_fields((VertexDim, KDim), self.u_vert, self.v_vert) - self._wait(res) + res = self._sync_fields(VertexDim, self.u_vert, self.v_vert) + self._wait(res, VertexDim) # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 @@ -754,8 +754,8 @@ def _do_diffusion_step( log.debug("running fused stencil fused stencil 02_03: end") # # # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH - res = self._sync_fields((EdgeDim, KDim), self.z_nabla2_e) - self._wait(res) + res = self._sync_fields(EdgeDim, self.z_nabla2_e) + self._wait(res, EdgeDim) # # 5. CALL rbf_vec_interpol_vertex_wp log.debug("rbf interpolation: start") @@ -774,9 +774,8 @@ def _do_diffusion_step( log.debug("rbf interpolation: end") # # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult # # - res = self._sync_fields((VertexDim, KDim), self.u_vert, self.v_vert) - self._wait(res) - + res = self._sync_fields(VertexDim, self.u_vert, self.v_vert) + self._wait(res, VertexDim) # # # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 # # # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 log.debug("running fused stencil 04_05_06: start") @@ -893,8 +892,8 @@ def _do_diffusion_step( # ) log.debug("running fused stencil 07_08_09_10: end") # # 8. HALO EXCHANGE: CALL sync_patch_array - comm_res = self._sync_fields((EdgeDim, KDim), prognostic_state.vn) - self._wait(comm_res) + comm_res = self._sync_fields(EdgeDim, prognostic_state.vn) + self._wait(comm_res, EdgeDim) # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, # # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 @@ -978,20 +977,21 @@ def _do_diffusion_step( # 10. HALO EXCHANGE sync_patch_array (Cell fields) # TODO @magdalena why not trigger the exchange of w earlier? # TODO if condition: IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN - # TODO magdalena: why does this crash? - res_1 = self._sync_fields( - (CellDim, KDim), + res = self._sync_fields( + CellDim, prognostic_state.theta_v, prognostic_state.exner_pressure, + prognostic_state.w, ) - self._wait(res_1) - res_w = self._sync_fields((CellDim, KHalfDim), prognostic_state.w) - self._wait(res_w) + self._wait(res, CellDim) - def _sync_fields(self, dim: tuple[Dimension, Dimension], *field): + def _sync_fields(self, dim: Dimension, *field): if self._exchange: return self._exchange.exchange(dim, *field) - def _wait(self, comm_handle): + def _wait(self, comm_handle, dim): if comm_handle: comm_handle.wait() + print( + f"rank={self._exchange._context.rank()}/{self._exchange._context.size()}:communication dim={dim} done" + ) diff --git a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py index 9a6cce873..95c903ba6 100644 --- a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py @@ -123,24 +123,13 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): self._context = context self._decomposition_info = domain_decomposition self._domain_descriptors = { - (CellDim, KDim): self._create_domain_descriptor( - CellDim, domain_decomposition.klevels + CellDim: self._create_domain_descriptor( + CellDim, ), - (CellDim, KHalfDim): self._create_domain_descriptor( - CellDim, domain_decomposition.klevels + 1 - ), - (VertexDim, KDim): self._create_domain_descriptor( - VertexDim, domain_decomposition.klevels - ), - (VertexDim, KHalfDim): self._create_domain_descriptor( - VertexDim, domain_decomposition.klevels + 1 - ), - (EdgeDim, KDim): self._create_domain_descriptor( - EdgeDim, domain_decomposition.klevels - ), - (EdgeDim, KHalfDim): self._create_domain_descriptor( - EdgeDim, domain_decomposition.klevels + 1 + VertexDim: self._create_domain_descriptor( + VertexDim, ), + EdgeDim: self._create_domain_descriptor(EdgeDim), } print( f"rank={self._context.rank()}/{self._context.size()} :domain descriptors initialized" @@ -163,16 +152,11 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): ) self._patterns = { - (CellDim, KDim): self._create_pattern(CellDim, KDim), - (CellDim, KHalfDim): self._create_pattern(CellDim, KHalfDim), - (VertexDim, KDim): self._create_pattern(VertexDim, KDim), - (VertexDim, KHalfDim): self._create_pattern(VertexDim, KHalfDim), - (EdgeDim, KDim): self._create_pattern(EdgeDim, KDim), - (EdgeDim, KHalfDim): self._create_pattern(EdgeDim, KHalfDim), + CellDim: self._create_pattern(CellDim), + VertexDim: self._create_pattern(VertexDim), + EdgeDim: self._create_pattern(EdgeDim), } - # TODO [magdalena] reuse communication object as much as possible. The pattern argument is needed for template args and should go away here - # if the pattern arg weren't there, we would need only as many communiciation objects as there are overlapping exchanges not as many as there are different patterns - self._comms = {k: ghex.make_co(context, v) for k, v in self._patterns.items()} + self._comms = {k: ghex.make_co(context) for k, v in self._patterns.items()} print( f"rank={self._context.rank()}/{self._context.size()} : patterns and communicators initialized " ) @@ -181,13 +165,13 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): ) def _domain_descriptor_info(self, descr): - return f" id={descr.domain_id()} levels = {descr.levels()}, size={descr.size()}, inner_size={descr.inner_size()}" + return f" id={descr.domain_id()}, size={descr.size()}, inner_size={descr.inner_size()}" def get_size(self): return self._context.size() # TODO [magdalena] is the tolist() necessary? - def _create_domain_descriptor(self, dim: Dimension, levels: int): + def _create_domain_descriptor(self, dim: Dimension): all_global = self._decomposition_info.global_index( dim, DecompositionInfo.EntryType.ALL ) @@ -198,7 +182,7 @@ def _create_domain_descriptor(self, dim: Dimension, levels: int): # on the contrary it is safer if those are different for all domain descriptors (otherwise the system deadlocks if 2 parallel exchanges are done # with the same domain-id domain_desc = ghex.domain_descriptor( - self._context.rank(), all_global.tolist(), local_halo.tolist(), levels + self._context.rank(), all_global.tolist(), local_halo.tolist() ) print( f"rank={self._context.rank()}/{self._context.size()}: domain descriptor for dim {dim} with properties {self._domain_descriptor_info(domain_desc)}" @@ -206,7 +190,7 @@ def _create_domain_descriptor(self, dim: Dimension, levels: int): return domain_desc - def _create_pattern(self, horizontal_dim: Dimension, vertical_dim: Dimension): + def _create_pattern(self, horizontal_dim: Dimension): halo_generator = ghex.halo_generator_with_gids( self._decomposition_info.global_index( horizontal_dim, DecompositionInfo.EntryType.HALO @@ -215,27 +199,27 @@ def _create_pattern(self, horizontal_dim: Dimension, vertical_dim: Dimension): print( f"rank={self._context.rank()}/{self._context.size()}: halo generator for dim={horizontal_dim} created" ) - dimensions = (horizontal_dim, vertical_dim) pattern = ghex.make_pattern( - self._context, halo_generator, [self._domain_descriptors[dimensions]] + self._context, halo_generator, [self._domain_descriptors[horizontal_dim]] ) print( - f"rank={self._context.rank()}/{self._context.size()}: pattern for dim={dimensions} and {self._domain_descriptors[dimensions]} created" + f"rank={self._context.rank()}/{self._context.size()}: pattern for dim={horizontal_dim} and {self._domain_descriptors[horizontal_dim]} created" ) return pattern - def exchange(self, dims: tuple[Dimension, Dimension], *fields): - assert dims[0] in [CellDim, EdgeDim, VertexDim] - horizontal_size = self._field_size[dims[0]] - pattern = self._patterns[dims] + def exchange(self, dim: Dimension, *fields): + assert dim in [CellDim, EdgeDim, VertexDim] + horizontal_size = self._field_size[dim] + pattern = self._patterns[dim] assert pattern is not None fields = [np.asarray(f)[:horizontal_size, :] for f in fields] + shapes = list(map(lambda f: f.shape, fields)) print( - f"rank = {self._context.rank()}/{self._context.size()}: communicating fields of dim = {dims} (shape = {fields[0].shape})" + f"rank = {self._context.rank()}/{self._context.size()}: communicating fields of dim = {dim} : shapes = {shapes}" ) - communicator = self._comms[dims] + patterns_of_field = [ - pattern(ghex.field_descriptor(self._domain_descriptors[dims], f)) + pattern(ghex.field_descriptor(self._domain_descriptors[dim], f)) for f in fields ] - return communicator.exchange(patterns_of_field) + return self._comms[dim].exchange(patterns_of_field) diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index e66ac002d..16d341879 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -30,7 +30,6 @@ from icon4py.driver.parallel_setup import ( DecompositionInfo, Exchange, - finalize_mpi, get_processor_properties, ) @@ -165,8 +164,8 @@ def test_decomposition_info_matches_gridsize(datapath, caplog): @pytest.mark.mpi -def test_parallel_diffusion(r04b09_diffusion_config, step_date_init): - # caplog.set_level(logging.DEBUG) +def test_parallel_diffusion(r04b09_diffusion_config, step_date_init, caplog): + caplog.set_level(logging.DEBUG) props = get_processor_properties() num_nodes = props.comm_size @@ -192,7 +191,7 @@ def test_parallel_diffusion(r04b09_diffusion_config, step_date_init): print( f"rank={props.rank}/{props.comm_size}: GHEX context setup: from {props.comm_name} with {props.comm_size} nodes" ) - #assert context.size() == 2 + # assert context.size() == 2 icon_grid = read_icon_grid(path, rank=props.rank) print( @@ -247,11 +246,13 @@ def test_parallel_diffusion(r04b09_diffusion_config, step_date_init): diffusion_savepoint_exit = IconSerialDataProvider( "icon_pydycore", str(path), True, mpi_rank=props.rank ).from_savepoint_diffusion_init(linit=initial_run, date=step_date_init) - verify_fields(diffusion_savepoint_exit, diagnostic_state, prognostic_state, - diffusion.metric_state.mask_hdiff) + # verify_fields(diffusion_savepoint_exit, diagnostic_state, prognostic_state, + # diffusion.metric_state.mask_hdiff) -def verify_fields(diffusion_savepoint_exit, diagnostic_state, prognostic_state, steep_points): +def verify_fields( + diffusion_savepoint_exit, diagnostic_state, prognostic_state, steep_points +): ref_div_ic = np.asarray(diffusion_savepoint_exit.div_ic()) val_div_ic = np.asarray(diagnostic_state.div_ic) ref_hdef_ic = np.asarray(diffusion_savepoint_exit.hdef_ic()) From 0787dbaee0abe3814df4a0c2bb05fb8b43e7b644 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Mon, 26 Jun 2023 22:56:48 +0200 Subject: [PATCH 147/263] fix validation of stencil 15 --- atm_dyn_iconam/tests/test_diffusion.py | 33 ++++++++++++------- .../tests/test_utils/serialbox_utils.py | 28 +++++++++++----- 2 files changed, 42 insertions(+), 19 deletions(-) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 0980a90c6..c23aeecf6 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -405,7 +405,14 @@ def test_verify_diffusion_init_against_other_regular_savepoint( @pytest.mark.datatest -# @pytest.mark.parametrize("step_date_init, step_date_exit", [("2021-06-20T12:00:10.000", "2021-06-20T12:00:10.000"), ("2021-06-20T12:00:20.000", "2021-06-20T12:00:20.000"), ("2021-06-20T12:01:00.000", "2021-06-20T12:01:00.000")] ) +@pytest.mark.parametrize( + "step_date_init, step_date_exit", + [ + ("2021-06-20T12:00:10.000", "2021-06-20T12:00:10.000"), + ("2021-06-20T12:00:20.000", "2021-06-20T12:00:20.000"), + ("2021-06-20T12:01:00.000", "2021-06-20T12:01:00.000"), + ], +) def test_run_diffusion_single_step( diffusion_savepoint_init, diffusion_savepoint_exit, @@ -479,11 +486,8 @@ def test_run_diffusion_single_step( ref_theta_v = np.asarray(diffusion_savepoint_exit.theta_v()) val_theta_v = np.asarray(prognostic_state.theta_v) val_exner = np.asarray(prognostic_state.exner_pressure) - steep_points = np.asarray(diffusion.metric_state.mask_hdiff) - assert np.allclose(ref_theta_v[~steep_points], val_theta_v[~steep_points]) - assert np.allclose(ref_exner[~steep_points], val_exner[~steep_points]) - # assert np.allclose(ref_theta_v[steep_points], val_theta_v[steep_points]) - # assert np.allclose(ref_exner[steep_points], val_exner[steep_points]) + assert np.allclose(ref_theta_v, val_theta_v) + assert np.allclose(ref_exner, val_exner) @pytest.mark.datatest @@ -561,8 +565,15 @@ def test_run_diffusion_initial_step( ref_theta_v = np.asarray(diffusion_savepoint_exit.theta_v()) val_theta_v = np.asarray(prognostic_state.theta_v) val_exner = np.asarray(prognostic_state.exner_pressure) - steep_points = np.asarray(diffusion.metric_state.mask_hdiff) - assert np.allclose(ref_theta_v[~steep_points], val_theta_v[~steep_points]) - assert np.allclose(ref_exner[~steep_points], val_exner[~steep_points]) - # assert np.allclose(ref_theta_v[steep_points], val_theta_v[steep_points]) - # assert np.allclose(ref_exner[steep_points], val_exner[steep_points]) + assert np.allclose(ref_theta_v, val_theta_v) + assert np.allclose(ref_exner, val_exner) + + +def test_verify_stencil15_field_manipulation(interpolation_savepoint, icon_grid): + geofac_n2s = np.asarray(interpolation_savepoint.geofac_n2s()) + int_state = interpolation_savepoint.construct_interpolation_state() + geofac_c = np.asarray(int_state.geofac_n2s_c) + geofac_nbh = np.asarray(int_state.geofac_n2s_nbh) + cec_table = icon_grid.get_c2cec_connectivity().table + assert np.allclose(geofac_c, geofac_n2s[:, 0]) + assert np.allclose(geofac_nbh[cec_table], geofac_n2s[:, 1:]) diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index 19f767704..fa3fbf238 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -49,6 +49,9 @@ from .helpers import as_1D_sparse_field +C2E2C_DIM = 3 + + class IconSavepoint: def __init__(self, sp: ser.Savepoint, ser: ser.Serializer): self.savepoint = sp @@ -311,19 +314,28 @@ def zd_diffcoef(self): return self._get_field("zd_diffcoef", CellDim, KDim) def zd_intcoef(self): - return self._from_cell_c2e2c_to_cec("vcoef") - - def _from_cell_c2e2c_to_cec(self, field_name: str, offset: int = 0): - ser_input = ( - np.squeeze(self.serializer.read(field_name, self.savepoint)) + offset + ser_input = np.moveaxis( + (np.squeeze(self.serializer.read("vcoef", self.savepoint))), 1, -1 ) - old_shape = ser_input.shape + return self._linearize_first_2dims(ser_input, sparse_size=3) + + def _linearize_first_2dims(self, data: np.ndarray, sparse_size): + old_shape = data.shape + assert old_shape[1] == sparse_size return np_as_located_field(CECDim, KDim)( - ser_input.reshape(old_shape[0] * old_shape[1], old_shape[2]) + data.reshape(old_shape[0] * old_shape[1], old_shape[2]) ) def zd_vertoffset(self): - return self._from_cell_c2e2c_to_cec("zd_vertoffset", 0) + ser_input = np.squeeze(self.serializer.read("zd_vertoffset", self.savepoint)) + ser_input = np.moveaxis(ser_input, 1, -1) + return self._linearize_first_2dims(ser_input, sparse_size=3) + + def zd_vertidx(self): + return np.squeeze(self.serializer.read("zd_vertidx", self.savepoint)) + + def zd_indlist(self): + return np.squeeze(self.serializer.read("zd_indlist", self.savepoint)) def theta_ref_mc(self): return self._get_field("theta_ref_mc", CellDim, KDim) From 238b4694133bb7a1d6830c817567e3f158a84d8d Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 29 Jun 2023 10:02:26 +0200 Subject: [PATCH 148/263] refactor tests --- atm_dyn_iconam/tests/test_diffusion.py | 43 ++++++++------------------ 1 file changed, 13 insertions(+), 30 deletions(-) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index c23aeecf6..336efcd77 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -461,14 +461,20 @@ def test_run_diffusion_single_step( cell_areas=cell_geometry.area, ) assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v + _verify_diffusion_exit_fields( + diagnostic_state, prognostic_state, diffusion_savepoint_exit + ) + + +def _verify_diffusion_exit_fields( + diagnostic_state, prognostic_state, diffusion_savepoint_exit +): ref_div_ic = np.asarray(diffusion_savepoint_exit.div_ic()) val_div_ic = np.asarray(diagnostic_state.div_ic) ref_hdef_ic = np.asarray(diffusion_savepoint_exit.hdef_ic()) val_hdef_ic = np.asarray(diagnostic_state.hdef_ic) - assert np.allclose(ref_div_ic, val_div_ic) assert np.allclose(ref_hdef_ic, val_hdef_ic) - ref_w = np.asarray(diffusion_savepoint_exit.w()) val_w = np.asarray(prognostic_state.w) ref_dwdx = np.asarray(diffusion_savepoint_exit.dwdx()) @@ -477,7 +483,6 @@ def test_run_diffusion_single_step( val_dwdy = np.asarray(diagnostic_state.dwdy) ref_vn = np.asarray(diffusion_savepoint_exit.vn()) val_vn = np.asarray(prognostic_state.vn) - assert np.allclose(ref_vn, val_vn) assert np.allclose(ref_dwdx, val_dwdx) assert np.allclose(ref_dwdy, val_dwdy) @@ -540,33 +545,11 @@ def test_run_diffusion_initial_step( cell_areas=cell_geometry.area, ) assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v - ref_div_ic = np.asarray(diffusion_savepoint_exit.div_ic()) - val_div_ic = np.asarray(diagnostic_state.div_ic) - ref_hdef_ic = np.asarray(diffusion_savepoint_exit.hdef_ic()) - val_hdef_ic = np.asarray(diagnostic_state.hdef_ic) - - assert np.allclose(ref_div_ic, val_div_ic) - assert np.allclose(ref_hdef_ic, val_hdef_ic) - - ref_w = np.asarray(diffusion_savepoint_exit.w()) - val_w = np.asarray(prognostic_state.w) - ref_dwdx = np.asarray(diffusion_savepoint_exit.dwdx()) - val_dwdx = np.asarray(diagnostic_state.dwdx) - ref_dwdy = np.asarray(diffusion_savepoint_exit.dwdy()) - val_dwdy = np.asarray(diagnostic_state.dwdy) - ref_vn = np.asarray(diffusion_savepoint_exit.vn()) - val_vn = np.asarray(prognostic_state.vn) - - assert np.allclose(ref_vn, val_vn) - assert np.allclose(ref_dwdx, val_dwdx) - assert np.allclose(ref_dwdy, val_dwdy) - assert np.allclose(ref_w, val_w) - ref_exner = np.asarray(diffusion_savepoint_exit.exner()) - ref_theta_v = np.asarray(diffusion_savepoint_exit.theta_v()) - val_theta_v = np.asarray(prognostic_state.theta_v) - val_exner = np.asarray(prognostic_state.exner_pressure) - assert np.allclose(ref_theta_v, val_theta_v) - assert np.allclose(ref_exner, val_exner) + _verify_diffusion_exit_fields( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + diffusion_savepoint_exit=diffusion_savepoint_exit, + ) def test_verify_stencil15_field_manipulation(interpolation_savepoint, icon_grid): From 8a9ac4927b71cfa2b90ea7de4f5d90cb8a346cb6 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 29 Jun 2023 13:25:32 +0200 Subject: [PATCH 149/263] switch back to fused stencil 07, 08, 09, 10 --- .../src/icon4py/diffusion/diffusion.py | 173 ++++++------------ atm_dyn_iconam/src/icon4py/diffusion/utils.py | 4 + 2 files changed, 57 insertions(+), 120 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 865955937..14293b030 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -25,26 +25,20 @@ run_gtfn_cached, ) -from icon4py.atm_dyn_iconam.apply_nabla2_to_w import apply_nabla2_to_w -from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import ( - apply_nabla2_to_w_in_upper_damping_layer, -) +from icon4py.atm_dyn_iconam.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance + from icon4py.atm_dyn_iconam.calculate_diagnostic_quantities_for_turbulence import ( calculate_diagnostic_quantities_for_turbulence, ) from icon4py.atm_dyn_iconam.calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools import ( calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools, ) -from icon4py.atm_dyn_iconam.calculate_horizontal_gradients_for_turbulence import ( - calculate_horizontal_gradients_for_turbulence, -) from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( calculate_nabla2_and_smag_coefficients_for_vn, ) from icon4py.atm_dyn_iconam.calculate_nabla2_for_theta import ( calculate_nabla2_for_theta, ) -from icon4py.atm_dyn_iconam.calculate_nabla2_for_w import calculate_nabla2_for_w from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import ( fused_mo_nh_diffusion_stencil_04_05_06, ) @@ -71,9 +65,8 @@ init_diffusion_local_fields_for_regular_timestep, init_nabla2_factor_in_upper_damping_zone, scale_k, - set_zero_v_k, setup_fields_for_initial_step, - zero_field, + zero_field, copy_field, ) @@ -474,6 +467,7 @@ def _index_field(dim: Dimension, size=None): self.vertical_index = _index_field(KDim, self.grid.n_lev() + 1) self.horizontal_cell_index = _index_field(CellDim) self.horizontal_edge_index = _index_field(EdgeDim) + self.w_tmp = np_as_located_field(CellDim, KDim)(np.zeros((self.grid.num_cells(), self.grid.n_lev() + 1), dtype=float)) def initial_step( self, @@ -659,18 +653,15 @@ def _do_diffusion_step( HorizontalMarkerIndex.local(VertexDim) - 1, ) - # 0b call timer start - # - # 0c. dtime dependent stuff: enh_smag_factor, + # dtime dependent: enh_smag_factor, scale_k.with_backend(backend)( self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={} ) # TODO: @magdalena is this needed?, if not remove - set_zero_v_k.with_backend(backend)(self.u_vert, offset_provider={}) - set_zero_v_k.with_backend(backend)(self.v_vert, offset_provider={}) + #set_zero_v_k.with_backend(backend)(self.u_vert, offset_provider={}) + #set_zero_v_k.with_backend(backend)(self.v_vert, offset_provider={}) log.debug("rbf interpolation: start") - # # 1. CALL rbf_vec_interpol_vertex mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( p_e_in=prognostic_state.vn, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, @@ -684,10 +675,10 @@ def _do_diffusion_step( offset_provider={"V2E": self.grid.get_v2e_connectivity()}, ) log.debug("rbf interpolation: end") - # 2. HALO EXCHANGE -- CALL sync_patch_array_mult - # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 - log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: start") + # HALO EXCHANGE -- CALL sync_patch_array_mult + + log.debug("running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start") calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend)( diff_multfac_smag=self.diff_multfac_smag, tangent_orientation=tangent_orientation, @@ -714,8 +705,8 @@ def _do_diffusion_step( "E2ECV": self.grid.get_e2ecv_connectivity(), }, ) - log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: end") - log.debug("running fused stencil fused stencil 02_03: start") + log.debug("running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end") + log.debug("running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start") calculate_diagnostic_quantities_for_turbulence.with_backend(backend)( kh_smag_ec=self.kh_smag_ec, vn=prognostic_state.vn, @@ -735,11 +726,10 @@ def _do_diffusion_step( "Koff": KDim, }, ) - log.debug("running fused stencil fused stencil 02_03: end") - # - # # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH - # - # # 5. CALL rbf_vec_interpol_vertex_wp + log.debug("running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end") + + # HALO EXCHANGE IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH + log.debug("rbf interpolation: start") mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( p_e_in=self.z_nabla2_e, @@ -754,13 +744,10 @@ def _do_diffusion_step( offset_provider={"V2E": self.grid.get_v2e_connectivity()}, ) log.debug("rbf interpolation: end") - # # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult - # # - # # # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 - # # # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 - # # - # - log.debug("running fused stencil 04_05_06: start") + + # 6. HALO EXCHANGE -- CALL sync_patch_array_mult + + log.debug("running stencil 04 05 06: start") fused_mo_nh_diffusion_stencil_04_05_06.with_backend(backend)( u_vert=self.u_vert, v_vert=self.v_vert, @@ -787,98 +774,44 @@ def _do_diffusion_step( "E2ECV": self.grid.get_e2ecv_connectivity(), }, ) - log.debug("running fused stencil 04_05_06: end") - # # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, - # # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 + log.debug("runningstencils 04 05 06: end") - log.debug("running stencils 07_08_09_10: start") - calculate_horizontal_gradients_for_turbulence.with_backend(backend)( - w=prognostic_state.w, + log.debug("running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start") + copy_field.with_backend(backend)(prognostic_state.w, self.w_tmp, offset_provider={}) + # TODO @magdalena: fix this python offset problem (+1): int(self.vertical_params.index_of_damping_layer + 1) + apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.with_backend(backend)( + area=cell_areas, + geofac_n2s=self.interpolation_state.geofac_n2s, geofac_grg_x=self.interpolation_state.geofac_grg_x, geofac_grg_y=self.interpolation_state.geofac_grg_y, + w_old=self.w_tmp, + w=prognostic_state.w, dwdx=diagnostic_state.dwdx, dwdy=diagnostic_state.dwdy, - vertical_start=1, - vertical_end=klevels, - horizontal_start=cell_start_nudging, - horizontal_end=cell_end_halo, - offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, - ) - - z_nabla2_c = zero_field(self.grid, CellDim, KDim, dtype=float) - - calculate_nabla2_for_w.with_backend(backend)( - w=prognostic_state.w, - geofac_n2s=self.interpolation_state.geofac_n2s, - z_nabla2_c=z_nabla2_c, - vertical_start=0, - vertical_end=klevels, + diff_multfac_w=self.diff_multfac_w, + diff_multfac_n2w=self.diff_multfac_n2w, + vert_idx=self.vertical_index, + horz_idx=self.horizontal_cell_index, + nrdmax=int32(self.vertical_params.index_of_damping_layer +1) , + interior_idx=int32( + cell_start_interior + ), + halo_idx=int32( + cell_end_local + ), horizontal_start=cell_start_nudging, horizontal_end=cell_end_halo, - offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, - ) - apply_nabla2_to_w.with_backend(backend)( - area=cell_areas, - z_nabla2_c=z_nabla2_c, - w=prognostic_state.w, - diff_multfac_w=self.diff_multfac_w, - geofac_n2s=self.interpolation_state.geofac_n2s, vertical_start=0, vertical_end=klevels, - horizontal_start=cell_start_interior, - horizontal_end=cell_end_local, - offset_provider={"C2E2CO": self.grid.get_c2e2co_connectivity()}, + offset_provider={ + "C2E2CO": self.grid.get_c2e2co_connectivity(), + }, ) + log.debug("running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end") + # HALO EXCHANGE: CALL sync_patch_array - # TODO @magdalena: fix this python offset problem (+1): int(self.vertical_params.index_of_damping_layer + 1) - apply_nabla2_to_w_in_upper_damping_layer.with_backend(backend)( - w=prognostic_state.w, - diff_multfac_n2w=self.diff_multfac_n2w, - cell_area=cell_areas, - z_nabla2_c=z_nabla2_c, - vertical_start=1, - vertical_end=int(self.vertical_params.index_of_damping_layer + 1), - horizontal_start=int(cell_start_interior), - horizontal_end=int(cell_end_local), - offset_provider={}, - ) - # w_old = prognostic_state.w - # apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.with_backend(backend)( - # area=cell_areas, - # geofac_n2s=self.interpolation_state.geofac_n2s, - # geofac_grg_x=self.interpolation_state.geofac_grg_x, - # geofac_grg_y=self.interpolation_state.geofac_grg_y, - # w_old=w_old, - # w=prognostic_state.w, - # dwdx=diagnostic_state.dwdx, - # dwdy=diagnostic_state.dwdy, - # diff_multfac_w=self.diff_multfac_w, - # diff_multfac_n2w=self.diff_multfac_n2w, - # vert_idx=self.vertical_index, - # horz_idx=self.horizontal_cell_index, - # nrdmax=int32(self.vertical_params.index_of_damping_layer +1) , - # interior_idx=int32( - # cell_start_interior -1 - # ), # h end index for stencil_09 and stencil_10 # TODO: in ICON: start_interior_idx_c -1 ?? - # halo_idx=int32( - # cell_end_local - # ), # h end index for stencil_09 and stencil_10, - # horizontal_start=cell_start_nudging, # h start index for stencil_07 and stencil_08 - # horizontal_end=cell_end_halo, # h end index for stencil_07 and stencil_08 - # vertical_start=0, - # vertical_end=klevels, - # offset_provider={ - # "C2E2CO": self.grid.get_c2e2co_connectivity(), - # }, - # ) - log.debug("running fused stencil 07_08_09_10: end") - # # 8. HALO EXCHANGE: CALL sync_patch_array - # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, - # # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 - # # # TODO @magdalena check: kh_smag_e is an out field, should not be calculated in init? - # - log.debug("running fused stencil 11_12: start") + log.debug("running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start") calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.with_backend( backend )( @@ -895,8 +828,8 @@ def _do_diffusion_step( "C2E2C": self.grid.get_c2e2c_connectivity(), }, ) - log.debug("running fused stencil 11_12: end") - log.debug("running fused stencil 13_14: start") + log.debug("running stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): end") + log.debug("running stencils 13 14 (calculate_nabla2_for_theta): start") calculate_nabla2_for_theta.with_backend(backend)( kh_smag_e=self.kh_smag_e, inv_dual_edge_length=inverse_dual_edge_length, @@ -913,8 +846,8 @@ def _do_diffusion_step( "C2CE": self.grid.get_c2ce_connectivity(), }, ) - log.debug("running fused stencil 13_14: end") - log.debug("running fused stencil 15: start") + log.debug("running stencils 13_14 (calculate_nabla2_for_theta): end") + log.debug("running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start") truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_backend( backend )( @@ -937,8 +870,8 @@ def _do_diffusion_step( }, ) - log.debug("running fused stencil 15: end") - log.debug("running fused stencil update_theta_and_exner: start") + log.debug("running fused stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end") + log.debug("running fused stencil 16 (update_theta_and_exner): start") update_theta_and_exner.with_backend(backend)( z_temp=self.z_temp, area=cell_areas, @@ -951,5 +884,5 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={}, ) - log.debug("running fused stencil update_theta_and_exner: end") + log.debug("running stencil 16 (update_theta_and_exner): end") # 10. HALO EXCHANGE sync_patch_array diff --git a/atm_dyn_iconam/src/icon4py/diffusion/utils.py b/atm_dyn_iconam/src/icon4py/diffusion/utils.py index 674a5db7a..077c8b1ac 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/utils.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/utils.py @@ -33,6 +33,10 @@ def _identity_c_k( ) -> Field[[CellDim, KDim], float]: return field +@program +def copy_field(old_f:Field[[CellDim, KDim], float], new_f:Field[[CellDim, KDim], float]): + _identity_c_k(old_f, out=new_f) + @field_operator def _identity_e_k( From 90b3bd8117c2880500db509eee751827a278bb04 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 29 Jun 2023 14:38:44 +0200 Subject: [PATCH 150/263] unify construction of NeighborTableOffsetProviders for 1D sparse fields --- .../src/icon4py/diffusion/diffusion.py | 7 +++--- .../src/icon4py/diffusion/icon_grid.py | 24 +++++++++---------- 2 files changed, 16 insertions(+), 15 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 14293b030..8ccf848f2 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -66,7 +66,7 @@ init_nabla2_factor_in_upper_damping_zone, scale_k, setup_fields_for_initial_step, - zero_field, copy_field, + zero_field, copy_field, set_zero_v_k, ) @@ -111,6 +111,7 @@ def __init__( max_nudging_coeff: float = 0.02, nudging_decay_rate: float = 2.0, ): + """ Set the diffusion configuration parameters with the ICON default values.""" # parameters from namelist diffusion_nml self.diffusion_type: int = diffusion_type @@ -659,8 +660,8 @@ def _do_diffusion_step( ) # TODO: @magdalena is this needed?, if not remove - #set_zero_v_k.with_backend(backend)(self.u_vert, offset_provider={}) - #set_zero_v_k.with_backend(backend)(self.v_vert, offset_provider={}) + set_zero_v_k.with_backend(backend)(self.u_vert, offset_provider={}) + set_zero_v_k.with_backend(backend)(self.v_vert, offset_provider={}) log.debug("rbf interpolation: start") mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( p_e_in=prognostic_state.vn, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index b40e95d2e..5d2feae3f 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -136,7 +136,7 @@ def num_edges(self): def get_indices_from_to( self, dim: Dimension, start_marker: int, end_marker: int - ) -> Tuple[int32, int32]: + ) -> tuple[int32, int32]: """ Use to specify domains of a field for field_operator. @@ -181,23 +181,23 @@ def get_v2e_connectivity(self): return NeighborTableOffsetProvider(table, VertexDim, EdgeDim, table.shape[1]) def get_e2ecv_connectivity(self): - old_shape = self.connectivities["e2c2v"].shape - v2ecv_table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) + return self._neighbortable_offset_provider_for_1d_sparse_fields( + self.connectivities["e2c2v"].shape, EdgeDim, ECVDim) + + def _neighbortable_offset_provider_for_1d_sparse_fields(self, old_shape:tuple[int, int], origin_axis:Dimension, neighbor_axis:Dimension ): + table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) return NeighborTableOffsetProvider( - v2ecv_table, EdgeDim, ECVDim, v2ecv_table.shape[1] + table, origin_axis, neighbor_axis, table.shape[1] ) def get_c2cec_connectivity(self): - old_shape = self.connectivities["c2e2c"].shape - c2cec_table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) - return NeighborTableOffsetProvider( - c2cec_table, CellDim, CECDim, c2cec_table.shape[1] - ) + return self._neighbortable_offset_provider_for_1d_sparse_fields( + self.connectivities["c2e2c"].shape, CellDim, CECDim) def get_c2ce_connectivity(self): - old_shape = self.connectivities["c2e"].shape - table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) - return NeighborTableOffsetProvider(table, CellDim, CEDim, table.shape[1]) + return self._neighbortable_offset_provider_for_1d_sparse_fields( + self.connectivities["c2e"].shape, CellDim, CEDim) + class VerticalModelParams: From d61eec7d78c1835fc4710dadac084e80df39b824 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 29 Jun 2023 14:39:49 +0200 Subject: [PATCH 151/263] remove unused param --- atm_dyn_iconam/src/icon4py/diffusion/diffusion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 8ccf848f2..557fafed6 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -360,7 +360,7 @@ def _diffusion_type_5_smagorinski_factor(config: DiffusionConfig): class Diffusion: """Class that configures diffusion and does one diffusion step.""" - def __init__(self, run_program=True): + def __init__(self): self._initialized = False self.rd_o_cvd: float = GAS_CONSTANT_DRY_AIR / (CPD - GAS_CONSTANT_DRY_AIR) From 46810571335a9e576738a646911a73bc3367895c Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 29 Jun 2023 16:16:34 +0200 Subject: [PATCH 152/263] fix dummy_driver to manually import testutils --- .../src/icon4py/diffusion/diffusion.py | 6 ++--- .../src/icon4py/driver/dycore_driver.py | 10 ++++--- .../src/icon4py/driver/icon_configuration.py | 4 +-- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 26 ++++++++++++++----- atm_dyn_iconam/tests/test_diffusion.py | 2 +- .../tests/test_utils/serialbox_utils.py | 3 --- 6 files changed, 31 insertions(+), 20 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 557fafed6..799f65382 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -345,7 +345,7 @@ def determine_smagorinski_factor(self, config: DiffusionConfig): @staticmethod def _diffusion_type_5_smagorinski_factor(config: DiffusionConfig): """ - Initialize smagorinski factors used in diffusion type 5. + Initialize Smagorinski factors used in diffusion type 5. The calculation and magic numbers are taken from mo_diffusion_nml.f90 """ @@ -392,7 +392,7 @@ def init( """ Initialize Diffusion granule with configuration. - calculates all local fields that are used in diffusion within the time loop + calculates all local fields that are used in diffusion within the time loop. """ self.config: DiffusionConfig = config self.params: DiffusionParams = params @@ -470,7 +470,7 @@ def _index_field(dim: Dimension, size=None): self.horizontal_edge_index = _index_field(EdgeDim) self.w_tmp = np_as_located_field(CellDim, KDim)(np.zeros((self.grid.num_cells(), self.grid.n_lev() + 1), dtype=float)) - def initial_step( + def initial_run( self, diagnostic_state: DiagnosticState, prognostic_state: PrognosticState, diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index c551d5264..54bf386ed 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -20,7 +20,6 @@ from devtools import Timer from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn -from atm_dyn_iconam.tests.test_utils.serialbox_utils import IconSerialDataProvider from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.diffusion.horizontal import CellParams, EdgeParams @@ -33,15 +32,18 @@ read_geometry_fields, read_icon_grid, read_initial_state, - read_static_fields, + read_static_fields, import_testutils, ) +helpers = import_testutils() +from helpers import serialbox_utils as sb_utils + log = logging.getLogger(__name__) class DummyAtmoNonHydro: - def __init__(self, data_provider: IconSerialDataProvider): + def __init__(self, data_provider: sb_utils.IconSerialDataProvider): self.config = None self.data_provider = data_provider self.simulation_date = datetime.fromisoformat(SIMULATION_START_DATE) @@ -197,7 +199,7 @@ def initialize(n_time_steps, file_path: Path): log.info("initializing dycore") diffusion_params = DiffusionParams(config.diffusion_config) - diffusion = Diffusion(run_program=False) + diffusion = Diffusion() diffusion.init( icon_grid, config.diffusion_config, diff --git a/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py b/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py index 0c55ae88f..d0c14aa5c 100644 --- a/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py +++ b/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py @@ -16,7 +16,7 @@ from icon4py.diffusion.diffusion import DiffusionConfig - +n_substeps_reduced = 2 @dataclass class IconRunConfig: n_time_steps: int = 5 @@ -35,7 +35,6 @@ class IconConfig: dycore_config: AtmoNonHydroConfig -# TODO @magdalena move to io_utils? def read_config(experiment: Optional[str], n_time_steps: int) -> IconConfig: def _default_run_config(n_steps: int): if n_steps > 5: @@ -46,6 +45,7 @@ def mch_ch_r04b09_diffusion_config(): return DiffusionConfig( diffusion_type=5, hdiff_w=True, + n_substeps=n_substeps_reduced, hdiff_vn=True, type_t_diffu=2, type_vn_diffu=1, diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index f5f15ba09..d11095077 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -10,13 +10,13 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import importlib import logging +import sys from datetime import datetime from enum import Enum from pathlib import Path -from atm_dyn_iconam.tests.test_utils import serialbox_utils -from atm_dyn_iconam.tests.test_utils.serialbox_utils import IconSerialDataProvider from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.horizontal import CellParams, EdgeParams from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams @@ -29,6 +29,18 @@ log = logging.getLogger(__name__) +def import_testutils(): + testutils = Path(__file__).parent.__str__() + "/../../../tests/test_utils/__init__.py" + spec = importlib.util.spec_from_file_location("helpers", testutils) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + +helpers = import_testutils() + +from helpers import serialbox_utils as sb + class SerializationType(str, Enum): SB = "serialbox" NC = "netcdf" @@ -46,7 +58,7 @@ def read_icon_grid(path: Path, ser_type=SerializationType.SB) -> IconGrid: """ if ser_type == SerializationType.SB: return ( - serialbox_utils.IconSerialDataProvider( + sb.IconSerialDataProvider( "icon_pydycore", str(path.absolute()), False ) .from_savepoint_grid() @@ -58,7 +70,7 @@ def read_icon_grid(path: Path, ser_type=SerializationType.SB) -> IconGrid: def read_initial_state( gridfile_path: Path, -) -> tuple[IconSerialDataProvider, DiagnosticState, PrognosticState]: +) -> tuple[sb.IconSerialDataProvider, DiagnosticState, PrognosticState]: """ Read prognostic and diagnostic state from serialized data. @@ -70,7 +82,7 @@ def read_initial_state( read from within the dummy timeloop """ - data_provider = serialbox_utils.IconSerialDataProvider( + data_provider = sb.IconSerialDataProvider( "icon_pydycore", str(gridfile_path), False ) init_savepoint = data_provider.from_savepoint_diffusion_init( @@ -95,7 +107,7 @@ def read_geometry_fields( the data is originally obtained from the grid file (horizontal fields) or some special input files. """ if ser_type == SerializationType.SB: - sp = serialbox_utils.IconSerialDataProvider( + sp = sb.IconSerialDataProvider( "icon_pydycore", str(path.absolute()), False ).from_savepoint_grid() edge_geometry = sp.construct_edge_geometry() @@ -123,7 +135,7 @@ def read_static_fields( """ if ser_type == SerializationType.SB: - dataprovider = serialbox_utils.IconSerialDataProvider( + dataprovider = sb.IconSerialDataProvider( "icon_pydycore", str(path.absolute()), False ) interpolation_state = ( diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 336efcd77..16781912c 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -531,7 +531,7 @@ def test_run_diffusion_initial_step( metric_state=metric_state, interpolation_state=interpolation_state, ) - diffusion.initial_step( + diffusion.initial_run( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime, diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index fa3fbf238..94b02fb98 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -49,9 +49,6 @@ from .helpers import as_1D_sparse_field -C2E2C_DIM = 3 - - class IconSavepoint: def __init__(self, sp: ser.Savepoint, ser: ser.Serializer): self.savepoint = sp From a5dc6461a99327cbadced01aad89a5e312c1b8dd Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 29 Jun 2023 16:48:40 +0200 Subject: [PATCH 153/263] fix some todos --- .../src/icon4py/diffusion/diffusion.py | 78 ++++++++++++------- .../src/icon4py/diffusion/icon_grid.py | 19 +++-- atm_dyn_iconam/src/icon4py/diffusion/utils.py | 5 +- .../src/icon4py/driver/dycore_driver.py | 6 +- .../src/icon4py/driver/icon_configuration.py | 3 + atm_dyn_iconam/src/icon4py/driver/io_utils.py | 12 +-- atm_dyn_iconam/tests/test_vertical.py | 3 +- .../py2f/wrappers/diffusion_wrapper.py | 3 +- 8 files changed, 84 insertions(+), 45 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 799f65382..5aa746122 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -25,8 +25,9 @@ run_gtfn_cached, ) -from icon4py.atm_dyn_iconam.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance - +from icon4py.atm_dyn_iconam.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( + apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance, +) from icon4py.atm_dyn_iconam.calculate_diagnostic_quantities_for_turbulence import ( calculate_diagnostic_quantities_for_turbulence, ) @@ -62,11 +63,13 @@ from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.prognostic_state import PrognosticState from icon4py.diffusion.utils import ( + copy_field, init_diffusion_local_fields_for_regular_timestep, init_nabla2_factor_in_upper_damping_zone, scale_k, + set_zero_v_k, setup_fields_for_initial_step, - zero_field, copy_field, set_zero_v_k, + zero_field, ) @@ -111,7 +114,7 @@ def __init__( max_nudging_coeff: float = 0.02, nudging_decay_rate: float = 2.0, ): - """ Set the diffusion configuration parameters with the ICON default values.""" + """Set the diffusion configuration parameters with the ICON default values.""" # parameters from namelist diffusion_nml self.diffusion_type: int = diffusion_type @@ -468,7 +471,9 @@ def _index_field(dim: Dimension, size=None): self.vertical_index = _index_field(KDim, self.grid.n_lev() + 1) self.horizontal_cell_index = _index_field(CellDim) self.horizontal_edge_index = _index_field(EdgeDim) - self.w_tmp = np_as_located_field(CellDim, KDim)(np.zeros((self.grid.num_cells(), self.grid.n_lev() + 1), dtype=float)) + self.w_tmp = np_as_located_field(CellDim, KDim)( + np.zeros((self.grid.num_cells(), self.grid.n_lev() + 1), dtype=float) + ) def initial_run( self, @@ -659,7 +664,6 @@ def _do_diffusion_step( self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={} ) - # TODO: @magdalena is this needed?, if not remove set_zero_v_k.with_backend(backend)(self.u_vert, offset_provider={}) set_zero_v_k.with_backend(backend)(self.v_vert, offset_provider={}) log.debug("rbf interpolation: start") @@ -679,7 +683,9 @@ def _do_diffusion_step( # HALO EXCHANGE -- CALL sync_patch_array_mult - log.debug("running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start") + log.debug( + "running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start" + ) calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend)( diff_multfac_smag=self.diff_multfac_smag, tangent_orientation=tangent_orientation, @@ -706,8 +712,12 @@ def _do_diffusion_step( "E2ECV": self.grid.get_e2ecv_connectivity(), }, ) - log.debug("running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end") - log.debug("running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start") + log.debug( + "running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end" + ) + log.debug( + "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start" + ) calculate_diagnostic_quantities_for_turbulence.with_backend(backend)( kh_smag_ec=self.kh_smag_ec, vn=prognostic_state.vn, @@ -727,7 +737,9 @@ def _do_diffusion_step( "Koff": KDim, }, ) - log.debug("running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end") + log.debug( + "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end" + ) # HALO EXCHANGE IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH @@ -777,10 +789,15 @@ def _do_diffusion_step( ) log.debug("runningstencils 04 05 06: end") - log.debug("running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start") - copy_field.with_backend(backend)(prognostic_state.w, self.w_tmp, offset_provider={}) - # TODO @magdalena: fix this python offset problem (+1): int(self.vertical_params.index_of_damping_layer + 1) - apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.with_backend(backend)( + log.debug( + "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" + ) + copy_field.with_backend(backend)( + prognostic_state.w, self.w_tmp, offset_provider={} + ) + apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.with_backend( + backend + )( area=cell_areas, geofac_n2s=self.interpolation_state.geofac_n2s, geofac_grg_x=self.interpolation_state.geofac_grg_x, @@ -793,13 +810,11 @@ def _do_diffusion_step( diff_multfac_n2w=self.diff_multfac_n2w, vert_idx=self.vertical_index, horz_idx=self.horizontal_cell_index, - nrdmax=int32(self.vertical_params.index_of_damping_layer +1) , - interior_idx=int32( - cell_start_interior - ), - halo_idx=int32( - cell_end_local - ), + nrdmax=int32( + self.vertical_params.index_of_damping_layer + 1 + ), # +1 since Fortran includes boundaries + interior_idx=int32(cell_start_interior), + halo_idx=int32(cell_end_local), horizontal_start=cell_start_nudging, horizontal_end=cell_end_halo, vertical_start=0, @@ -808,11 +823,14 @@ def _do_diffusion_step( "C2E2CO": self.grid.get_c2e2co_connectivity(), }, ) - log.debug("running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end") + log.debug( + "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end" + ) # HALO EXCHANGE: CALL sync_patch_array - # # TODO @magdalena check: kh_smag_e is an out field, should not be calculated in init? - log.debug("running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start") + log.debug( + "running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" + ) calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.with_backend( backend )( @@ -829,7 +847,9 @@ def _do_diffusion_step( "C2E2C": self.grid.get_c2e2c_connectivity(), }, ) - log.debug("running stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): end") + log.debug( + "running stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): end" + ) log.debug("running stencils 13 14 (calculate_nabla2_for_theta): start") calculate_nabla2_for_theta.with_backend(backend)( kh_smag_e=self.kh_smag_e, @@ -848,7 +868,9 @@ def _do_diffusion_step( }, ) log.debug("running stencils 13_14 (calculate_nabla2_for_theta): end") - log.debug("running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start") + log.debug( + "running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start" + ) truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_backend( backend )( @@ -871,7 +893,9 @@ def _do_diffusion_step( }, ) - log.debug("running fused stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end") + log.debug( + "running fused stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end" + ) log.debug("running fused stencil 16 (update_theta_and_exner): start") update_theta_and_exner.with_backend(backend)( z_temp=self.z_temp, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index 5d2feae3f..53b8eb823 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from typing import Dict, Tuple +from typing import Dict import numpy as np from gt4py.next.common import Dimension, DimensionKind, Field @@ -182,9 +182,15 @@ def get_v2e_connectivity(self): def get_e2ecv_connectivity(self): return self._neighbortable_offset_provider_for_1d_sparse_fields( - self.connectivities["e2c2v"].shape, EdgeDim, ECVDim) + self.connectivities["e2c2v"].shape, EdgeDim, ECVDim + ) - def _neighbortable_offset_provider_for_1d_sparse_fields(self, old_shape:tuple[int, int], origin_axis:Dimension, neighbor_axis:Dimension ): + def _neighbortable_offset_provider_for_1d_sparse_fields( + self, + old_shape: tuple[int, int], + origin_axis: Dimension, + neighbor_axis: Dimension, + ): table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) return NeighborTableOffsetProvider( table, origin_axis, neighbor_axis, table.shape[1] @@ -192,12 +198,13 @@ def _neighbortable_offset_provider_for_1d_sparse_fields(self, old_shape:tuple[in def get_c2cec_connectivity(self): return self._neighbortable_offset_provider_for_1d_sparse_fields( - self.connectivities["c2e2c"].shape, CellDim, CECDim) + self.connectivities["c2e2c"].shape, CellDim, CECDim + ) def get_c2ce_connectivity(self): return self._neighbortable_offset_provider_for_1d_sparse_fields( - self.connectivities["c2e"].shape, CellDim, CEDim) - + self.connectivities["c2e"].shape, CellDim, CEDim + ) class VerticalModelParams: diff --git a/atm_dyn_iconam/src/icon4py/diffusion/utils.py b/atm_dyn_iconam/src/icon4py/diffusion/utils.py index 077c8b1ac..073a7915c 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/utils.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/utils.py @@ -33,8 +33,11 @@ def _identity_c_k( ) -> Field[[CellDim, KDim], float]: return field + @program -def copy_field(old_f:Field[[CellDim, KDim], float], new_f:Field[[CellDim, KDim], float]): +def copy_field( + old_f: Field[[CellDim, KDim], float], new_f: Field[[CellDim, KDim], float] +): _identity_c_k(old_f, out=new_f) diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index 54bf386ed..55687db98 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -29,14 +29,16 @@ from icon4py.driver.io_utils import ( SIMULATION_START_DATE, configure_logging, + import_testutils, read_geometry_fields, read_icon_grid, read_initial_state, - read_static_fields, import_testutils, + read_static_fields, ) + helpers = import_testutils() -from helpers import serialbox_utils as sb_utils +from helpers import serialbox_utils as sb_utils # noqa log = logging.getLogger(__name__) diff --git a/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py b/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py index d0c14aa5c..7e75ebf63 100644 --- a/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py +++ b/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py @@ -16,7 +16,10 @@ from icon4py.diffusion.diffusion import DiffusionConfig + n_substeps_reduced = 2 + + @dataclass class IconRunConfig: n_time_steps: int = 5 diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index d11095077..07ac5f5e1 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -30,16 +30,20 @@ def import_testutils(): - testutils = Path(__file__).parent.__str__() + "/../../../tests/test_utils/__init__.py" + testutils = ( + Path(__file__).parent.__str__() + "/../../../tests/test_utils/__init__.py" + ) spec = importlib.util.spec_from_file_location("helpers", testutils) module = importlib.util.module_from_spec(spec) sys.modules[spec.name] = module spec.loader.exec_module(module) return module + helpers = import_testutils() -from helpers import serialbox_utils as sb +from helpers import serialbox_utils as sb # noqa + class SerializationType(str, Enum): SB = "serialbox" @@ -58,9 +62,7 @@ def read_icon_grid(path: Path, ser_type=SerializationType.SB) -> IconGrid: """ if ser_type == SerializationType.SB: return ( - sb.IconSerialDataProvider( - "icon_pydycore", str(path.absolute()), False - ) + sb.IconSerialDataProvider("icon_pydycore", str(path.absolute()), False) .from_savepoint_grid() .construct_icon_grid() ) diff --git a/atm_dyn_iconam/tests/test_vertical.py b/atm_dyn_iconam/tests/test_vertical.py index 7cd3335dc..597ee5acf 100644 --- a/atm_dyn_iconam/tests/test_vertical.py +++ b/atm_dyn_iconam/tests/test_vertical.py @@ -34,10 +34,9 @@ def test_nrdmax_calculation(max_h, damping, delta): @pytest.mark.datatest -def test_nrdmax_calculation_from_icon_input(icon_grid, grid_savepoint): +def test_nrdmax_calculation_from_icon_input(icon_grid, grid_savepoint, damping_height): a = grid_savepoint.vct_a() nrdmax = grid_savepoint.nrdmax() - damping_height = 12500 vertical_params = VerticalModelParams( rayleigh_damping_height=damping_height, vct_a=a ) diff --git a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py index fb55f4411..017b65a49 100644 --- a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py +++ b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py @@ -39,7 +39,7 @@ from icon4py.py2f.cffi_utils import CffiMethod -diffusion: Diffusion(run_program=True) +diffusion: Diffusion() @CffiMethod.register @@ -154,5 +154,4 @@ def diffusion_run( class DuplicateInitializationException(Exception): """Raised if the component is already initilalized.""" - pass From dc58678dda968465d6a45ae93ea1a3be9fe220c2 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 29 Jun 2023 22:48:22 +0200 Subject: [PATCH 154/263] change diffusion interface: pass geometry fields to init --- .../src/icon4py/diffusion/diffusion.py | 101 +++++++----------- .../src/icon4py/driver/dycore_driver.py | 25 +---- atm_dyn_iconam/tests/test_diffusion.py | 51 +++++---- .../py2f/wrappers/diffusion_wrapper.py | 7 +- 4 files changed, 70 insertions(+), 114 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 5aa746122..4675c0b26 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -57,7 +57,11 @@ ) from icon4py.common.dimension import CellDim, ECVDim, EdgeDim, KDim, VertexDim from icon4py.diffusion.diagnostic_state import DiagnosticState -from icon4py.diffusion.horizontal import HorizontalMarkerIndex +from icon4py.diffusion.horizontal import ( + CellParams, + EdgeParams, + HorizontalMarkerIndex, +) from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState @@ -382,6 +386,8 @@ def __init__(self): self.fac_bdydiff_v: Optional[float] = None self.bdy_diff: Optional[float] = None self.nudgezone_diff: Optional[float] = None + self.edge_params: Optional[EdgeParams] = None + self.cell_params: Optional[CellParams] = None def init( self, @@ -391,11 +397,23 @@ def init( vertical_params: VerticalModelParams, metric_state: MetricState, interpolation_state: InterpolationState, + edge_params: EdgeParams, + cell_params: CellParams, ): """ Initialize Diffusion granule with configuration. calculates all local fields that are used in diffusion within the time loop. + + Args: + grid: + config: + params: + vertical_params: + metric_state: + interpolation_state: + edge_params: + cell_params: """ self.config: DiffusionConfig = config self.params: DiffusionParams = params @@ -403,6 +421,8 @@ def init( self.vertical_params = vertical_params self.metric_state: MetricState = metric_state self.interpolation_state: InterpolationState = interpolation_state + self.edge_params = edge_params + self.cell_params = cell_params self._allocate_local_fields() @@ -480,14 +500,6 @@ def initial_run( diagnostic_state: DiagnosticState, prognostic_state: PrognosticState, dtime: float, - tangent_orientation: Field[[EdgeDim], float], - inverse_primal_edge_lengths: Field[[EdgeDim], float], - inverse_dual_edge_length: Field[[EdgeDim], float], - inverse_vert_vert_lengths: Field[[EdgeDim], float], - primal_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], - dual_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], - edge_areas: Field[[EdgeDim], float], - cell_areas: Field[[CellDim], float], ): """ Calculate initial diffusion step. @@ -514,14 +526,6 @@ def initial_run( diagnostic_state, prognostic_state, dtime, - tangent_orientation, - inverse_primal_edge_lengths, - inverse_dual_edge_length, - inverse_vert_vert_lengths, - primal_normal_vert, - dual_normal_vert, - edge_areas, - cell_areas, diff_multfac_vn, smag_limit, 0.0, @@ -532,14 +536,6 @@ def run( diagnostic_state: DiagnosticState, prognostic_state: PrognosticState, dtime: float, - tangent_orientation: Field[[EdgeDim], float], - inverse_primal_edge_lengths: Field[[EdgeDim], float], - inverse_dual_edge_length: Field[[EdgeDim], float], - inverse_vert_vert_lengths: Field[[EdgeDim], float], - primal_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], - dual_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], - edge_areas: Field[[EdgeDim], float], - cell_areas: Field[[CellDim], float], ): """ Do one diffusion step within regular time loop. @@ -551,33 +547,16 @@ def run( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime, - tangent_orientation=tangent_orientation, - inverse_primal_edge_lengths=inverse_primal_edge_lengths, - inverse_dual_edge_length=inverse_dual_edge_length, - inverse_vertex_vertex_lengths=inverse_vert_vert_lengths, - primal_normal_vert=primal_normal_vert, - dual_normal_vert=dual_normal_vert, - edge_areas=edge_areas, - cell_areas=cell_areas, diff_multfac_vn=self.diff_multfac_vn, smag_limit=self.smag_limit, smag_offset=self.smag_offset, ) - log.info("diffusion program: end") def _do_diffusion_step( self, diagnostic_state: DiagnosticState, prognostic_state: PrognosticState, dtime: float, - tangent_orientation: Field[[EdgeDim], float], - inverse_primal_edge_lengths: Field[[EdgeDim], float], - inverse_dual_edge_length: Field[[EdgeDim], float], - inverse_vertex_vertex_lengths: Field[[EdgeDim], float], - primal_normal_vert: Tuple[Field[[ECVDim], float], Field[[ECVDim], float]], - dual_normal_vert: Tuple[Field[[ECVDim], float], Field[[ECVDim], float]], - edge_areas: Field[[EdgeDim], float], - cell_areas: Field[[CellDim], float], diff_multfac_vn: Field[[KDim], float], smag_limit: Field[[KDim], float], smag_offset: float, @@ -589,14 +568,6 @@ def _do_diffusion_step( diagnostic_state: output argument, data class that contains diagnostic variables prognostic_state: output argument, data class that contains prognostic variables dtime: the time step, - tangent_orientation: - inverse_primal_edge_lengths: - inverse_dual_edge_length: - inverse_vertex_vertex_lengths: - primal_normal_vert: - dual_normal_vert: - edge_areas: - cell_areas: diff_multfac_vn: smag_limit: smag_offset: @@ -688,15 +659,15 @@ def _do_diffusion_step( ) calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend)( diff_multfac_smag=self.diff_multfac_smag, - tangent_orientation=tangent_orientation, - inv_primal_edge_length=inverse_primal_edge_lengths, - inv_vert_vert_length=inverse_vertex_vertex_lengths, + tangent_orientation=self.edge_params.tangent_orientation, + inv_primal_edge_length=self.edge_params.inverse_primal_edge_lengths, + inv_vert_vert_length=self.edge_params.inverse_vertex_vertex_lengths, u_vert=self.u_vert, v_vert=self.v_vert, - primal_normal_vert_x=primal_normal_vert[0], - primal_normal_vert_y=primal_normal_vert[1], - dual_normal_vert_x=dual_normal_vert[0], - dual_normal_vert_y=dual_normal_vert[1], + primal_normal_vert_x=self.edge_params.primal_normal_vert[0], + primal_normal_vert_y=self.edge_params.primal_normal_vert[1], + dual_normal_vert_x=self.edge_params.dual_normal_vert[0], + dual_normal_vert_y=self.edge_params.dual_normal_vert[1], vn=prognostic_state.vn, smag_limit=smag_limit, kh_smag_e=self.kh_smag_e, @@ -764,12 +735,12 @@ def _do_diffusion_step( fused_mo_nh_diffusion_stencil_04_05_06.with_backend(backend)( u_vert=self.u_vert, v_vert=self.v_vert, - primal_normal_vert_v1=primal_normal_vert[0], - primal_normal_vert_v2=primal_normal_vert[1], + primal_normal_vert_v1=self.edge_params.primal_normal_vert[0], + primal_normal_vert_v2=self.edge_params.primal_normal_vert[1], z_nabla2_e=self.z_nabla2_e, - inv_vert_vert_length=inverse_vertex_vertex_lengths, - inv_primal_edge_length=inverse_primal_edge_lengths, - area_edge=edge_areas, + inv_vert_vert_length=self.edge_params.inverse_vertex_vertex_lengths, + inv_primal_edge_length=self.edge_params.inverse_primal_edge_lengths, + area_edge=self.edge_params.edge_areas, kh_smag_e=self.kh_smag_e, diff_multfac_vn=diff_multfac_vn, nudgecoeff_e=self.interpolation_state.nudgecoeff_e, @@ -798,7 +769,7 @@ def _do_diffusion_step( apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.with_backend( backend )( - area=cell_areas, + area=self.cell_params.area, geofac_n2s=self.interpolation_state.geofac_n2s, geofac_grg_x=self.interpolation_state.geofac_grg_x, geofac_grg_y=self.interpolation_state.geofac_grg_y, @@ -853,7 +824,7 @@ def _do_diffusion_step( log.debug("running stencils 13 14 (calculate_nabla2_for_theta): start") calculate_nabla2_for_theta.with_backend(backend)( kh_smag_e=self.kh_smag_e, - inv_dual_edge_length=inverse_dual_edge_length, + inv_dual_edge_length=self.edge_params.inverse_dual_edge_lengths, theta_v=prognostic_state.theta_v, geofac_div=self.interpolation_state.geofac_div, z_temp=self.z_temp, @@ -899,7 +870,7 @@ def _do_diffusion_step( log.debug("running fused stencil 16 (update_theta_and_exner): start") update_theta_and_exner.with_backend(backend)( z_temp=self.z_temp, - area=cell_areas, + area=self.cell_params.area, theta_v=prognostic_state.theta_v, exner=prognostic_state.exner_pressure, rd_o_cvd=self.rd_o_cvd, diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index 55687db98..89baa998f 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -22,7 +22,6 @@ from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.diffusion import Diffusion, DiffusionParams -from icon4py.diffusion.horizontal import CellParams, EdgeParams from icon4py.diffusion.prognostic_state import PrognosticState from icon4py.diffusion.utils import copy_diagnostic_and_prognostics from icon4py.driver.icon_configuration import IconRunConfig, read_config @@ -105,14 +104,10 @@ def __init__( config: IconRunConfig, diffusion: Diffusion, atmo_non_hydro: DummyAtmoNonHydro, - edge_geometry: EdgeParams, - cell_geometry: CellParams, ): self.config = config self.diffusion = diffusion self.atmo_non_hydro = atmo_non_hydro - self.edges = edge_geometry - self.cells = cell_geometry def _full_name(self, func: Callable): return ":".join((self.__class__.__name__, func.__name__)) @@ -130,14 +125,6 @@ def _timestep( diagnostic_state, prognostic_state, self.config.dtime, - self.edges.tangent_orientation, - self.edges.inverse_primal_edge_lengths, - self.edges.inverse_dual_edge_lengths, - self.edges.inverse_vertex_vertex_lengths, - self.edges.primal_normal_vert, - self.edges.dual_normal_vert, - self.edges.edge_areas, - self.cells.area, ) def __call__( @@ -153,14 +140,6 @@ def __call__( diagnostic_state, prognostic_state, self.config.dtime, - self.edges.tangent_orientation, - self.edges.inverse_primal_edge_lengths, - self.edges.inverse_dual_edge_lengths, - self.edges.inverse_vertex_vertex_lengths, - self.edges.primal_normal_vert, - self.edges.dual_normal_vert, - self.edges.edge_areas, - self.cells.area, ) log.info( f"starting real time loop for dtime={self.config.dtime} n_timesteps={self.config.n_time_steps}" @@ -209,6 +188,8 @@ def initialize(n_time_steps, file_path: Path): vertical_geometry, metric_state, interpolation_state, + edge_geometry, + cell_geometry, ) data_provider, diagnostic_state, prognostic_state = read_initial_state(file_path) @@ -220,8 +201,6 @@ def initialize(n_time_steps, file_path: Path): config=config.run_config, diffusion=diffusion, atmo_non_hydro=atmo_non_hydro, - edge_geometry=edge_geometry, - cell_geometry=cell_geometry, ) return tl, diagnostic_state, prognostic_state diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 16781912c..1a7c17ea6 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -243,15 +243,19 @@ def test_diffusion_init( interpolation_state = interpolation_savepoint.construct_interpolation_state() metric_state = metrics_savepoint.construct_metric_state() + edge_params = grid_savepoint.construct_edge_geometry() + cell_params = grid_savepoint.construct_cell_geometry() diffusion = Diffusion() diffusion.init( grid=icon_grid, - vertical_params=vertical_params, config=config, params=additional_parameters, + vertical_params=vertical_params, metric_state=metric_state, interpolation_state=interpolation_state, + edge_params=edge_params, + cell_params=cell_params, ) assert diffusion.diff_multfac_w == min( @@ -355,18 +359,22 @@ def test_verify_diffusion_init_against_first_regular_savepoint( config.ndyn_substeps = datarun_reduced_substeps additional_parameters = DiffusionParams(config) vct_a = grid_savepoint.vct_a() + cell_geometry = grid_savepoint.construct_cell_geometry() + edge_geometry = grid_savepoint.construct_edge_geometry() interpolation_state = interpolation_savepoint.construct_interpolation_state() metric_state = metrics_savepoint.construct_metric_state() diffusion = Diffusion() diffusion.init( - config=config, grid=icon_grid, + config=config, params=additional_parameters, vertical_params=VerticalModelParams(vct_a, damping_height), metric_state=metric_state, interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry, ) _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) @@ -390,6 +398,8 @@ def test_verify_diffusion_init_against_other_regular_savepoint( vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) interpolation_state = interpolation_savepoint.construct_interpolation_state() metric_state = metrics_savepoint.construct_metric_state() + edge_params = grid_savepoint.construct_edge_geometry() + cell_params = grid_savepoint.construct_cell_geometry() diffusion = Diffusion() diffusion.init( @@ -399,6 +409,8 @@ def test_verify_diffusion_init_against_other_regular_savepoint( vertical_params, metric_state, interpolation_state, + edge_params, + cell_params, ) _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) @@ -446,27 +458,24 @@ def test_run_diffusion_single_step( vertical_params=vertical_params, metric_state=metric_state, interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry, + ) + _verify_diffusion_fields( + diagnostic_state, prognostic_state, diffusion_savepoint_init ) + assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v diffusion.run( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime, - tangent_orientation=edge_geometry.tangent_orientation, - inverse_primal_edge_lengths=edge_geometry.inverse_primal_edge_lengths, - inverse_dual_edge_length=edge_geometry.inverse_dual_edge_lengths, - inverse_vert_vert_lengths=edge_geometry.inverse_vertex_vertex_lengths, - primal_normal_vert=edge_geometry.primal_normal_vert, - dual_normal_vert=edge_geometry.dual_normal_vert, - edge_areas=edge_geometry.edge_areas, - cell_areas=cell_geometry.area, ) - assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v - _verify_diffusion_exit_fields( + _verify_diffusion_fields( diagnostic_state, prognostic_state, diffusion_savepoint_exit ) -def _verify_diffusion_exit_fields( +def _verify_diffusion_fields( diagnostic_state, prognostic_state, diffusion_savepoint_exit ): ref_div_ic = np.asarray(diffusion_savepoint_exit.div_ic()) @@ -530,22 +539,18 @@ def test_run_diffusion_initial_step( vertical_params=vertical_params, metric_state=metric_state, interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry, ) + assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v + diffusion.initial_run( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime, - tangent_orientation=edge_geometry.tangent_orientation, - inverse_primal_edge_lengths=edge_geometry.inverse_primal_edge_lengths, - inverse_dual_edge_length=edge_geometry.inverse_dual_edge_lengths, - inverse_vert_vert_lengths=edge_geometry.inverse_vertex_vertex_lengths, - primal_normal_vert=edge_geometry.primal_normal_vert, - dual_normal_vert=edge_geometry.dual_normal_vert, - edge_areas=edge_geometry.edge_areas, - cell_areas=cell_geometry.area, ) - assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v - _verify_diffusion_exit_fields( + + _verify_diffusion_fields( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, diffusion_savepoint_exit=diffusion_savepoint_exit, diff --git a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py index 017b65a49..821081046 100644 --- a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py +++ b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py @@ -81,7 +81,7 @@ def diffusion_init( """ grid = IconGrid() # TODO where to get this from - edges_params = EdgeParams( + edge_params = EdgeParams( tangent_orientation=tangent_orientation, primal_edge_lengths=primal_edge_lengths, inverse_primal_edge_lengths=inverse_primal_edge_lengths, @@ -112,13 +112,13 @@ def diffusion_init( diffusion.init( grid=grid, - cell_params=cell_params, - edges_params=edges_params, config=config, params=derived_diffusion_params, vertical_params=vertical_params, metric_state=metric_state, interpolation_state=interpolation_state, + edge_params=edge_params, + cell_params=cell_params, ) @@ -154,4 +154,5 @@ def diffusion_run( class DuplicateInitializationException(Exception): """Raised if the component is already initilalized.""" + pass From b0cf692536017a34d6a08e1c345967960d10f7a3 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 30 Jun 2023 09:02:12 +0200 Subject: [PATCH 155/263] make setup_icon_data test for specific datafolder --- atm_dyn_iconam/tests/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 4f5cd0cc6..626214535 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -38,7 +38,7 @@ def setup_icon_data(): Session scoped fixture which is a prerequisite of all the other fixtures in this file. """ data_path.mkdir(parents=True, exist_ok=True) - if not any(data_path.iterdir()): + if not data_path.joinpath(extracted_path).exists(): print( f"directory {data_path} is empty: downloading data from {data_uri} and extracting" ) From 1354599a8314aa9992c393a667b4f96c264ece2f Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 30 Jun 2023 11:49:20 +0200 Subject: [PATCH 156/263] move parallel_setup.py --- .../src/icon4py/{driver => decomposition}/parallel_setup.py | 0 atm_dyn_iconam/src/icon4py/diffusion/diffusion.py | 4 ++-- atm_dyn_iconam/src/icon4py/driver/dycore_driver.py | 2 +- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 2 +- atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py | 2 +- atm_dyn_iconam/tests/test_utils/serialbox_utils.py | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) rename atm_dyn_iconam/src/icon4py/{driver => decomposition}/parallel_setup.py (100%) diff --git a/atm_dyn_iconam/src/icon4py/driver/parallel_setup.py b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py similarity index 100% rename from atm_dyn_iconam/src/icon4py/driver/parallel_setup.py rename to atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 8465149e2..acda67cde 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -14,7 +14,7 @@ import math import sys from collections import namedtuple -from typing import Final, Optional, Tuple +from typing import Final, Optional import numpy as np from gt4py.next.common import Dimension @@ -76,7 +76,7 @@ setup_fields_for_initial_step, zero_field, ) -from icon4py.driver.parallel_setup import Exchange +from icon4py.decomposition.parallel_setup import Exchange # flake8: noqa diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index cfc986f75..8f10357f2 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -35,7 +35,7 @@ read_initial_state, read_static_fields, ) -from icon4py.driver.parallel_setup import get_processor_properties +from icon4py.decomposition.parallel_setup import get_processor_properties helpers = import_testutils() diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index a4c349df1..feeb8fa4c 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -24,7 +24,7 @@ from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.prognostic_state import PrognosticState -from icon4py.driver.parallel_setup import DecompositionInfo +from icon4py.decomposition.parallel_setup import DecompositionInfo SERIALBOX_ONLY_MSG = "Only ser_type='sb (Serialbox)' is implemented so far." diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index 16d341879..2e437a625 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -27,7 +27,7 @@ read_icon_grid, read_static_fields, ) -from icon4py.driver.parallel_setup import ( +from icon4py.decomposition.parallel_setup import ( DecompositionInfo, Exchange, get_processor_properties, diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index 0c93af6d4..33885ed62 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -46,7 +46,7 @@ from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.prognostic_state import PrognosticState -from icon4py.driver.parallel_setup import DecompositionInfo +from icon4py.decomposition.parallel_setup import DecompositionInfo from .helpers import as_1D_sparse_field From 07de07d5ba562491b08f92acf66d4635862999e8 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 30 Jun 2023 11:49:31 +0200 Subject: [PATCH 157/263] readme for ghex build --- atm_dyn_iconam/tests/mpi_tests/README.md | 71 ++++++++++++++++++++++++ 1 file changed, 71 insertions(+) create mode 100644 atm_dyn_iconam/tests/mpi_tests/README.md diff --git a/atm_dyn_iconam/tests/mpi_tests/README.md b/atm_dyn_iconam/tests/mpi_tests/README.md new file mode 100644 index 000000000..133097620 --- /dev/null +++ b/atm_dyn_iconam/tests/mpi_tests/README.md @@ -0,0 +1,71 @@ +## running parallel version of diffusion + +### installation +The parallelized code uses [GHEX](https://github.com/ghex-org/GHEX) with MPI for halo exchanges. +GHEX has a CMake build but no setup script for pip, so it needs to be installed manually: + +1. You need a running MPI installation in the system. +2. clone GHEX +```bash +cd {icon4py}/_external_src +git clone --recursive -b refactoring2 git@github.com:boeschf/GHEX.git +``` +3. build GHEX +``` +cd GHEX +mkdir build +cd build +cmake .. -DCMAKE_BUILD_TYPE=Debug \ +-DGHEX_GIT_SUBMODULE=OFF \ +-DGHEX_USE_BUNDLED_LIBS=ON \ +-DGHEX_USE_BUNDLED_GRIDTOOLS=ON \ +-DBUILD_TESTING=OFF \ +-DGT_BUILD_TESTING=OFF \ +-DGT_INSTALL_EXAMPLES=OFF \ +-DGHEX_USE_BUNDLED_OOMPH=ON \ +-DGHEX_TRANSPORT_BACKEND=MPI \ +-DGHEX_USE_XPMEM=OFF \ +-DGHEX_WITH_TESTING=ON \ +-DGHEX_BUILD_PYTHON_BINDINGS=ON \ +-DGHEX_BUILD_BENCHMARKS=ON \ +-DGHEX_BUILD_FORTRAN=ON \ +-DGHEX_ENABLE_ATLAS_BINDINGS=OFF \ +-DGHEX_ENABLE_PARMETIS_BINDINGS=ON \ +-DMETIS_INCLUDE_DIR=/opt/metis/include \ +-DMETIS_LIB_DIR=/opt/metis/lib \ +-DPARMETIS_INCLUDE_DIR=/opt/parmetis/include \ +-DPARMETIS_LIB_DIR=/opt/parmetis/lib \ +-DMPIEXEC_PREFLAGS=--oversubscribe \ +-DGHEX_USE_GPU=OFF +make +make test ## runs the C++ tests +``` +#### generating python bindings +```bash +cmake -DGHEX_BUILD_PYTHON_BINDINGS=ON . # turns on python bindings, need pybind11 wo we install it in the +``` +builds GHEX including python bindings. If the build fails with an error message that `pybind11` was not found, you can either install it +int the Python `.venv` in GHEX build folder or set the `pybind11_DIR` variable to some other location. +```bash +cd pyghex_venv/ +source bin/activate +pip install pybind11 +export pybind11_DIR=./pyghex_venv/lib/python3.10/site-packages/pybind11 +make +make test ## will now run the python tests +``` +#### use GHEX bindings from icon4py +simply create a sym link the the installation above: +``` +cd {icon4py}/.venv/lib/python3.10/site-packages +ln -s ../../../../_external_src/GHEX/build/bindings/python/ghex ghex +``` +### run parallel test +all MPI related tests are in the folder +`icon4py/atm_dyn_iconam/tests/mpi_tests` tests depending on MPI are marked with `@pytest.mark.mpi` are +therefore skipped by any serial pytest run. + +run them with +```bash +mpirun -np 2 pytest -v --with-mpi tests/mpi_tests/ +``` From 1c9d8c07d26e595e984dd231e75ea4d32c0c0fd8 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 30 Jun 2023 12:19:50 +0200 Subject: [PATCH 158/263] fix diffusion calls in test_parallel_setup.py --- .../tests/mpi_tests/test_parallel_setup.py | 59 ++++--------------- atm_dyn_iconam/tests/test_diffusion.py | 8 +-- 2 files changed, 15 insertions(+), 52 deletions(-) diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index 2e437a625..1af8a283e 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -17,6 +17,7 @@ import numpy as np import pytest +from atm_dyn_iconam.tests.test_diffusion import verify_diffusion_fields from atm_dyn_iconam.tests.test_utils.serialbox_utils import IconSerialDataProvider from icon4py.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.diffusion.diffusion import Diffusion, DiffusionParams @@ -191,7 +192,7 @@ def test_parallel_diffusion(r04b09_diffusion_config, step_date_init, caplog): print( f"rank={props.rank}/{props.comm_size}: GHEX context setup: from {props.comm_name} with {props.comm_size} nodes" ) - # assert context.size() == 2 + assert context.size() == 2 icon_grid = read_icon_grid(path, rank=props.rank) print( @@ -218,12 +219,14 @@ def test_parallel_diffusion(r04b09_diffusion_config, step_date_init, caplog): diffusion = Diffusion(exchange) diffusion.init( - icon_grid, - r04b09_diffusion_config, - diffusion_params, - vertical_geometry, - metric_state, - interpolation_state, + grid=icon_grid, + config=r04b09_diffusion_config, + params=diffusion_params, + vertical_params=vertical_geometry, + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry ) print(f"rank={props.rank}/{props.comm_size}: diffusion initialized ") diagnostic_state = diffusion_initial_data.construct_diagnostics() @@ -232,50 +235,10 @@ def test_parallel_diffusion(r04b09_diffusion_config, step_date_init, caplog): diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime, - tangent_orientation=edge_geometry.tangent_orientation, - inverse_primal_edge_lengths=edge_geometry.inverse_primal_edge_lengths, - inverse_dual_edge_length=edge_geometry.inverse_dual_edge_lengths, - inverse_vert_vert_lengths=edge_geometry.inverse_vertex_vertex_lengths, - primal_normal_vert=edge_geometry.primal_normal_vert, - dual_normal_vert=edge_geometry.dual_normal_vert, - edge_areas=edge_geometry.edge_areas, - cell_areas=cell_geometry.area, ) print(f"rank={props.rank}/{props.comm_size}: diffusion run ") diffusion_savepoint_exit = IconSerialDataProvider( "icon_pydycore", str(path), True, mpi_rank=props.rank ).from_savepoint_diffusion_init(linit=initial_run, date=step_date_init) - # verify_fields(diffusion_savepoint_exit, diagnostic_state, prognostic_state, - # diffusion.metric_state.mask_hdiff) - - -def verify_fields( - diffusion_savepoint_exit, diagnostic_state, prognostic_state, steep_points -): - ref_div_ic = np.asarray(diffusion_savepoint_exit.div_ic()) - val_div_ic = np.asarray(diagnostic_state.div_ic) - ref_hdef_ic = np.asarray(diffusion_savepoint_exit.hdef_ic()) - val_hdef_ic = np.asarray(diagnostic_state.hdef_ic) - assert np.allclose(ref_div_ic, val_div_ic) - assert np.allclose(ref_hdef_ic, val_hdef_ic) - ref_w = np.asarray(diffusion_savepoint_exit.w()) - val_w = np.asarray(prognostic_state.w) - ref_dwdx = np.asarray(diffusion_savepoint_exit.dwdx()) - val_dwdx = np.asarray(prognostic_state.dwdx) - ref_dwdy = np.asarray(diffusion_savepoint_exit.dwdy()) - val_dwdy = np.asarray(prognostic_state.dwdy) - ref_vn = np.asarray(diffusion_savepoint_exit.vn()) - val_vn = np.asarray(prognostic_state.vn) - assert np.allclose(ref_vn, val_vn) - assert np.allclose(ref_dwdx, val_dwdx) - assert np.allclose(ref_dwdy, val_dwdy) - assert np.allclose(ref_w, val_w) - ref_exner = np.asarray(diffusion_savepoint_exit.exner()) - ref_theta_v = np.asarray(diffusion_savepoint_exit.theta_v()) - val_theta_v = np.asarray(prognostic_state.theta_v) - val_exner = np.asarray(prognostic_state.exner_pressure) - assert np.allclose(ref_theta_v[~steep_points], val_theta_v[~steep_points]) - assert np.allclose(ref_exner[~steep_points], val_exner[~steep_points]) - # assert np.allclose(ref_theta_v[steep_points], val_theta_v[steep_points]) - # assert np.allclose(ref_exner[steep_points], val_exner[steep_points]) + #verify_diffusion_fields(diffusion_savepoint_exit, diagnostic_state, prognostic_state) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 1a7c17ea6..38e49c8c7 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -461,7 +461,7 @@ def test_run_diffusion_single_step( edge_params=edge_geometry, cell_params=cell_geometry, ) - _verify_diffusion_fields( + verify_diffusion_fields( diagnostic_state, prognostic_state, diffusion_savepoint_init ) assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v @@ -470,12 +470,12 @@ def test_run_diffusion_single_step( prognostic_state=prognostic_state, dtime=dtime, ) - _verify_diffusion_fields( + verify_diffusion_fields( diagnostic_state, prognostic_state, diffusion_savepoint_exit ) -def _verify_diffusion_fields( +def verify_diffusion_fields( diagnostic_state, prognostic_state, diffusion_savepoint_exit ): ref_div_ic = np.asarray(diffusion_savepoint_exit.div_ic()) @@ -550,7 +550,7 @@ def test_run_diffusion_initial_step( dtime=dtime, ) - _verify_diffusion_fields( + verify_diffusion_fields( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, diffusion_savepoint_exit=diffusion_savepoint_exit, From c719af44f2ad386b25da53286d7ab41c29f2537d Mon Sep 17 00:00:00 2001 From: Magdalena Date: Fri, 30 Jun 2023 16:55:59 +0200 Subject: [PATCH 159/263] Greenline diffusion verification (#232) Changes from making the diffusion verify. - bug fixes in diffusion - diffusion interface change: pass all static fields to init - test updates from changed sparse fields and newly added domain boundary args - deleted duplicate fused stencils --- .../atm_dyn_iconam/apply_nabla2_to_w.py | 19 +- ...pply_nabla2_to_w_in_upper_damping_layer.py | 16 +- ...te_diagnostic_quantities_for_turbulence.py | 10 +- ...ate_horizontal_gradients_for_turbulence.py | 15 +- .../atm_dyn_iconam/calculate_nabla2_for_w.py | 16 +- .../fused_mo_nh_diffusion_stencil_02_03.py | 69 --- ...sed_mo_nh_diffusion_stencil_07_08_09_10.py | 126 ----- .../fused_mo_nh_diffusion_stencil_11_12.py | 63 --- .../fused_mo_nh_diffusion_stencil_13_14.py | 58 -- ...orary_fields_for_turbulence_diagnostics.py | 24 +- ...fusion_nabla_of_theta_over_steep_points.py | 8 + .../src/icon4py/diffusion/diffusion.py | 529 ++++++------------ .../icon4py/diffusion/diffusion_program.py | 314 ----------- .../src/icon4py/diffusion/horizontal.py | 1 - .../src/icon4py/diffusion/icon_grid.py | 39 +- .../icon4py/diffusion/interpolation_state.py | 32 +- .../src/icon4py/diffusion/metric_state.py | 7 +- atm_dyn_iconam/src/icon4py/diffusion/utils.py | 7 + .../src/icon4py/driver/dycore_driver.py | 39 +- .../src/icon4py/driver/icon_configuration.py | 5 +- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 32 +- atm_dyn_iconam/tests/conftest.py | 8 +- .../tests/test_apply_nabla2_to_w.py | 12 +- ...pply_nabla2_to_w_in_upper_damping_layer.py | 6 + ...ate_horizontal_gradients_for_turbulence.py | 5 + .../tests/test_calculate_nabla2_for_w.py | 5 + atm_dyn_iconam/tests/test_diffusion.py | 168 +++--- .../tests/test_interpolation_state.py | 27 + atm_dyn_iconam/tests/test_io_utils.py | 4 +- ...orary_fields_for_turbulence_diagnostics.py | 13 +- ...fusion_nabla_of_theta_over_steep_points.py | 4 + .../tests/test_utils/serialbox_utils.py | 54 +- .../tests/test_utils/simple_mesh.py | 12 +- atm_dyn_iconam/tests/test_vertical.py | 10 +- base-requirements-dev.txt | 2 +- .../py2f/wrappers/diffusion_wrapper.py | 12 +- 36 files changed, 593 insertions(+), 1178 deletions(-) delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py delete mode 100644 atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py create mode 100644 atm_dyn_iconam/tests/test_interpolation_state.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py index 532191cf2..3e40b54f3 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w.py @@ -13,7 +13,7 @@ from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, neighbor_sum +from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim @@ -39,5 +39,20 @@ def apply_nabla2_to_w( geofac_n2s: Field[[CellDim, C2E2CODim], float], w: Field[[CellDim, KDim], float], diff_multfac_w: float, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): - _apply_nabla2_to_w(area, z_nabla2_c, geofac_n2s, w, diff_multfac_w, out=w) + _apply_nabla2_to_w( + area, + z_nabla2_c, + geofac_n2s, + w, + diff_multfac_w, + out=w, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py index 9c704b6ca..2eaef9f4b 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_to_w_in_upper_damping_layer.py @@ -13,7 +13,7 @@ from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.common.dimension import CellDim, KDim @@ -35,7 +35,19 @@ def apply_nabla2_to_w_in_upper_damping_layer( diff_multfac_n2w: Field[[KDim], float], cell_area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _apply_nabla2_to_w_in_upper_damping_layer( - w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w + w, + diff_multfac_n2w, + cell_area, + z_nabla2_c, + out=w, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py index 9db5202cc..593c95fa5 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py @@ -20,15 +20,15 @@ from icon4py.atm_dyn_iconam.temporary_fields_for_turbulence_diagnostics import ( _temporary_fields_for_turbulence_diagnostics, ) -from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim +from icon4py.common.dimension import CEDim, CellDim, EdgeDim, KDim @field_operator def _calculate_diagnostic_quantities_for_turbulence( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], + e_bln_c_s: Field[[CEDim], float], + geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], wgtfac_c: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: @@ -43,8 +43,8 @@ def _calculate_diagnostic_quantities_for_turbulence( def calculate_diagnostic_quantities_for_turbulence( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], + e_bln_c_s: Field[[CEDim], float], + geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], wgtfac_c: Field[[CellDim, KDim], float], div_ic: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py index f3b88d633..05ccae179 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_horizontal_gradients_for_turbulence.py @@ -13,7 +13,7 @@ from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, neighbor_sum +from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim @@ -36,7 +36,18 @@ def calculate_horizontal_gradients_for_turbulence( geofac_grg_y: Field[[CellDim, C2E2CODim], float], dwdx: Field[[CellDim, KDim], float], dwdy: Field[[CellDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _calculate_horizontal_gradients_for_turbulence( - w, geofac_grg_x, geofac_grg_y, out=(dwdx, dwdy) + w, + geofac_grg_x, + geofac_grg_y, + out=(dwdx, dwdy), + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py index 1e6acc248..5fd0d0f2e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_w.py @@ -13,7 +13,7 @@ from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, neighbor_sum +from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim @@ -31,5 +31,17 @@ def calculate_nabla2_for_w( w: Field[[CellDim, KDim], float], geofac_n2s: Field[[CellDim, C2E2CODim], float], z_nabla2_c: Field[[CellDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): - _calculate_nabla2_for_w(w, geofac_n2s, out=z_nabla2_c) + _calculate_nabla2_for_w( + w, + geofac_n2s, + out=z_nabla2_c, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py deleted file mode 100644 index 3bbca2b79..000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ /dev/null @@ -1,69 +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 - -from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field - -from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulence import ( - _calculate_diagnostics_for_turbulence, -) -from icon4py.atm_dyn_iconam.temporary_fields_for_turbulence_diagnostics import ( - _temporary_fields_for_turbulence_diagnostics, -) -from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim - - -@field_operator -def _fused_mo_nh_diffusion_stencil_02_03( - kh_smag_ec: Field[[EdgeDim, KDim], float], - vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], - diff_multfac_smag: Field[[KDim], float], - wgtfac_c: Field[[CellDim, KDim], float], -) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - kh_c, div = _temporary_fields_for_turbulence_diagnostics( - kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag - ) - div_ic, hdef_ic = _calculate_diagnostics_for_turbulence(div, kh_c, wgtfac_c) - return div_ic, hdef_ic - - -@program -def fused_mo_nh_diffusion_stencil_02_03( - kh_smag_ec: Field[[EdgeDim, KDim], float], - vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], - diff_multfac_smag: Field[[KDim], float], - wgtfac_c: Field[[CellDim, KDim], float], - div_ic: Field[[CellDim, KDim], float], - hdef_ic: Field[[CellDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, -): - _fused_mo_nh_diffusion_stencil_02_03( - kh_smag_ec, - vn, - e_bln_c_s, - geofac_div, - diff_multfac_smag, - wgtfac_c, - out=(div_ic, hdef_ic), - domain={ - CellDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, - ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py deleted file mode 100644 index b83b00dfd..000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ /dev/null @@ -1,126 +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 - -from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where - -from icon4py.atm_dyn_iconam.apply_nabla2_to_w import _apply_nabla2_to_w -from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import ( - _apply_nabla2_to_w_in_upper_damping_layer, -) -from icon4py.atm_dyn_iconam.calculate_horizontal_gradients_for_turbulence import ( - _calculate_horizontal_gradients_for_turbulence, -) -from icon4py.atm_dyn_iconam.calculate_nabla2_for_w import _calculate_nabla2_for_w -from icon4py.common.dimension import C2E2CODim, CellDim, KDim - - -@field_operator -def _fused_mo_nh_diffusion_stencil_07_08_09_10( - area: Field[[CellDim], float], - geofac_n2s: Field[[CellDim, C2E2CODim], float], - geofac_grg_x: Field[[CellDim, C2E2CODim], float], - geofac_grg_y: Field[[CellDim, C2E2CODim], float], - w_old: Field[[CellDim, KDim], float], - w: Field[[CellDim, KDim], float], - dwdx: Field[[CellDim, KDim], float], - dwdy: Field[[CellDim, KDim], float], - diff_multfac_w: float, - diff_multfac_n2w: Field[[KDim], float], - vert_idx: Field[[KDim], int32], - horz_idx: Field[[CellDim], int32], - nrdmax: int32, - interior_idx: int32, - halo_idx: int32, -) -> tuple[ - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], -]: - - vert_idx = broadcast(vert_idx, (CellDim, KDim)) - - dwdx, dwdy = where( - vert_idx > int32(0), - _calculate_horizontal_gradients_for_turbulence( - w_old, geofac_grg_x, geofac_grg_y - ), - (dwdx, dwdy), - ) - - z_nabla2_c = _calculate_nabla2_for_w(w_old, geofac_n2s) - - w = where( - (horz_idx >= interior_idx) & (horz_idx < halo_idx), - _apply_nabla2_to_w(area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w), - w_old, - ) - - w = where( - (vert_idx > int32(0)) - & (vert_idx < nrdmax) - & (horz_idx >= interior_idx) - & (horz_idx < halo_idx), - _apply_nabla2_to_w_in_upper_damping_layer( - w, diff_multfac_n2w, area, z_nabla2_c - ), - w, - ) - - return w, dwdx, dwdy - - -@program -def fused_mo_nh_diffusion_stencil_07_08_09_10( - area: Field[[CellDim], float], - geofac_n2s: Field[[CellDim, C2E2CODim], float], - geofac_grg_x: Field[[CellDim, C2E2CODim], float], - geofac_grg_y: Field[[CellDim, C2E2CODim], float], - w_old: Field[[CellDim, KDim], float], - w: Field[[CellDim, KDim], float], - dwdx: Field[[CellDim, KDim], float], - dwdy: Field[[CellDim, KDim], float], - diff_multfac_w: float, - diff_multfac_n2w: Field[[KDim], float], - vert_idx: Field[[KDim], int32], - horz_idx: Field[[CellDim], int32], - nrdmax: int32, - interior_idx: int32, - halo_idx: int32, - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, -): - _fused_mo_nh_diffusion_stencil_07_08_09_10( - area, - geofac_n2s, - geofac_grg_x, - geofac_grg_y, - w_old, - w, - dwdx, - dwdy, - diff_multfac_w, - diff_multfac_n2w, - vert_idx, - horz_idx, - nrdmax, - interior_idx, - halo_idx, - out=(w, dwdx, dwdy), - domain={ - CellDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, - ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py deleted file mode 100644 index a9fb3dbd2..000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py +++ /dev/null @@ -1,63 +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 - -from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field - -from icon4py.atm_dyn_iconam.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( - _enhance_diffusion_coefficient_for_grid_point_cold_pools, -) -from icon4py.atm_dyn_iconam.temporary_field_for_grid_point_cold_pools_enhancement import ( - _temporary_field_for_grid_point_cold_pools_enhancement, -) -from icon4py.common.dimension import CellDim, EdgeDim, KDim - - -@field_operator -def _fused_mo_nh_diffusion_stencil_11_12( - theta_v: Field[[CellDim, KDim], float], - theta_ref_mc: Field[[CellDim, KDim], float], - thresh_tdiff: float, - kh_smag_e: Field[[EdgeDim, KDim], float], -) -> Field[[EdgeDim, KDim], float]: - enh_diffu_3d = _temporary_field_for_grid_point_cold_pools_enhancement( - theta_v, theta_ref_mc, thresh_tdiff - ) - kh_smag_e = _enhance_diffusion_coefficient_for_grid_point_cold_pools( - kh_smag_e, enh_diffu_3d - ) - return kh_smag_e - - -@program -def fused_mo_nh_diffusion_stencil_11_12( - theta_v: Field[[CellDim, KDim], float], - theta_ref_mc: Field[[CellDim, KDim], float], - thresh_tdiff: float, - kh_smag_e: Field[[EdgeDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, -): - _fused_mo_nh_diffusion_stencil_11_12( - theta_v, - theta_ref_mc, - thresh_tdiff, - kh_smag_e, - out=kh_smag_e, - domain={ - EdgeDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, - ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py deleted file mode 100644 index d13d7c8e1..000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ /dev/null @@ -1,58 +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 - -from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field - -from icon4py.atm_dyn_iconam.calculate_nabla2_for_z import _calculate_nabla2_for_z -from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( - _calculate_nabla2_of_theta, -) -from icon4py.common.dimension import CEDim, CellDim, EdgeDim, KDim - - -@field_operator -def _fused_mo_nh_diffusion_stencil_13_14( - kh_smag_e: Field[[EdgeDim, KDim], float], - inv_dual_edge_length: Field[[EdgeDim], float], - theta_v: Field[[CellDim, KDim], float], - geofac_div: Field[[CEDim], float], -) -> Field[[CellDim, KDim], float]: - z_nabla2_e = _calculate_nabla2_for_z(kh_smag_e, inv_dual_edge_length, theta_v) - z_temp = _calculate_nabla2_of_theta(z_nabla2_e, geofac_div) - return z_temp - - -@program -def fused_mo_nh_diffusion_stencil_13_14( - kh_smag_e: Field[[EdgeDim, KDim], float], - inv_dual_edge_length: Field[[EdgeDim], float], - theta_v: Field[[CellDim, KDim], float], - geofac_div: Field[[CEDim], float], - z_temp: Field[[CellDim, KDim], float], - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, -): - _fused_mo_nh_diffusion_stencil_13_14( - kh_smag_e, - inv_dual_edge_length, - theta_v, - geofac_div, - out=z_temp, - domain={ - CellDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, - ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_fields_for_turbulence_diagnostics.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_fields_for_turbulence_diagnostics.py index c2ffe9734..c26add794 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_fields_for_turbulence_diagnostics.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/temporary_fields_for_turbulence_diagnostics.py @@ -15,19 +15,29 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim +from icon4py.common.dimension import ( + C2CE, + C2E, + C2EDim, + CEDim, + CellDim, + EdgeDim, + KDim, +) @field_operator def _temporary_fields_for_turbulence_diagnostics( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], + e_bln_c_s: Field[[CEDim], float], + geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - kh_c = neighbor_sum(kh_smag_ec(C2E) * e_bln_c_s, axis=C2EDim) / diff_multfac_smag - div = neighbor_sum(vn(C2E) * geofac_div, axis=C2EDim) + kh_c = ( + neighbor_sum(kh_smag_ec(C2E) * e_bln_c_s(C2CE), axis=C2EDim) / diff_multfac_smag + ) + div = neighbor_sum(vn(C2E) * geofac_div(C2CE), axis=C2EDim) return kh_c, div @@ -35,8 +45,8 @@ def _temporary_fields_for_turbulence_diagnostics( def temporary_fields_for_turbulence_diagnostics( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], + e_bln_c_s: Field[[CEDim], float], + geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], kh_c: Field[[CellDim, KDim], float], div: Field[[CellDim, KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index ab0574e40..663fa0efe 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -75,6 +75,10 @@ def truly_horizontal_diffusion_nabla_of_theta_over_steep_points( vcoef: Field[[CECDim, KDim], float], theta_v: Field[[CellDim, KDim], float], z_temp: Field[[CellDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _truly_horizontal_diffusion_nabla_of_theta_over_steep_points( mask, @@ -86,4 +90,8 @@ def truly_horizontal_diffusion_nabla_of_theta_over_steep_points( theta_v, z_temp, out=z_temp, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index d99d39f83..4675c0b26 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -20,9 +20,11 @@ from gt4py.next.common import Dimension from gt4py.next.ffront.fbuiltins import Field, int32 from gt4py.next.iterator.embedded import np_as_located_field -from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn +from gt4py.next.program_processors.runners.gtfn_cpu import ( + run_gtfn, + run_gtfn_cached, +) -import icon4py.diffusion.diffusion_program as diff_prog from icon4py.atm_dyn_iconam.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance, ) @@ -44,32 +46,28 @@ from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) +from icon4py.atm_dyn_iconam.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( + truly_horizontal_diffusion_nabla_of_theta_over_steep_points, +) from icon4py.atm_dyn_iconam.update_theta_and_exner import update_theta_and_exner from icon4py.common.constants import ( CPD, DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO, GAS_CONSTANT_DRY_AIR, ) -from icon4py.common.dimension import ( - C2E2CDim, - C2E2CODim, - C2EDim, - CellDim, - E2C2VDim, - E2CDim, - ECVDim, - EdgeDim, - KDim, - V2EDim, - VertexDim, -) +from icon4py.common.dimension import CellDim, ECVDim, EdgeDim, KDim, VertexDim from icon4py.diffusion.diagnostic_state import DiagnosticState -from icon4py.diffusion.horizontal import HorizontalMarkerIndex +from icon4py.diffusion.horizontal import ( + CellParams, + EdgeParams, + HorizontalMarkerIndex, +) from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams from icon4py.diffusion.interpolation_state import InterpolationState from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.prognostic_state import PrognosticState from icon4py.diffusion.utils import ( + copy_field, init_diffusion_local_fields_for_regular_timestep, init_nabla2_factor_in_upper_damping_zone, scale_k, @@ -84,6 +82,10 @@ VectorTuple = namedtuple("VectorTuple", "x y") +cached_backend = run_gtfn_cached +compiled_backend = run_gtfn +backend = compiled_backend # + class DiffusionConfig: """ @@ -116,6 +118,7 @@ def __init__( max_nudging_coeff: float = 0.02, nudging_decay_rate: float = 2.0, ): + """Set the diffusion configuration parameters with the ICON default values.""" # parameters from namelist diffusion_nml self.diffusion_type: int = diffusion_type @@ -349,7 +352,7 @@ def determine_smagorinski_factor(self, config: DiffusionConfig): @staticmethod def _diffusion_type_5_smagorinski_factor(config: DiffusionConfig): """ - Initialize smagorinski factors used in diffusion type 5. + Initialize Smagorinski factors used in diffusion type 5. The calculation and magic numbers are taken from mo_diffusion_nml.f90 """ @@ -364,26 +367,27 @@ def _diffusion_type_5_smagorinski_factor(config: DiffusionConfig): class Diffusion: """Class that configures diffusion and does one diffusion step.""" - def __init__(self, run_program=True): + def __init__(self): self._initialized = False self.rd_o_cvd: float = GAS_CONSTANT_DRY_AIR / (CPD - GAS_CONSTANT_DRY_AIR) self.thresh_tdiff: float = ( -5.0 ) # threshold temperature deviation from neighboring grid points hat activates extra diffusion against runaway cooling - self._run_program = run_program self.grid: Optional[IconGrid] = None self.config: Optional[DiffusionConfig] = None self.params: Optional[DiffusionParams] = None self.vertical_params: Optional[VerticalModelParams] = None - self.interpolation_state = None - self.metric_state = None + self.interpolation_state: InterpolationState = None + self.metric_state: MetricState = None self.diff_multfac_w: Optional[float] = None self.diff_multfac_n2w: Field[[KDim], float] = None self.smag_offset: Optional[float] = None self.fac_bdydiff_v: Optional[float] = None self.bdy_diff: Optional[float] = None self.nudgezone_diff: Optional[float] = None + self.edge_params: Optional[EdgeParams] = None + self.cell_params: Optional[CellParams] = None def init( self, @@ -393,11 +397,23 @@ def init( vertical_params: VerticalModelParams, metric_state: MetricState, interpolation_state: InterpolationState, + edge_params: EdgeParams, + cell_params: CellParams, ): """ Initialize Diffusion granule with configuration. - calculates all local fields that are used in diffusion within the time loop + calculates all local fields that are used in diffusion within the time loop. + + Args: + grid: + config: + params: + vertical_params: + metric_state: + interpolation_state: + edge_params: + cell_params: """ self.config: DiffusionConfig = config self.params: DiffusionParams = params @@ -405,6 +421,8 @@ def init( self.vertical_params = vertical_params self.metric_state: MetricState = metric_state self.interpolation_state: InterpolationState = interpolation_state + self.edge_params = edge_params + self.cell_params = cell_params self._allocate_local_fields() @@ -426,7 +444,7 @@ def init( 1.0 / 48.0, params.K4W * config.substep_as_float() ) - init_diffusion_local_fields_for_regular_timestep.with_backend(run_gtfn)( + init_diffusion_local_fields_for_regular_timestep.with_backend(backend)( params.K4, config.substep_as_float(), *params.smagorinski_factor, @@ -473,20 +491,15 @@ def _index_field(dim: Dimension, size=None): self.vertical_index = _index_field(KDim, self.grid.n_lev() + 1) self.horizontal_cell_index = _index_field(CellDim) self.horizontal_edge_index = _index_field(EdgeDim) + self.w_tmp = np_as_located_field(CellDim, KDim)( + np.zeros((self.grid.num_cells(), self.grid.n_lev() + 1), dtype=float) + ) - def initial_step( + def initial_run( self, diagnostic_state: DiagnosticState, prognostic_state: PrognosticState, dtime: float, - tangent_orientation: Field[[EdgeDim], float], - inverse_primal_edge_lengths: Field[[EdgeDim], float], - inverse_dual_edge_length: Field[[EdgeDim], float], - inverse_vert_vert_lengths: Field[[EdgeDim], float], - primal_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], - dual_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], - edge_areas: Field[[EdgeDim], float], - cell_areas: Field[[CellDim], float], ): """ Calculate initial diffusion step. @@ -502,7 +515,7 @@ def initial_step( diff_multfac_vn = zero_field(self.grid, KDim) smag_limit = zero_field(self.grid, KDim) - setup_fields_for_initial_step.with_backend(run_gtfn)( + setup_fields_for_initial_step.with_backend(backend)( self.params.K4, self.config.hdiff_efdt_ratio, diff_multfac_vn, @@ -513,32 +526,16 @@ def initial_step( diagnostic_state, prognostic_state, dtime, - tangent_orientation, - inverse_primal_edge_lengths, - inverse_dual_edge_length, - inverse_vert_vert_lengths, - primal_normal_vert, - dual_normal_vert, - edge_areas, - cell_areas, diff_multfac_vn, smag_limit, 0.0, ) - def time_step( + def run( self, diagnostic_state: DiagnosticState, prognostic_state: PrognosticState, dtime: float, - tangent_orientation: Field[[EdgeDim], float], - inverse_primal_edge_lengths: Field[[EdgeDim], float], - inverse_dual_edge_length: Field[[EdgeDim], float], - inverse_vert_vert_lengths: Field[[EdgeDim], float], - primal_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], - dual_normal_vert: VectorTuple[Field[[ECVDim], float], Field[[ECVDim], float]], - edge_areas: Field[[EdgeDim], float], - cell_areas: Field[[CellDim], float], ): """ Do one diffusion step within regular time loop. @@ -546,182 +543,20 @@ def time_step( runs a diffusion step for the parameter linit=False, within regular time loop. """ - if not self._run_program: - self._do_diffusion_step( - diagnostic_state=diagnostic_state, - prognostic_state=prognostic_state, - dtime=dtime, - tangent_orientation=tangent_orientation, - inverse_primal_edge_lengths=inverse_primal_edge_lengths, - inverse_dual_edge_length=inverse_dual_edge_length, - inverse_vertex_vertex_lengths=inverse_vert_vert_lengths, - primal_normal_vert=primal_normal_vert, - dual_normal_vert=dual_normal_vert, - edge_areas=edge_areas, - cell_areas=cell_areas, - diff_multfac_vn=self.diff_multfac_vn, - smag_limit=self.smag_limit, - smag_offset=self.smag_offset, - ) - else: - log.info("running diffusion_program") - ( - cell_startindex_nudging, - cell_endindex_local, - ) = self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.local(CellDim), - ) - - ( - cell_startindex_interior, - cell_endindex_local_plus1, - ) = self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.interior(CellDim), - HorizontalMarkerIndex.local(CellDim) - 1, - ) - - ( - edge_startindex_nudging_plus1, - edge_endindex_local, - ) = self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim) + 1, - HorizontalMarkerIndex.local(EdgeDim), - ) - - ( - edge_startindex_nudging_minus1, - edge_endindex_local_minus2, - ) = self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim) - 1, - HorizontalMarkerIndex.local(EdgeDim) - 2, - ) - - ( - vertex_startindex_lb_plus3, - vertex_endindex_local, - ) = self.grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, - HorizontalMarkerIndex.local(VertexDim), - ) - - ( - vertex_startindex_lb_plus1, - vertex_endindex_local_minus1, - ) = self.grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, - HorizontalMarkerIndex.local(VertexDim) - 1, - ) - edge_start_lb_plus4, _ = self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, - ) - log.info("diffusion program: start") - diff_prog.diffusion_run.with_backend(run_gtfn)( - diagnostic_hdef_ic=diagnostic_state.hdef_ic, - diagnostic_div_ic=diagnostic_state.div_ic, - diagnostic_dwdx=diagnostic_state.dwdx, - diagnostic_dwdy=diagnostic_state.dwdy, - prognostic_w=prognostic_state.w, - prognostic_vn=prognostic_state.vn, - prognostic_exner_pressure=prognostic_state.exner_pressure, - prognostic_theta_v=prognostic_state.theta_v, - metric_theta_ref_mc=self.metric_state.theta_ref_mc, - metric_wgtfac_c=self.metric_state.wgtfac_c, - metric_mask_hdiff=self.metric_state.mask_hdiff, - metric_zd_vertidx=self.metric_state.zd_vertidx, - metric_zd_diffcoef=self.metric_state.zd_diffcoef, - metric_zd_intcoef=self.metric_state.zd_intcoef, - interpolation_e_bln_c_s=self.interpolation_state.e_bln_c_s, - interpolation_rbf_coeff_1=self.interpolation_state.rbf_coeff_1, - interpolation_rbf_coeff_2=self.interpolation_state.rbf_coeff_2, - interpolation_geofac_div=self.interpolation_state._geofac_div, - interpolation_geofac_grg_x=self.interpolation_state.geofac_grg_x, - interpolation_geofac_grg_y=self.interpolation_state.geofac_grg_y, - interpolation_nudgecoeff_e=self.interpolation_state.nudgecoeff_e, - interpolation_geofac_n2s=self.interpolation_state.geofac_n2s, - interpolation_geofac_n2s_c=self.interpolation_state.geofac_n2s_c, - interpolation_geofac_n2s_nbh=self.interpolation_state.geofac_n2s_nbh, - tangent_orientation=tangent_orientation, - inverse_primal_edge_lengths=inverse_primal_edge_lengths, - inverse_dual_edge_lengths=inverse_dual_edge_length, - inverse_vert_vert_lengths=inverse_vert_vert_lengths, - primal_normal_vert_1=primal_normal_vert[0], - primal_normal_vert_2=primal_normal_vert[1], - dual_normal_vert_1=dual_normal_vert[0], - dual_normal_vert_2=dual_normal_vert[1], - edge_areas=edge_areas, - cell_areas=cell_areas, - diff_multfac_vn=self.diff_multfac_vn, - dtime=dtime, - rd_o_cvd=self.rd_o_cvd, - local_thresh_tdiff=self.thresh_tdiff, - local_smag_limit=self.smag_limit, - local_u_vert=self.u_vert, - local_v_vert=self.v_vert, - local_enh_smag_fac=self.enh_smag_fac, - local_kh_smag_e=self.kh_smag_e, - local_kh_smag_ec=self.kh_smag_ec, - local_z_nabla2_e=self.z_nabla2_e, - local_z_temp=self.z_temp, - local_diff_multfac_smag=self.diff_multfac_smag, - local_diff_multfac_n2w=self.diff_multfac_n2w, - local_smag_offset=self.smag_offset, - local_nudgezone_diff=self.nudgezone_diff, - local_fac_bdydiff_v=self.fac_bdydiff_v, - local_diff_multfac_w=self.diff_multfac_w, - local_vertical_index=self.vertical_index, - local_horizontal_cell_index=self.horizontal_cell_index, - local_horizontal_edge_index=self.horizontal_edge_index, - cell_startindex_interior=int32(cell_startindex_interior), - cell_startindex_nudging=cell_startindex_nudging, - cell_endindex_local_plus1=cell_endindex_local_plus1, - cell_endindex_local=cell_endindex_local, - cell_halo_idx=int32(cell_endindex_local), - edge_startindex_nudging_plus1=edge_startindex_nudging_plus1, - edge_startindex_nudging_minus1=int32(edge_startindex_nudging_minus1), - edge_endindex_local=edge_endindex_local, - edge_endindex_local_minus2=edge_endindex_local_minus2, - vertex_startindex_lb_plus3=vertex_startindex_lb_plus3, - vertex_startindex_lb_plus1=vertex_startindex_lb_plus1, - vertex_endindex_local=vertex_endindex_local, - vertex_endindex_local_minus1=vertex_endindex_local_minus1, - index_of_damping_height=self.vertical_params.index_of_damping_layer, - nlev=self.grid.n_lev(), - boundary_diffusion_start_index_edges=edge_start_lb_plus4, - offset_provider={ - "V2E": self.grid.get_v2e_connectivity(), - "E2C2V": self.grid.get_e2c2v_connectivity(), - "E2ECV": self.grid.get_e2ecv_connectivity(), - "C2E": self.grid.get_c2e_connectivity(), - "E2C": self.grid.get_e2c_connectivity(), - "C2E2C": self.grid.get_c2e2c_connectivity(), - "C2E2CO": self.grid.get_c2e2co_connectivity(), - "Koff": KDim, - }, - ) - log.info("diffusion program: end") + self._do_diffusion_step( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + diff_multfac_vn=self.diff_multfac_vn, + smag_limit=self.smag_limit, + smag_offset=self.smag_offset, + ) def _do_diffusion_step( self, diagnostic_state: DiagnosticState, prognostic_state: PrognosticState, dtime: float, - tangent_orientation: Field[[EdgeDim], float], - inverse_primal_edge_lengths: Field[[EdgeDim], float], - inverse_dual_edge_length: Field[[EdgeDim], float], - inverse_vertex_vertex_lengths: Field[[EdgeDim], float], - primal_normal_vert: Tuple[Field[[ECVDim], float], Field[[ECVDim], float]], - dual_normal_vert: Tuple[Field[[ECVDim], float], Field[[ECVDim], float]], - edge_areas: Field[[EdgeDim], float], - cell_areas: Field[[CellDim], float], diff_multfac_vn: Field[[KDim], float], smag_limit: Field[[KDim], float], smag_offset: float, @@ -733,38 +568,22 @@ def _do_diffusion_step( diagnostic_state: output argument, data class that contains diagnostic variables prognostic_state: output argument, data class that contains prognostic variables dtime: the time step, - tangent_orientation: - inverse_primal_edge_lengths: - inverse_dual_edge_length: - inverse_vertex_vertex_lengths: - primal_normal_vert: - dual_normal_vert: - edge_areas: - cell_areas: diff_multfac_vn: smag_limit: smag_offset: """ klevels = self.grid.n_lev() - k_start_end_minus2 = klevels - 2 - - cell_start_nudging_minus1, cell_end_local_plus1 = self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim) - 1, - HorizontalMarkerIndex.local(CellDim) - 1, - ) - cell_start_interior, cell_end_local = self.grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.interior(CellDim), HorizontalMarkerIndex.local(CellDim), ) - cell_start_nudging, _ = self.grid.get_indices_from_to( + cell_start_nudging, cell_end_halo = self.grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.local(CellDim), + HorizontalMarkerIndex.halo(CellDim), ) edge_start_nudging_plus_one, edge_end_local = self.grid.get_indices_from_to( @@ -773,6 +592,12 @@ def _do_diffusion_step( HorizontalMarkerIndex.local(EdgeDim), ) + edge_start_nudging, edge_end_halo = self.grid.get_indices_from_to( + EdgeDim, + HorizontalMarkerIndex.nudging(EdgeDim), + HorizontalMarkerIndex.halo(EdgeDim), + ) + edge_start_lb_plus4, _ = self.grid.get_indices_from_to( EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, @@ -797,7 +622,7 @@ def _do_diffusion_step( HorizontalMarkerIndex.local(VertexDim), ) ( - vertex_start_local_boundary_plus1, + vertex_start_lb_plus1, vertex_end_local_minus1, ) = self.grid.get_indices_from_to( VertexDim, @@ -805,46 +630,44 @@ def _do_diffusion_step( HorizontalMarkerIndex.local(VertexDim) - 1, ) - # 0b call timer start - # - # 0c. dtime dependent stuff: enh_smag_factor, - scale_k.with_backend(run_gtfn)( + # dtime dependent: enh_smag_factor, + scale_k.with_backend(backend)( self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={} ) - # TODO: @magdalena is this needed?, if not remove - set_zero_v_k.with_backend(run_gtfn)(self.u_vert, offset_provider={}) - set_zero_v_k.with_backend(run_gtfn)(self.v_vert, offset_provider={}) + set_zero_v_k.with_backend(backend)(self.u_vert, offset_provider={}) + set_zero_v_k.with_backend(backend)(self.v_vert, offset_provider={}) log.debug("rbf interpolation: start") - # # 1. CALL rbf_vec_interpol_vertex - mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(run_gtfn)( + mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( p_e_in=prognostic_state.vn, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, ptr_coeff_2=self.interpolation_state.rbf_coeff_2, p_u_out=self.u_vert, p_v_out=self.v_vert, - horizontal_start=vertex_start_local_boundary_plus3, - horizontal_end=vertex_end_local_minus1, + horizontal_start=vertex_start_lb_plus1, + horizontal_end=vertex_end_local, vertical_start=0, vertical_end=klevels, offset_provider={"V2E": self.grid.get_v2e_connectivity()}, ) log.debug("rbf interpolation: end") - # 2. HALO EXCHANGE -- CALL sync_patch_array_mult - # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 - log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: start") - calculate_nabla2_and_smag_coefficients_for_vn.with_backend(run_gtfn)( + # HALO EXCHANGE -- CALL sync_patch_array_mult + + log.debug( + "running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start" + ) + calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend)( diff_multfac_smag=self.diff_multfac_smag, - tangent_orientation=tangent_orientation, - inv_primal_edge_length=inverse_primal_edge_lengths, - inv_vert_vert_length=inverse_vertex_vertex_lengths, + tangent_orientation=self.edge_params.tangent_orientation, + inv_primal_edge_length=self.edge_params.inverse_primal_edge_lengths, + inv_vert_vert_length=self.edge_params.inverse_vertex_vertex_lengths, u_vert=self.u_vert, v_vert=self.v_vert, - primal_normal_vert_x=primal_normal_vert[0], - primal_normal_vert_y=primal_normal_vert[1], - dual_normal_vert_x=dual_normal_vert[0], - dual_normal_vert_y=dual_normal_vert[1], + primal_normal_vert_x=self.edge_params.primal_normal_vert[0], + primal_normal_vert_y=self.edge_params.primal_normal_vert[1], + dual_normal_vert_x=self.edge_params.dual_normal_vert[0], + dual_normal_vert_y=self.edge_params.dual_normal_vert[1], vn=prognostic_state.vn, smag_limit=smag_limit, kh_smag_e=self.kh_smag_e, @@ -860,61 +683,64 @@ def _do_diffusion_step( "E2ECV": self.grid.get_e2ecv_connectivity(), }, ) - log.debug("running calculate_nabla2_and_smag_coefficients_for_vn: end") - log.debug("running fused stencil fused stencil 02_03: start") - calculate_diagnostic_quantities_for_turbulence.with_backend(run_gtfn)( + log.debug( + "running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end" + ) + log.debug( + "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start" + ) + calculate_diagnostic_quantities_for_turbulence.with_backend(backend)( kh_smag_ec=self.kh_smag_ec, vn=prognostic_state.vn, e_bln_c_s=self.interpolation_state.e_bln_c_s, - geofac_div=self.interpolation_state._geofac_div, + geofac_div=self.interpolation_state.geofac_div, diff_multfac_smag=self.diff_multfac_smag, wgtfac_c=self.metric_state.wgtfac_c, div_ic=diagnostic_state.div_ic, hdef_ic=diagnostic_state.hdef_ic, horizontal_start=cell_start_nudging, horizontal_end=cell_end_local, - vertical_start=0, + vertical_start=1, vertical_end=klevels, offset_provider={ "C2E": self.grid.get_c2e_connectivity(), + "C2CE": self.grid.get_c2ce_connectivity(), "Koff": KDim, }, ) - log.debug("running fused stencil fused stencil 02_03: end") - # - # # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH - # - # # 5. CALL rbf_vec_interpol_vertex_wp + log.debug( + "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end" + ) + + # HALO EXCHANGE IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH + log.debug("rbf interpolation: start") - mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(run_gtfn)( + mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( p_e_in=self.z_nabla2_e, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, ptr_coeff_2=self.interpolation_state.rbf_coeff_2, p_u_out=self.u_vert, p_v_out=self.v_vert, - horizontal_start=vertex_start_local_boundary_plus3, + horizontal_start=vertex_start_lb_plus1, horizontal_end=vertex_end_local, vertical_start=0, vertical_end=klevels, offset_provider={"V2E": self.grid.get_v2e_connectivity()}, ) log.debug("rbf interpolation: end") - # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult - # - # # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 - # # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 - # - - log.debug("running fused stencil 04_05_06: start") - fused_mo_nh_diffusion_stencil_04_05_06.with_backend(run_gtfn)( + + # 6. HALO EXCHANGE -- CALL sync_patch_array_mult + + log.debug("running stencil 04 05 06: start") + fused_mo_nh_diffusion_stencil_04_05_06.with_backend(backend)( u_vert=self.u_vert, v_vert=self.v_vert, - primal_normal_vert_v1=primal_normal_vert[0], - primal_normal_vert_v2=primal_normal_vert[1], + primal_normal_vert_v1=self.edge_params.primal_normal_vert[0], + primal_normal_vert_v2=self.edge_params.primal_normal_vert[1], z_nabla2_e=self.z_nabla2_e, - inv_vert_vert_length=inverse_vertex_vertex_lengths, - inv_primal_edge_length=inverse_primal_edge_lengths, - area_edge=edge_areas, + inv_vert_vert_length=self.edge_params.inverse_vertex_vertex_lengths, + inv_primal_edge_length=self.edge_params.inverse_primal_edge_lengths, + area_edge=self.edge_params.edge_areas, kh_smag_e=self.kh_smag_e, diff_multfac_vn=diff_multfac_vn, nudgecoeff_e=self.interpolation_state.nudgecoeff_e, @@ -922,8 +748,8 @@ def _do_diffusion_step( horz_idx=self.horizontal_edge_index, nudgezone_diff=self.nudgezone_diff, fac_bdydiff_v=self.fac_bdydiff_v, - start_2nd_nudge_line_idx_e=int32(edge_start_nudging_minus1), - horizontal_start=edge_start_nudging_plus_one, + start_2nd_nudge_line_idx_e=int32(edge_start_nudging_plus_one), + horizontal_start=edge_start_lb_plus4, horizontal_end=edge_end_local, vertical_start=0, vertical_end=klevels, @@ -932,20 +758,22 @@ def _do_diffusion_step( "E2ECV": self.grid.get_e2ecv_connectivity(), }, ) - # # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, - # # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 - log.debug("running fused stencil 04_05_06: end") + log.debug("runningstencils 04 05 06: end") - log.debug("running fused stencil 07_08_09_10: start") - w_old = prognostic_state.w + log.debug( + "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" + ) + copy_field.with_backend(backend)( + prognostic_state.w, self.w_tmp, offset_provider={} + ) apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.with_backend( - run_gtfn + backend )( - area=cell_areas, + area=self.cell_params.area, geofac_n2s=self.interpolation_state.geofac_n2s, geofac_grg_x=self.interpolation_state.geofac_grg_x, geofac_grg_y=self.interpolation_state.geofac_grg_y, - w_old=w_old, + w_old=self.w_tmp, w=prognostic_state.w, dwdx=diagnostic_state.dwdx, dwdy=diagnostic_state.dwdy, @@ -953,50 +781,50 @@ def _do_diffusion_step( diff_multfac_n2w=self.diff_multfac_n2w, vert_idx=self.vertical_index, horz_idx=self.horizontal_cell_index, - nrdmax=self.vertical_params.index_of_damping_layer, - interior_idx=int32( - cell_start_interior - ), # h end index for stencil_09 and stencil_10 - halo_idx=int32( - cell_end_local - ), # h end index for stencil_09 and stencil_10, - horizontal_start=cell_start_nudging, # h start index for stencil_07 and stencil_08 - horizontal_end=cell_end_local_plus1, # h end index for stencil_07 and stencil_08 + nrdmax=int32( + self.vertical_params.index_of_damping_layer + 1 + ), # +1 since Fortran includes boundaries + interior_idx=int32(cell_start_interior), + halo_idx=int32(cell_end_local), + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_halo, vertical_start=0, vertical_end=klevels, offset_provider={ "C2E2CO": self.grid.get_c2e2co_connectivity(), }, ) - log.debug("running fused stencil 07_08_09_10: start") - # # 8. HALO EXCHANGE: CALL sync_patch_array - # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, - # # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 - # - # # TODO @magdalena check: kh_smag_e is an out field, should not be calculated in init? - # - log.debug("running fused stencil 11_12: start") + log.debug( + "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end" + ) + # HALO EXCHANGE: CALL sync_patch_array + + log.debug( + "running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" + ) calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.with_backend( - run_gtfn + backend )( theta_v=prognostic_state.theta_v, theta_ref_mc=self.metric_state.theta_ref_mc, thresh_tdiff=self.thresh_tdiff, kh_smag_e=self.kh_smag_e, - horizontal_start=edge_start_nudging_plus_one, - horizontal_end=edge_end_local, - vertical_start=k_start_end_minus2, + horizontal_start=edge_start_nudging, + horizontal_end=edge_end_halo, + vertical_start=(klevels - 2), vertical_end=klevels, offset_provider={ "E2C": self.grid.get_e2c_connectivity(), "C2E2C": self.grid.get_c2e2c_connectivity(), }, ) - log.debug("running fused stencil 11_12: end") - log.debug("running fused stencil 13_14: start") - calculate_nabla2_for_theta.with_backend(run_gtfn)( + log.debug( + "running stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): end" + ) + log.debug("running stencils 13 14 (calculate_nabla2_for_theta): start") + calculate_nabla2_for_theta.with_backend(backend)( kh_smag_e=self.kh_smag_e, - inv_dual_edge_length=inverse_dual_edge_length, + inv_dual_edge_length=self.edge_params.inverse_dual_edge_lengths, theta_v=prognostic_state.theta_v, geofac_div=self.interpolation_state.geofac_div, z_temp=self.z_temp, @@ -1007,33 +835,42 @@ def _do_diffusion_step( offset_provider={ "C2E": self.grid.get_c2e_connectivity(), "E2C": self.grid.get_e2c_connectivity(), + "C2CE": self.grid.get_c2ce_connectivity(), }, ) - log.debug("running fused stencil 13_14: end") - log.debug("running fused stencil 15: start") - # truly_horizontal_diffusion_nabla_of_theta_over_steep_points( - # self.metric_state.mask_hdiff, - # self.metric_state.zd_vertidx, - # self.metric_state.zd_diffcoef, - # self.interpolation_state.geofac_n2s_c, - # self.interpolation_state.geofac_n2s_nbh, - # self.metric_state.zd_intcoef, - # prognostic_state.theta_v, - # self.z_temp, - # cell_start_nudging, - # cell_end_local, - # 0, - # klevels, - # offset_provider={ - # "C2E2C": self.grid.get_c2e2c_connectivity(), - # "Koff": KDim, - # }, - # ) - log.debug("running fused stencil 15: end") - log.debug("running fused stencil update_theta_and_exner: start") - update_theta_and_exner.with_backend(run_gtfn)( + log.debug("running stencils 13_14 (calculate_nabla2_for_theta): end") + log.debug( + "running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start" + ) + truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_backend( + backend + )( + mask=self.metric_state.mask_hdiff, + zd_vertoffset=self.metric_state.zd_vertoffset, + zd_diffcoef=self.metric_state.zd_diffcoef, + geofac_n2s_c=self.interpolation_state.geofac_n2s_c, + geofac_n2s_nbh=self.interpolation_state.geofac_n2s_nbh, + vcoef=self.metric_state.zd_intcoef, + theta_v=prognostic_state.theta_v, + z_temp=self.z_temp, + horizontal_start=cell_start_nudging, + horizontal_end=cell_end_local, + vertical_start=0, + vertical_end=klevels, + offset_provider={ + "C2CEC": self.grid.get_c2cec_connectivity(), + "C2E2C": self.grid.get_c2e2c_connectivity(), + "Koff": KDim, + }, + ) + + log.debug( + "running fused stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end" + ) + log.debug("running fused stencil 16 (update_theta_and_exner): start") + update_theta_and_exner.with_backend(backend)( z_temp=self.z_temp, - area=cell_areas, + area=self.cell_params.area, theta_v=prognostic_state.theta_v, exner=prognostic_state.exner_pressure, rd_o_cvd=self.rd_o_cvd, @@ -1043,5 +880,5 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={}, ) - log.debug("running fused stencil update_theta_and_exner: end") + log.debug("running stencil 16 (update_theta_and_exner): end") # 10. HALO EXCHANGE sync_patch_array diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py deleted file mode 100644 index a4f74c4c7..000000000 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_program.py +++ /dev/null @@ -1,314 +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 - -# flake8: noqa - -from gt4py.next.common import Field -from gt4py.next.ffront.decorator import program -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.atm_dyn_iconam.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( - _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance, -) -from icon4py.atm_dyn_iconam.calculate_diagnostic_quantities_for_turbulence import ( - _calculate_diagnostic_quantities_for_turbulence, -) -from icon4py.atm_dyn_iconam.calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools import ( - _calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools, -) -from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( - _calculate_nabla2_and_smag_coefficients_for_vn, -) -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import ( - _fused_mo_nh_diffusion_stencil_04_05_06, -) -from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( - _mo_intp_rbf_rbf_vec_interpol_vertex, -) -from icon4py.atm_dyn_iconam.update_theta_and_exner import _update_theta_and_exner -from icon4py.common.dimension import ( - C2E2CDim, - C2E2CODim, - C2EDim, - CellDim, - ECVDim, - EdgeDim, - KDim, - V2EDim, - VertexDim, -) -from icon4py.diffusion.utils import _scale_k, _set_zero_v_k - - -@program -def diffusion_run( - diagnostic_hdef_ic: Field[[CellDim, KDim], float], - diagnostic_div_ic: Field[[CellDim, KDim], float], - diagnostic_dwdx: Field[[CellDim, KDim], float], - diagnostic_dwdy: Field[[CellDim, KDim], float], - prognostic_w: Field[[CellDim, KDim], float], - prognostic_vn: Field[[EdgeDim, KDim], float], - prognostic_exner_pressure: Field[[CellDim, KDim], float], - prognostic_theta_v: Field[[CellDim, KDim], float], - metric_theta_ref_mc: Field[[CellDim, KDim], float], - metric_wgtfac_c: Field[[CellDim, KDim], float], - metric_mask_hdiff: Field[[CellDim, KDim], int], - metric_zd_vertidx: Field[[CellDim, C2E2CDim, KDim], int], - metric_zd_diffcoef: Field[[CellDim, KDim], float], - metric_zd_intcoef: Field[[CellDim, C2E2CDim, KDim], float], - interpolation_e_bln_c_s: Field[[CellDim, C2EDim], float], - interpolation_rbf_coeff_1: Field[[VertexDim, V2EDim], float], - interpolation_rbf_coeff_2: Field[[VertexDim, V2EDim], float], - interpolation_geofac_div: Field[[CellDim, C2EDim], float], - interpolation_geofac_grg_x: Field[[CellDim, C2E2CODim], float], - interpolation_geofac_grg_y: Field[[CellDim, C2E2CODim], float], - interpolation_nudgecoeff_e: Field[[EdgeDim], float], - interpolation_geofac_n2s: Field[[CellDim, C2E2CODim], float], - interpolation_geofac_n2s_c: Field[[CellDim], float], - interpolation_geofac_n2s_nbh: Field[[CellDim, C2E2CDim], float], - tangent_orientation: Field[[EdgeDim], float], - inverse_primal_edge_lengths: Field[[EdgeDim], float], - inverse_dual_edge_lengths: Field[[EdgeDim], float], - inverse_vert_vert_lengths: Field[[EdgeDim], float], - primal_normal_vert_1: Field[[ECVDim], float], - primal_normal_vert_2: Field[[ECVDim], float], - dual_normal_vert_1: Field[[ECVDim], float], - dual_normal_vert_2: Field[[ECVDim], float], - edge_areas: Field[[EdgeDim], float], - cell_areas: Field[[CellDim], float], - diff_multfac_vn: Field[[KDim], float], - dtime: float, - rd_o_cvd: float, - local_thresh_tdiff: float, - local_smag_limit: Field[[KDim], float], - local_u_vert: Field[[VertexDim, KDim], float], - local_v_vert: Field[[VertexDim, KDim], float], - local_enh_smag_fac: Field[[KDim], float], - local_kh_smag_e: Field[[EdgeDim, KDim], float], - local_kh_smag_ec: Field[[EdgeDim, KDim], float], - local_z_nabla2_e: Field[[EdgeDim, KDim], float], - local_z_temp: Field[[CellDim, KDim], float], - local_diff_multfac_smag: Field[[KDim], float], - local_diff_multfac_n2w: Field[[KDim], float], - local_smag_offset: float, - local_nudgezone_diff: float, - local_fac_bdydiff_v: float, - local_diff_multfac_w: float, - local_vertical_index: Field[[KDim], int32], - local_horizontal_cell_index: Field[[CellDim], int32], - local_horizontal_edge_index: Field[[EdgeDim], int32], - cell_startindex_interior: int32, - cell_halo_idx: int32, - cell_startindex_nudging: int32, - cell_endindex_local_plus1: int32, - cell_endindex_local: int32, - edge_startindex_nudging_plus1: int32, - edge_startindex_nudging_minus1: int32, - edge_endindex_local: int32, - edge_endindex_local_minus2: int32, - vertex_startindex_lb_plus3: int32, - vertex_startindex_lb_plus1: int32, - vertex_endindex_local: int32, - vertex_endindex_local_minus1: int32, - index_of_damping_height: int32, - nlev: int32, - boundary_diffusion_start_index_edges: int32, -): - _scale_k(local_enh_smag_fac, dtime, out=local_diff_multfac_smag) - - # TODO: @magdalena is this needed?, if not remove - _set_zero_v_k(out=local_u_vert) - _set_zero_v_k(out=local_v_vert) - - # # 1. CALL rbf_vec_interpol_vertex - _mo_intp_rbf_rbf_vec_interpol_vertex( - prognostic_vn, - interpolation_rbf_coeff_1, - interpolation_rbf_coeff_2, - out=(local_u_vert, local_v_vert), - domain={ - VertexDim: (vertex_startindex_lb_plus1, vertex_endindex_local_minus1), - KDim: (0, nlev), - }, - ) - - # 2. HALO EXCHANGE -- CALL sync_patch_array_mult - # 3. mo_nh_diffusion_stencil_01, mo_nh_diffusion_stencil_02, mo_nh_diffusion_stencil_03 - - _calculate_nabla2_and_smag_coefficients_for_vn( - local_diff_multfac_smag, - tangent_orientation, - inverse_primal_edge_lengths, - inverse_vert_vert_lengths, - local_u_vert, - local_v_vert, - primal_normal_vert_1, - primal_normal_vert_2, - dual_normal_vert_1, - dual_normal_vert_2, - prognostic_vn, - local_smag_limit, - local_smag_offset, - out=(local_kh_smag_e, local_kh_smag_ec, local_z_nabla2_e), - domain={ - # TODO @magdalena wrong start index?? - EdgeDim: (boundary_diffusion_start_index_edges, edge_endindex_local_minus2), - KDim: (0, nlev), - }, - ) - - _calculate_diagnostic_quantities_for_turbulence( - local_kh_smag_ec, - prognostic_vn, - interpolation_e_bln_c_s, - interpolation_geofac_div, - local_diff_multfac_smag, - metric_wgtfac_c, - out=( - diagnostic_div_ic, - diagnostic_hdef_ic, - ), - domain={ - CellDim: (cell_startindex_nudging, cell_endindex_local), - KDim: (0, nlev), - }, - ) - # - # # 4. IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH - # - # # 5. CALL rbf_vec_interpol_vertex_wp - _mo_intp_rbf_rbf_vec_interpol_vertex( - local_z_nabla2_e, - interpolation_rbf_coeff_1, - interpolation_rbf_coeff_2, - out=(local_u_vert, local_v_vert), - domain={ - VertexDim: (vertex_startindex_lb_plus3, vertex_endindex_local), - KDim: (0, nlev), - }, - ) - # # 6. HALO EXCHANGE -- CALL sync_patch_array_mult - # - # # 7. mo_nh_diffusion_stencil_04, mo_nh_diffusion_stencil_05 - # # 7a. IF (l_limited_area .OR. jg > 1) mo_nh_diffusion_stencil_06 - # - _fused_mo_nh_diffusion_stencil_04_05_06( - local_u_vert, - local_v_vert, - primal_normal_vert_1, - primal_normal_vert_2, - local_z_nabla2_e, - inverse_vert_vert_lengths, - inverse_primal_edge_lengths, - edge_areas, - local_kh_smag_e, - diff_multfac_vn, - interpolation_nudgecoeff_e, - prognostic_vn, - local_horizontal_edge_index, - local_nudgezone_diff, - local_fac_bdydiff_v, - edge_startindex_nudging_minus1, - out=prognostic_vn, - domain={ - EdgeDim: (edge_startindex_nudging_plus1, edge_endindex_local), - KDim: (0, nlev), - }, - ) - # # 7b. mo_nh_diffusion_stencil_07, mo_nh_diffusion_stencil_08, - # # mo_nh_diffusion_stencil_09, mo_nh_diffusion_stencil_10 - - _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( - cell_areas, - interpolation_geofac_n2s, - interpolation_geofac_grg_x, - interpolation_geofac_grg_y, - prognostic_w, - diagnostic_dwdx, - diagnostic_dwdy, - local_diff_multfac_w, - local_diff_multfac_n2w, - local_vertical_index, - local_horizontal_cell_index, - index_of_damping_height, - cell_startindex_interior, - cell_halo_idx, - out=( - prognostic_w, - diagnostic_dwdx, - diagnostic_dwdy, - ), - domain={ - CellDim: (cell_startindex_nudging, cell_endindex_local_plus1), - KDim: (0, nlev), - }, - ) - # # 8. HALO EXCHANGE: CALL sync_patch_array - # # 9. mo_nh_diffusion_stencil_11, mo_nh_diffusion_stencil_12, mo_nh_diffusion_stencil_13, - # # mo_nh_diffusion_stencil_14, mo_nh_diffusion_stencil_15, mo_nh_diffusion_stencil_16 - # - # # TODO @magdalena check: kh_smag_e is an out field, should not be calculated in init? - # - - _calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( - prognostic_theta_v, - metric_theta_ref_mc, - local_thresh_tdiff, - local_kh_smag_e, - out=local_kh_smag_e, - domain={ - EdgeDim: (edge_startindex_nudging_plus1, edge_endindex_local), - KDim: (nlev - 2, nlev), - }, - ) - # _calculate_nabla2_for_theta( - # local_kh_smag_e, - # inverse_dual_edge_lengths, - # prognostic_theta_v, - # interpolation_geofac_div, - # out=local_z_temp, - # domain={ - # CellDim: (cell_startindex_nudging, cell_endindex_local), - # KDim: (0, nlev), - # }, - # ) - - # MO_NH_DIFFUSION_STENCIL_15: as_offset index fields! - # _mo_nh_diffusion_stencil_15( - # metric_mask_hdiff, - # metric_zd_vertidx, - # metric_zd_diffcoef, - # interpolation_geofac_n2s_c, - # interpolation_geofac_n2s_nbh, - # metric_zd_intcoef, - # prognostic_theta_v, - # local_z_temp, - # domain = { - # CellDim:(cell_startindex_nudging, cell_endindex_local), - # KDim: (0,nlev) - # } - # ) - - _update_theta_and_exner( - local_z_temp, - cell_areas, - prognostic_theta_v, - prognostic_exner_pressure, - rd_o_cvd, - out=(prognostic_theta_v, prognostic_exner_pressure), - domain={ - CellDim: (cell_startindex_nudging, cell_endindex_local), - KDim: (0, nlev), - }, - ) - # 10. HALO EXCHANGE sync_patch_array diff --git a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py index aa59ec458..7c41c7331 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py @@ -88,7 +88,6 @@ def local(cls, dim: Dimension) -> int: @classmethod def halo(cls, dim: Dimension) -> int: - """Indicate the halo points.""" match (dim): case (dimension.CellDim): return cls._HALO_CELLS diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index 2cc3dc56f..53b8eb823 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -11,14 +11,22 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from typing import Dict, Tuple +from typing import Dict import numpy as np from gt4py.next.common import Dimension, DimensionKind, Field from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import NeighborTableOffsetProvider -from icon4py.common.dimension import CellDim, ECVDim, EdgeDim, KDim, VertexDim +from icon4py.common.dimension import ( + CECDim, + CEDim, + CellDim, + ECVDim, + EdgeDim, + KDim, + VertexDim, +) from icon4py.diffusion.horizontal import HorizontalMeshSize @@ -128,7 +136,7 @@ def num_edges(self): def get_indices_from_to( self, dim: Dimension, start_marker: int, end_marker: int - ) -> Tuple[int32, int32]: + ) -> tuple[int32, int32]: """ Use to specify domains of a field for field_operator. @@ -173,10 +181,29 @@ def get_v2e_connectivity(self): return NeighborTableOffsetProvider(table, VertexDim, EdgeDim, table.shape[1]) def get_e2ecv_connectivity(self): - old_shape = self.connectivities["e2c2v"].shape - v2ecv_table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) + return self._neighbortable_offset_provider_for_1d_sparse_fields( + self.connectivities["e2c2v"].shape, EdgeDim, ECVDim + ) + + def _neighbortable_offset_provider_for_1d_sparse_fields( + self, + old_shape: tuple[int, int], + origin_axis: Dimension, + neighbor_axis: Dimension, + ): + table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) return NeighborTableOffsetProvider( - v2ecv_table, EdgeDim, ECVDim, v2ecv_table.shape[1] + table, origin_axis, neighbor_axis, table.shape[1] + ) + + def get_c2cec_connectivity(self): + return self._neighbortable_offset_provider_for_1d_sparse_fields( + self.connectivities["c2e2c"].shape, CellDim, CECDim + ) + + def get_c2ce_connectivity(self): + return self._neighbortable_offset_provider_for_1d_sparse_fields( + self.connectivities["c2e"].shape, CellDim, CEDim ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py index 0c8241062..98586b90c 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py @@ -18,9 +18,8 @@ from gt4py.next.iterator.embedded import np_as_located_field from icon4py.common.dimension import ( - C2E2CDim, C2E2CODim, - C2EDim, + CECDim, CEDim, CellDim, EdgeDim, @@ -38,7 +37,7 @@ class InterpolationState: """ e_bln_c_s: Field[ - [CellDim, C2EDim], float + [CEDim], float ] # coefficent for bilinear interpolation from edge to cell () rbf_coeff_1: Field[ [VertexDim, V2EDim], float @@ -47,32 +46,29 @@ class InterpolationState: [VertexDim, V2EDim], float ] # rbf_vec_coeff_v_2(nproma, rbf_vec_dim_v, nblks_v) - _geofac_div: Field[ - [CellDim, C2EDim], float + geofac_div: Field[ + [CEDim], float ] # factor for divergence (nproma,cell_type,nblks_c) geofac_n2s: Field[ [CellDim, C2E2CODim], float ] # factor for nabla2-scalar (nproma,cell_type+1,nblks_c) - geofac_grg_x: Field[ + geofac_grg_x: Field[[CellDim, C2E2CODim], float] + geofac_grg_y: Field[ [CellDim, C2E2CODim], float - ] # factor for green gauss gradient (nproma,4,nblks_c,2) - geofac_grg_y: Field[[CellDim, C2E2CODim], float] + ] # factors for green gauss gradient (nproma,4,nblks_c,2) nudgecoeff_e: Field[[EdgeDim], float] # Nudgeing coeffients for edges - @property - def geofac_div(self): - old_shape = np.asarray(self._geofac_div).shape - return np_as_located_field(CEDim)( - np.asarray(self._geofac_div).reshape((old_shape[0] * old_shape[1],)) - ) - @property def geofac_n2s_c(self) -> Field[[CellDim], float]: return np_as_located_field(CellDim)(np.asarray(self.geofac_n2s)[:, 0]) @property - def geofac_n2s_nbh(self) -> Field[[CellDim, C2E2CDim], float]: - return np_as_located_field(CellDim, C2E2CDim)( - np.asarray(self.geofac_n2s)[:, 1:] + def geofac_n2s_nbh(self) -> Field[[CECDim], float]: + geofac_nbh_ar = np.asarray(self.geofac_n2s)[:, 1:] + old_shape = geofac_nbh_ar.shape + return np_as_located_field(CECDim)( + geofac_nbh_ar.reshape( + old_shape[0] * old_shape[1], + ) ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py b/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py index 069863583..4fcdb0830 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py @@ -14,8 +14,9 @@ from dataclasses import dataclass from gt4py.next.common import Field +from numpy import int32 -from icon4py.common.dimension import C2E2CDim, CellDim, KDim +from icon4py.common.dimension import CECDim, CellDim, KDim @dataclass @@ -25,6 +26,6 @@ class MetricState: [CellDim, KDim], float ] # weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_c) mask_hdiff: Field[[CellDim, KDim], bool] - zd_vertidx: Field[[CellDim, C2E2CDim, KDim], int] + zd_vertoffset: Field[[CECDim, KDim], int32] zd_diffcoef: Field[[CellDim, KDim], float] - zd_intcoef: Field[[CellDim, C2E2CDim, KDim], float] + zd_intcoef: Field[[CECDim, KDim], float] diff --git a/atm_dyn_iconam/src/icon4py/diffusion/utils.py b/atm_dyn_iconam/src/icon4py/diffusion/utils.py index 674a5db7a..073a7915c 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/utils.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/utils.py @@ -34,6 +34,13 @@ def _identity_c_k( return field +@program +def copy_field( + old_f: Field[[CellDim, KDim], float], new_f: Field[[CellDim, KDim], float] +): + _identity_c_k(old_f, out=new_f) + + @field_operator def _identity_e_k( field: Field[[EdgeDim, KDim], float] diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index d8ead5a3a..89baa998f 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -20,16 +20,15 @@ from devtools import Timer from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn -from atm_dyn_iconam.tests.test_utils.serialbox_utils import IconSerialDataProvider from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.diffusion import Diffusion, DiffusionParams -from icon4py.diffusion.horizontal import CellParams, EdgeParams from icon4py.diffusion.prognostic_state import PrognosticState from icon4py.diffusion.utils import copy_diagnostic_and_prognostics from icon4py.driver.icon_configuration import IconRunConfig, read_config from icon4py.driver.io_utils import ( SIMULATION_START_DATE, configure_logging, + import_testutils, read_geometry_fields, read_icon_grid, read_initial_state, @@ -37,11 +36,15 @@ ) +helpers = import_testutils() +from helpers import serialbox_utils as sb_utils # noqa + + log = logging.getLogger(__name__) class DummyAtmoNonHydro: - def __init__(self, data_provider: IconSerialDataProvider): + def __init__(self, data_provider: sb_utils.IconSerialDataProvider): self.config = None self.data_provider = data_provider self.simulation_date = datetime.fromisoformat(SIMULATION_START_DATE) @@ -101,14 +104,10 @@ def __init__( config: IconRunConfig, diffusion: Diffusion, atmo_non_hydro: DummyAtmoNonHydro, - edge_geometry: EdgeParams, - cell_geometry: CellParams, ): self.config = config self.diffusion = diffusion self.atmo_non_hydro = atmo_non_hydro - self.edges = edge_geometry - self.cells = cell_geometry def _full_name(self, func: Callable): return ":".join((self.__class__.__name__, func.__name__)) @@ -122,18 +121,10 @@ def _timestep( self.atmo_non_hydro.do_dynamics_substepping( self.config.dtime, diagnostic_state, prognostic_state ) - self.diffusion.time_step( + self.diffusion.run( diagnostic_state, prognostic_state, self.config.dtime, - self.edges.tangent_orientation, - self.edges.inverse_primal_edge_lengths, - self.edges.inverse_dual_edge_lengths, - self.edges.inverse_vertex_vertex_lengths, - self.edges.primal_normal_vert, - self.edges.dual_normal_vert, - self.edges.edge_areas, - self.cells.area, ) def __call__( @@ -145,18 +136,10 @@ def __call__( f"starting time loop for dtime={self.config.dtime} n_timesteps={self.config.n_time_steps}" ) log.info("running initial step to diffuse fields before timeloop starts") - self.diffusion.initial_step( + self.diffusion.initial_run( diagnostic_state, prognostic_state, self.config.dtime, - self.edges.tangent_orientation, - self.edges.inverse_primal_edge_lengths, - self.edges.inverse_dual_edge_lengths, - self.edges.inverse_vertex_vertex_lengths, - self.edges.primal_normal_vert, - self.edges.dual_normal_vert, - self.edges.edge_areas, - self.cells.area, ) log.info( f"starting real time loop for dtime={self.config.dtime} n_timesteps={self.config.n_time_steps}" @@ -197,7 +180,7 @@ def initialize(n_time_steps, file_path: Path): log.info("initializing dycore") diffusion_params = DiffusionParams(config.diffusion_config) - diffusion = Diffusion(run_program=False) + diffusion = Diffusion() diffusion.init( icon_grid, config.diffusion_config, @@ -205,6 +188,8 @@ def initialize(n_time_steps, file_path: Path): vertical_geometry, metric_state, interpolation_state, + edge_geometry, + cell_geometry, ) data_provider, diagnostic_state, prognostic_state = read_initial_state(file_path) @@ -216,8 +201,6 @@ def initialize(n_time_steps, file_path: Path): config=config.run_config, diffusion=diffusion, atmo_non_hydro=atmo_non_hydro, - edge_geometry=edge_geometry, - cell_geometry=cell_geometry, ) return tl, diagnostic_state, prognostic_state diff --git a/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py b/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py index 0c55ae88f..7e75ebf63 100644 --- a/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py +++ b/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py @@ -17,6 +17,9 @@ from icon4py.diffusion.diffusion import DiffusionConfig +n_substeps_reduced = 2 + + @dataclass class IconRunConfig: n_time_steps: int = 5 @@ -35,7 +38,6 @@ class IconConfig: dycore_config: AtmoNonHydroConfig -# TODO @magdalena move to io_utils? def read_config(experiment: Optional[str], n_time_steps: int) -> IconConfig: def _default_run_config(n_steps: int): if n_steps > 5: @@ -46,6 +48,7 @@ def mch_ch_r04b09_diffusion_config(): return DiffusionConfig( diffusion_type=5, hdiff_w=True, + n_substeps=n_substeps_reduced, hdiff_vn=True, type_t_diffu=2, type_vn_diffu=1, diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index f5f15ba09..07ac5f5e1 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -10,13 +10,13 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import importlib import logging +import sys from datetime import datetime from enum import Enum from pathlib import Path -from atm_dyn_iconam.tests.test_utils import serialbox_utils -from atm_dyn_iconam.tests.test_utils.serialbox_utils import IconSerialDataProvider from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.horizontal import CellParams, EdgeParams from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams @@ -29,6 +29,22 @@ log = logging.getLogger(__name__) +def import_testutils(): + testutils = ( + Path(__file__).parent.__str__() + "/../../../tests/test_utils/__init__.py" + ) + spec = importlib.util.spec_from_file_location("helpers", testutils) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +helpers = import_testutils() + +from helpers import serialbox_utils as sb # noqa + + class SerializationType(str, Enum): SB = "serialbox" NC = "netcdf" @@ -46,9 +62,7 @@ def read_icon_grid(path: Path, ser_type=SerializationType.SB) -> IconGrid: """ if ser_type == SerializationType.SB: return ( - serialbox_utils.IconSerialDataProvider( - "icon_pydycore", str(path.absolute()), False - ) + sb.IconSerialDataProvider("icon_pydycore", str(path.absolute()), False) .from_savepoint_grid() .construct_icon_grid() ) @@ -58,7 +72,7 @@ def read_icon_grid(path: Path, ser_type=SerializationType.SB) -> IconGrid: def read_initial_state( gridfile_path: Path, -) -> tuple[IconSerialDataProvider, DiagnosticState, PrognosticState]: +) -> tuple[sb.IconSerialDataProvider, DiagnosticState, PrognosticState]: """ Read prognostic and diagnostic state from serialized data. @@ -70,7 +84,7 @@ def read_initial_state( read from within the dummy timeloop """ - data_provider = serialbox_utils.IconSerialDataProvider( + data_provider = sb.IconSerialDataProvider( "icon_pydycore", str(gridfile_path), False ) init_savepoint = data_provider.from_savepoint_diffusion_init( @@ -95,7 +109,7 @@ def read_geometry_fields( the data is originally obtained from the grid file (horizontal fields) or some special input files. """ if ser_type == SerializationType.SB: - sp = serialbox_utils.IconSerialDataProvider( + sp = sb.IconSerialDataProvider( "icon_pydycore", str(path.absolute()), False ).from_savepoint_grid() edge_geometry = sp.construct_edge_geometry() @@ -123,7 +137,7 @@ def read_static_fields( """ if ser_type == SerializationType.SB: - dataprovider = serialbox_utils.IconSerialDataProvider( + dataprovider = sb.IconSerialDataProvider( "icon_pydycore", str(path.absolute()), False ) interpolation_state = ( diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 86da046cb..4f5cd0cc6 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -91,8 +91,8 @@ def diffusion_savepoint_init(data_provider, linit, step_date_init): """ Load data from ICON savepoint at start of diffusion module. - date of the timestamp to be selected can be set seperately by overriding the 'step_data' - fixture, passing 'step_data=' + date of the timestamp to be selected can be set seperately by overriding the 'step_date_init' + fixture, passing 'step_date_init=' linit flag can be set by overriding the 'linit' fixture """ @@ -100,14 +100,14 @@ def diffusion_savepoint_init(data_provider, linit, step_date_init): @pytest.fixture -def diffusion_savepoint_exit(data_provider, step_date_exit): +def diffusion_savepoint_exit(data_provider, linit, step_date_exit): """ Load data from ICON savepoint at exist of diffusion module. date of the timestamp to be selected can be set seperately by overriding the 'step_data' fixture, passing 'step_data=' """ - sp = data_provider.from_savepoint_diffusion_exit(linit=False, date=step_date_exit) + sp = data_provider.from_savepoint_diffusion_exit(linit=linit, date=step_date_exit) return sp diff --git a/atm_dyn_iconam/tests/test_apply_nabla2_to_w.py b/atm_dyn_iconam/tests/test_apply_nabla2_to_w.py index ff3fe64c9..cbeee1ae1 100644 --- a/atm_dyn_iconam/tests/test_apply_nabla2_to_w.py +++ b/atm_dyn_iconam/tests/test_apply_nabla2_to_w.py @@ -13,6 +13,7 @@ import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 from icon4py.atm_dyn_iconam.apply_nabla2_to_w import apply_nabla2_to_w from icon4py.common.dimension import C2E2CODim, CellDim, KDim @@ -32,7 +33,8 @@ def reference( z_nabla2_c: np.array, geofac_n2s: np.array, w: np.array, - diff_multfac_w, + diff_multfac_w: float, + **kwargs, ) -> np.array: geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) area = np.expand_dims(area, axis=-1) @@ -47,12 +49,14 @@ def input_data(self, mesh): z_nabla2_c = random_field(mesh, CellDim, KDim) geofac_n2s = random_field(mesh, CellDim, C2E2CODim) w = random_field(mesh, CellDim, KDim) - diff_multfac_w = 5.0 - return dict( area=area, z_nabla2_c=z_nabla2_c, geofac_n2s=geofac_n2s, w=w, - diff_multfac_w=diff_multfac_w, + diff_multfac_w=5.0, + horizontal_start=int32(0), + horizontal_end=int32(mesh.n_cells), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), ) diff --git a/atm_dyn_iconam/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py b/atm_dyn_iconam/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py index c6cba862f..c9a0a86df 100644 --- a/atm_dyn_iconam/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py +++ b/atm_dyn_iconam/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py @@ -13,6 +13,7 @@ import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import ( apply_nabla2_to_w_in_upper_damping_layer, @@ -39,6 +40,10 @@ def input_data(self, mesh): diff_multfac_n2w=diff_multfac_n2w, cell_area=cell_area, z_nabla2_c=z_nabla2_c, + horizontal_start=int32(0), + horizontal_end=int(mesh.n_cells), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), ) @staticmethod @@ -48,6 +53,7 @@ def reference( diff_multfac_n2w: np.array, cell_area: np.array, z_nabla2_c: np.array, + **kwargs, ) -> np.array: cell_area = np.expand_dims(cell_area, axis=-1) w = w + diff_multfac_n2w * cell_area * z_nabla2_c diff --git a/atm_dyn_iconam/tests/test_calculate_horizontal_gradients_for_turbulence.py b/atm_dyn_iconam/tests/test_calculate_horizontal_gradients_for_turbulence.py index 0f5e29c62..4590b91f4 100644 --- a/atm_dyn_iconam/tests/test_calculate_horizontal_gradients_for_turbulence.py +++ b/atm_dyn_iconam/tests/test_calculate_horizontal_gradients_for_turbulence.py @@ -13,6 +13,7 @@ import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 from icon4py.atm_dyn_iconam.calculate_horizontal_gradients_for_turbulence import ( calculate_horizontal_gradients_for_turbulence, @@ -52,4 +53,8 @@ def input_data(self, mesh): geofac_grg_y=geofac_grg_y, dwdx=dwdx, dwdy=dwdy, + horizontal_start=int32(0), + horizontal_end=int32(mesh.n_cells), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), ) diff --git a/atm_dyn_iconam/tests/test_calculate_nabla2_for_w.py b/atm_dyn_iconam/tests/test_calculate_nabla2_for_w.py index 42a971996..f8eb5d3bd 100644 --- a/atm_dyn_iconam/tests/test_calculate_nabla2_for_w.py +++ b/atm_dyn_iconam/tests/test_calculate_nabla2_for_w.py @@ -13,6 +13,7 @@ import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 from icon4py.atm_dyn_iconam.calculate_nabla2_for_w import calculate_nabla2_for_w from icon4py.common.dimension import C2E2CODim, CellDim, KDim @@ -41,4 +42,8 @@ def input_data(self, mesh): w=w, geofac_n2s=geofac_n2s, z_nabla2_c=z_nabla2_c, + horizontal_start=int32(0), + horizontal_end=int32(mesh.n_cells), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), ) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index a62688dcb..1a7c17ea6 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -243,15 +243,19 @@ def test_diffusion_init( interpolation_state = interpolation_savepoint.construct_interpolation_state() metric_state = metrics_savepoint.construct_metric_state() + edge_params = grid_savepoint.construct_edge_geometry() + cell_params = grid_savepoint.construct_cell_geometry() diffusion = Diffusion() diffusion.init( grid=icon_grid, - vertical_params=vertical_params, config=config, params=additional_parameters, + vertical_params=vertical_params, metric_state=metric_state, interpolation_state=interpolation_state, + edge_params=edge_params, + cell_params=cell_params, ) assert diffusion.diff_multfac_w == min( @@ -355,18 +359,22 @@ def test_verify_diffusion_init_against_first_regular_savepoint( config.ndyn_substeps = datarun_reduced_substeps additional_parameters = DiffusionParams(config) vct_a = grid_savepoint.vct_a() + cell_geometry = grid_savepoint.construct_cell_geometry() + edge_geometry = grid_savepoint.construct_edge_geometry() interpolation_state = interpolation_savepoint.construct_interpolation_state() metric_state = metrics_savepoint.construct_metric_state() diffusion = Diffusion() diffusion.init( - config=config, grid=icon_grid, + config=config, params=additional_parameters, vertical_params=VerticalModelParams(vct_a, damping_height), metric_state=metric_state, interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry, ) _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) @@ -390,6 +398,8 @@ def test_verify_diffusion_init_against_other_regular_savepoint( vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) interpolation_state = interpolation_savepoint.construct_interpolation_state() metric_state = metrics_savepoint.construct_metric_state() + edge_params = grid_savepoint.construct_edge_geometry() + cell_params = grid_savepoint.construct_cell_geometry() diffusion = Diffusion() diffusion.init( @@ -399,16 +409,23 @@ def test_verify_diffusion_init_against_other_regular_savepoint( vertical_params, metric_state, interpolation_state, + edge_params, + cell_params, ) _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) -@pytest.mark.skip("fix: diffusion_stencil_15") -@pytest.mark.parametrize("run_with_program", [True, False]) @pytest.mark.datatest +@pytest.mark.parametrize( + "step_date_init, step_date_exit", + [ + ("2021-06-20T12:00:10.000", "2021-06-20T12:00:10.000"), + ("2021-06-20T12:00:20.000", "2021-06-20T12:00:20.000"), + ("2021-06-20T12:01:00.000", "2021-06-20T12:01:00.000"), + ], +) def test_run_diffusion_single_step( - run_with_program, diffusion_savepoint_init, diffusion_savepoint_exit, interpolation_savepoint, @@ -433,7 +450,7 @@ def test_run_diffusion_single_step( config.ndyn_substeps = datarun_reduced_substeps additional_parameters = DiffusionParams(config) - diffusion = Diffusion(run_program=run_with_program) + diffusion = Diffusion() diffusion.init( grid=icon_grid, config=config, @@ -441,49 +458,63 @@ def test_run_diffusion_single_step( vertical_params=vertical_params, metric_state=metric_state, interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry, ) - diffusion.time_step( + _verify_diffusion_fields( + diagnostic_state, prognostic_state, diffusion_savepoint_init + ) + assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v + diffusion.run( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime, - tangent_orientation=edge_geometry.tangent_orientation, - inverse_primal_edge_lengths=edge_geometry.inverse_primal_edge_lengths, - inverse_dual_edge_length=edge_geometry.inverse_dual_edge_lengths, - inverse_vert_vert_lengths=edge_geometry.inverse_vertex_vertex_lengths, - primal_normal_vert=edge_geometry.primal_normal_vert, - dual_normal_vert=edge_geometry.dual_normal_vert, - edge_areas=edge_geometry.edge_areas, - cell_areas=cell_geometry.area, + ) + _verify_diffusion_fields( + diagnostic_state, prognostic_state, diffusion_savepoint_exit ) - icon_result_exner = diffusion_savepoint_exit.exner() - icon_result_vn = diffusion_savepoint_exit.vn() - icon_result_w = diffusion_savepoint_exit.w() - icon_result_theta_w = diffusion_savepoint_exit.theta_v() - assert np.allclose(icon_result_w, np.asarray(prognostic_state.w)) - assert np.allclose(np.asarray(icon_result_vn), np.asarray(prognostic_state.vn)) - assert np.allclose( - np.asarray(icon_result_theta_w), np.asarray(prognostic_state.theta_v) - ) - assert np.allclose( - np.asarray(icon_result_exner), np.asarray(prognostic_state.exner_pressure) - ) +def _verify_diffusion_fields( + diagnostic_state, prognostic_state, diffusion_savepoint_exit +): + ref_div_ic = np.asarray(diffusion_savepoint_exit.div_ic()) + val_div_ic = np.asarray(diagnostic_state.div_ic) + ref_hdef_ic = np.asarray(diffusion_savepoint_exit.hdef_ic()) + val_hdef_ic = np.asarray(diagnostic_state.hdef_ic) + assert np.allclose(ref_div_ic, val_div_ic) + assert np.allclose(ref_hdef_ic, val_hdef_ic) + ref_w = np.asarray(diffusion_savepoint_exit.w()) + val_w = np.asarray(prognostic_state.w) + ref_dwdx = np.asarray(diffusion_savepoint_exit.dwdx()) + val_dwdx = np.asarray(diagnostic_state.dwdx) + ref_dwdy = np.asarray(diffusion_savepoint_exit.dwdy()) + val_dwdy = np.asarray(diagnostic_state.dwdy) + ref_vn = np.asarray(diffusion_savepoint_exit.vn()) + val_vn = np.asarray(prognostic_state.vn) + assert np.allclose(ref_vn, val_vn) + assert np.allclose(ref_dwdx, val_dwdx) + assert np.allclose(ref_dwdy, val_dwdy) + assert np.allclose(ref_w, val_w) + ref_exner = np.asarray(diffusion_savepoint_exit.exner()) + ref_theta_v = np.asarray(diffusion_savepoint_exit.theta_v()) + val_theta_v = np.asarray(prognostic_state.theta_v) + val_exner = np.asarray(prognostic_state.exner_pressure) + assert np.allclose(ref_theta_v, val_theta_v) + assert np.allclose(ref_exner, val_exner) -@pytest.mark.skip("fix: diffusion_stencil_15") @pytest.mark.datatest -def test_diffusion_five_steps( - damping_height, - r04b09_diffusion_config, - icon_grid, - grid_savepoint, - interpolation_savepoint, - metrics_savepoint, +@pytest.mark.parametrize("linit", [True]) +def test_run_diffusion_initial_step( diffusion_savepoint_init, diffusion_savepoint_exit, - linit=True, - step_date_exit="2021-06-20T12:01:00.000", + interpolation_savepoint, + metrics_savepoint, + grid_savepoint, + icon_grid, + r04b09_diffusion_config, + damping_height, ): dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() @@ -492,14 +523,14 @@ def test_diffusion_five_steps( metric_state = metrics_savepoint.construct_metric_state() diagnostic_state = diffusion_savepoint_init.construct_diagnostics() prognostic_state = diffusion_savepoint_init.construct_prognostics() - + vct_a = grid_savepoint.vct_a() vertical_params = VerticalModelParams( - vct_a=grid_savepoint.vct_a(), rayleigh_damping_height=damping_height + vct_a=vct_a, rayleigh_damping_height=damping_height ) - - additional_parameters = DiffusionParams(r04b09_diffusion_config) config = r04b09_diffusion_config config.ndyn_substeps = datarun_reduced_substeps + additional_parameters = DiffusionParams(config) + diffusion = Diffusion() diffusion.init( grid=icon_grid, @@ -508,44 +539,29 @@ def test_diffusion_five_steps( vertical_params=vertical_params, metric_state=metric_state, interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry, ) - diffusion.initial_step( + assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v + + diffusion.initial_run( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime, - tangent_orientation=edge_geometry.tangent_orientation, - inverse_primal_edge_lengths=edge_geometry.inverse_primal_edge_lengths, - inverse_dual_edge_length=edge_geometry.inverse_dual_edge_lengths, - inverse_vert_vert_lengths=edge_geometry.inverse_vertex_vertex_lengths, - primal_normal_vert=edge_geometry.primal_normal_vert, - dual_normal_vert=edge_geometry.dual_normal_vert, - edge_areas=edge_geometry.edge_areas, - cell_areas=cell_geometry.area, ) - for _ in range(4): - diffusion.time_step( - diagnostic_state=diagnostic_state, - prognostic_state=prognostic_state, - dtime=dtime, - tangent_orientation=edge_geometry.tangent_orientation, - inverse_primal_edge_lengths=edge_geometry.inverse_primal_edge_lengths, - inverse_dual_edge_length=edge_geometry.inverse_dual_edge_lengths, - inverse_vert_vert_lengths=edge_geometry.inverse_vertex_vertex_lengths, - primal_normal_vert=edge_geometry.primal_normal_vert, - dual_normal_vert=edge_geometry.dual_normal_vert, - edge_areas=edge_geometry.edge_areas, - cell_areas=cell_geometry.area, - ) - - icon_result_exner = diffusion_savepoint_exit.exner() - icon_result_vn = diffusion_savepoint_exit.vn() - icon_result_w = diffusion_savepoint_exit.w() - icon_result_theta_w = diffusion_savepoint_exit.theta_v() - assert np.allclose(icon_result_w, np.asarray(prognostic_state.w)) - assert np.allclose(np.asarray(icon_result_vn), np.asarray(prognostic_state.vn)) - assert np.allclose( - np.asarray(icon_result_theta_w), np.asarray(prognostic_state.theta_v) - ) - assert np.allclose( - np.asarray(icon_result_exner), np.asarray(prognostic_state.exner_pressure) + + _verify_diffusion_fields( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + diffusion_savepoint_exit=diffusion_savepoint_exit, ) + + +def test_verify_stencil15_field_manipulation(interpolation_savepoint, icon_grid): + geofac_n2s = np.asarray(interpolation_savepoint.geofac_n2s()) + int_state = interpolation_savepoint.construct_interpolation_state() + geofac_c = np.asarray(int_state.geofac_n2s_c) + geofac_nbh = np.asarray(int_state.geofac_n2s_nbh) + cec_table = icon_grid.get_c2cec_connectivity().table + assert np.allclose(geofac_c, geofac_n2s[:, 0]) + assert np.allclose(geofac_nbh[cec_table], geofac_n2s[:, 1:]) diff --git a/atm_dyn_iconam/tests/test_interpolation_state.py b/atm_dyn_iconam/tests/test_interpolation_state.py new file mode 100644 index 000000000..c49c73019 --- /dev/null +++ b/atm_dyn_iconam/tests/test_interpolation_state.py @@ -0,0 +1,27 @@ +# 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 + + +@pytest.mark.datatest +def test_cecdim(interpolation_savepoint, icon_grid): + interpolation_fields = interpolation_savepoint.construct_interpolation_state() + geofac_n2s = np.asarray(interpolation_fields.geofac_n2s) + geofac_n2s_nbh = np.asarray(interpolation_fields.geofac_n2s_nbh) + assert np.count_nonzero(geofac_n2s_nbh) > 0 + c2cec = icon_grid.get_c2cec_connectivity().table + ported = geofac_n2s_nbh[c2cec] + assert ported.shape == geofac_n2s[:, 1:].shape + assert np.allclose(ported, geofac_n2s[:, 1:]) diff --git a/atm_dyn_iconam/tests/test_io_utils.py b/atm_dyn_iconam/tests/test_io_utils.py index 932af8646..c31ce40d8 100644 --- a/atm_dyn_iconam/tests/test_io_utils.py +++ b/atm_dyn_iconam/tests/test_io_utils.py @@ -101,7 +101,7 @@ def assert_metric_state_fields(metric_state): assert metric_state.zd_diffcoef assert metric_state.theta_ref_mc assert metric_state.mask_hdiff - assert metric_state.zd_vertidx + assert metric_state.zd_vertoffset def assert_interpolation_state_fields(interpolation_state): @@ -109,7 +109,7 @@ def assert_interpolation_state_fields(interpolation_state): assert interpolation_state.e_bln_c_s assert interpolation_state.nudgecoeff_e assert interpolation_state.geofac_n2s_nbh - assert interpolation_state._geofac_div + assert interpolation_state.geofac_div assert interpolation_state.geofac_grg_y assert interpolation_state.geofac_grg_x assert interpolation_state.rbf_coeff_2 diff --git a/atm_dyn_iconam/tests/test_temporary_fields_for_turbulence_diagnostics.py b/atm_dyn_iconam/tests/test_temporary_fields_for_turbulence_diagnostics.py index 46fb11cd9..acda21fb3 100644 --- a/atm_dyn_iconam/tests/test_temporary_fields_for_turbulence_diagnostics.py +++ b/atm_dyn_iconam/tests/test_temporary_fields_for_turbulence_diagnostics.py @@ -17,9 +17,9 @@ from icon4py.atm_dyn_iconam.temporary_fields_for_turbulence_diagnostics import ( temporary_fields_for_turbulence_diagnostics, ) -from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim +from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from .test_utils.helpers import random_field, zero_field +from .test_utils.helpers import as_1D_sparse_field, random_field, zero_field from .test_utils.stencil_test import StencilTest @@ -38,12 +38,11 @@ def reference( **kwargs, ) -> dict: geofac_div = np.expand_dims(geofac_div, axis=-1) - vn_geofac = vn[mesh.c2e] * geofac_div + vn_geofac = vn[mesh.c2e] * geofac_div[mesh.get_c2ce_offset_provider().table] div = np.sum(vn_geofac, axis=1) - e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) diff_multfac_smag = np.expand_dims(diff_multfac_smag, axis=0) - mul = kh_smag_ec[mesh.c2e] * e_bln_c_s + mul = kh_smag_ec[mesh.c2e] * e_bln_c_s[mesh.get_c2ce_offset_provider().table] summed = np.sum(mul, axis=1) kh_c = summed / diff_multfac_smag @@ -52,9 +51,9 @@ def reference( @pytest.fixture def input_data(self, mesh): vn = random_field(mesh, EdgeDim, KDim) - geofac_div = random_field(mesh, CellDim, C2EDim) + geofac_div = as_1D_sparse_field(random_field(mesh, CellDim, C2EDim), CEDim) kh_smag_ec = random_field(mesh, EdgeDim, KDim) - e_bln_c_s = random_field(mesh, CellDim, C2EDim) + e_bln_c_s = as_1D_sparse_field(random_field(mesh, CellDim, C2EDim), CEDim) diff_multfac_smag = random_field(mesh, KDim) kh_c = zero_field(mesh, CellDim, KDim) diff --git a/atm_dyn_iconam/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/atm_dyn_iconam/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index 925e04d55..2e989dba4 100644 --- a/atm_dyn_iconam/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/atm_dyn_iconam/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -115,6 +115,10 @@ def test_mo_nh_diffusion_stencil_15(): vcoef_new, theta_v, z_temp, + horizontal_start=int32(0), + horizontal_end=int32(mesh.n_cells), + vertical_start=int32(0), + vertical_end=int32(mesh.k_level), offset_provider={ "C2E2C": mesh.get_c2e2c_offset_provider(), "C2CEC": StridedNeighborOffsetProvider(CellDim, CECDim, mesh.n_c2e2c), diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index a40807b87..94b02fb98 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -18,11 +18,12 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import np_as_located_field -from atm_dyn_iconam.tests.test_utils.helpers import as_1D_sparse_field from icon4py.common.dimension import ( C2E2CDim, C2E2CODim, C2EDim, + CECDim, + CEDim, CellDim, E2C2VDim, E2CDim, @@ -45,6 +46,8 @@ from icon4py.diffusion.metric_state import MetricState from icon4py.diffusion.prognostic_state import PrognosticState +from .helpers import as_1D_sparse_field + class IconSavepoint: def __init__(self, sp: ser.Savepoint, ser: ser.Serializer): @@ -183,6 +186,9 @@ def e2c2v(self): def v2e(self): return self._get_connectiviy_array("v2e") + def nrdmax(self): + return self._get_connectiviy_array("nrdmax") + def construct_icon_grid(self) -> IconGrid: sp_meta = self.get_metadata( "nproma", "nlev", "num_vert", "num_cells", "num_edges" @@ -279,10 +285,10 @@ def nudgecoeff_e(self): def construct_interpolation_state(self) -> InterpolationState: grg = self.geofac_grg() return InterpolationState( - e_bln_c_s=self.e_bln_c_s(), + e_bln_c_s=as_1D_sparse_field(self.e_bln_c_s(), CEDim), rbf_coeff_1=self.rbf_vec_coeff_v1(), rbf_coeff_2=self.rbf_vec_coeff_v2(), - _geofac_div=self.geofac_div(), + geofac_div=as_1D_sparse_field(self.geofac_div(), CEDim), geofac_n2s=self.geofac_n2s(), geofac_grg_x=grg[0], geofac_grg_y=grg[1], @@ -297,7 +303,7 @@ def construct_metric_state(self) -> MetricState: theta_ref_mc=self.theta_ref_mc(), wgtfac_c=self.wgtfac_c(), zd_intcoef=self.zd_intcoef(), - zd_vertidx=self.zd_vertoffset(), + zd_vertoffset=self.zd_vertoffset(), zd_diffcoef=self.zd_diffcoef(), ) @@ -305,13 +311,28 @@ def zd_diffcoef(self): return self._get_field("zd_diffcoef", CellDim, KDim) def zd_intcoef(self): - return self._get_field("vcoef", CellDim, C2E2CDim, KDim) + ser_input = np.moveaxis( + (np.squeeze(self.serializer.read("vcoef", self.savepoint))), 1, -1 + ) + return self._linearize_first_2dims(ser_input, sparse_size=3) - def zd_vertidx(self): - return self._get_field("zd_vertidx", CellDim, C2E2CDim, dtype=int) + def _linearize_first_2dims(self, data: np.ndarray, sparse_size): + old_shape = data.shape + assert old_shape[1] == sparse_size + return np_as_located_field(CECDim, KDim)( + data.reshape(old_shape[0] * old_shape[1], old_shape[2]) + ) def zd_vertoffset(self): - return self._get_field("zd_vertoffset", CellDim, C2E2CDim, KDim, dtype=int) + ser_input = np.squeeze(self.serializer.read("zd_vertoffset", self.savepoint)) + ser_input = np.moveaxis(ser_input, 1, -1) + return self._linearize_first_2dims(ser_input, sparse_size=3) + + def zd_vertidx(self): + return np.squeeze(self.serializer.read("zd_vertidx", self.savepoint)) + + def zd_indlist(self): + return np.squeeze(self.serializer.read("zd_indlist", self.savepoint)) def theta_ref_mc(self): return self._get_field("theta_ref_mc", CellDim, KDim) @@ -323,7 +344,7 @@ def wgtfac_e(self): return self._get_field("wgtfac_e", EdgeDim, KDim) def mask_diff(self): - return self._get_field("mask_hdiff", CellDim, KDim, dtype=int) + return self._get_field("mask_hdiff", CellDim, KDim, dtype=bool) class IconDiffusionInitSavepoint(IconSavepoint): @@ -405,9 +426,24 @@ def theta_v(self): def w(self): return self._get_field("x_w", CellDim, KDim) + def dwdx(self): + return self._get_field("x_dwdx", CellDim, KDim) + + def dwdy(self): + return self._get_field("x_dwdy", CellDim, KDim) + def exner(self): return self._get_field("x_exner", CellDim, KDim) + def z_temp(self): + return self._get_field("x_z_temp", CellDim, KDim) + + def div_ic(self): + return self._get_field("x_div_ic", CellDim, KDim) + + def hdef_ic(self): + return self._get_field("x_hdef_ic", CellDim, KDim) + class IconSerialDataProvider: def __init__(self, fname_prefix, path=".", do_print=False): diff --git a/atm_dyn_iconam/tests/test_utils/simple_mesh.py b/atm_dyn_iconam/tests/test_utils/simple_mesh.py index 82cc9e44f..873a017b0 100644 --- a/atm_dyn_iconam/tests/test_utils/simple_mesh.py +++ b/atm_dyn_iconam/tests/test_utils/simple_mesh.py @@ -428,9 +428,16 @@ def get_e2c2v_offset_provider(self) -> NeighborTableOffsetProvider: def get_e2ecv_offset_provider(self): old_shape = self.e2c2v.shape - v2ecv_table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) + e2ecv_table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) return NeighborTableOffsetProvider( - v2ecv_table, EdgeDim, ECVDim, v2ecv_table.shape[1] + e2ecv_table, EdgeDim, ECVDim, e2ecv_table.shape[1] + ) + + def get_c2ce_offset_provider(self): + old_shape = self.c2e.shape + c2ce_table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) + return NeighborTableOffsetProvider( + c2ce_table, CellDim, CEDim, c2ce_table.shape[1] ) def get_offset_provider(self): @@ -445,5 +452,6 @@ def get_offset_provider(self): "E2C": self.get_e2c_offset_provider(), "E2V": self.get_e2v_offset_provider(), "E2C2V": self.get_e2c2v_offset_provider(), + "C2CE": self.get_c2ce_offset_provider(), "Koff": KDim, } diff --git a/atm_dyn_iconam/tests/test_vertical.py b/atm_dyn_iconam/tests/test_vertical.py index ccf0b7dbf..597ee5acf 100644 --- a/atm_dyn_iconam/tests/test_vertical.py +++ b/atm_dyn_iconam/tests/test_vertical.py @@ -34,13 +34,13 @@ def test_nrdmax_calculation(max_h, damping, delta): @pytest.mark.datatest -def test_nrdmax_calculation_from_icon_input(icon_grid, grid_savepoint): +def test_nrdmax_calculation_from_icon_input(icon_grid, grid_savepoint, damping_height): a = grid_savepoint.vct_a() - damping_height = 12500 + nrdmax = grid_savepoint.nrdmax() vertical_params = VerticalModelParams( rayleigh_damping_height=damping_height, vct_a=a ) - assert 9 == vertical_params.index_of_damping_layer + assert nrdmax == vertical_params.index_of_damping_layer a_array = np.asarray(a) - assert a_array[9] > damping_height - assert a_array[10] < damping_height + assert a_array[nrdmax] > damping_height + assert a_array[nrdmax + 1] < damping_height diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index e5dd89a97..4ee67b578 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -16,7 +16,7 @@ flake8-mutable>=1.2.0 flake8-pyproject>=1.2.2 isort~=5.10 mypy>=0.942 -typing-extensions>=4.5 +typing-extensions==4.5.0 pre-commit~=2.15 pytest>=6.1 pytest-benchmark>=4.0.0 diff --git a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py index 09e1dfce8..821081046 100644 --- a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py +++ b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py @@ -39,7 +39,7 @@ from icon4py.py2f.cffi_utils import CffiMethod -diffusion: Diffusion(run_program=True) +diffusion: Diffusion() @CffiMethod.register @@ -81,7 +81,7 @@ def diffusion_init( """ grid = IconGrid() # TODO where to get this from - edges_params = EdgeParams( + edge_params = EdgeParams( tangent_orientation=tangent_orientation, primal_edge_lengths=primal_edge_lengths, inverse_primal_edge_lengths=inverse_primal_edge_lengths, @@ -112,13 +112,13 @@ def diffusion_init( diffusion.init( grid=grid, - cell_params=cell_params, - edges_params=edges_params, config=config, params=derived_diffusion_params, vertical_params=vertical_params, metric_state=metric_state, interpolation_state=interpolation_state, + edge_params=edge_params, + cell_params=cell_params, ) @@ -143,13 +143,13 @@ def diffusion_run( theta_v=theta_v, ) if linit: - diffusion.initial_step( + diffusion.initial_run( diagnostic_state, prognostic_state, dtime, ) else: - diffusion.time_step(diagnostic_state, prognostic_state, dtime) + diffusion.run(diagnostic_state, prognostic_state, dtime) class DuplicateInitializationException(Exception): From 11b3466d0b1c6e858e3e78b45417295f6d37412a Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 30 Jun 2023 17:39:50 +0200 Subject: [PATCH 160/263] (fix) datatests: - add missing datatest marker in test_diffusion.py - use fixtures in test_io_utils.py --- atm_dyn_iconam/tests/conftest.py | 7 ++++++- atm_dyn_iconam/tests/test_diffusion.py | 2 +- atm_dyn_iconam/tests/test_io_utils.py | 22 +++++++++++----------- 3 files changed, 18 insertions(+), 13 deletions(-) diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 4f5cd0cc6..28a17a533 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -30,6 +30,11 @@ data_file = data_path.joinpath("mch_ch_r04b09_dsl_v2.tar.gz").name +@pytest.fixture +def get_data_path(setup_icon_data): + return extracted_path + + @pytest.fixture(scope="session") def setup_icon_data(): """ @@ -139,7 +144,7 @@ def grid_savepoint(data_provider): @pytest.fixture -def r04b09_diffusion_config(setup_icon_data) -> DiffusionConfig: +def r04b09_diffusion_config() -> DiffusionConfig: """ Create DiffusionConfig matching MCH_CH_r04b09_dsl. diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 1a7c17ea6..dab84d613 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -556,7 +556,7 @@ def test_run_diffusion_initial_step( diffusion_savepoint_exit=diffusion_savepoint_exit, ) - +@pytest.mark.datatest def test_verify_stencil15_field_manipulation(interpolation_savepoint, icon_grid): geofac_n2s = np.asarray(interpolation_savepoint.geofac_n2s()) int_state = interpolation_savepoint.construct_interpolation_state() diff --git a/atm_dyn_iconam/tests/test_io_utils.py b/atm_dyn_iconam/tests/test_io_utils.py index c31ce40d8..d2c9e1100 100644 --- a/atm_dyn_iconam/tests/test_io_utils.py +++ b/atm_dyn_iconam/tests/test_io_utils.py @@ -25,17 +25,17 @@ ) -test_data_path = pathlib.Path(__file__).parent.joinpath( - "./ser_icondata/mch_ch_r04b09_dsl/ser_data" -) + @pytest.mark.parametrize( "read_fun", (read_geometry_fields, read_static_fields, read_icon_grid) ) -def test_read_geometry_fields_not_implemented_type(read_fun): +@pytest.mark.datatest +def test_read_geometry_fields_not_implemented_type(read_fun, get_data_path): + path = get_data_path with pytest.raises(NotImplementedError, match=r"Only ser_type='sb'"): - read_fun(path=test_data_path, ser_type=SerializationType.NC) + read_fun(path=path, ser_type=SerializationType.NC) def assert_grid_size_and_connectivities(grid): @@ -52,24 +52,24 @@ def assert_grid_size_and_connectivities(grid): @pytest.mark.datatest -def test_read_icon_grid_for_type_sb(): - grid = read_icon_grid(test_data_path, ser_type=SerializationType.SB) +def test_read_icon_grid_for_type_sb(get_data_path): + grid = read_icon_grid(get_data_path, ser_type=SerializationType.SB) assert_grid_size_and_connectivities(grid) @pytest.mark.datatest -def test_read_static_fields_for_type_sb(): +def test_read_static_fields_for_type_sb(get_data_path): metric_state, interpolation_state = read_static_fields( - test_data_path, ser_type=SerializationType.SB + get_data_path, ser_type=SerializationType.SB ) assert_metric_state_fields(metric_state) assert_interpolation_state_fields(interpolation_state) @pytest.mark.datatest -def test_read_geometry_fields_for_type_sb(): +def test_read_geometry_fields_for_type_sb(get_data_path): edge_geometry, cell_geometry, vertical_geometry = read_geometry_fields( - test_data_path, ser_type=SerializationType.SB + get_data_path, ser_type=SerializationType.SB ) assert_edge_geometry_fields(edge_geometry) assert_cell_geometry_fields(cell_geometry) From 254d8ac00c9a39a6f0add8bd7a6110d94997e820 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 30 Jun 2023 17:58:59 +0200 Subject: [PATCH 161/263] (fix) need boost if we want to build serialbox --- .github/workflows/icon4pytools-qa.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/icon4pytools-qa.yml b/.github/workflows/icon4pytools-qa.yml index 3539d285c..7f122b0f7 100644 --- a/.github/workflows/icon4pytools-qa.yml +++ b/.github/workflows/icon4pytools-qa.yml @@ -20,6 +20,10 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v2 + - name: Install boost + run: | + sudo apt-get update + sudo apt-get install libboost-all-dev - name: Set up Python uses: actions/setup-python@v3 with: From c9b6b96044c30709c8d9aef779ef45233010a964 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 30 Jun 2023 18:12:27 +0200 Subject: [PATCH 162/263] (fix) remove erroneous changes in tools. --- tools/src/icon4pytools/liskov/codegen/exceptions.py | 0 tools/tests/icon4pygen/test_field_rendering.py | 2 -- 2 files changed, 2 deletions(-) delete mode 100644 tools/src/icon4pytools/liskov/codegen/exceptions.py diff --git a/tools/src/icon4pytools/liskov/codegen/exceptions.py b/tools/src/icon4pytools/liskov/codegen/exceptions.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tools/tests/icon4pygen/test_field_rendering.py b/tools/tests/icon4pygen/test_field_rendering.py index d1db3252d..9caa19683 100644 --- a/tools/tests/icon4pygen/test_field_rendering.py +++ b/tools/tests/icon4pygen/test_field_rendering.py @@ -10,7 +10,6 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -import pytest from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum from icon4py.common.dimension import E2CDim, EdgeDim, KDim @@ -73,7 +72,6 @@ def identity_prog(field: Field[[EdgeDim, KDim], float], out: Field[[EdgeDim, KDi ) -@pytest.mark.skip("new lowering: dims in offset provider") def test_vertical_sparse_field_sid_rendering(): @field_operator def reduction(nb_field: Field[[EdgeDim, E2CDim, KDim], float]) -> Field[[EdgeDim, KDim], float]: From 6b8bcb6399d6b4e7e7eaff12cbb28df79d97d409 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 30 Jun 2023 18:13:21 +0200 Subject: [PATCH 163/263] (fix) remove erroneous changes in tools. --- tools/tests/icon4pygen/test_field_rendering.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tools/tests/icon4pygen/test_field_rendering.py b/tools/tests/icon4pygen/test_field_rendering.py index 9caa19683..2ace9a6e2 100644 --- a/tools/tests/icon4pygen/test_field_rendering.py +++ b/tools/tests/icon4pygen/test_field_rendering.py @@ -10,6 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later + from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum from icon4py.common.dimension import E2CDim, EdgeDim, KDim From e7d9e3849445a9cadbc6fc2e3f955d86126e3719 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 30 Jun 2023 18:30:56 +0200 Subject: [PATCH 164/263] (fix) remove dimensions from offset providers list --- ...fused_mo_nh_diffusion_stencil_01_02_03_rbf | 143 ------------------ atm_dyn_iconam/tests/test_diffusion.py | 1 + atm_dyn_iconam/tests/test_io_utils.py | 4 - .../test_mo_solve_nonhydro_stencil_39.py | 6 +- .../test_mo_solve_nonhydro_stencil_40.py | 6 +- .../test_mo_velocity_advection_stencil_19.py | 1 - .../test_mo_velocity_advection_stencil_20.py | 2 - 7 files changed, 3 insertions(+), 160 deletions(-) delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf deleted file mode 100644 index 88ecc9645..000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf +++ /dev/null @@ -1,143 +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 - -from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field - -from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( - _mo_intp_rbf_rbf_vec_interpol_vertex, -) -from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( - _calculate_nabla2_and_smag_coefficients_for_vn, -) -from icon4py.atm_dyn_iconam.temporary_fields_for_turbulance_diagnostics import ( - _temporary_fields_for_turbulance_diagnostics, -) -from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulance import ( - _calculate_diagnostics_for_turbulance, -) -from icon4py.common.dimension import ( - C2EDim, - CellDim, - ECVDim, - EdgeDim, - KDim, - V2EDim, - VertexDim, -) - - -@field_operator -def _fused_mo_nh_diffusion_stencil_01_02_03_rbf( - diff_multfac_smag: Field[[KDim], float], - tangent_orientation: Field[[EdgeDim], float], - inv_primal_edge_length: Field[[EdgeDim], float], - inv_vert_vert_length: Field[[EdgeDim], float], - u_vert_old: Field[[VertexDim, KDim], float], - v_vert_old: Field[[VertexDim, KDim], float], - primal_normal_vert_x: Field[[ECVDim], float], - primal_normal_vert_y: Field[[ECVDim], float], - dual_normal_vert_x: Field[[ECVDim], float], - dual_normal_vert_y: Field[[ECVDim], float], - vn: Field[[EdgeDim, KDim], float], - smag_limit: Field[[KDim], float], - smag_offset: float, - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], - wgtfac_c: Field[[CellDim, KDim], float], - ptr_coeff_1: Field[[VertexDim, V2EDim], float], - ptr_coeff_2: Field[[VertexDim, V2EDim], float], -) -> tuple[ - Field[[EdgeDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[VertexDim, KDim], float], - Field[[VertexDim, KDim], float], -]: - - kh_smag_e, kh_smag_ec, z_nabla2_e = _calculate_nabla2_and_smag_coefficients_for_vn( - diff_multfac_smag, - tangent_orientation, - inv_primal_edge_length, - inv_vert_vert_length, - u_vert_old, - v_vert_old, - primal_normal_vert_x, - primal_normal_vert_y, - dual_normal_vert_x, - dual_normal_vert_y, - vn, - smag_limit, - smag_offset, - ) - - kh_c, div = _temporary_fields_for_turbulance_diagnostics( - kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag - ) - - div_ic, hdef_ic = _calculate_diagnostics_for_turbulance(div, kh_c, wgtfac_c) - - u_vert, v_vert = _mo_intp_rbf_rbf_vec_interpol_vertex( - z_nabla2_e, ptr_coeff_1, ptr_coeff_2 - ) - - return kh_smag_e, div_ic, hdef_ic, u_vert, v_vert - - -@program -def fused_mo_nh_diffusion_stencil_01_02_03_rbf( - diff_multfac_smag: Field[[KDim], float], - tangent_orientation: Field[[EdgeDim], float], - inv_primal_edge_length: Field[[EdgeDim], float], - inv_vert_vert_length: Field[[EdgeDim], float], - u_vert_old: Field[[VertexDim, KDim], float], - v_vert_old: Field[[VertexDim, KDim], float], - primal_normal_vert_x: Field[[ECVDim], float], - primal_normal_vert_y: Field[[ECVDim], float], - dual_normal_vert_x: Field[[ECVDim], float], - dual_normal_vert_y: Field[[ECVDim], float], - vn: Field[[EdgeDim, KDim], float], - smag_limit: Field[[KDim], float], - smag_offset: float, - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], - wgtfac_c: Field[[CellDim, KDim], float], - ptr_coeff_1: Field[[VertexDim, V2EDim], float], - ptr_coeff_2: Field[[VertexDim, V2EDim], float], - kh_smag_e: Field[[EdgeDim, KDim], float], - div_ic: Field[[CellDim, KDim], float], - hdef_ic: Field[[CellDim, KDim], float], - u_vert: Field[[VertexDim, KDim], float], - v_vert: Field[[VertexDim, KDim], float], -): - _fused_mo_nh_diffusion_stencil_01_02_03_rbf( - diff_multfac_smag, - tangent_orientation, - inv_primal_edge_length, - inv_vert_vert_length, - u_vert_old, - v_vert_old, - primal_normal_vert_x, - primal_normal_vert_y, - dual_normal_vert_x, - dual_normal_vert_y, - vn, - smag_limit, - smag_offset, - e_bln_c_s, - geofac_div, - wgtfac_c, - ptr_coeff_1, - ptr_coeff_2, - out=(kh_smag_e, div_ic, hdef_ic, u_vert, v_vert), - ) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index dab84d613..0c3988098 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -556,6 +556,7 @@ def test_run_diffusion_initial_step( diffusion_savepoint_exit=diffusion_savepoint_exit, ) + @pytest.mark.datatest def test_verify_stencil15_field_manipulation(interpolation_savepoint, icon_grid): geofac_n2s = np.asarray(interpolation_savepoint.geofac_n2s()) diff --git a/atm_dyn_iconam/tests/test_io_utils.py b/atm_dyn_iconam/tests/test_io_utils.py index d2c9e1100..bc8b379f7 100644 --- a/atm_dyn_iconam/tests/test_io_utils.py +++ b/atm_dyn_iconam/tests/test_io_utils.py @@ -11,7 +11,6 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import pathlib import pytest @@ -25,9 +24,6 @@ ) - - - @pytest.mark.parametrize( "read_fun", (read_geometry_fields, read_static_fields, read_icon_grid) ) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_39.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_39.py index 73ab3e89e..1fbef1a76 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_39.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_39.py @@ -57,11 +57,7 @@ def test_mo_solve_nonhydro_stencil_39(): z_w_concorr_me, wgtfac_c, w_concorr_c, - offset_provider={ - "Koff": KDim, - "C2E": mesh.get_c2e_offset_provider(), - "C2EDim": C2EDim, - }, + offset_provider={"Koff": KDim, "C2E": mesh.get_c2e_offset_provider()}, ) assert np.allclose(w_concorr_c[:, 1:], w_concorr_c_ref[:, 1:]) diff --git a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_40.py b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_40.py index 436c6a363..ab0fce797 100644 --- a/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_40.py +++ b/atm_dyn_iconam/tests/test_mo_solve_nonhydro_stencil_40.py @@ -64,11 +64,7 @@ def test_mo_solve_nonhydro_stencil_40(): z_w_concorr_me, wgtfacq_c, w_concorr_c, - offset_provider={ - "Koff": KDim, - "C2E": mesh.get_c2e_offset_provider(), - "C2EDim": C2EDim, - }, + offset_provider={"Koff": KDim, "C2E": mesh.get_c2e_offset_provider()}, ) assert np.allclose(w_concorr_c[:, 3:], w_concorr_c_ref[:, 3:]) diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_19.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_19.py index 4061ddef5..d0c0a2839 100644 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_19.py +++ b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_19.py @@ -121,7 +121,6 @@ def test_mo_velocity_advection_stencil_19(): offset_provider={ "E2V": mesh.get_e2v_offset_provider(), "E2C": mesh.get_e2c_offset_provider(), - "E2CDim": E2CDim, "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, mesh.n_e2c), "Koff": KDim, }, diff --git a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_20.py b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_20.py index 0446e05e0..101628edb 100644 --- a/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_20.py +++ b/atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_20.py @@ -141,9 +141,7 @@ def test_mo_velocity_advection_stencil_20(): offset_provider={ "Koff": KDim, "E2C": mesh.get_e2c_offset_provider(), - "E2CDim": E2CDim, "E2C2EO": mesh.get_e2c2eO_offset_provider(), - "E2C2EODim": E2C2EODim, "E2V": mesh.get_e2v_offset_provider(), }, ) From adf0b44b3e40dfa2ccc5f1b7d6274fd68e5e4c82 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 30 Jun 2023 19:03:36 +0200 Subject: [PATCH 165/263] (fix) ignore common and py2f in tools pre-commit config for now. --- tools/.pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/.pre-commit-config.yaml b/tools/.pre-commit-config.yaml index 6c78812df..a200a7100 100644 --- a/tools/.pre-commit-config.yaml +++ b/tools/.pre-commit-config.yaml @@ -5,7 +5,7 @@ default_language_version: # Remove frozen version once we migrated away from tsa node: 17.9.1 minimum_pre_commit_version: 2.20.0 -exclude: "atm_dyn_iconam/.*" +exclude: "atm_dyn_iconam/.*|common/.*|py2f/.*" repos: - repo: meta From 5710554d0fc21e6faaea7ef189c1588b010750cc Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 7 Jul 2023 16:00:27 +0200 Subject: [PATCH 166/263] remove duplicated global stencil --- .../apply_nabla2_and_nabla4_to_vn.py | 28 ------------------- 1 file changed, 28 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py index 432d783be..98513ebed 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py @@ -32,34 +32,6 @@ def _mo_nh_diffusion_stencil_05_global_mode( return vn + z_d_vn_hdf -@program -def mo_nh_diffusion_stencil_05_global_mode( - area_edge: Field[[EdgeDim], float], - kh_smag_e: Field[[EdgeDim, KDim], float], - z_nabla2_e: Field[[EdgeDim, KDim], float], - z_nabla4_e2: Field[[EdgeDim, KDim], float], - diff_multfac_vn: Field[[KDim], float], - vn: Field[[EdgeDim, KDim], float], - horizontal_start: int32, - horizontal_end: int32, - vertical_start: int32, - vertical_end: int32, -): - _mo_nh_diffusion_stencil_05_global_mode( - area_edge, - kh_smag_e, - z_nabla2_e, - z_nabla4_e2, - diff_multfac_vn, - vn, - out=vn, - domain={ - EdgeDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, - ) - - @field_operator def _apply_nabla2_and_nabla4_to_vn( area_edge: Field[[EdgeDim], float], From 364ef8fc05bfd387f55d4103f48119005e231e5f Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 7 Jul 2023 16:09:58 +0200 Subject: [PATCH 167/263] convert staticmethod to module level function --- .../src/icon4py/diffusion/diffusion.py | 27 ++++++++++--------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 4675c0b26..e486ce036 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -332,7 +332,7 @@ def determine_smagorinski_factor(self, config: DiffusionConfig): ( smagorinski_factor, smagorinski_height, - ) = self._diffusion_type_5_smagorinski_factor(config) + ) = diffusion_type_5_smagorinski_factor(config) case 4: # according to mo_nh_diffusion.f90 this isn't used anywhere the factor is only # used for diffusion_type (3,5) but the defaults are only defined for iequations=3 @@ -349,19 +349,20 @@ def determine_smagorinski_factor(self, config: DiffusionConfig): pass return smagorinski_factor, smagorinski_height - @staticmethod - def _diffusion_type_5_smagorinski_factor(config: DiffusionConfig): - """ - Initialize Smagorinski factors used in diffusion type 5. - The calculation and magic numbers are taken from mo_diffusion_nml.f90 - """ - magic_sqrt = math.sqrt(1600.0 * (1600 + 50000.0)) - magic_fac2_value = 2e-6 * (1600.0 + 25000.0 + magic_sqrt) - magic_z2 = 1600.0 + 50000.0 + magic_sqrt - factor = (config.smagorinski_scaling_factor, magic_fac2_value, 0.0, 1.0) - heights = (32500.0, magic_z2, 50000.0, 90000.0) - return factor, heights + +def diffusion_type_5_smagorinski_factor(config: DiffusionConfig): + """ + Initialize Smagorinski factors used in diffusion type 5. + + The calculation and magic numbers are taken from mo_diffusion_nml.f90 + """ + magic_sqrt = math.sqrt(1600.0 * (1600 + 50000.0)) + magic_fac2_value = 2e-6 * (1600.0 + 25000.0 + magic_sqrt) + magic_z2 = 1600.0 + 50000.0 + magic_sqrt + factor = (config.smagorinski_scaling_factor, magic_fac2_value, 0.0, 1.0) + heights = (32500.0, magic_z2, 50000.0, 90000.0) + return factor, heights class Diffusion: From 17cd0737341d7aa2b1c391cefea34e4efa5554ad Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 7 Jul 2023 17:15:44 +0200 Subject: [PATCH 168/263] make DiffusionParams a frozen data class rename _allocate_local_fields --- .../apply_nabla2_and_nabla4_to_vn.py | 15 +----- .../src/icon4py/diffusion/diffusion.py | 49 +++++++++++-------- 2 files changed, 30 insertions(+), 34 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py index 98513ebed..8d3c665a4 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py @@ -12,24 +12,11 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, int32, maximum +from gt4py.next.ffront.fbuiltins import Field, maximum from icon4py.common.dimension import EdgeDim, KDim -@field_operator -def _mo_nh_diffusion_stencil_05_global_mode( - area_edge: Field[[EdgeDim], float], - kh_smag_e: Field[[EdgeDim, KDim], float], - z_nabla2_e: Field[[EdgeDim, KDim], float], - z_nabla_e2: Field[[EdgeDim, KDim], float], - diff_multfac_vn: Field[[KDim], float], - vn: Field[[EdgeDim, KDim], float], -) -> Field[[EdgeDim, KDim], float]: - z_d_vn_hdf = area_edge * ( - kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla_e2 * area_edge - ) - return vn + z_d_vn_hdf @field_operator diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index e486ce036..be376f4ee 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -14,6 +14,7 @@ import math import sys from collections import namedtuple +from dataclasses import dataclass, field, InitVar from typing import Final, Optional, Tuple import numpy as np @@ -290,36 +291,44 @@ def _validate(self): def substep_as_float(self): return float(self.ndyn_substeps) - +@dataclass(frozen=True) class DiffusionParams: """Calculates derived quantities depending on the diffusion config.""" - - def __init__(self, config: DiffusionConfig): - - self.K2: Final[float] = ( + config: InitVar[DiffusionConfig] + K2: Final[float] = field(init=False) + K4: Final[float] = field(init=False) + K6: Final[float] = field(init=False) + K4W: Final[float]= field(init=False) + smagorinski_factor: Final[float] = field(init=False) + smagorinski_height: Final[float] = field(init=False) + scaled_nudge_max_coeff: Final[float] = field(init=False) + + + def __post_init__(self, config): + object.__setattr__(self, 'K2', ( 1.0 / (config.hdiff_efdt_ratio * 8.0) if config.hdiff_efdt_ratio > 0.0 else 0.0 - ) - self.K4: Final[float] = self.K2 / 8.0 - self.K6: Final[float] = self.K2 / 64.0 - - self.K4W: Final[float] = ( + )) + object.__setattr__(self, 'K4', self.K2 / 8.0) + object.__setattr__(self, 'K6', self.K2 / 64.0) + object.__setattr__(self, 'K4W', ( 1.0 / (config.hdiff_w_efdt_ratio * 36.0) if config.hdiff_w_efdt_ratio > 0 else 0.0 - ) + ) ) ( - self.smagorinski_factor, - self.smagorinski_height, - ) = self.determine_smagorinski_factor(config) + smagorinski_factor, + smagorinski_height, + ) = self._determine_smagorinski_factor(config) + object.__setattr__(self, 'smagorinski_factor', smagorinski_factor) + object.__setattr__(self, 'smagorinski_height', smagorinski_height) # see mo_interpol_nml.f90: - self.scaled_nudge_max_coeff = ( - config.nudge_max_coeff * DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO - ) + object.__setattr__(self, 'scaled_nudge_max_coeff', config.nudge_max_coeff * DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO) + - def determine_smagorinski_factor(self, config: DiffusionConfig): + def _determine_smagorinski_factor(self, config: DiffusionConfig): """Enhanced Smagorinsky diffusion factor. Smagorinsky diffusion factor is defined as a profile in height @@ -425,7 +434,7 @@ def init( self.edge_params = edge_params self.cell_params = cell_params - self._allocate_local_fields() + self._allocate_temporary_fields() self.nudgezone_diff: float = 0.04 / ( params.scaled_nudge_max_coeff + sys.float_info.epsilon @@ -469,7 +478,7 @@ def init( def initialized(self): return self._initialized - def _allocate_local_fields(self): + def _allocate_temporary_fields(self): def _allocate(*dims: Dimension): return zero_field(self.grid, *dims) From 9ac647f7b56e5d30b32e74f4983775dd79e69963 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 11 Jul 2023 14:31:59 +0200 Subject: [PATCH 169/263] TODO format --- atm_dyn_iconam/src/icon4py/diffusion/diffusion.py | 6 +++--- atm_dyn_iconam/src/icon4py/diffusion/horizontal.py | 2 +- atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py | 2 +- atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py | 2 +- atm_dyn_iconam/src/icon4py/diffusion/utils.py | 5 +++-- 5 files changed, 9 insertions(+), 8 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index be376f4ee..37bf1365a 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -95,9 +95,9 @@ class DiffusionConfig: Encapsulates namelist parameters and derived parameters. Values should be read from configuration. Default values are taken from the defaults in the corresponding ICON Fortran namelist files. - TODO: @magdalena to be read from config - TODO: @magdalena handle dependencies on other namelists (see below...) """ + # TODO(Magdalena): to be read from config + # TODO(Magdalena): handle dependencies on other namelists (see below...) def __init__( self, @@ -497,7 +497,7 @@ def _index_field(dim: Dimension, size=None): self.z_nabla2_e = _allocate(EdgeDim, KDim) self.z_temp = _allocate(CellDim, KDim) self.diff_multfac_smag = _allocate(KDim) - # TODO @magdalena this is KHalfDim + # TODO(Magdalena): this is KHalfDim self.vertical_index = _index_field(KDim, self.grid.n_lev() + 1) self.horizontal_cell_index = _index_field(CellDim) self.horizontal_edge_index = _index_field(EdgeDim) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py index 7c41c7331..def62ab14 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py @@ -136,7 +136,7 @@ class HorizontalMeshSize: num_cells: int -# TODO [@Magdalena] allow initialization with only partial values +# TODO(Magdalena): allow initialization with only partial values # (becomes tedious for testing otherwise): hence this should # that should not be a data class class EdgeParams: diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index 53b8eb823..29acd613c 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -30,7 +30,7 @@ from icon4py.diffusion.horizontal import HorizontalMeshSize -# TODO @magdalena keep naming grid vs mesh consistent +# TODO(Magdalena): keep naming grid vs mesh consistent class VerticalMeshConfig: def __init__(self, num_lev: int): self._num_lev = num_lev diff --git a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py index 98586b90c..4a22668d7 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py @@ -33,8 +33,8 @@ class InterpolationState: """ represents the ICON interpolation state. - TODO [magdalena]: keep? does this state make sense at all? """ + # TODO(Magdalena): keep? does this state make sense at all? e_bln_c_s: Field[ [CEDim], float diff --git a/atm_dyn_iconam/src/icon4py/diffusion/utils.py b/atm_dyn_iconam/src/icon4py/diffusion/utils.py index 073a7915c..981aa76c7 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/utils.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/utils.py @@ -21,7 +21,7 @@ from icon4py.common.dimension import CellDim, EdgeDim, KDim, Koff, VertexDim -# TODO [@Magdalena] fix duplication: duplicated from test testutils/utils.py +# TODO(Magdalena): fix duplication: duplicated from test testutils/utils.py def zero_field(mesh, *dims: Dimension, dtype=float): shapex = tuple(map(lambda x: mesh.size[x], dims)) return np_as_located_field(*dims)(np.zeros(shapex, dtype=dtype)) @@ -248,7 +248,6 @@ def init_nabla2_factor_in_upper_damping_zone( Calculate diff_multfac_n2w. numpy version, since gt4py does not allow non-constant indexing into fields - TODO: @magdalena fix this once IndexedFields are implemented Args k_size: number of vertical levels @@ -256,6 +255,8 @@ def init_nabla2_factor_in_upper_damping_zone( nshift: physcial_heights: vector of physical heights [m] of the height levels """ + # TODO(Magdalena): fix with as_offset in gt4py + buffer = np.zeros(k_size) buffer[1 : nrdmax + 1] = ( 1.0 From 61aeefa880fe9360efd98caa85b8a2096aac1485 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 11 Jul 2023 15:05:51 +0200 Subject: [PATCH 170/263] use dict instead of switch case to lookup marker indices --- .../apply_nabla2_and_nabla4_to_vn.py | 2 - .../src/icon4py/diffusion/diffusion.py | 52 ++++++----- .../src/icon4py/diffusion/horizontal.py | 87 +++++++++---------- .../icon4py/diffusion/interpolation_state.py | 4 +- atm_dyn_iconam/tests/test_icon_grid.py | 7 +- 5 files changed, 76 insertions(+), 76 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py index 8d3c665a4..3993a1f25 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_nabla2_and_nabla4_to_vn.py @@ -17,8 +17,6 @@ from icon4py.common.dimension import EdgeDim, KDim - - @field_operator def _apply_nabla2_and_nabla4_to_vn( area_edge: Field[[EdgeDim], float], diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 37bf1365a..b13045d03 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -14,7 +14,7 @@ import math import sys from collections import namedtuple -from dataclasses import dataclass, field, InitVar +from dataclasses import InitVar, dataclass, field from typing import Final, Optional, Tuple import numpy as np @@ -96,6 +96,7 @@ class DiffusionConfig: Values should be read from configuration. Default values are taken from the defaults in the corresponding ICON Fortran namelist files. """ + # TODO(Magdalena): to be read from config # TODO(Magdalena): handle dependencies on other namelists (see below...) @@ -291,42 +292,54 @@ def _validate(self): def substep_as_float(self): return float(self.ndyn_substeps) + @dataclass(frozen=True) class DiffusionParams: """Calculates derived quantities depending on the diffusion config.""" + config: InitVar[DiffusionConfig] K2: Final[float] = field(init=False) K4: Final[float] = field(init=False) K6: Final[float] = field(init=False) - K4W: Final[float]= field(init=False) + K4W: Final[float] = field(init=False) smagorinski_factor: Final[float] = field(init=False) smagorinski_height: Final[float] = field(init=False) scaled_nudge_max_coeff: Final[float] = field(init=False) - def __post_init__(self, config): - object.__setattr__(self, 'K2', ( - 1.0 / (config.hdiff_efdt_ratio * 8.0) - if config.hdiff_efdt_ratio > 0.0 - else 0.0 - )) - object.__setattr__(self, 'K4', self.K2 / 8.0) - object.__setattr__(self, 'K6', self.K2 / 64.0) - object.__setattr__(self, 'K4W', ( - 1.0 / (config.hdiff_w_efdt_ratio * 36.0) - if config.hdiff_w_efdt_ratio > 0 - else 0.0 - ) ) + object.__setattr__( + self, + "K2", + ( + 1.0 / (config.hdiff_efdt_ratio * 8.0) + if config.hdiff_efdt_ratio > 0.0 + else 0.0 + ), + ) + object.__setattr__(self, "K4", self.K2 / 8.0) + object.__setattr__(self, "K6", self.K2 / 64.0) + object.__setattr__( + self, + "K4W", + ( + 1.0 / (config.hdiff_w_efdt_ratio * 36.0) + if config.hdiff_w_efdt_ratio > 0 + else 0.0 + ), + ) ( smagorinski_factor, smagorinski_height, ) = self._determine_smagorinski_factor(config) - object.__setattr__(self, 'smagorinski_factor', smagorinski_factor) - object.__setattr__(self, 'smagorinski_height', smagorinski_height) + object.__setattr__(self, "smagorinski_factor", smagorinski_factor) + object.__setattr__(self, "smagorinski_height", smagorinski_height) # see mo_interpol_nml.f90: - object.__setattr__(self, 'scaled_nudge_max_coeff', config.nudge_max_coeff * DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO) - + object.__setattr__( + self, + "scaled_nudge_max_coeff", + config.nudge_max_coeff * DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO, + ) def _determine_smagorinski_factor(self, config: DiffusionConfig): """Enhanced Smagorinsky diffusion factor. @@ -359,7 +372,6 @@ def _determine_smagorinski_factor(self, config: DiffusionConfig): return smagorinski_factor, smagorinski_height - def diffusion_type_5_smagorinski_factor(config: DiffusionConfig): """ Initialize Smagorinski factors used in diffusion type 5. diff --git a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py index def62ab14..1245fb8bc 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py @@ -39,94 +39,89 @@ class HorizontalMarkerIndex: _MIN_RL_EDGE: Final[int] = _MIN_RL_EDGE_INT - (2 * NUM_GHOST_ROWS + 1) _MAX_RL_EDGE: Final[int] = 2 * _MAX_RL_CELL - _LOCAL_BOUNDARY_EDGES: Final[int] = 1 + _ICON_INDEX_OFFSET_EDGES + _LATERAL_BOUNDARY_EDGES: Final[int] = 1 + _ICON_INDEX_OFFSET_EDGES _INTERIOR_EDGES: Final[int] = _ICON_INDEX_OFFSET_EDGES _NUDGING_EDGES: Final[int] = _GRF_BOUNDARY_WIDTH_EDGES + _ICON_INDEX_OFFSET_EDGES _HALO_EDGES: Final[int] = _MIN_RL_EDGE_INT - 1 + _ICON_INDEX_OFFSET_EDGES _LOCAL_EDGES: Final[int] = _MIN_RL_EDGE_INT + _ICON_INDEX_OFFSET_EDGES _END_EDGES: Final[int] = 0 - _LOCAL_BOUNDARY_CELLS: Final[int] = 1 + _ICON_INDEX_OFFSET_CELLS + _LATERAL_BOUNDARY_CELLS: Final[int] = 1 + _ICON_INDEX_OFFSET_CELLS _INTERIOR_CELLS: Final[int] = _ICON_INDEX_OFFSET_CELLS _NUDGING_CELLS: Final[int] = _GRF_BOUNDARY_WIDTH_CELL + 1 + _ICON_INDEX_OFFSET_CELLS _HALO_CELLS: Final[int] = _MIN_RL_CELL_INT - 1 + _ICON_INDEX_OFFSET_CELLS _LOCAL_CELLS: Final[int] = _MIN_RL_CELL_INT + _ICON_INDEX_OFFSET_CELLS _END_CELLS: Final[int] = 0 - _LOCAL_BOUNDARY_VERTICES = 1 + _ICON_INDEX_OFFSET_VERTEX + _LATERAL_BOUNDARY_VERTICES = 1 + _ICON_INDEX_OFFSET_VERTEX _INTERIOR_VERTICES: Final[int] = _ICON_INDEX_OFFSET_VERTEX _NUDGING_VERTICES: Final[int] = 0 _HALO_VERTICES: Final[int] = _MIN_RL_VERTEX_INT - 1 + _ICON_INDEX_OFFSET_VERTEX _LOCAL_VERTICES: Final[int] = _MIN_RL_VERTEX_INT + _ICON_INDEX_OFFSET_VERTEX _END_VERTICES: Final[int] = 0 + _lateral_boundary = { + dimension.CellDim: _LATERAL_BOUNDARY_CELLS, + dimension.EdgeDim: _LATERAL_BOUNDARY_EDGES, + dimension.VertexDim: _LATERAL_BOUNDARY_VERTICES, + } + _local = { + dimension.CellDim: _LOCAL_CELLS, + dimension.EdgeDim: _LOCAL_EDGES, + dimension.VertexDim: _LOCAL_VERTICES, + } + _halo = { + dimension.CellDim: _HALO_CELLS, + dimension.EdgeDim: _HALO_EDGES, + dimension.VertexDim: _HALO_VERTICES, + } + _interior = { + dimension.CellDim: _INTERIOR_CELLS, + dimension.EdgeDim: _INTERIOR_EDGES, + dimension.VertexDim: _INTERIOR_VERTICES, + } + _nudging = { + dimension.CellDim: _NUDGING_CELLS, + dimension.EdgeDim: _NUDGING_EDGES, + dimension.VertexDim: _NUDGING_VERTICES, + } + _end = { + dimension.CellDim: _END_CELLS, + dimension.EdgeDim: _END_EDGES, + dimension.VertexDim: _END_VERTICES, + } + @classmethod def lateral_boundary(cls, dim: Dimension) -> int: """Indicate lateral boundary. These points correspond to the sorted points in ICON, the marker can be incremented in order - to accesss higher boundary lines + to access higher order boundary lines """ - match (dim): - case (dimension.CellDim): - return cls._LOCAL_BOUNDARY_CELLS - case (dimension.EdgeDim): - return cls._LOCAL_BOUNDARY_EDGES - case (dimension.VertexDim): - return cls._LOCAL_BOUNDARY_VERTICES + return cls._lateral_boundary[dim] @classmethod def local(cls, dim: Dimension) -> int: """Indicate points that are owned by the processing unit, i.e. no halo points.""" - match (dim): - case (dimension.CellDim): - return cls._LOCAL_CELLS - case (dimension.EdgeDim): - return cls._LOCAL_EDGES - case (dimension.VertexDim): - return cls._LOCAL_VERTICES + return cls._local[dim] @classmethod def halo(cls, dim: Dimension) -> int: - match (dim): - case (dimension.CellDim): - return cls._HALO_CELLS - case (dimension.EdgeDim): - return cls._HALO_EDGES - case (dimension.VertexDim): - return cls._HALO_VERTICES + return cls._halo[dim] @classmethod def nudging(cls, dim: Dimension) -> int: """Indicate the nudging zone.""" - match (dim): - case (dimension.CellDim): - return cls._NUDGING_CELLS - case (dimension.EdgeDim): - return cls._NUDGING_EDGES - case (dimension.VertexDim): - return cls._NUDGING_VERTICES + return cls._nudging[dim] @classmethod def interior(cls, dim: Dimension) -> int: """Indicate interior i.e. unordered prognostic cells in ICON.""" - match (dim): - case (dimension.CellDim): - return cls._INTERIOR_CELLS - case (dimension.EdgeDim): - return cls._INTERIOR_EDGES - case (dimension.VertexDim): - return cls._INTERIOR_VERTICES + return cls._interior[dim] @classmethod def end(cls, dim: Dimension) -> int: - match (dim): - case (dimension.CellDim): - return cls._END_CELLS - case (dimension.EdgeDim): - return cls._END_EDGES - case (dimension.VertexDim): - return cls._END_VERTICES + return cls._end[dim] @dataclass(frozen=True) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py index 4a22668d7..625f177a0 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py @@ -30,10 +30,8 @@ @dataclass class InterpolationState: - """ - represents the ICON interpolation state. + """Represents the ICON interpolation state.""" - """ # TODO(Magdalena): keep? does this state make sense at all? e_bln_c_s: Field[ diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/atm_dyn_iconam/tests/test_icon_grid.py index 3c04aeb69..e62ed01a0 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/atm_dyn_iconam/tests/test_icon_grid.py @@ -260,12 +260,9 @@ def test_horizontal_vertex_markers( @pytest.mark.datatest def test_cross_check_marker_equivalences(icon_grid): - """ - Check actual equivalences of calculated markers. + """Check actual equivalences of calculated markers.""" + # TODO(Magdalena): This should go away once we refactor these markers in a good way, such that no calculation need to be done with them anymore. - TODO [magdalena] This should go away once we refactor these markers in a good way, such that no - calculation need to be done with them anymore. - """ assert icon_grid.get_indices_from_to( CellDim, HorizontalMarkerIndex.local(CellDim) - 1, From c03be01a43efd8c987ff3fa9b84c552b98be602a Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 11 Jul 2023 15:20:29 +0200 Subject: [PATCH 171/263] convert mesh config classes to dataclasses --- .../src/icon4py/diffusion/diffusion.py | 2 +- .../src/icon4py/diffusion/icon_grid.py | 90 ++++++++----------- 2 files changed, 39 insertions(+), 53 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index b13045d03..75538d6da 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -85,7 +85,7 @@ cached_backend = run_gtfn_cached compiled_backend = run_gtfn -backend = compiled_backend # +backend = compiled_backend class DiffusionConfig: diff --git a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py index 29acd613c..056673f06 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/icon_grid.py @@ -10,8 +10,8 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - -from typing import Dict +from dataclasses import dataclass, field +from typing import Dict, Final import numpy as np from gt4py.next.common import Dimension, DimensionKind, Field @@ -31,50 +31,35 @@ # TODO(Magdalena): keep naming grid vs mesh consistent +@dataclass(frozen=True) class VerticalMeshConfig: - def __init__(self, num_lev: int): - self._num_lev = num_lev - - @property - def num_lev(self) -> int: - return self._num_lev + num_lev: int +@dataclass( + frozen=True, +) class MeshConfig: - def __init__( - self, - horizontal_config: HorizontalMeshSize, - vertical_config: VerticalMeshConfig, - limited_area=True, - ): - self._vertical = vertical_config - self._n_shift_total = 0 - self._limited_area = limited_area - self._horizontal = horizontal_config - - @property - def limited_area(self): - return self._limited_area + horizontal_config: HorizontalMeshSize + vertical_config: VerticalMeshConfig + limited_area: bool = True + n_shift_total: int = 0 @property def num_k_levels(self): - return self._vertical.num_lev - - @property - def n_shift_total(self): - return self._n_shift_total + return self.vertical_config.num_lev @property def num_vertices(self): - return self._horizontal.num_vertices + return self.horizontal_config.num_vertices @property def num_edges(self): - return self._horizontal.num_edges + return self.horizontal_config.num_edges @property def num_cells(self): - return self._horizontal.num_cells + return self.horizontal_config.num_cells def builder(func): @@ -207,31 +192,32 @@ def get_c2ce_connectivity(self): ) +@dataclass(frozen=True) class VerticalModelParams: - def __init__(self, vct_a: Field[[KDim], float], rayleigh_damping_height: float): - """ - Contains vertical physical parameters defined on the grid. - - Args: - vct_a: field containing the physical heights of the k level - rayleigh_damping_height: height of rayleigh damping in [m] mo_nonhydro_nml - """ - self._rayleigh_damping_height = rayleigh_damping_height - self._vct_a = vct_a - self._index_of_damping_height = int32( - np.argmax( - np.where(np.asarray(self._vct_a) >= self._rayleigh_damping_height) - ) + """ + Contains vertical physical parameters defined on the grid. + + vct_a: field containing the physical heights of the k level + rayleigh_damping_height: height of rayleigh damping in [m] mo_nonhydro_nml + """ + + vct_a: Field[[KDim], float] + rayleigh_damping_height: Final[float] + index_of_damping_layer: Final[int32] = field(init=False) + + def __post_init__(self): + object.__setattr__( + self, + "index_of_damping_layer", + self._determine_damping_height_index( + np.asarray(self.vct_a), self.rayleigh_damping_height + ), ) - @property - def index_of_damping_layer(self): - return self._index_of_damping_height + @classmethod + def _determine_damping_height_index(cls, vct_a: np.ndarray, damping_height: float): + return int32(np.argmax(np.where(vct_a >= damping_height))) @property def physical_heights(self) -> Field[[KDim], float]: - return self._vct_a - - @property - def rayleigh_damping_height(self): - return self._rayleigh_damping_height + return self.vct_a From feb22bbf7a2bb9042d02dbe72f95d7b2578fbf52 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 11 Jul 2023 16:30:52 +0200 Subject: [PATCH 172/263] use functools.cached_property for geofac fields --- .../src/icon4py/diffusion/interpolation_state.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py index 625f177a0..e57a27ce6 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - +import functools from dataclasses import dataclass import numpy as np @@ -28,7 +28,7 @@ ) -@dataclass +@dataclass(frozen=True) class InterpolationState: """Represents the ICON interpolation state.""" @@ -57,11 +57,11 @@ class InterpolationState: ] # factors for green gauss gradient (nproma,4,nblks_c,2) nudgecoeff_e: Field[[EdgeDim], float] # Nudgeing coeffients for edges - @property + @functools.cached_property def geofac_n2s_c(self) -> Field[[CellDim], float]: return np_as_located_field(CellDim)(np.asarray(self.geofac_n2s)[:, 0]) - @property + @functools.cached_property def geofac_n2s_nbh(self) -> Field[[CECDim], float]: geofac_nbh_ar = np.asarray(self.geofac_n2s)[:, 1:] old_shape = geofac_nbh_ar.shape From 5f9fbc16a8199fc5cf1dd74299a67c6ea0b2a644 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 11 Jul 2023 16:45:39 +0200 Subject: [PATCH 173/263] added TODO for import of testutils --- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index 07ac5f5e1..cc1d81193 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -29,6 +29,9 @@ log = logging.getLogger(__name__) +# TODO(Magdalena): for preliminary version of the driver we need serialbox data which is in +# testutils, since that is no proper package we need to import it by hand here. +# Hence: Turn testutils into a package again? def import_testutils(): testutils = ( Path(__file__).parent.__str__() + "/../../../tests/test_utils/__init__.py" From 6776d12204488a19d6fe5bae3db4700e8b4c5834 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 11 Jul 2023 16:56:37 +0200 Subject: [PATCH 174/263] fixes in io_utils.py: type annotations, docstrings --- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 43 +++++++++++-------- 1 file changed, 24 insertions(+), 19 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index cc1d81193..f80a27fa1 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -25,6 +25,8 @@ from icon4py.diffusion.prognostic_state import PrognosticState +SB_ONLY_MSG = "Only ser_type='sb' is implemented so far." + SIMULATION_START_DATE = "2021-06-20T12:00:10.000" log = logging.getLogger(__name__) @@ -53,15 +55,17 @@ class SerializationType(str, Enum): NC = "netcdf" -def read_icon_grid(path: Path, ser_type=SerializationType.SB) -> IconGrid: +def read_icon_grid( + path: Path, ser_type: SerializationType = SerializationType.SB +) -> IconGrid: """ - Return IconGrid parsed from a given input type. - - Factory method that returns an icon grid dependeing on the ser_type. + Read icon grid. Args: - path: str - path where to find the input data - ser_type: str - type of input data. Currently only 'sb (serialbox)' is supported. It reads from ppser serialized test data + path: path where to find the input data + ser_type: type of input data. Currently only 'sb (serialbox)' is supported. It reads + from ppser serialized test data + Returns: IconGrid parsed from a given input type. """ if ser_type == SerializationType.SB: return ( @@ -70,7 +74,7 @@ def read_icon_grid(path: Path, ser_type=SerializationType.SB) -> IconGrid: .construct_icon_grid() ) else: - raise NotImplementedError("Only ser_type='sb' is implemented so far.") + raise NotImplementedError(SB_ONLY_MSG) def read_initial_state( @@ -80,11 +84,11 @@ def read_initial_state( Read prognostic and diagnostic state from serialized data. Args: - gridfile_path: path the the serialized input data + gridfile_path: path the serialized input data Returns: a tuple containing the data_provider, the initial diagnostic and prognostic state. - The data_provider is returned such that further timesteps of diagnostics and prognostics can be - read from within the dummy timeloop + The data_provider is returned such that further timesteps of diagnostics and prognostics + can be read from within the dummy timeloop """ data_provider = sb.IconSerialDataProvider( @@ -99,14 +103,14 @@ def read_initial_state( def read_geometry_fields( - path: Path, ser_type=SerializationType.SB + path: Path, ser_type: SerializationType = SerializationType.SB ) -> tuple[EdgeParams, CellParams, VerticalModelParams]: """ Read fields containing grid properties. Args: path: path to the serialized input data - ser_type: (optional) defualt so SB=serialbox, type of input data to be read + ser_type: (optional) defaults to SB=serialbox, type of input data to be read Returns: a tuple containing fields describing edges, cells, vertical properties of the model the data is originally obtained from the grid file (horizontal fields) or some special input files. @@ -122,21 +126,22 @@ def read_geometry_fields( ) return edge_geometry, cell_geometry, vertical_geometry else: - raise NotImplementedError("Only ser_type='sb' is implemented so far.") + raise NotImplementedError(SB_ONLY_MSG) def read_static_fields( - path: Path, ser_type=SerializationType.SB + path: Path, ser_type: SerializationType = SerializationType.SB ) -> tuple[MetricState, InterpolationState]: """ Read fields for metric and interpolation state. Args: path: path to the serialized input data - ser_type: (optional) defualt so SB=serialbox, type of input data to be read + ser_type: (optional) defaults to SB=serialbox, type of input data to be read - Returns: a tuple containing the metric_state and interpolation state, - the fields are precalculated in the icon setup. + Returns: + a tuple containing the metric_state and interpolation state, + the fields are precalculated in the icon setup. """ if ser_type == SerializationType.SB: @@ -149,10 +154,10 @@ def read_static_fields( metric_state = dataprovider.from_metrics_savepoint().construct_metric_state() return metric_state, interpolation_state else: - raise NotImplementedError("Only ser_type='sb' is implemented so far.") + raise NotImplementedError(SB_ONLY_MSG) -def configure_logging(run_path: str, start_time): +def configure_logging(run_path: str, start_time) -> None: """ Configure logging. From 3542e36b1a9eb0eb2fbb651acf7e38569cf8fb5c Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 11 Jul 2023 17:24:42 +0200 Subject: [PATCH 175/263] - remove Annotated type from constants.py - use raw docstring in horizontal.py --- .flake8 | 1 - .../src/icon4py/diffusion/horizontal.py | 4 +-- common/src/icon4py/common/constants.py | 31 ++++++++++--------- 3 files changed, 18 insertions(+), 18 deletions(-) diff --git a/.flake8 b/.flake8 index 8d46210ae..b37b750f9 100644 --- a/.flake8 +++ b/.flake8 @@ -43,4 +43,3 @@ rst-roles = per-file-ignores = tools/src/icon4pytools/icon4pygen/icochainsize.py:E800 - atm_dyn_iconam/src/icon4py/diffusion/horizontal.py:W605 diff --git a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py index 1245fb8bc..301792a0c 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/horizontal.py @@ -151,7 +151,7 @@ def __init__( ): self.tangent_orientation: Field[[EdgeDim], float] = tangent_orientation - """ + r""" Orientation of vector product of the edge and the adjacent cell centers v3 / \ @@ -205,7 +205,7 @@ def __init__( self.inverse_vertex_vertex_lengths: Field[ [EdgeDim], float ] = inverse_vertex_vertex_lengths - """ + r""" Inverse distance between outer vertices of adjacent cells. v1-------- diff --git a/common/src/icon4py/common/constants.py b/common/src/icon4py/common/constants.py index 8aa751bde..1c64d9f04 100644 --- a/common/src/icon4py/common/constants.py +++ b/common/src/icon4py/common/constants.py @@ -11,20 +11,21 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from typing import Annotated +from typing import Final -GAS_CONSTANT_DRY_AIR: Annotated[ - float, - "gas constant for dry air [J/K/kg], called 'rd' in ICON (mo_physical_constants.f90), https://glossary.ametsoc.org/wiki/Gas_constant", -] = 287.04 -CPD: Annotated[float, "specific heat at constant pressure [J/K/kg]"] = 1004.64 -GAS_CONSTANT_WATER_VAPOR: Annotated[ - float, "gas constant for water vapor [J/K/kg], rv in Icon" -] = 461.51 -GRAVITATIONAL_ACCELERATION: Annotated[ - float, "av. gravitational acceleration [m/s^2]" -] = 9.8066 -DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO: Annotated[ - float, "default physics to dynamics time step ratio" -] = 5.0 +#: Gas constant for dry air [J/K/kg], called 'rd' in ICON (mo_physical_constants.f90), +#: see https://glossary.ametsoc.org/wiki/Gas_constant. +GAS_CONSTANT_DRY_AIR: Final[float] = 287.04 + +#: Specific heat at constant pressure [J/K/kg] +CPD: Final[float] = 1004.64 + +#: Gas constant for water vapor [J/K/kg], rv in ICON. +GAS_CONSTANT_WATER_VAPOR: Final[float] = 461.51 + +#: Av. gravitational acceleration [m/s^2] +GRAVITATIONAL_ACCELERATION: Final[float] = 9.8066 + +#: Default physics to dynamics time step ratio +DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO: Final[float] = 5.0 From eb2bd5b80fb9765875f667af1f3f6cd64dabdebe Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 12 Jul 2023 17:55:26 +0200 Subject: [PATCH 176/263] - introduce diffusiontype enum - change format of attribute docstrings --- .../src/icon4py/diffusion/diffusion.py | 156 ++++++------------ 1 file changed, 54 insertions(+), 102 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 75538d6da..cd45edbcb 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -15,6 +15,7 @@ import sys from collections import namedtuple from dataclasses import InitVar, dataclass, field +from enum import Enum from typing import Final, Optional, Tuple import numpy as np @@ -87,6 +88,18 @@ compiled_backend = run_gtfn backend = compiled_backend +class DiffusionType(int, Enum): + """ + Order of nabla operator for diffusion. + Note: Called `hdiff_order` in `mo_diffusion_nml.f90`. + Note: We currently only support type 5. + """ + NO_DIFFUSION = -1 #: no diffusion + LINEAR_2ND_ORDER = 2 #: 2nd order linear diffusion on all vertical levels + SMAGORINSKY_NO_BACKGROUND = 3 #: Smagorinsky diffusion without background diffusion + LINEAR_4TH_ORDER = 4 #: 4th order linear diffusion on all vertical levels + SMAGORINSKY_4TH_ORDER = 5 #: Smagorinsky diffusion with fourth-order background diffusion + class DiffusionConfig: """ @@ -102,7 +115,7 @@ class DiffusionConfig: def __init__( self, - diffusion_type: int = 5, + diffusion_type: DiffusionType = DiffusionType.SMAGORINSKY_4TH_ORDER, hdiff_w=True, hdiff_vn=True, hdiff_temp=True, @@ -121,150 +134,89 @@ def __init__( nudging_decay_rate: float = 2.0, ): """Set the diffusion configuration parameters with the ICON default values.""" - # parameters from namelist diffusion_nml + self.diffusion_type: int = diffusion_type - """ - Order of nabla operator for diffusion. - - Called `hdiff_order` in mo_diffusion_nml.f90. - Possible values are: - - -1: no diffusion - - 2: 2nd order linear diffusion on all vertical levels - - 3: Smagorinsky diffusion without background diffusion - - 4: 4th order linear diffusion on all vertical levels - - 5: Smagorinsky diffusion with fourth-order background diffusion - - We only support type 5. - TODO: [ml] use enum - """ + #: If True, apply diffusion on the vertical wind field + #: Called `lhdiff_w` in mo_diffusion_nml.f90 self.apply_to_vertical_wind: bool = hdiff_w - """ - If True, apply diffusion on the vertical wind field - - Called `lhdiff_w` in in mo_diffusion_nml.f90 - """ + #: True apply diffusion on the horizontal wind field, is ONLY used in mo_nh_stepping.f90 + #: Called `lhdiff_vn` in mo_diffusion_nml.f90 self.apply_to_horizontal_wind = hdiff_vn - """ - If True apply diffusion on the horizontal wind field, is ONLY used in mo_nh_stepping.f90 - - Called `lhdiff_vn` in in mo_diffusion_nml.f90 - """ + #: If True, apply horizontal diffusion to temperature field + #: Called `lhdiff_temp` in mo_diffusion_nml.f90 self.apply_to_temperature: bool = hdiff_temp - """ - If True, apply horizontal diffusion to temperature field - - Called `lhdiff_temp` in in mo_diffusion_nml.f90 - """ + #: If True, compute 3D Smagorinsky diffusion coefficient + #: Called `lsmag_3d` in mo_diffusion_nml.f90 self.compute_3d_smag_coeff: bool = smag_3d - """ - If True, compute 3D Smagorinsky diffusion coefficient. - - Called `lsmag_3d` in in mo_diffusion_nml.f90 - """ + #: Options for discretizing the Smagorinsky momentum diffusion + #: Called `itype_vn_diffu` in mo_diffusion_nml.f90 self.type_vn_diffu: int = type_vn_diffu - """ - Options for discretizing the Smagorinsky momentum diffusion. - - Called `itype_vn_diffu` in in mo_diffusion_nml.f90 - """ + #: Options for discretizing the Smagorinsky temperature diffusion + #: Called `itype_t_diffu` inmo_diffusion_nml.f90 self.type_t_diffu = type_t_diffu - """ - Options for discretizing the Smagorinsky temperature diffusion. - - Called `itype_t_diffu` in in mo_diffusion_nml.f90 - """ + #: Ratio of e-folding time to (2*)time step + #: Called `hdiff_efdt_ratio` inmo_diffusion_nml.f90 self.hdiff_efdt_ratio: float = hdiff_efdt_ratio - """ - Ratio of e-folding time to (2*)time step. - - Called `hdiff_efdt_ratio` in in mo_diffusion_nml.f90. - """ + #: Ratio of e-folding time to time step for w diffusion (NH only) + #: Called `hdiff_w_efdt_ratio` inmo_diffusion_nml.f90. self.hdiff_w_efdt_ratio: float = hdiff_w_efdt_ratio - """ - Ratio of e-folding time to time step for w diffusion (NH only). - - Called `hdiff_w_efdt_ratio` in in mo_diffusion_nml.f90. - """ + #: Scaling factor for Smagorinsky diffusion at height hdiff_smag_z and below + #: Called `hdiff_smag_fac` inmo_diffusion_nml.f90 self.smagorinski_scaling_factor: float = smagorinski_scaling_factor - """ - Scaling factor for Smagorinsky diffusion at height hdiff_smag_z and below. - - Called `hdiff_smag_fac` in in mo_diffusion_nml.f90. - """ + #: If True, apply truly horizontal temperature diffusion over steep slopes + #: Called 'l_zdiffu_t' in mo_nonhydrostatic_nml.f90 self.apply_zdiffusion_t: bool = zdiffu_t - """ - If True, apply truly horizontal temperature diffusion over steep slopes. - - From parent namelist mo_nonhydrostatic_nml.f90, but is only used in diffusion, - and in mo_vertical_grid.f90>prepare_zdiffu. - Called 'l_zdiffu_t' in mo_nonhydrostatic_nml.f90. - """ - - # from other namelists + # from other namelists: # from parent namelist mo_nonhydrostatic_nml - self.ndyn_substeps: int = n_substeps - """ - Number of dynamics substeps per fast-physics step. - Called 'ndyn_substeps' in mo_nonhydrostatic_nml.f90. - """ + #: Number of dynamics substeps per fast-physics step + #: Called 'ndyn_substeps' in mo_nonhydrostatic_nml.f90 + self.ndyn_substeps: int = n_substeps + #: If True, compute horizontal diffusion only at the large time step + #: Called 'lhdiff_rcf' in mo_nonhydrostatic_nml.f90 self.lhdiff_rcf: bool = hdiff_rcf - """ - If True, compute horizontal diffusion only at the large time step. - - Called 'lhdiff_rcf' in mo_nonhydrostatic_nml.f90. - """ # namelist mo_gridref_nml.f90 + + #: Denominator for temperature boundary diffusion + #: Called 'denom_diffu_t' in mo_gridref_nml.f90 self.temperature_boundary_diffusion_denominator: float = ( temperature_boundary_diffusion_denom ) - """ - Denominator for temperature boundary diffusion. - - Called 'denom_diffu_t' in mo_gridref_nml.f90. - """ + #: Denominator for velocity boundary diffusion + #: Called 'denom_diffu_v' in mo_gridref_nml.f90 self.velocity_boundary_diffusion_denominator: float = ( velocity_boundary_diffusion_denom ) - """ - Denominator for velocity boundary diffusion. - - Called 'denom_diffu_v' in mo_gridref_nml.f90. - """ # parameters from namelist: mo_interpol_nml.f90 + + #: Parameter describing the lateral boundary nudging in limited area mode. + #: + #: Maximal value of the nudging coefficients used cell row bordering the boundary interpolation zone, + #: from there nudging coefficients decay exponentially with `nudge_efold_width` in units of cell rows. + #: Called `nudge_max_coeff` in mo_interpol_nml.f90 self.nudge_max_coeff: float = max_nudging_coeff - """ - Parameter describing the lateral boundary nudging in limited area mode. - Maximal value of the nudging coefficients used cell row bordering the boundary - interpolation zone, from there nudging coefficients decay exponentially with - `nudge_efold_width` in units of cell rows. - Called `nudge_max_coeff` in mo_interpol_nml.f90 - """ + #: Exponential decay rate (in units of cell rows) of the lateral boundary nudging coefficients + #: Called `nudge_efold_width` in mo_interpol_nml.f90 self.nudge_efold_width: float = nudging_decay_rate - """ - Exponential decay rate (in units of cell rows) of the lateral boundary nudging coefficients. - - Called `nudge_efold_width` in mo_interpol_nml.f90 - """ self._validate() From 4f2b6df38952997d42a15e64e4eee5836bd3d85b Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 12 Jul 2023 21:48:53 +0200 Subject: [PATCH 177/263] - remove set_zero calls for u_vert and v_vert --- atm_dyn_iconam/src/icon4py/diffusion/diffusion.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index cd45edbcb..ffceb2340 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -609,8 +609,6 @@ def _do_diffusion_step( self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={} ) - set_zero_v_k.with_backend(backend)(self.u_vert, offset_provider={}) - set_zero_v_k.with_backend(backend)(self.v_vert, offset_provider={}) log.debug("rbf interpolation: start") mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( p_e_in=prognostic_state.vn, From de2c6df5be62df187056e931602b82ef57a15160 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 13 Jul 2023 12:58:43 +0200 Subject: [PATCH 178/263] merge together the state modules --- .../src/icon4py/diffusion/diagnostic_state.py | 36 ------------ .../src/icon4py/diffusion/diffusion.py | 21 ++++--- .../src/icon4py/diffusion/metric_state.py | 31 ----------- .../src/icon4py/diffusion/prognostic_state.py | 33 ----------- ...{interpolation_state.py => state_utils.py} | 55 ++++++++++++++++++- .../src/icon4py/driver/dycore_driver.py | 5 +- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 14 +++-- atm_dyn_iconam/tests/test_diffusion.py | 26 ++++++--- ...rpolation_state.py => test_state_utils.py} | 4 +- .../tests/test_utils/serialbox_utils.py | 14 +++-- .../py2f/wrappers/diffusion_wrapper.py | 10 ++-- 11 files changed, 109 insertions(+), 140 deletions(-) delete mode 100644 atm_dyn_iconam/src/icon4py/diffusion/diagnostic_state.py delete mode 100644 atm_dyn_iconam/src/icon4py/diffusion/metric_state.py delete mode 100644 atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py rename atm_dyn_iconam/src/icon4py/diffusion/{interpolation_state.py => state_utils.py} (56%) rename atm_dyn_iconam/tests/{test_interpolation_state.py => test_state_utils.py} (90%) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diagnostic_state.py b/atm_dyn_iconam/src/icon4py/diffusion/diagnostic_state.py deleted file mode 100644 index b74384af8..000000000 --- a/atm_dyn_iconam/src/icon4py/diffusion/diagnostic_state.py +++ /dev/null @@ -1,36 +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 - -from dataclasses import dataclass - -from gt4py.next.common import Field - -from icon4py.common.dimension import CellDim, KDim - - -@dataclass -class DiagnosticState: - # fields for 3D elements in turbdiff - hdef_ic: Field[ - [CellDim, KDim], float - ] # ! divergence at half levels(nproma,nlevp1,nblks_c) [1/s] - div_ic: Field[ - [CellDim, KDim], float - ] # ! horizontal wind field deformation (nproma,nlevp1,nblks_c) [1/s^2] - dwdx: Field[ - [CellDim, KDim], float - ] # zonal gradient of vertical wind speed (nproma,nlevp1,nblks_c) [1/s] - - dwdy: Field[ - [CellDim, KDim], float - ] # meridional gradient of vertical wind speed (nproma,nlevp1,nblks_c) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index ffceb2340..84550e89a 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -57,23 +57,24 @@ DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO, GAS_CONSTANT_DRY_AIR, ) -from icon4py.common.dimension import CellDim, ECVDim, EdgeDim, KDim, VertexDim -from icon4py.diffusion.diagnostic_state import DiagnosticState +from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim from icon4py.diffusion.horizontal import ( CellParams, EdgeParams, HorizontalMarkerIndex, ) from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams -from icon4py.diffusion.interpolation_state import InterpolationState -from icon4py.diffusion.metric_state import MetricState -from icon4py.diffusion.prognostic_state import PrognosticState +from icon4py.diffusion.state_utils import ( + DiagnosticState, + InterpolationState, + MetricState, + PrognosticState, +) from icon4py.diffusion.utils import ( copy_field, init_diffusion_local_fields_for_regular_timestep, init_nabla2_factor_in_upper_damping_zone, scale_k, - set_zero_v_k, setup_fields_for_initial_step, zero_field, ) @@ -88,17 +89,21 @@ compiled_backend = run_gtfn backend = compiled_backend + class DiffusionType(int, Enum): """ Order of nabla operator for diffusion. Note: Called `hdiff_order` in `mo_diffusion_nml.f90`. Note: We currently only support type 5. """ + NO_DIFFUSION = -1 #: no diffusion LINEAR_2ND_ORDER = 2 #: 2nd order linear diffusion on all vertical levels SMAGORINSKY_NO_BACKGROUND = 3 #: Smagorinsky diffusion without background diffusion LINEAR_4TH_ORDER = 4 #: 4th order linear diffusion on all vertical levels - SMAGORINSKY_4TH_ORDER = 5 #: Smagorinsky diffusion with fourth-order background diffusion + SMAGORINSKY_4TH_ORDER = ( + 5 #: Smagorinsky diffusion with fourth-order background diffusion + ) class DiffusionConfig: @@ -212,8 +217,6 @@ def __init__( #: Called `nudge_max_coeff` in mo_interpol_nml.f90 self.nudge_max_coeff: float = max_nudging_coeff - - #: Exponential decay rate (in units of cell rows) of the lateral boundary nudging coefficients #: Called `nudge_efold_width` in mo_interpol_nml.f90 self.nudge_efold_width: float = nudging_decay_rate diff --git a/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py b/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py deleted file mode 100644 index 4fcdb0830..000000000 --- a/atm_dyn_iconam/src/icon4py/diffusion/metric_state.py +++ /dev/null @@ -1,31 +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 - -from dataclasses import dataclass - -from gt4py.next.common import Field -from numpy import int32 - -from icon4py.common.dimension import CECDim, CellDim, KDim - - -@dataclass -class MetricState: - theta_ref_mc: Field[[CellDim, KDim], float] - wgtfac_c: Field[ - [CellDim, KDim], float - ] # weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_c) - mask_hdiff: Field[[CellDim, KDim], bool] - zd_vertoffset: Field[[CECDim, KDim], int32] - zd_diffcoef: Field[[CellDim, KDim], float] - zd_intcoef: Field[[CECDim, KDim], float] diff --git a/atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py b/atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py deleted file mode 100644 index f659eaca9..000000000 --- a/atm_dyn_iconam/src/icon4py/diffusion/prognostic_state.py +++ /dev/null @@ -1,33 +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 - -from dataclasses import dataclass - -from gt4py.next.common import Field - -from icon4py.common.dimension import CellDim, EdgeDim, KDim - - -@dataclass -class PrognosticState: - """Class that contains the prognostic state. - - corresponds to ICON t_nh_prog - """ - - w: Field[ - [CellDim, KDim], float - ] # vertical_wind field, w(nproma, nlevp1, nblks_c) [m/s] - vn: Field[[EdgeDim, KDim], float] # vn(nproma, nlev, nblks_e) [m/s] - exner_pressure: Field[[CellDim, KDim], float] # exner(nrpoma, nlev, nblks_c) - theta_v: Field[[CellDim, KDim], float] # (nproma, nlev, nlbks_c) [K] diff --git a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py b/atm_dyn_iconam/src/icon4py/diffusion/state_utils.py similarity index 56% rename from atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py rename to atm_dyn_iconam/src/icon4py/diffusion/state_utils.py index e57a27ce6..8caf8be52 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/state_utils.py @@ -15,6 +15,7 @@ import numpy as np from gt4py.next.common import Field +from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import np_as_located_field from icon4py.common.dimension import ( @@ -23,16 +24,49 @@ CEDim, CellDim, EdgeDim, + KDim, V2EDim, VertexDim, ) +@dataclass +class DiagnosticState: + """Represent diagnostic fields needed in diffusion.""" + + # fields for 3D elements in turbdiff + hdef_ic: Field[ + [CellDim, KDim], float + ] # ! divergence at half levels(nproma,nlevp1,nblks_c) [1/s] + div_ic: Field[ + [CellDim, KDim], float + ] # ! horizontal wind field deformation (nproma,nlevp1,nblks_c) [1/s^2] + dwdx: Field[ + [CellDim, KDim], float + ] # zonal gradient of vertical wind speed (nproma,nlevp1,nblks_c) [1/s] + + dwdy: Field[ + [CellDim, KDim], float + ] # meridional gradient of vertical wind speed (nproma,nlevp1,nblks_c) + + +@dataclass +class MetricState: + """Represents the metric state fields needed in diffusion.""" + + theta_ref_mc: Field[[CellDim, KDim], float] + wgtfac_c: Field[ + [CellDim, KDim], float + ] # weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_c) + mask_hdiff: Field[[CellDim, KDim], bool] + zd_vertoffset: Field[[CECDim, KDim], int32] + zd_diffcoef: Field[[CellDim, KDim], float] + zd_intcoef: Field[[CECDim, KDim], float] + + @dataclass(frozen=True) class InterpolationState: - """Represents the ICON interpolation state.""" - - # TODO(Magdalena): keep? does this state make sense at all? + """Represents the ICON interpolation state needed in diffusion.""" e_bln_c_s: Field[ [CEDim], float @@ -70,3 +104,18 @@ def geofac_n2s_nbh(self) -> Field[[CECDim], float]: old_shape[0] * old_shape[1], ) ) + + +@dataclass +class PrognosticState: + """Class that contains the prognostic state. + + corresponds to ICON t_nh_prog + """ + + w: Field[ + [CellDim, KDim], float + ] # vertical_wind field, w(nproma, nlevp1, nblks_c) [m/s] + vn: Field[[EdgeDim, KDim], float] # vn(nproma, nlev, nblks_e) [m/s] + exner_pressure: Field[[CellDim, KDim], float] # exner(nrpoma, nlev, nblks_c) + theta_v: Field[[CellDim, KDim], float] # (nproma, nlev, nlbks_c) [K] diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index 89baa998f..a3a8d50c9 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -20,9 +20,8 @@ from devtools import Timer from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn -from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.diffusion import Diffusion, DiffusionParams -from icon4py.diffusion.prognostic_state import PrognosticState +from icon4py.diffusion.state_utils import DiagnosticState, PrognosticState from icon4py.diffusion.utils import copy_diagnostic_and_prognostics from icon4py.driver.icon_configuration import IconRunConfig, read_config from icon4py.driver.io_utils import ( @@ -72,7 +71,7 @@ def do_dynamics_substepping( linit=False, date=self.simulation_date.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] ) new_p = sp.construct_prognostics() - new_d = sp.construct_diagnostics() + new_d = sp.construct_diagnostics_for_diffusion() copy_diagnostic_and_prognostics.with_backend(run_gtfn)( new_d.hdef_ic, diagnostic_state.hdef_ic, diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index f80a27fa1..a2cd858ff 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -17,12 +17,14 @@ from enum import Enum from pathlib import Path -from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.horizontal import CellParams, EdgeParams from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams -from icon4py.diffusion.interpolation_state import InterpolationState -from icon4py.diffusion.metric_state import MetricState -from icon4py.diffusion.prognostic_state import PrognosticState +from icon4py.diffusion.state_utils import ( + DiagnosticState, + InterpolationState, + MetricState, + PrognosticState, +) SB_ONLY_MSG = "Only ser_type='sb' is implemented so far." @@ -98,7 +100,7 @@ def read_initial_state( linit=True, date=SIMULATION_START_DATE ) prognostic_state = init_savepoint.construct_prognostics() - diagnostic_state = init_savepoint.construct_diagnostics() + diagnostic_state = init_savepoint.construct_diagnostics_for_diffusion() return data_provider, diagnostic_state, prognostic_state @@ -149,7 +151,7 @@ def read_static_fields( "icon_pydycore", str(path.absolute()), False ) interpolation_state = ( - dataprovider.from_interpolation_savepoint().construct_interpolation_state() + dataprovider.from_interpolation_savepoint().construct_interpolation_state_for_diffusion() ) metric_state = dataprovider.from_metrics_savepoint().construct_metric_state() return metric_state, interpolation_state diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 0c3988098..6fda2cd17 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -241,7 +241,9 @@ def test_diffusion_init( assert meta["linit"] is False assert meta["date"] == step_date_init - interpolation_state = interpolation_savepoint.construct_interpolation_state() + interpolation_state = ( + interpolation_savepoint.construct_interpolation_state_for_diffusion() + ) metric_state = metrics_savepoint.construct_metric_state() edge_params = grid_savepoint.construct_edge_geometry() cell_params = grid_savepoint.construct_cell_geometry() @@ -362,7 +364,9 @@ def test_verify_diffusion_init_against_first_regular_savepoint( cell_geometry = grid_savepoint.construct_cell_geometry() edge_geometry = grid_savepoint.construct_edge_geometry() - interpolation_state = interpolation_savepoint.construct_interpolation_state() + interpolation_state = ( + interpolation_savepoint.construct_interpolation_state_for_diffusion() + ) metric_state = metrics_savepoint.construct_metric_state() diffusion = Diffusion() @@ -396,7 +400,9 @@ def test_verify_diffusion_init_against_other_regular_savepoint( additional_parameters = DiffusionParams(config) vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) - interpolation_state = interpolation_savepoint.construct_interpolation_state() + interpolation_state = ( + interpolation_savepoint.construct_interpolation_state_for_diffusion() + ) metric_state = metrics_savepoint.construct_metric_state() edge_params = grid_savepoint.construct_edge_geometry() cell_params = grid_savepoint.construct_cell_geometry() @@ -438,9 +444,11 @@ def test_run_diffusion_single_step( dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() - interpolation_state = interpolation_savepoint.construct_interpolation_state() + interpolation_state = ( + interpolation_savepoint.construct_interpolation_state_for_diffusion() + ) metric_state = metrics_savepoint.construct_metric_state() - diagnostic_state = diffusion_savepoint_init.construct_diagnostics() + diagnostic_state = diffusion_savepoint_init.construct_diagnostics_for_diffusion() prognostic_state = diffusion_savepoint_init.construct_prognostics() vct_a = grid_savepoint.vct_a() vertical_params = VerticalModelParams( @@ -519,9 +527,11 @@ def test_run_diffusion_initial_step( dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() - interpolation_state = interpolation_savepoint.construct_interpolation_state() + interpolation_state = ( + interpolation_savepoint.construct_interpolation_state_for_diffusion() + ) metric_state = metrics_savepoint.construct_metric_state() - diagnostic_state = diffusion_savepoint_init.construct_diagnostics() + diagnostic_state = diffusion_savepoint_init.construct_diagnostics_for_diffusion() prognostic_state = diffusion_savepoint_init.construct_prognostics() vct_a = grid_savepoint.vct_a() vertical_params = VerticalModelParams( @@ -560,7 +570,7 @@ def test_run_diffusion_initial_step( @pytest.mark.datatest def test_verify_stencil15_field_manipulation(interpolation_savepoint, icon_grid): geofac_n2s = np.asarray(interpolation_savepoint.geofac_n2s()) - int_state = interpolation_savepoint.construct_interpolation_state() + int_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() geofac_c = np.asarray(int_state.geofac_n2s_c) geofac_nbh = np.asarray(int_state.geofac_n2s_nbh) cec_table = icon_grid.get_c2cec_connectivity().table diff --git a/atm_dyn_iconam/tests/test_interpolation_state.py b/atm_dyn_iconam/tests/test_state_utils.py similarity index 90% rename from atm_dyn_iconam/tests/test_interpolation_state.py rename to atm_dyn_iconam/tests/test_state_utils.py index c49c73019..8e6ca355f 100644 --- a/atm_dyn_iconam/tests/test_interpolation_state.py +++ b/atm_dyn_iconam/tests/test_state_utils.py @@ -17,7 +17,9 @@ @pytest.mark.datatest def test_cecdim(interpolation_savepoint, icon_grid): - interpolation_fields = interpolation_savepoint.construct_interpolation_state() + interpolation_fields = ( + interpolation_savepoint.construct_interpolation_state_for_diffusion() + ) geofac_n2s = np.asarray(interpolation_fields.geofac_n2s) geofac_n2s_nbh = np.asarray(interpolation_fields.geofac_n2s_nbh) assert np.count_nonzero(geofac_n2s_nbh) > 0 diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index 94b02fb98..231be5ffb 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -34,7 +34,6 @@ V2EDim, VertexDim, ) -from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.diffusion import VectorTuple from icon4py.diffusion.horizontal import ( CellParams, @@ -42,9 +41,12 @@ HorizontalMeshSize, ) from icon4py.diffusion.icon_grid import IconGrid, MeshConfig, VerticalMeshConfig -from icon4py.diffusion.interpolation_state import InterpolationState -from icon4py.diffusion.metric_state import MetricState -from icon4py.diffusion.prognostic_state import PrognosticState +from icon4py.diffusion.state_utils import ( + DiagnosticState, + InterpolationState, + MetricState, + PrognosticState, +) from .helpers import as_1D_sparse_field @@ -282,7 +284,7 @@ def rbf_vec_coeff_v2(self): def nudgecoeff_e(self): return self._get_field("nudgecoeff_e", EdgeDim) - def construct_interpolation_state(self) -> InterpolationState: + def construct_interpolation_state_for_diffusion(self) -> InterpolationState: grg = self.geofac_grg() return InterpolationState( e_bln_c_s=as_1D_sparse_field(self.e_bln_c_s(), CEDim), @@ -407,7 +409,7 @@ def construct_prognostics(self) -> PrognosticState: theta_v=self.theta_v(), ) - def construct_diagnostics(self) -> DiagnosticState: + def construct_diagnostics_for_diffusion(self) -> DiagnosticState: return DiagnosticState( hdef_ic=self.hdef_ic(), div_ic=self.div_ic(), diff --git a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py index 821081046..007a901f2 100644 --- a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py +++ b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py @@ -25,7 +25,6 @@ V2EDim, VertexDim, ) -from icon4py.diffusion.diagnostic_state import DiagnosticState from icon4py.diffusion.diffusion import ( Diffusion, DiffusionConfig, @@ -33,9 +32,12 @@ ) from icon4py.diffusion.horizontal import CellParams, EdgeParams from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams -from icon4py.diffusion.interpolation_state import InterpolationState -from icon4py.diffusion.metric_state import MetricState -from icon4py.diffusion.prognostic_state import PrognosticState +from icon4py.diffusion.state_utils import ( + DiagnosticState, + InterpolationState, + MetricState, + PrognosticState, +) from icon4py.py2f.cffi_utils import CffiMethod From 8efae753fbdcefeef72b9b633c45b3a30a41bbfc Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 13 Jul 2023 14:01:01 +0200 Subject: [PATCH 179/263] clean up tests: - move utils tests from test_diffusion.py to test_diffusion_utils.py - add fixture for dynamical substeps --- .../src/icon4py/diffusion/diffusion.py | 16 +- .../{utils.py => diffusion_utils.py} | 0 .../src/icon4py/diffusion/state_utils.py | 10 +- .../src/icon4py/driver/dycore_driver.py | 2 +- atm_dyn_iconam/tests/conftest.py | 14 +- atm_dyn_iconam/tests/test_diffusion.py | 189 ++---------------- atm_dyn_iconam/tests/test_diffusion_utils.py | 176 ++++++++++++++++ atm_dyn_iconam/tests/test_dycore_driver.py | 2 +- 8 files changed, 218 insertions(+), 191 deletions(-) rename atm_dyn_iconam/src/icon4py/diffusion/{utils.py => diffusion_utils.py} (100%) create mode 100644 atm_dyn_iconam/tests/test_diffusion_utils.py diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 84550e89a..3b2756117 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -58,6 +58,14 @@ GAS_CONSTANT_DRY_AIR, ) from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim +from icon4py.diffusion.diffusion_utils import ( + copy_field, + init_diffusion_local_fields_for_regular_timestep, + init_nabla2_factor_in_upper_damping_zone, + scale_k, + setup_fields_for_initial_step, + zero_field, +) from icon4py.diffusion.horizontal import ( CellParams, EdgeParams, @@ -70,14 +78,6 @@ MetricState, PrognosticState, ) -from icon4py.diffusion.utils import ( - copy_field, - init_diffusion_local_fields_for_regular_timestep, - init_nabla2_factor_in_upper_damping_zone, - scale_k, - setup_fields_for_initial_step, - zero_field, -) # flake8: noqa diff --git a/atm_dyn_iconam/src/icon4py/diffusion/utils.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_utils.py similarity index 100% rename from atm_dyn_iconam/src/icon4py/diffusion/utils.py rename to atm_dyn_iconam/src/icon4py/diffusion/diffusion_utils.py diff --git a/atm_dyn_iconam/src/icon4py/diffusion/state_utils.py b/atm_dyn_iconam/src/icon4py/diffusion/state_utils.py index 8caf8be52..80672dc1b 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/state_utils.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/state_utils.py @@ -30,9 +30,9 @@ ) -@dataclass +@dataclass(frozen=True) class DiagnosticState: - """Represent diagnostic fields needed in diffusion.""" + """Represents the diagnostic fields needed in diffusion.""" # fields for 3D elements in turbdiff hdef_ic: Field[ @@ -50,7 +50,7 @@ class DiagnosticState: ] # meridional gradient of vertical wind speed (nproma,nlevp1,nblks_c) -@dataclass +@dataclass(frozen=True) class MetricState: """Represents the metric state fields needed in diffusion.""" @@ -106,11 +106,11 @@ def geofac_n2s_nbh(self) -> Field[[CECDim], float]: ) -@dataclass +@dataclass(frozen=True) class PrognosticState: """Class that contains the prognostic state. - corresponds to ICON t_nh_prog + Corresponds to ICON t_nh_prog """ w: Field[ diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index a3a8d50c9..aa6703bfc 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -21,8 +21,8 @@ from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn from icon4py.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.diffusion.diffusion_utils import copy_diagnostic_and_prognostics from icon4py.diffusion.state_utils import DiagnosticState, PrognosticState -from icon4py.diffusion.utils import copy_diagnostic_and_prognostics from icon4py.driver.icon_configuration import IconRunConfig, read_config from icon4py.driver.io_utils import ( SIMULATION_START_DATE, diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 28a17a533..9a3de4c6d 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -35,6 +35,17 @@ def get_data_path(setup_icon_data): return extracted_path +@pytest.fixture +def ndyn_substeps(): + """ + Return number of dynamical substeps. + + Serialized data uses a reduced number (2 instead of the default 5) in order to reduce the amount + of data generated. + """ + return 2 + + @pytest.fixture(scope="session") def setup_icon_data(): """ @@ -144,7 +155,7 @@ def grid_savepoint(data_provider): @pytest.fixture -def r04b09_diffusion_config() -> DiffusionConfig: +def r04b09_diffusion_config(ndyn_substeps) -> DiffusionConfig: """ Create DiffusionConfig matching MCH_CH_r04b09_dsl. @@ -163,6 +174,7 @@ def r04b09_diffusion_config() -> DiffusionConfig: zdiffu_t=True, velocity_boundary_diffusion_denom=150.0, max_nudging_coeff=0.075, + n_substeps=ndyn_substeps, ) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 6fda2cd17..f1c7db477 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -17,146 +17,16 @@ from atm_dyn_iconam.tests.test_utils.serialbox_utils import ( IconDiffusionInitSavepoint, ) -from icon4py.common.dimension import KDim, VertexDim from icon4py.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.diffusion.diffusion_utils import scale_k from icon4py.diffusion.horizontal import CellParams, EdgeParams from icon4py.diffusion.icon_grid import VerticalModelParams -from icon4py.diffusion.utils import ( - _en_smag_fac_for_zero_nshift, - _setup_runtime_diff_multfac_vn, - _setup_smag_limit, - scale_k, - set_zero_v_k, - setup_fields_for_initial_step, -) - -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh - - -datarun_reduced_substeps = 2 - - -def test_scale_k(): - mesh = SimpleMesh() - field = random_field(mesh, KDim) - scaled_field = zero_field(mesh, KDim) - factor = 2.0 - scale_k(field, factor, scaled_field, offset_provider={}) - assert np.allclose(factor * np.asarray(field), scaled_field) - - -def initial_diff_multfac_vn_numpy(shape, k4, hdiff_efdt_ratio): - return k4 * hdiff_efdt_ratio / 3.0 * np.ones(shape) - - -def smag_limit_numpy(func, *args): - return 0.125 - 4.0 * func(*args) - - -def test_diff_multfac_vn_and_smag_limit_for_initial_step(): - mesh = SimpleMesh() - diff_multfac_vn_init = zero_field(mesh, KDim) - smag_limit_init = zero_field(mesh, KDim) - k4 = 1.0 - efdt_ratio = 24.0 - shape = np.asarray(diff_multfac_vn_init).shape - - expected_diff_multfac_vn_init = initial_diff_multfac_vn_numpy(shape, k4, efdt_ratio) - expected_smag_limit_init = smag_limit_numpy( - initial_diff_multfac_vn_numpy, shape, k4, efdt_ratio - ) - - setup_fields_for_initial_step( - k4, efdt_ratio, diff_multfac_vn_init, smag_limit_init, offset_provider={} - ) - - assert np.allclose(expected_diff_multfac_vn_init, diff_multfac_vn_init) - assert np.allclose(expected_smag_limit_init, smag_limit_init) - - -def test_diff_multfac_vn_smag_limit_for_time_step_with_const_value(): - mesh = SimpleMesh() - diff_multfac_vn = zero_field(mesh, KDim) - smag_limit = zero_field(mesh, KDim) - k4 = 1.0 - substeps = 5.0 - efdt_ratio = 24.0 - shape = np.asarray(diff_multfac_vn).shape - - expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps) - expected_smag_limit = smag_limit_numpy(diff_multfac_vn_numpy, shape, k4, substeps) - - _setup_runtime_diff_multfac_vn( - k4, efdt_ratio, out=diff_multfac_vn, offset_provider={} - ) - _setup_smag_limit(diff_multfac_vn, out=smag_limit, offset_provider={}) - - assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) - assert np.allclose(expected_smag_limit, smag_limit) - - -def test_diff_multfac_vn_smag_limit_for_loop_run_with_k4_substeps(): - mesh = SimpleMesh() - diff_multfac_vn = zero_field(mesh, KDim) - smag_limit = zero_field(mesh, KDim) - k4 = 0.003 - substeps = 1.0 - - shape = np.asarray(diff_multfac_vn).shape - expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps) - expected_smag_limit = smag_limit_numpy(diff_multfac_vn_numpy, shape, k4, substeps) - _setup_runtime_diff_multfac_vn( - k4, substeps, out=diff_multfac_vn, offset_provider={} - ) - _setup_smag_limit(diff_multfac_vn, out=smag_limit, offset_provider={}) - - assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) - assert np.allclose(expected_smag_limit, smag_limit) - - -def test_init_enh_smag_fac(): - mesh = SimpleMesh() - enh_smag_fac = zero_field(mesh, KDim) - a_vec = random_field(mesh, KDim, low=1.0, high=10.0, extend={KDim: 1}) - fac = (0.67, 0.5, 1.3, 0.8) - z = (0.1, 0.2, 0.3, 0.4) - - enhanced_smag_fac_np = enhanced_smagorinski_factor_numpy(fac, z, np.asarray(a_vec)) - - _en_smag_fac_for_zero_nshift( - a_vec, *fac, *z, out=enh_smag_fac, offset_provider={"Koff": KDim} - ) - assert np.allclose(enhanced_smag_fac_np, np.asarray(enh_smag_fac)) - - -def diff_multfac_vn_numpy(shape, k4, substeps): - factor = min(1.0 / 128.0, k4 * substeps / 3.0) - return factor * np.ones(shape) - - -def enhanced_smagorinski_factor_numpy(factor_in, heigths_in, a_vec): - alin = (factor_in[1] - factor_in[0]) / (heigths_in[1] - heigths_in[0]) - df32 = factor_in[2] - factor_in[1] - df42 = factor_in[3] - factor_in[1] - dz32 = heigths_in[2] - heigths_in[1] - dz42 = heigths_in[3] - heigths_in[1] - bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32)) - aqdr = df32 / dz32 - bqdr * dz32 - zf = 0.5 * (a_vec[:-1] + a_vec[1:]) - max0 = np.maximum(0.0, zf - heigths_in[0]) - dzlin = np.minimum(heigths_in[1] - heigths_in[0], max0) - max1 = np.maximum(0.0, zf - heigths_in[1]) - dzqdr = np.minimum(heigths_in[3] - heigths_in[1], max1) - return factor_in[0] + dzlin * alin + dzqdr * (aqdr + dzqdr * bqdr) - - -def test_set_zero_vertex_k(): - mesh = SimpleMesh() - f = random_field(mesh, VertexDim, KDim) - set_zero_v_k(f, offset_provider={}) - assert np.allclose(0.0, f) +from .test_diffusion_utils import ( + diff_multfac_vn_numpy, + enhanced_smagorinski_factor_numpy, + smag_limit_numpy, +) def test_diffusion_coefficients_with_hdiff_efdt_ratio(r04b09_diffusion_config): config = r04b09_diffusion_config @@ -231,7 +101,6 @@ def test_diffusion_init( damping_height, ): config = r04b09_diffusion_config - config.ndyn_substeps = datarun_reduced_substeps additional_parameters = DiffusionParams(config) vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) @@ -261,7 +130,7 @@ def test_diffusion_init( ) assert diffusion.diff_multfac_w == min( - 1.0 / 48.0, additional_parameters.K4W * config.substep_as_float() + 1.0 / 48.0, additional_parameters.K4W * config.substep_as_float ) assert np.allclose(0.0, np.asarray(diffusion.v_vert)) @@ -274,17 +143,17 @@ def test_diffusion_init( diff_multfac_vn_numpy, shape_k, additional_parameters.K4, - config.substep_as_float(), + config.substep_as_float, ) assert ( diffusion.smag_offset - == 0.25 * additional_parameters.K4 * config.substep_as_float() + == 0.25 * additional_parameters.K4 * config.substep_as_float ) assert np.allclose(expected_smag_limit, diffusion.smag_limit) expected_diff_multfac_vn = diff_multfac_vn_numpy( - shape_k, additional_parameters.K4, config.substep_as_float() + shape_k, additional_parameters.K4, config.substep_as_float ) assert np.allclose(expected_diff_multfac_vn, diffusion.diff_multfac_vn) expected_enh_smag_fac = enhanced_smagorinski_factor_numpy( @@ -319,34 +188,6 @@ def _verify_init_values_against_savepoint( assert np.allclose(savepoint.diff_multfac_vn(), diffusion.diff_multfac_vn) -@pytest.mark.datatest -@pytest.mark.parametrize("linit", [True]) -def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( - diffusion_savepoint_init, r04b09_diffusion_config, icon_grid -): - savepoint = diffusion_savepoint_init - config = r04b09_diffusion_config - config.ndyn_substeps = datarun_reduced_substeps - - params = DiffusionParams(config) - expected_diff_multfac_vn = savepoint.diff_multfac_vn() - expected_smag_limit = savepoint.smag_limit() - exptected_smag_offset = savepoint.smag_offset() - - diff_multfac_vn = zero_field(icon_grid, KDim) - smag_limit = zero_field(icon_grid, KDim) - setup_fields_for_initial_step( - params.K4, - config.hdiff_efdt_ratio, - diff_multfac_vn, - smag_limit, - offset_provider={}, - ) - assert np.allclose(expected_smag_limit, smag_limit) - assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) - assert exptected_smag_offset == 0.0 - - @pytest.mark.datatest def test_verify_diffusion_init_against_first_regular_savepoint( diffusion_savepoint_init, @@ -358,7 +199,6 @@ def test_verify_diffusion_init_against_first_regular_savepoint( damping_height, ): config = r04b09_diffusion_config - config.ndyn_substeps = datarun_reduced_substeps additional_parameters = DiffusionParams(config) vct_a = grid_savepoint.vct_a() cell_geometry = grid_savepoint.construct_cell_geometry() @@ -396,7 +236,6 @@ def test_verify_diffusion_init_against_other_regular_savepoint( damping_height, ): config = r04b09_diffusion_config - config.ndyn_substeps = datarun_reduced_substeps additional_parameters = DiffusionParams(config) vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) @@ -455,7 +294,6 @@ def test_run_diffusion_single_step( vct_a=vct_a, rayleigh_damping_height=damping_height ) config = r04b09_diffusion_config - config.ndyn_substeps = datarun_reduced_substeps additional_parameters = DiffusionParams(config) diffusion = Diffusion() @@ -484,8 +322,9 @@ def test_run_diffusion_single_step( def _verify_diffusion_fields( - diagnostic_state, prognostic_state, diffusion_savepoint_exit + diagnostic_state, prognostic_state, diffusion_savepoint_exit, diffusion_savepoint_init = None ): + start_nudging = 3317 ref_div_ic = np.asarray(diffusion_savepoint_exit.div_ic()) val_div_ic = np.asarray(diagnostic_state.div_ic) ref_hdef_ic = np.asarray(diffusion_savepoint_exit.hdef_ic()) @@ -494,6 +333,7 @@ def _verify_diffusion_fields( assert np.allclose(ref_hdef_ic, val_hdef_ic) ref_w = np.asarray(diffusion_savepoint_exit.w()) val_w = np.asarray(prognostic_state.w) + init_w = np.asarray(diffusion_savepoint_init.w()) if diffusion_savepoint_init else None ref_dwdx = np.asarray(diffusion_savepoint_exit.dwdx()) val_dwdx = np.asarray(diagnostic_state.dwdx) ref_dwdy = np.asarray(diffusion_savepoint_exit.dwdy()) @@ -538,7 +378,6 @@ def test_run_diffusion_initial_step( vct_a=vct_a, rayleigh_damping_height=damping_height ) config = r04b09_diffusion_config - config.ndyn_substeps = datarun_reduced_substeps additional_parameters = DiffusionParams(config) diffusion = Diffusion() @@ -563,7 +402,7 @@ def test_run_diffusion_initial_step( _verify_diffusion_fields( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, - diffusion_savepoint_exit=diffusion_savepoint_exit, + diffusion_savepoint_exit=diffusion_savepoint_exit, diffusion_savepoint_init= diffusion_savepoint_init ) diff --git a/atm_dyn_iconam/tests/test_diffusion_utils.py b/atm_dyn_iconam/tests/test_diffusion_utils.py new file mode 100644 index 000000000..53142b8b6 --- /dev/null +++ b/atm_dyn_iconam/tests/test_diffusion_utils.py @@ -0,0 +1,176 @@ +# 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 atm_dyn_iconam.tests.test_utils.helpers import random_field, zero_field +from atm_dyn_iconam.tests.test_utils.simple_mesh import SimpleMesh +from icon4py.common.dimension import KDim, VertexDim +from icon4py.diffusion.diffusion import DiffusionParams +from icon4py.diffusion.diffusion_utils import ( + _en_smag_fac_for_zero_nshift, + _setup_runtime_diff_multfac_vn, + _setup_smag_limit, + scale_k, + set_zero_v_k, + setup_fields_for_initial_step, +) + + +def diff_multfac_vn_numpy(shape, k4, substeps): + factor = min(1.0 / 128.0, k4 * substeps / 3.0) + return factor * np.ones(shape) + + +def initial_diff_multfac_vn_numpy(shape, k4, hdiff_efdt_ratio): + return k4 * hdiff_efdt_ratio / 3.0 * np.ones(shape) + + +def smag_limit_numpy(func, *args): + return 0.125 - 4.0 * func(*args) + + +def test_scale_k(): + mesh = SimpleMesh() + field = random_field(mesh, KDim) + scaled_field = zero_field(mesh, KDim) + factor = 2.0 + scale_k(field, factor, scaled_field, offset_provider={}) + assert np.allclose(factor * np.asarray(field), scaled_field) + + +def test_diff_multfac_vn_and_smag_limit_for_initial_step(): + mesh = SimpleMesh() + diff_multfac_vn_init = zero_field(mesh, KDim) + smag_limit_init = zero_field(mesh, KDim) + k4 = 1.0 + efdt_ratio = 24.0 + shape = np.asarray(diff_multfac_vn_init).shape + + expected_diff_multfac_vn_init = initial_diff_multfac_vn_numpy(shape, k4, efdt_ratio) + expected_smag_limit_init = smag_limit_numpy( + initial_diff_multfac_vn_numpy, shape, k4, efdt_ratio + ) + + setup_fields_for_initial_step( + k4, efdt_ratio, diff_multfac_vn_init, smag_limit_init, offset_provider={} + ) + + assert np.allclose(expected_diff_multfac_vn_init, diff_multfac_vn_init) + assert np.allclose(expected_smag_limit_init, smag_limit_init) + + +def test_diff_multfac_vn_smag_limit_for_time_step_with_const_value(): + mesh = SimpleMesh() + diff_multfac_vn = zero_field(mesh, KDim) + smag_limit = zero_field(mesh, KDim) + k4 = 1.0 + substeps = 5.0 + efdt_ratio = 24.0 + shape = np.asarray(diff_multfac_vn).shape + + expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps) + expected_smag_limit = smag_limit_numpy(diff_multfac_vn_numpy, shape, k4, substeps) + + _setup_runtime_diff_multfac_vn( + k4, efdt_ratio, out=diff_multfac_vn, offset_provider={} + ) + _setup_smag_limit(diff_multfac_vn, out=smag_limit, offset_provider={}) + + assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) + assert np.allclose(expected_smag_limit, smag_limit) + + +def test_diff_multfac_vn_smag_limit_for_loop_run_with_k4_substeps(): + mesh = SimpleMesh() + diff_multfac_vn = zero_field(mesh, KDim) + smag_limit = zero_field(mesh, KDim) + k4 = 0.003 + substeps = 1.0 + + shape = np.asarray(diff_multfac_vn).shape + expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps) + expected_smag_limit = smag_limit_numpy(diff_multfac_vn_numpy, shape, k4, substeps) + _setup_runtime_diff_multfac_vn( + k4, substeps, out=diff_multfac_vn, offset_provider={} + ) + _setup_smag_limit(diff_multfac_vn, out=smag_limit, offset_provider={}) + + assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) + assert np.allclose(expected_smag_limit, smag_limit) + + +def enhanced_smagorinski_factor_numpy(factor_in, heigths_in, a_vec): + alin = (factor_in[1] - factor_in[0]) / (heigths_in[1] - heigths_in[0]) + df32 = factor_in[2] - factor_in[1] + df42 = factor_in[3] - factor_in[1] + dz32 = heigths_in[2] - heigths_in[1] + dz42 = heigths_in[3] - heigths_in[1] + bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32)) + aqdr = df32 / dz32 - bqdr * dz32 + zf = 0.5 * (a_vec[:-1] + a_vec[1:]) + max0 = np.maximum(0.0, zf - heigths_in[0]) + dzlin = np.minimum(heigths_in[1] - heigths_in[0], max0) + max1 = np.maximum(0.0, zf - heigths_in[1]) + dzqdr = np.minimum(heigths_in[3] - heigths_in[1], max1) + return factor_in[0] + dzlin * alin + dzqdr * (aqdr + dzqdr * bqdr) + + +def test_init_enh_smag_fac(): + mesh = SimpleMesh() + enh_smag_fac = zero_field(mesh, KDim) + a_vec = random_field(mesh, KDim, low=1.0, high=10.0, extend={KDim: 1}) + fac = (0.67, 0.5, 1.3, 0.8) + z = (0.1, 0.2, 0.3, 0.4) + + enhanced_smag_fac_np = enhanced_smagorinski_factor_numpy(fac, z, np.asarray(a_vec)) + + _en_smag_fac_for_zero_nshift( + a_vec, *fac, *z, out=enh_smag_fac, offset_provider={"Koff": KDim} + ) + assert np.allclose(enhanced_smag_fac_np, np.asarray(enh_smag_fac)) + + +def test_set_zero_vertex_k(): + mesh = SimpleMesh() + f = random_field(mesh, VertexDim, KDim) + set_zero_v_k(f, offset_provider={}) + assert np.allclose(0.0, f) + + +@pytest.mark.datatest +@pytest.mark.parametrize("linit", [True]) +def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( + diffusion_savepoint_init, r04b09_diffusion_config, icon_grid, ndyn_substeps +): + savepoint = diffusion_savepoint_init + config = r04b09_diffusion_config + + params = DiffusionParams(config) + expected_diff_multfac_vn = savepoint.diff_multfac_vn() + expected_smag_limit = savepoint.smag_limit() + exptected_smag_offset = savepoint.smag_offset() + + diff_multfac_vn = zero_field(icon_grid, KDim) + smag_limit = zero_field(icon_grid, KDim) + setup_fields_for_initial_step( + params.K4, + config.hdiff_efdt_ratio, + diff_multfac_vn, + smag_limit, + offset_provider={}, + ) + assert np.allclose(expected_smag_limit, smag_limit) + assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) + assert exptected_smag_offset == 0.0 diff --git a/atm_dyn_iconam/tests/test_dycore_driver.py b/atm_dyn_iconam/tests/test_dycore_driver.py index b9fae7a81..b261dcd6e 100644 --- a/atm_dyn_iconam/tests/test_dycore_driver.py +++ b/atm_dyn_iconam/tests/test_dycore_driver.py @@ -14,7 +14,7 @@ import numpy as np from icon4py.common.dimension import CellDim, EdgeDim, KDim -from icon4py.diffusion.utils import copy_diagnostic_and_prognostics +from icon4py.diffusion.diffusion_utils import copy_diagnostic_and_prognostics from .test_utils.helpers import random_field, zero_field from .test_utils.simple_mesh import SimpleMesh From 74c1736ee5eeb051a36cc22e6a42810a8b8d3078 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 13 Jul 2023 16:28:54 +0200 Subject: [PATCH 180/263] make DiffusionConfig.substep_as_float a property --- atm_dyn_iconam/src/icon4py/diffusion/diffusion.py | 10 ++++++---- atm_dyn_iconam/src/icon4py/diffusion/state_utils.py | 2 +- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 3b2756117..0eb302e46 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -10,6 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import functools import logging import math import sys @@ -244,6 +245,7 @@ def _validate(self): "zdiffu_t = False is not implemented (leaves out stencil_15)" ) + @functools.cached_property def substep_as_float(self): return float(self.ndyn_substeps) @@ -410,20 +412,20 @@ def init( params.scaled_nudge_max_coeff + sys.float_info.epsilon ) self.fac_bdydiff_v: float = ( - math.sqrt(config.substep_as_float()) + math.sqrt(config.substep_as_float) / config.velocity_boundary_diffusion_denominator if config.lhdiff_rcf else 1.0 / config.velocity_boundary_diffusion_denominator ) - self.smag_offset: float = 0.25 * params.K4 * config.substep_as_float() + self.smag_offset: float = 0.25 * params.K4 * config.substep_as_float self.diff_multfac_w: float = min( - 1.0 / 48.0, params.K4W * config.substep_as_float() + 1.0 / 48.0, params.K4W * config.substep_as_float ) init_diffusion_local_fields_for_regular_timestep.with_backend(backend)( params.K4, - config.substep_as_float(), + config.substep_as_float, *params.smagorinski_factor, *params.smagorinski_height, self.vertical_params.physical_heights, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/state_utils.py b/atm_dyn_iconam/src/icon4py/diffusion/state_utils.py index 80672dc1b..dc1bc9944 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/state_utils.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/state_utils.py @@ -106,7 +106,7 @@ def geofac_n2s_nbh(self) -> Field[[CECDim], float]: ) -@dataclass(frozen=True) +@dataclass class PrognosticState: """Class that contains the prognostic state. From c4f75a9458181cca2c3d69831f49dc8681eba1af Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 13 Jul 2023 16:44:29 +0200 Subject: [PATCH 181/263] fix test_diffusion.py --- atm_dyn_iconam/tests/test_diffusion.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index f1c7db477..438f8fdd0 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -28,6 +28,7 @@ smag_limit_numpy, ) + def test_diffusion_coefficients_with_hdiff_efdt_ratio(r04b09_diffusion_config): config = r04b09_diffusion_config config.hdiff_efdt_ratio = 1.0 @@ -322,9 +323,10 @@ def test_run_diffusion_single_step( def _verify_diffusion_fields( - diagnostic_state, prognostic_state, diffusion_savepoint_exit, diffusion_savepoint_init = None + diagnostic_state, + prognostic_state, + diffusion_savepoint_exit, ): - start_nudging = 3317 ref_div_ic = np.asarray(diffusion_savepoint_exit.div_ic()) val_div_ic = np.asarray(diagnostic_state.div_ic) ref_hdef_ic = np.asarray(diffusion_savepoint_exit.hdef_ic()) @@ -333,7 +335,6 @@ def _verify_diffusion_fields( assert np.allclose(ref_hdef_ic, val_hdef_ic) ref_w = np.asarray(diffusion_savepoint_exit.w()) val_w = np.asarray(prognostic_state.w) - init_w = np.asarray(diffusion_savepoint_init.w()) if diffusion_savepoint_init else None ref_dwdx = np.asarray(diffusion_savepoint_exit.dwdx()) val_dwdx = np.asarray(diagnostic_state.dwdx) ref_dwdy = np.asarray(diffusion_savepoint_exit.dwdy()) @@ -402,7 +403,7 @@ def test_run_diffusion_initial_step( _verify_diffusion_fields( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, - diffusion_savepoint_exit=diffusion_savepoint_exit, diffusion_savepoint_init= diffusion_savepoint_init + diffusion_savepoint_exit=diffusion_savepoint_exit, ) From 4b84d836da0a62ae7d42138937df73d4e49cb72b Mon Sep 17 00:00:00 2001 From: Magdalena Date: Fri, 14 Jul 2023 15:37:49 +0200 Subject: [PATCH 182/263] diffusion for non limited area runs (#241) incorporate changes for non limited area runs in diffusion --- .../atm_dyn_iconam/apply_diffusion_to_vn.py | 8 +++++ .../src/icon4py/diffusion/diffusion.py | 29 ++++++++++++++----- 2 files changed, 29 insertions(+), 8 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py index 64f79b4d8..679a8eca1 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py @@ -112,6 +112,10 @@ def apply_diffusion_to_vn( nudgezone_diff: float, fac_bdydiff_v: float, start_2nd_nudge_line_idx_e: int32, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, limited_area: bool, ): _apply_diffusion_to_vn( @@ -133,4 +137,8 @@ def apply_diffusion_to_vn( start_2nd_nudge_line_idx_e, limited_area, out=vn, + domain={ + EdgeDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 0eb302e46..c8e78fca1 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -17,7 +17,7 @@ from collections import namedtuple from dataclasses import InitVar, dataclass, field from enum import Enum -from typing import Final, Optional, Tuple +from typing import Final, Optional import numpy as np from gt4py.next.common import Dimension @@ -28,6 +28,7 @@ run_gtfn_cached, ) +from icon4py.atm_dyn_iconam.apply_diffusion_to_vn import apply_diffusion_to_vn from icon4py.atm_dyn_iconam.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance, ) @@ -43,9 +44,6 @@ from icon4py.atm_dyn_iconam.calculate_nabla2_for_theta import ( calculate_nabla2_for_theta, ) -from icon4py.atm_dyn_iconam.fused_mo_nh_diffusion_stencil_04_05_06 import ( - fused_mo_nh_diffusion_stencil_04_05_06, -) from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) @@ -367,6 +365,7 @@ def __init__(self): self.nudgezone_diff: Optional[float] = None self.edge_params: Optional[EdgeParams] = None self.cell_params: Optional[CellParams] = None + self._horizontal_start_index_w_diffusion: int32 = 0 def init( self, @@ -405,6 +404,17 @@ def init( self._allocate_temporary_fields() + def _get_start_index_for_w_diffusion() -> int32: + marker = ( + HorizontalMarkerIndex.nudging(CellDim) + if self.grid.limited_area() + else HorizontalMarkerIndex.interior(CellDim) + ) + + return self.grid.get_indices_from_to( + CellDim, marker, HorizontalMarkerIndex.interior(CellDim) + )[0] + self.nudgezone_diff: float = 0.04 / ( params.scaled_nudge_max_coeff + sys.float_info.epsilon ) @@ -441,6 +451,7 @@ def init( physical_heights=np.asarray(self.vertical_params.physical_heights), nrdmax=self.vertical_params.index_of_damping_layer, ) + self._horizontal_start_index_w_diffusion = _get_start_index_for_w_diffusion() self._initialized = True @property @@ -466,6 +477,7 @@ def _index_field(dim: Dimension, size=None): self.z_nabla2_e = _allocate(EdgeDim, KDim) self.z_temp = _allocate(CellDim, KDim) self.diff_multfac_smag = _allocate(KDim) + self.z_nabla4_e2 = _allocate(EdgeDim, KDim) # TODO(Magdalena): this is KHalfDim self.vertical_index = _index_field(KDim, self.grid.n_lev() + 1) self.horizontal_cell_index = _index_field(CellDim) @@ -708,8 +720,8 @@ def _do_diffusion_step( # 6. HALO EXCHANGE -- CALL sync_patch_array_mult - log.debug("running stencil 04 05 06: start") - fused_mo_nh_diffusion_stencil_04_05_06.with_backend(backend)( + log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): start") + apply_diffusion_to_vn.with_backend(backend)( u_vert=self.u_vert, v_vert=self.v_vert, primal_normal_vert_v1=self.edge_params.primal_normal_vert[0], @@ -726,6 +738,7 @@ def _do_diffusion_step( nudgezone_diff=self.nudgezone_diff, fac_bdydiff_v=self.fac_bdydiff_v, start_2nd_nudge_line_idx_e=int32(edge_start_nudging_plus_one), + limited_area=self.grid.limited_area(), horizontal_start=edge_start_lb_plus4, horizontal_end=edge_end_local, vertical_start=0, @@ -735,7 +748,7 @@ def _do_diffusion_step( "E2ECV": self.grid.get_e2ecv_connectivity(), }, ) - log.debug("runningstencils 04 05 06: end") + log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): end") log.debug( "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" @@ -763,7 +776,7 @@ def _do_diffusion_step( ), # +1 since Fortran includes boundaries interior_idx=int32(cell_start_interior), halo_idx=int32(cell_end_local), - horizontal_start=cell_start_nudging, + horizontal_start=self._horizontal_start_index_w_diffusion, horizontal_end=cell_end_halo, vertical_start=0, vertical_end=klevels, From 88e304eeb1ea60834123e36b89e948954840c2f7 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 14 Jul 2023 15:47:10 +0200 Subject: [PATCH 183/263] - remove unused stencil - consolicate duplicated test --- .../fused_mo_nh_diffusion_stencil_04_05_06.py | 122 ------------------ atm_dyn_iconam/tests/test_diffusion.py | 11 -- atm_dyn_iconam/tests/test_state_utils.py | 21 ++- 3 files changed, 10 insertions(+), 144 deletions(-) delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py deleted file mode 100644 index 93bcc055e..000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py +++ /dev/null @@ -1,122 +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 - -from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, int32, where - -from icon4py.atm_dyn_iconam.apply_nabla2_and_nabla4_to_vn import ( - _apply_nabla2_and_nabla4_to_vn, -) -from icon4py.atm_dyn_iconam.apply_nabla2_to_vn_in_lateral_boundary import ( - _apply_nabla2_to_vn_in_lateral_boundary, -) -from icon4py.atm_dyn_iconam.calculate_nabla4 import _calculate_nabla4 -from icon4py.common.dimension import ECVDim, EdgeDim, KDim, VertexDim - - -@field_operator -def _fused_mo_nh_diffusion_stencil_04_05_06( - u_vert: Field[[VertexDim, KDim], float], - v_vert: Field[[VertexDim, KDim], float], - primal_normal_vert_v1: Field[[ECVDim], float], - primal_normal_vert_v2: Field[[ECVDim], float], - z_nabla2_e: Field[[EdgeDim, KDim], float], - inv_vert_vert_length: Field[[EdgeDim], float], - inv_primal_edge_length: Field[[EdgeDim], float], - area_edge: Field[[EdgeDim], float], - kh_smag_e: Field[[EdgeDim, KDim], float], - diff_multfac_vn: Field[[KDim], float], - nudgecoeff_e: Field[[EdgeDim], float], - vn: Field[[EdgeDim, KDim], float], - horz_idx: Field[[EdgeDim], int32], - nudgezone_diff: float, - fac_bdydiff_v: float, - start_2nd_nudge_line_idx_e: int32, -) -> Field[[EdgeDim, KDim], float]: - - z_nabla4_e2 = _calculate_nabla4( - u_vert, - v_vert, - primal_normal_vert_v1, - primal_normal_vert_v2, - z_nabla2_e, - inv_vert_vert_length, - inv_primal_edge_length, - ) - - vn = where( - horz_idx >= start_2nd_nudge_line_idx_e, - _apply_nabla2_and_nabla4_to_vn( - area_edge, - kh_smag_e, - z_nabla2_e, - z_nabla4_e2, - diff_multfac_vn, - nudgecoeff_e, - vn, - nudgezone_diff, - ), - _apply_nabla2_to_vn_in_lateral_boundary( - z_nabla2_e, area_edge, vn, fac_bdydiff_v - ), - ) - - return vn - - -@program -def fused_mo_nh_diffusion_stencil_04_05_06( - u_vert: Field[[VertexDim, KDim], float], - v_vert: Field[[VertexDim, KDim], float], - primal_normal_vert_v1: Field[[ECVDim], float], - primal_normal_vert_v2: Field[[ECVDim], float], - z_nabla2_e: Field[[EdgeDim, KDim], float], - inv_vert_vert_length: Field[[EdgeDim], float], - inv_primal_edge_length: Field[[EdgeDim], float], - area_edge: Field[[EdgeDim], float], - kh_smag_e: Field[[EdgeDim, KDim], float], - diff_multfac_vn: Field[[KDim], float], - nudgecoeff_e: Field[[EdgeDim], float], - vn: Field[[EdgeDim, KDim], float], - horz_idx: Field[[EdgeDim], int32], - nudgezone_diff: float, - fac_bdydiff_v: float, - start_2nd_nudge_line_idx_e: int32, - horizontal_start: int32, - horizontal_end: int32, - vertical_start: int32, - vertical_end: int32, -): - _fused_mo_nh_diffusion_stencil_04_05_06( - u_vert, - v_vert, - primal_normal_vert_v1, - primal_normal_vert_v2, - z_nabla2_e, - inv_vert_vert_length, - inv_primal_edge_length, - area_edge, - kh_smag_e, - diff_multfac_vn, - nudgecoeff_e, - vn, - horz_idx, - nudgezone_diff, - fac_bdydiff_v, - start_2nd_nudge_line_idx_e, - out=vn, - domain={ - EdgeDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, - ) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 438f8fdd0..ad538db76 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -405,14 +405,3 @@ def test_run_diffusion_initial_step( prognostic_state=prognostic_state, diffusion_savepoint_exit=diffusion_savepoint_exit, ) - - -@pytest.mark.datatest -def test_verify_stencil15_field_manipulation(interpolation_savepoint, icon_grid): - geofac_n2s = np.asarray(interpolation_savepoint.geofac_n2s()) - int_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() - geofac_c = np.asarray(int_state.geofac_n2s_c) - geofac_nbh = np.asarray(int_state.geofac_n2s_nbh) - cec_table = icon_grid.get_c2cec_connectivity().table - assert np.allclose(geofac_c, geofac_n2s[:, 0]) - assert np.allclose(geofac_nbh[cec_table], geofac_n2s[:, 1:]) diff --git a/atm_dyn_iconam/tests/test_state_utils.py b/atm_dyn_iconam/tests/test_state_utils.py index 8e6ca355f..243ab01e2 100644 --- a/atm_dyn_iconam/tests/test_state_utils.py +++ b/atm_dyn_iconam/tests/test_state_utils.py @@ -16,14 +16,13 @@ @pytest.mark.datatest -def test_cecdim(interpolation_savepoint, icon_grid): - interpolation_fields = ( - interpolation_savepoint.construct_interpolation_state_for_diffusion() - ) - geofac_n2s = np.asarray(interpolation_fields.geofac_n2s) - geofac_n2s_nbh = np.asarray(interpolation_fields.geofac_n2s_nbh) - assert np.count_nonzero(geofac_n2s_nbh) > 0 - c2cec = icon_grid.get_c2cec_connectivity().table - ported = geofac_n2s_nbh[c2cec] - assert ported.shape == geofac_n2s[:, 1:].shape - assert np.allclose(ported, geofac_n2s[:, 1:]) +def test_verify_geofac_n2s_field_manipulation(interpolation_savepoint, icon_grid): + geofac_n2s = np.asarray(interpolation_savepoint.geofac_n2s()) + int_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() + geofac_c = np.asarray(int_state.geofac_n2s_c) + geofac_nbh = np.asarray(int_state.geofac_n2s_nbh) + assert np.count_nonzero(geofac_nbh) > 0 + cec_table = icon_grid.get_c2cec_connectivity().table + assert np.allclose(geofac_c, geofac_n2s[:, 0]) + assert geofac_nbh[cec_table].shape == geofac_n2s[:, 1:].shape + assert np.allclose(geofac_nbh[cec_table], geofac_n2s[:, 1:]) From 3d5c29e6a8048b2ca8aba0d662c79e6be2d72c78 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 18 Jul 2023 15:31:22 +0200 Subject: [PATCH 184/263] add vertical.py --- .../src/icon4py/diffusion/diffusion.py | 3 +- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 3 +- atm_dyn_iconam/tests/test_diffusion.py | 6 +- atm_dyn_iconam/tests/test_io_utils.py | 2 +- atm_dyn_iconam/tests/test_utils/fixtures.py | 1 + .../tests/test_utils/serialbox_utils.py | 4 +- common/src/icon4py/grid/grid_manager.py | 5 +- common/src/icon4py/grid/icon_grid.py | 40 +------------ common/src/icon4py/grid/vertical.py | 56 +++++++++++++++++++ common/tests/conftest.py | 4 +- common/tests/test_gridmanager.py | 12 ++-- common/tests/test_vertical.py | 5 +- .../py2f/wrappers/diffusion_wrapper.py | 5 +- 13 files changed, 84 insertions(+), 62 deletions(-) create mode 100644 common/src/icon4py/grid/vertical.py diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index f2da9bc21..7b8eae4aa 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -74,7 +74,8 @@ PrognosticState, ) from icon4py.grid.horizontal import CellParams, EdgeParams, HorizontalMarkerIndex -from icon4py.grid.icon_grid import IconGrid, VerticalModelParams +from icon4py.grid.icon_grid import IconGrid +from icon4py.grid.vertical import VerticalModelParams # flake8: noqa diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index 7c4bb5e07..c4548b5ec 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -24,7 +24,8 @@ PrognosticState, ) from icon4py.grid.horizontal import CellParams, EdgeParams -from icon4py.grid.icon_grid import IconGrid, VerticalModelParams +from icon4py.grid.vertical import VerticalModelParams +from icon4py.grid.icon_grid import IconGrid SB_ONLY_MSG = "Only ser_type='sb' is implemented so far." diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 9e3ba08f4..e6446daf2 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -19,9 +19,9 @@ ) from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.diffusion.diffusion_utils import scale_k -from icon4py.grid.grid_manager import ToGt4PyTransformation, GridManager +from icon4py.grid.grid_manager import GridManager, ToGt4PyTransformation from icon4py.grid.horizontal import CellParams, EdgeParams -from icon4py.grid.icon_grid import VerticalModelParams, VerticalGridConfig +from icon4py.grid.vertical import VerticalGridSize, VerticalModelParams from .test_diffusion_utils import ( diff_multfac_vn_numpy, @@ -92,7 +92,7 @@ def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): def init_grid_manager(fname, nlev): - config = VerticalGridConfig(nlev) + config = VerticalGridSize(nlev) gm = GridManager(ToGt4PyTransformation(), fname, config) gm.init() return gm diff --git a/atm_dyn_iconam/tests/test_io_utils.py b/atm_dyn_iconam/tests/test_io_utils.py index 6d99b1f69..7844fda5e 100644 --- a/atm_dyn_iconam/tests/test_io_utils.py +++ b/atm_dyn_iconam/tests/test_io_utils.py @@ -21,7 +21,7 @@ read_static_fields, ) from icon4py.grid.horizontal import CellParams, EdgeParams -from icon4py.grid.icon_grid import VerticalModelParams +from icon4py.grid.vertical import VerticalModelParams @pytest.mark.parametrize( diff --git a/atm_dyn_iconam/tests/test_utils/fixtures.py b/atm_dyn_iconam/tests/test_utils/fixtures.py index fd3e9cc8e..9774edbbf 100644 --- a/atm_dyn_iconam/tests/test_utils/fixtures.py +++ b/atm_dyn_iconam/tests/test_utils/fixtures.py @@ -18,6 +18,7 @@ from .data_handling import download_and_extract from .serialbox_utils import IconSerialDataProvider + data_uri = "https://polybox.ethz.ch/index.php/s/LcAbscZqnsx4WCf/download" data_path = Path(__file__).parent.joinpath("ser_icondata") extracted_path = data_path.joinpath("mch_ch_r04b09_dsl/ser_data") diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index f52c196a0..de7b8d0a6 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -43,7 +43,7 @@ PrognosticState, ) from icon4py.grid.horizontal import CellParams, EdgeParams, HorizontalGridSize -from icon4py.grid.icon_grid import GridConfig, IconGrid, VerticalGridConfig +from icon4py.grid.icon_grid import GridConfig, IconGrid, VerticalGridSize from .helpers import as_1D_sparse_field @@ -219,7 +219,7 @@ def construct_icon_grid(self) -> IconGrid: num_cells=sp_meta["nproma"], # or rather "num_cells" num_edges=sp_meta["nproma"], # or rather "num_edges" ), - vertical_config=VerticalGridConfig(num_lev=sp_meta["nlev"]), + vertical_config=VerticalGridSize(num_lev=sp_meta["nlev"]), ) c2e2c = self.c2e2c() c2e2c0 = np.column_stack(((np.asarray(range(c2e2c.shape[0]))), c2e2c)) diff --git a/common/src/icon4py/grid/grid_manager.py b/common/src/icon4py/grid/grid_manager.py index 53a379465..119ac80f0 100644 --- a/common/src/icon4py/grid/grid_manager.py +++ b/common/src/icon4py/grid/grid_manager.py @@ -36,8 +36,7 @@ VertexDim, ) from icon4py.grid.horizontal import HorizontalGridSize -from icon4py.grid.icon_grid import GridConfig, IconGrid -from icon4py.grid.icon_grid import VerticalGridConfig +from icon4py.grid.icon_grid import GridConfig, IconGrid, VerticalGridSize class GridFileName(str, Enum): @@ -156,7 +155,7 @@ def __init__( self, transformation: IndexTransformation, grid_file: str, - config: VerticalGridConfig, + config: VerticalGridSize, ): self._log = logging.getLogger(__name__) self._transformation = transformation diff --git a/common/src/icon4py/grid/icon_grid.py b/common/src/icon4py/grid/icon_grid.py index f22a94392..bdee4a1ab 100644 --- a/common/src/icon4py/grid/icon_grid.py +++ b/common/src/icon4py/grid/icon_grid.py @@ -28,11 +28,7 @@ VertexDim, ) from icon4py.grid.horizontal import HorizontalGridSize - - -@dataclass(frozen=True) -class VerticalGridConfig: - num_lev: int +from icon4py.grid.vertical import VerticalGridSize @dataclass( @@ -40,7 +36,7 @@ class VerticalGridConfig: ) class GridConfig: horizontal_config: HorizontalGridSize - vertical_config: VerticalGridConfig + vertical_config: VerticalGridSize limited_area: bool = True n_shift_total: int = 0 @@ -119,6 +115,7 @@ def num_vertices(self): def num_edges(self): return self.config.num_edges + # TODO(Magdalena) deprecate for after merge def get_indices_from_to( self, dim: Dimension, start_marker: int, end_marker: int ) -> tuple[int32, int32]: @@ -216,34 +213,3 @@ def get_c2ce_connectivity(self): return self._neighbortable_offset_provider_for_1d_sparse_fields( self.connectivities["c2e"].shape, CellDim, CEDim ) - - -@dataclass(frozen=True) -class VerticalModelParams: - """ - Contains vertical physical parameters defined on the grid. - - vct_a: field containing the physical heights of the k level - rayleigh_damping_height: height of rayleigh damping in [m] mo_nonhydro_nml - """ - - vct_a: Field[[KDim], float] - rayleigh_damping_height: Final[float] - index_of_damping_layer: Final[int32] = field(init=False) - - def __post_init__(self): - object.__setattr__( - self, - "index_of_damping_layer", - self._determine_damping_height_index( - np.asarray(self.vct_a), self.rayleigh_damping_height - ), - ) - - @classmethod - def _determine_damping_height_index(cls, vct_a: np.ndarray, damping_height: float): - return int32(np.argmax(np.where(vct_a >= damping_height))) - - @property - def physical_heights(self) -> Field[[KDim], float]: - return self.vct_a diff --git a/common/src/icon4py/grid/vertical.py b/common/src/icon4py/grid/vertical.py new file mode 100644 index 000000000..63e617089 --- /dev/null +++ b/common/src/icon4py/grid/vertical.py @@ -0,0 +1,56 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from dataclasses import Field, dataclass, field +from typing import Final + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.common.dimension import KDim + + +@dataclass(frozen=True) +class VerticalGridSize: + num_lev: int + + +@dataclass(frozen=True) +class VerticalModelParams: + """ + Contains vertical physical parameters defined on the grid. + + vct_a: field containing the physical heights of the k level + rayleigh_damping_height: height of rayleigh damping in [m] mo_nonhydro_nml + """ + + vct_a: Field[[KDim], float] + rayleigh_damping_height: Final[float] + index_of_damping_layer: Final[int32] = field(init=False) + + def __post_init__(self): + object.__setattr__( + self, + "index_of_damping_layer", + self._determine_damping_height_index( + np.asarray(self.vct_a), self.rayleigh_damping_height + ), + ) + + @classmethod + def _determine_damping_height_index(cls, vct_a: np.ndarray, damping_height: float): + return int32(np.argmax(np.where(vct_a >= damping_height))) + + @property + def physical_heights(self) -> Field[[KDim], float]: + return self.vct_a diff --git a/common/tests/conftest.py b/common/tests/conftest.py index ffa8adcd0..d2b5c5a91 100644 --- a/common/tests/conftest.py +++ b/common/tests/conftest.py @@ -17,7 +17,8 @@ def import_testutils(): testutils = ( - Path(__file__).parent.__str__() + "/../../atm_dyn_iconam/tests/test_utils/__init__.py" + Path(__file__).parent.__str__() + + "/../../atm_dyn_iconam/tests/test_utils/__init__.py" ) spec = importlib.util.spec_from_file_location("helpers", testutils) module = importlib.util.module_from_spec(spec) @@ -35,4 +36,3 @@ def import_testutils(): r04b09_dsl_gridfile, setup_icon_data, ) - diff --git a/common/tests/test_gridmanager.py b/common/tests/test_gridmanager.py index eb26eefb8..c2680a7fc 100644 --- a/common/tests/test_gridmanager.py +++ b/common/tests/test_gridmanager.py @@ -19,7 +19,6 @@ import pytest from netCDF4 import Dataset - from common.tests.conftest import import_testutils from icon4py.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.grid.grid_manager import ( @@ -30,14 +29,13 @@ ToGt4PyTransformation, ) from icon4py.grid.horizontal import HorizontalMarkerIndex -from icon4py.grid.icon_grid import VerticalGridConfig +from icon4py.grid.icon_grid import VerticalGridSize helpers = import_testutils() from helpers.simple_mesh import SimpleMesh - SIMPLE_MESH_NC = "./simple_mesh_grid.nc" @@ -413,7 +411,7 @@ def test_gridmanager_eval_c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): def init_grid_manager(fname): - gm = GridManager(ToGt4PyTransformation(), fname, VerticalGridConfig(65)) + gm = GridManager(ToGt4PyTransformation(), fname, VerticalGridSize(65)) gm.init() return gm @@ -422,7 +420,7 @@ def init_grid_manager(fname): def test_grid_manager_getsize(simple_mesh_data, simple_mesh_path, dim, size, caplog): caplog.set_level(logging.DEBUG) gm = GridManager( - IndexTransformation(), simple_mesh_path, VerticalGridConfig(num_lev=80) + IndexTransformation(), simple_mesh_path, VerticalGridSize(num_lev=80) ) gm.init() assert size == gm.get_size(dim) @@ -433,7 +431,7 @@ def test_grid_manager_diamond_offset(simple_mesh_path): gm = GridManager( IndexTransformation(), simple_mesh_path, - VerticalGridConfig(num_lev=mesh.k_level), + VerticalGridSize(num_lev=mesh.k_level), ) gm.init() assert np.allclose( @@ -444,7 +442,7 @@ def test_grid_manager_diamond_offset(simple_mesh_path): def test_gridmanager_given_file_not_found_then_abort(): fname = "./unknown_grid.nc" with pytest.raises(SystemExit) as error: - gm = GridManager(IndexTransformation(), fname, VerticalGridConfig(num_lev=80)) + gm = GridManager(IndexTransformation(), fname, VerticalGridSize(num_lev=80)) gm.init() assert error.type == SystemExit assert error.value == 1 diff --git a/common/tests/test_vertical.py b/common/tests/test_vertical.py index b09dcb178..6dce10272 100644 --- a/common/tests/test_vertical.py +++ b/common/tests/test_vertical.py @@ -16,8 +16,7 @@ import numpy as np import pytest -from icon4py.grid.icon_grid import VerticalModelParams - +from icon4py.grid.vertical import VerticalModelParams @pytest.mark.parametrize( @@ -33,7 +32,7 @@ def test_nrdmax_calculation(max_h, damping, delta): == vct_a.shape[0] - math.ceil(damping / delta) - 1 ) - +# TODO(Magdalena) fix for merge @pytest.mark.datatest @pytest.mark.skip("fix location of testdata - dont want to download it twice") def test_nrdmax_calculation_from_icon_input(icon_grid, grid_savepoint, damping_height): diff --git a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py index aae4895ad..626cff301 100644 --- a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py +++ b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py @@ -30,8 +30,9 @@ DiffusionConfig, DiffusionParams, ) -from icon4py.diffusion.horizontal import CellParams, EdgeParams -from icon4py.diffusion.icon_grid import IconGrid, VerticalModelParams +from icon4py.grid.horizontal import CellParams, EdgeParams +from icon4py.grid.icon_grid import IconGrid +from icon4py.grid.vertical import VerticalModelParams from icon4py.diffusion.state_utils import ( DiagnosticState, InterpolationState, From 6a112293b9e73a7220dd64aa49ca4648ecd1eb8b Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 18 Jul 2023 17:38:56 +0200 Subject: [PATCH 185/263] add deprecation decorator icon_grid.get_indices_from_to --- common/src/icon4py/grid/icon_grid.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/common/src/icon4py/grid/icon_grid.py b/common/src/icon4py/grid/icon_grid.py index bdee4a1ab..ba4f7f79d 100644 --- a/common/src/icon4py/grid/icon_grid.py +++ b/common/src/icon4py/grid/icon_grid.py @@ -17,6 +17,7 @@ from gt4py.next.common import Dimension, DimensionKind, Field from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import NeighborTableOffsetProvider +from typing_extensions import deprecated from icon4py.common.dimension import ( CECDim, @@ -115,7 +116,8 @@ def num_vertices(self): def num_edges(self): return self.config.num_edges - # TODO(Magdalena) deprecate for after merge + + @deprecated("use get_start_index and get_end_index instead, - should be removed after merge") def get_indices_from_to( self, dim: Dimension, start_marker: int, end_marker: int ) -> tuple[int32, int32]: From f6c5455cb266e1a8cd95adc3b6b310cfbe37617e Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 18 Jul 2023 17:42:56 +0200 Subject: [PATCH 186/263] pre-commit fixes --- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 2 +- common/src/icon4py/grid/grid_manager.py | 3 --- common/src/icon4py/grid/icon_grid.py | 11 ++++++----- common/tests/test_gridmanager.py | 3 +-- common/tests/test_vertical.py | 1 + py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py | 6 +++--- 6 files changed, 12 insertions(+), 14 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index c4548b5ec..4159bc15d 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -24,8 +24,8 @@ PrognosticState, ) from icon4py.grid.horizontal import CellParams, EdgeParams -from icon4py.grid.vertical import VerticalModelParams from icon4py.grid.icon_grid import IconGrid +from icon4py.grid.vertical import VerticalModelParams SB_ONLY_MSG = "Only ser_type='sb' is implemented so far." diff --git a/common/src/icon4py/grid/grid_manager.py b/common/src/icon4py/grid/grid_manager.py index 119ac80f0..430644e15 100644 --- a/common/src/icon4py/grid/grid_manager.py +++ b/common/src/icon4py/grid/grid_manager.py @@ -380,6 +380,3 @@ def construct_diamond_array(self, c2v: np.ndarray, e2c: np.ndarray): sh = expanded.shape flattened = expanded.reshape(sh[0], sh[1] * sh[2]) return np.apply_along_axis(np.unique, 1, flattened) - - def get_c2v_connectivity(self): - return self._grid.get_c2v_connectivity() diff --git a/common/src/icon4py/grid/icon_grid.py b/common/src/icon4py/grid/icon_grid.py index ba4f7f79d..435b77e13 100644 --- a/common/src/icon4py/grid/icon_grid.py +++ b/common/src/icon4py/grid/icon_grid.py @@ -10,11 +10,11 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -from dataclasses import dataclass, field -from typing import Dict, Final +from dataclasses import dataclass +from typing import Dict import numpy as np -from gt4py.next.common import Dimension, DimensionKind, Field +from gt4py.next.common import Dimension, DimensionKind from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import NeighborTableOffsetProvider from typing_extensions import deprecated @@ -116,8 +116,9 @@ def num_vertices(self): def num_edges(self): return self.config.num_edges - - @deprecated("use get_start_index and get_end_index instead, - should be removed after merge") + @deprecated( + "use get_start_index and get_end_index instead, - should be removed after merge" + ) def get_indices_from_to( self, dim: Dimension, start_marker: int, end_marker: int ) -> tuple[int32, int32]: diff --git a/common/tests/test_gridmanager.py b/common/tests/test_gridmanager.py index c2680a7fc..bd879dcbc 100644 --- a/common/tests/test_gridmanager.py +++ b/common/tests/test_gridmanager.py @@ -33,7 +33,7 @@ helpers = import_testutils() -from helpers.simple_mesh import SimpleMesh +from helpers.simple_mesh import SimpleMesh # noqa E402 SIMPLE_MESH_NC = "./simple_mesh_grid.nc" @@ -153,7 +153,6 @@ def simple_mesh_data(): GridFile.DimensionName.VERTEX_NAME, ), ) - # add_to_dataset(data, mesh.v2e2v, GridFile.Offsets.V2E2V, (GridFile.Dimension.V2E_SIZE, GridFile.Dimension.VERTEX_NAME)) _add_to_dataset( dataset, mesh.c2e2c, diff --git a/common/tests/test_vertical.py b/common/tests/test_vertical.py index 6dce10272..25bfef8df 100644 --- a/common/tests/test_vertical.py +++ b/common/tests/test_vertical.py @@ -32,6 +32,7 @@ def test_nrdmax_calculation(max_h, damping, delta): == vct_a.shape[0] - math.ceil(damping / delta) - 1 ) + # TODO(Magdalena) fix for merge @pytest.mark.datatest @pytest.mark.skip("fix location of testdata - dont want to download it twice") diff --git a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py index 626cff301..00158a122 100644 --- a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py +++ b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py @@ -30,15 +30,15 @@ DiffusionConfig, DiffusionParams, ) -from icon4py.grid.horizontal import CellParams, EdgeParams -from icon4py.grid.icon_grid import IconGrid -from icon4py.grid.vertical import VerticalModelParams from icon4py.diffusion.state_utils import ( DiagnosticState, InterpolationState, MetricState, PrognosticState, ) +from icon4py.grid.horizontal import CellParams, EdgeParams +from icon4py.grid.icon_grid import IconGrid +from icon4py.grid.vertical import VerticalModelParams from icon4py.py2f.cffi_utils import CffiMethod From abfd26298d83778d24bb59f9e81039ddd86a5979 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 19 Jul 2023 15:20:14 +0200 Subject: [PATCH 187/263] fix - consolidate location of test data --- atm_dyn_iconam/tests/conftest.py | 79 ++++--------------- atm_dyn_iconam/tests/test_utils/fixtures.py | 37 +++++++-- .../tests/test_utils/serialbox_utils.py | 2 +- common/tests/conftest.py | 2 + common/tests/test_gridmanager.py | 5 +- .../tests/test_icon_grid.py | 2 +- common/tests/test_vertical.py | 2 - 7 files changed, 50 insertions(+), 79 deletions(-) rename {atm_dyn_iconam => common}/tests/test_icon_grid.py (99%) diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index 9a3de4c6d..e23add07a 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -11,28 +11,23 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import tarfile -from pathlib import Path import pytest -import wget from gt4py.next.program_processors.runners.roundtrip import executor from atm_dyn_iconam.tests.test_utils.simple_mesh import SimpleMesh from icon4py.diffusion.diffusion import DiffusionConfig -from .test_utils.serialbox_utils import IconSerialDataProvider - - -data_uri = "https://polybox.ethz.ch/index.php/s/LcAbscZqnsx4WCf/download" -data_path = Path(__file__).parent.joinpath("ser_icondata") -extracted_path = data_path.joinpath("mch_ch_r04b09_dsl/ser_data") -data_file = data_path.joinpath("mch_ch_r04b09_dsl_v2.tar.gz").name - - -@pytest.fixture -def get_data_path(setup_icon_data): - return extracted_path +from .test_utils.fixtures import ( # noqa F401 + damping_height, + data_provider, + get_data_path, + get_grid_files, + grid_savepoint, + icon_grid, + r04b09_dsl_gridfile, + setup_icon_data, +) @pytest.fixture @@ -46,32 +41,6 @@ def ndyn_substeps(): return 2 -@pytest.fixture(scope="session") -def setup_icon_data(): - """ - Get the binary ICON data from a remote server. - - Session scoped fixture which is a prerequisite of all the other fixtures in this file. - """ - data_path.mkdir(parents=True, exist_ok=True) - if not any(data_path.iterdir()): - print( - f"directory {data_path} is empty: downloading data from {data_uri} and extracting" - ) - wget.download(data_uri, out=data_file) - # extract downloaded file - if not tarfile.is_tarfile(data_file): - raise NotImplementedError(f"{data_file} needs to be a valid tar file") - with tarfile.open(data_file, mode="r:*") as tf: - tf.extractall(path=data_path) - Path(data_file).unlink(missing_ok=True) - - -@pytest.fixture -def data_provider(setup_icon_data) -> IconSerialDataProvider: - return IconSerialDataProvider("icon_pydycore", str(extracted_path), True) - - @pytest.fixture def linit(): """ @@ -103,7 +72,7 @@ def step_date_exit(): @pytest.fixture -def diffusion_savepoint_init(data_provider, linit, step_date_init): +def diffusion_savepoint_init(data_provider, linit, step_date_init): # noqa F811 """ Load data from ICON savepoint at start of diffusion module. @@ -116,7 +85,7 @@ def diffusion_savepoint_init(data_provider, linit, step_date_init): @pytest.fixture -def diffusion_savepoint_exit(data_provider, linit, step_date_exit): +def diffusion_savepoint_exit(data_provider, linit, step_date_exit): # noqa F811 """ Load data from ICON savepoint at exist of diffusion module. @@ -128,32 +97,17 @@ def diffusion_savepoint_exit(data_provider, linit, step_date_exit): @pytest.fixture -def interpolation_savepoint(data_provider): +def interpolation_savepoint(data_provider): # noqa F811 """Load data from ICON interplation state savepoint.""" return data_provider.from_interpolation_savepoint() @pytest.fixture -def metrics_savepoint(data_provider): +def metrics_savepoint(data_provider): # noqa F811 """Load data from ICON mestric state savepoint.""" return data_provider.from_metrics_savepoint() -@pytest.fixture -def icon_grid(grid_savepoint): - """ - Load the icon grid from an ICON savepoint. - - Uses the special grid_savepoint that contains data from p_patch - """ - return grid_savepoint.construct_icon_grid() - - -@pytest.fixture -def grid_savepoint(data_provider): - return data_provider.from_savepoint_grid() - - @pytest.fixture def r04b09_diffusion_config(ndyn_substeps) -> DiffusionConfig: """ @@ -178,11 +132,6 @@ def r04b09_diffusion_config(ndyn_substeps) -> DiffusionConfig: ) -@pytest.fixture -def damping_height(): - return 12500 - - BACKENDS = {"embedded": executor} MESHES = {"simple_mesh": SimpleMesh()} diff --git a/atm_dyn_iconam/tests/test_utils/fixtures.py b/atm_dyn_iconam/tests/test_utils/fixtures.py index 9774edbbf..278e0aba4 100644 --- a/atm_dyn_iconam/tests/test_utils/fixtures.py +++ b/atm_dyn_iconam/tests/test_utils/fixtures.py @@ -19,21 +19,24 @@ from .serialbox_utils import IconSerialDataProvider -data_uri = "https://polybox.ethz.ch/index.php/s/LcAbscZqnsx4WCf/download" -data_path = Path(__file__).parent.joinpath("ser_icondata") -extracted_path = data_path.joinpath("mch_ch_r04b09_dsl/ser_data") -data_file = data_path.joinpath("mch_ch_r04b09_dsl_v2.tar.gz").name mch_ch_r04b09_dsl_grid_uri = ( "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download" ) - r02b04_global_grid_uri = "https://polybox.ethz.ch/index.php/s/0EM8O8U53GKGsst/download" -grids_path = Path(__file__).parent.joinpath("grids") +data_uri = "https://polybox.ethz.ch/index.php/s/LcAbscZqnsx4WCf/download" +data_path_old = Path(__file__).parent.joinpath("ser_icondata") +base_path = Path(__file__).parent.parent.parent.parent.joinpath("testdata") + +data_path = base_path.joinpath("ser_icondata") +extracted_path = data_path.joinpath("mpitask1/mch_ch_r04b09_dsl/ser_data") +data_file = data_path.joinpath("mch_ch_r04b09_dsl.tar.gz").name + +grids_path = base_path.joinpath("grids") r04b09_dsl_grid_path = grids_path.joinpath("mch_ch_r04b09_dsl") r04b09_dsl_data_file = r04b09_dsl_grid_path.joinpath( "mch_ch_r04b09_dsl_grids_v1.tar.gz" ).name -r02b04_global_grid_path = grids_path.joinpath("icon_r02b04_global") +r02b04_global_grid_path = grids_path.joinpath("r02b04_global") r02b04_global_data_file = r02b04_global_grid_path.joinpath( "icon_grid_0013_R02B04_G.tar.gz" ).name @@ -59,6 +62,21 @@ def grid_savepoint(data_provider): return data_provider.from_savepoint_grid() +@pytest.fixture +def icon_grid(grid_savepoint): + """ + Load the icon grid from an ICON savepoint. + + Uses the special grid_savepoint that contains data from p_patch + """ + return grid_savepoint.construct_icon_grid() + + +@pytest.fixture +def damping_height(): + return 12500 + + @pytest.fixture(scope="session") def get_grid_files(): """ @@ -77,3 +95,8 @@ def get_grid_files(): @pytest.fixture() def r04b09_dsl_gridfile(get_grid_files): return r04b09_dsl_grid_path.joinpath("grid.nc") + + +@pytest.fixture() +def get_data_path(): + return extracted_path diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index de7b8d0a6..4867a95df 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -201,7 +201,7 @@ def v2c(self): return self._get_connectivity_array("v2c") def nrdmax(self): - return self._get_connectiviy_array("nrdmax") + return self._get_connectivity_array("nrdmax") def construct_icon_grid(self) -> IconGrid: sp_meta = self.get_metadata( diff --git a/common/tests/conftest.py b/common/tests/conftest.py index d2b5c5a91..41ee0d276 100644 --- a/common/tests/conftest.py +++ b/common/tests/conftest.py @@ -30,9 +30,11 @@ def import_testutils(): helpers = import_testutils() from helpers.fixtures import ( # noqa F401 + damping_height, data_provider, get_grid_files, grid_savepoint, + icon_grid, r04b09_dsl_gridfile, setup_icon_data, ) diff --git a/common/tests/test_gridmanager.py b/common/tests/test_gridmanager.py index bd879dcbc..d9c45dd9e 100644 --- a/common/tests/test_gridmanager.py +++ b/common/tests/test_gridmanager.py @@ -10,16 +10,15 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - import logging from pathlib import Path from uuid import uuid4 +import conftest import numpy as np import pytest from netCDF4 import Dataset -from common.tests.conftest import import_testutils from icon4py.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.grid.grid_manager import ( GridFile, @@ -32,7 +31,7 @@ from icon4py.grid.icon_grid import VerticalGridSize -helpers = import_testutils() +helpers = conftest.import_testutils() from helpers.simple_mesh import SimpleMesh # noqa E402 diff --git a/atm_dyn_iconam/tests/test_icon_grid.py b/common/tests/test_icon_grid.py similarity index 99% rename from atm_dyn_iconam/tests/test_icon_grid.py rename to common/tests/test_icon_grid.py index 3657212fe..bff1bcdb8 100644 --- a/atm_dyn_iconam/tests/test_icon_grid.py +++ b/common/tests/test_icon_grid.py @@ -23,7 +23,7 @@ (HorizontalMarkerIndex.interior(CellDim), 4104), (HorizontalMarkerIndex.interior(CellDim) + 1, 0), (HorizontalMarkerIndex.local(CellDim) - 1, 20896), - (HorizontalMarkerIndex.local(CellDim), -1), # halo in icon is (1,20896) + (HorizontalMarkerIndex.local(CellDim), -1), (HorizontalMarkerIndex.nudging(CellDim), 3316), (HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 2511), (HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 1688), diff --git a/common/tests/test_vertical.py b/common/tests/test_vertical.py index 25bfef8df..da509a6f5 100644 --- a/common/tests/test_vertical.py +++ b/common/tests/test_vertical.py @@ -33,9 +33,7 @@ def test_nrdmax_calculation(max_h, damping, delta): ) -# TODO(Magdalena) fix for merge @pytest.mark.datatest -@pytest.mark.skip("fix location of testdata - dont want to download it twice") def test_nrdmax_calculation_from_icon_input(icon_grid, grid_savepoint, damping_height): a = grid_savepoint.vct_a() nrdmax = grid_savepoint.nrdmax() From abd640cf2fb1f315d2200dc068de631b705756e5 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 19 Jul 2023 15:41:19 +0200 Subject: [PATCH 188/263] fix - enable test with updated ser. data --- common/tests/test_gridmanager.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/common/tests/test_gridmanager.py b/common/tests/test_gridmanager.py index d9c45dd9e..c53fba2ba 100644 --- a/common/tests/test_gridmanager.py +++ b/common/tests/test_gridmanager.py @@ -221,7 +221,6 @@ def test_gridparser_dimension(simple_mesh_data): assert grid_parser.dimension(GridFile.DimensionName.EDGE_NAME) == mesh.n_edges -@pytest.mark.xfail def test_gridfile_vertex_cell_edge_dimensions(grid_savepoint, r04b09_dsl_gridfile): data = Dataset(r04b09_dsl_gridfile, "r") grid_file = GridFile(data) @@ -232,7 +231,6 @@ def test_gridfile_vertex_cell_edge_dimensions(grid_savepoint, r04b09_dsl_gridfil assert grid_file.dimension(GridFile.DimensionName.EDGE_NAME) == grid_savepoint.num( EdgeDim ) - # TODO: @magdalena fix in serialized data. it returns the num_cells assert grid_file.dimension( GridFile.DimensionName.VERTEX_NAME ) == grid_savepoint.num(VertexDim) From fc634a3b5c3a72f7737b31a40db1debe105f2c45 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 19 Jul 2023 15:41:39 +0200 Subject: [PATCH 189/263] fix - pre-commit warnings --- common/src/icon4py/common/dimension.py | 1 + common/src/icon4py/grid/grid_manager.py | 200 ++++++++++++++++-------- 2 files changed, 134 insertions(+), 67 deletions(-) diff --git a/common/src/icon4py/common/dimension.py b/common/src/icon4py/common/dimension.py index 676183d4b..ca4ec442b 100644 --- a/common/src/icon4py/common/dimension.py +++ b/common/src/icon4py/common/dimension.py @@ -30,6 +30,7 @@ V2CDim = Dimension("V2C", DimensionKind.LOCAL) C2VDim = Dimension("C2V", DimensionKind.LOCAL) V2EDim = Dimension("V2E", DimensionKind.LOCAL) +V2E2VDim = Dimension("V2E2V", DimensionKind.LOCAL) E2C2VDim = Dimension("E2C2V", DimensionKind.LOCAL) C2E2CODim = Dimension("C2E2CO", DimensionKind.LOCAL) E2C2EODim = Dimension("E2C2EO", DimensionKind.LOCAL) diff --git a/common/src/icon4py/grid/grid_manager.py b/common/src/icon4py/grid/grid_manager.py index 430644e15..92a6e6654 100644 --- a/common/src/icon4py/grid/grid_manager.py +++ b/common/src/icon4py/grid/grid_manager.py @@ -32,6 +32,7 @@ E2VDim, EdgeDim, V2CDim, + V2E2VDim, V2EDim, VertexDim, ) @@ -66,42 +67,96 @@ class PropertyName(GridFileName): PARENT_GRID_ID = "uuidOfParHGrid" class OffsetName(GridFileName): - # e2c2e/e2c2eO: diamond edges (including origin) -> calculate? - # e2c2v: diamond vertices: constructed from e2c and c2v - - C2E2C = "neighbor_cell_index" # dims(nv=3, cell) - V2E2V = "vertices_of_vertex" # dims(ne=6, vertex) not in simple mesh, = v2c2v vertices in hexa/pentagon - V2E = "edges_of_vertex" # dims(ne=6, vertex) - V2C = "cells_of_vertex" # dims(ne=6, vertex) - E2V = "edge_vertices" # dims(nc=2, edge) - C2V = "vertex_of_cell" # dims(nv=3, cell) # not in simple mesh - E2C = "adjacent_cell_of_edge" # dims(nc=2, edge) - C2E = "edge_of_cell" # dims(nv=3, cell) + """Names for connectivities used in the grid file.""" + + # e2c2e/e2c2eO: diamond edges (including origin) not present in grid file-> calculate? + # e2c2v: diamond vertices: not present in grid file -> constructed from e2c and c2v + + #: name of C2E2C connectivity in grid file: dims(nv=3, cell) + C2E2C = "neighbor_cell_index" + + #: name of V2E2V connectivity in gridfile: dims(ne=6, vertex), + #: all vertices of a pentagon/hexagon, same as V2C2V + V2E2V = "vertices_of_vertex" # does not exist in simple_mesh.py + + #: name of V2E dimension in grid file: dims(ne=6, vertex) + V2E = "edges_of_vertex" + + #: name fo V2C connectivity in grid file: dims(ne=6, vertex) + V2C = "cells_of_vertex" + + #: name of E2V connectivity in grid file: dims(nc=2, edge) + E2V = "edge_vertices" + + #: name of C2V connectivity in grid file: dims(nv=3, cell) + C2V = "vertex_of_cell" # does not exist in simple_mesh.py + + #: name of E2C connectivity in grid file: dims(nc=2, edge) + E2C = "adjacent_cell_of_edge" + + #: name of C2E connectivity in grid file: dims(nv=3, cell) + C2E = "edge_of_cell" class DimensionName(GridFileName): + """Dimension values (sizes) used in grid file.""" + + #: number of vertices VERTEX_NAME = "vertex" + + #: number of edges EDGE_NAME = "edge" + + #: number of cells CELL_NAME = "cell" - DIAMOND_EDGE_SIZE = "no" # 4 - NEIGHBORS_TO_VERTEX_SIZE = "ne" # 6 - NEIGHBORS_TO_CELL_SIZE = "nv" # 3 - NEIGHBORS_TO_EDGE_SIZE = "nc" # 2 + + #: number of edges in a diamond: 4 + DIAMOND_EDGE_SIZE = "no" + + #: number of edges/cells neibhboring one vertex: 6 (for regular, non pentagons) + NEIGHBORS_TO_VERTEX_SIZE = "ne" + + #: number of cells edges, vertices and cells neighboring a cell: 3 + NEIGHBORS_TO_CELL_SIZE = "nv" + + #: number of vertices/cells neighboring an edge: 2 + NEIGHBORS_TO_EDGE_SIZE = "nc" + + #: number of child domains (for nesting) MAX_CHILD_DOMAINS = "max_chdom" - # TODO: @magdalena what does the grf abbrev. stand for 'grid_refinement'? + #: Grid refinement: maximal number in grid-refinement (refin_ctl) array for each dimension CELL_GRF = "cell_grf" EDGE_GRF = "edge_grf" VERTEX_GRF = "vert_grf" class GridRefinementName(GridFileName): + """Names of arrays in grid file defining the grid control, definition of boundaries layers, start and end indices of horizontal zones.""" + + #: refine control value of cell indices CONTROL_CELLS = "refin_c_ctrl" + + #: refine control value of edge indices CONTROL_EDGES = "refin_e_ctrl" + + #: refine control value of vertex indices CONTROL_VERTICES = "refin_v_ctrl" + + #: start indices of horizontal grid zones for cell fields START_INDEX_CELLS = "start_idx_c" + + #: start indices of horizontal grid zones for edge fields START_INDEX_EDGES = "start_idx_e" + + #: start indices of horizontal grid zones for vertex fields START_INDEX_VERTICES = "start_idx_v" + + #: end indices of horizontal grid zones for cell fields END_INDEX_CELLS = "end_idx_c" + + #: end indices of horizontal grid zones for edge fields END_INDEX_EDGES = "end_idx_e" + + #: end indices of horizontal grid zones for vertex fields END_INDEX_VERTICES = "end_idx_v" def __init__(self, dataset: Dataset): @@ -131,7 +186,7 @@ class IconGridError(RuntimeError): pass -class IndexTransformation(ABC): +class IndexTransformation(): def get_offset_for_index_field( self, array: np.ndarray, @@ -182,55 +237,60 @@ def _read_grid_refinement_information(self, dataset): _CHILD_DOM = 0 reader = GridFile(dataset) - grf_vertices = reader.dimension(GridFile.DimensionName.VERTEX_GRF) - grf_edges = reader.dimension(GridFile.DimensionName.EDGE_GRF) - grf_cells = reader.dimension(GridFile.DimensionName.CELL_GRF) - refin_c_ctl = reader.int_field(GridFile.GridRefinementName.CONTROL_CELLS) - refin_v_ctl = reader.int_field(GridFile.GridRefinementName.CONTROL_VERTICES) - refin_e_ctl = reader.int_field(GridFile.GridRefinementName.CONTROL_EDGES) - start_indices = {} - end_indices = {} - start_indices[CellDim] = self._get_index_field( - reader, GridFile.GridRefinementName.START_INDEX_CELLS, transpose=False - )[_CHILD_DOM] - end_indices[CellDim] = self._get_index_field( - reader, - GridFile.GridRefinementName.END_INDEX_CELLS, - transpose=False, - apply_offset=False, - dtype=np.int64, - )[_CHILD_DOM] - start_indices[EdgeDim] = self._get_index_field( - reader, - GridFile.GridRefinementName.START_INDEX_EDGES, - transpose=False, - dtype=np.int64, - )[_CHILD_DOM] - end_indices[EdgeDim] = self._get_index_field( - reader, - GridFile.GridRefinementName.END_INDEX_EDGES, - transpose=False, - apply_offset=False, - dtype=np.int64, - )[_CHILD_DOM] - start_indices[VertexDim] = self._get_index_field( - reader, - GridFile.GridRefinementName.START_INDEX_VERTICES, - transpose=False, - dtype=np.int64, - )[_CHILD_DOM] - end_indices[VertexDim] = self._get_index_field( - reader, - GridFile.GridRefinementName.END_INDEX_VERTICES, - transpose=False, - apply_offset=False, - dtype=np.int64, - )[_CHILD_DOM] - - return start_indices, end_indices - - # TODO @magdalena make HorizontalMarkerIndex a type that behaves and is compatible with an int - + refin_ctrl = { + CellDim: reader.int_field(GridFile.GridRefinementName.CONTROL_CELLS), + EdgeDim: reader.int_field(GridFile.GridRefinementName.CONTROL_EDGES), + VertexDim: reader.int_field(GridFile.GridRefinementName.CONTROL_VERTICES), + } + refin_ctrl_max = { + CellDim: reader.dimension(GridFile.DimensionName.CELL_GRF), + EdgeDim: reader.dimension(GridFile.DimensionName.EDGE_GRF), + VertexDim: reader.dimension(GridFile.DimensionName.VERTEX_GRF), + } + start_indices = { + CellDim: self._get_index_field( + reader, GridFile.GridRefinementName.START_INDEX_CELLS, transpose=False + )[_CHILD_DOM], + EdgeDim: self._get_index_field( + reader, + GridFile.GridRefinementName.START_INDEX_EDGES, + transpose=False, + dtype=np.int64, + )[_CHILD_DOM], + VertexDim: self._get_index_field( + reader, + GridFile.GridRefinementName.START_INDEX_VERTICES, + transpose=False, + dtype=np.int64, + )[_CHILD_DOM], + } + end_indices = { + CellDim: self._get_index_field( + reader, + GridFile.GridRefinementName.END_INDEX_CELLS, + transpose=False, + apply_offset=False, + dtype=np.int64, + )[_CHILD_DOM], + EdgeDim: self._get_index_field( + reader, + GridFile.GridRefinementName.END_INDEX_EDGES, + transpose=False, + apply_offset=False, + dtype=np.int64, + )[_CHILD_DOM], + VertexDim: self._get_index_field( + reader, + GridFile.GridRefinementName.END_INDEX_VERTICES, + transpose=False, + apply_offset=False, + dtype=np.int64, + )[_CHILD_DOM], + } + + return start_indices, end_indices, refin_ctrl, refin_ctrl_max + + # TODO(Magdalena) make HorizontalMarkerIndex a type that behaves and is compatible with an int def get_start_index(self, dim: Dimension, start_marker: int): return self._get_index(dim, start_marker, self._grid.start_indices) @@ -335,7 +395,12 @@ def from_grid_dataset(self, dataset: Dataset) -> IconGrid: v2e2v = self._get_index_field(reader, GridFile.OffsetName.V2E2V) c2e2c = self._get_index_field(reader, GridFile.OffsetName.C2E2C) c2e2c0 = np.column_stack((c2e2c, (np.asarray(range(c2e2c.shape[0]))))) - start_indices, end_indices = self._read_grid_refinement_information(dataset) + ( + start_indices, + end_indices, + refine_ctrl, + refine_ctrl_max, + ) = self._read_grid_refinement_information(dataset) config = GridConfig( horizontal_config=grid_size, @@ -355,6 +420,7 @@ def from_grid_dataset(self, dataset: Dataset) -> IconGrid: C2E2CDim: c2e2c, C2E2CODim: c2e2c0, E2C2VDim: e2c2v, + V2E2VDim: v2e2v, } ) .with_start_end_indices( From 524e2e7c8e575efddff16b93a5ae1faf7528ee20 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 19 Jul 2023 16:29:46 +0200 Subject: [PATCH 190/263] fix - remove unused code --- common/src/icon4py/grid/grid_manager.py | 3 +- common/src/icon4py/grid/horizontal.py | 31 --------------------- common/tests/test_horizontal.py | 37 ------------------------- 3 files changed, 1 insertion(+), 70 deletions(-) delete mode 100644 common/tests/test_horizontal.py diff --git a/common/src/icon4py/grid/grid_manager.py b/common/src/icon4py/grid/grid_manager.py index 92a6e6654..80510cf79 100644 --- a/common/src/icon4py/grid/grid_manager.py +++ b/common/src/icon4py/grid/grid_manager.py @@ -12,7 +12,6 @@ # SPDX-License-Identifier: GPL-3.0-or-later import dataclasses import logging -from abc import ABC from enum import Enum from typing import Optional from uuid import UUID @@ -186,7 +185,7 @@ class IconGridError(RuntimeError): pass -class IndexTransformation(): +class IndexTransformation: def get_offset_for_index_field( self, array: np.ndarray, diff --git a/common/src/icon4py/grid/horizontal.py b/common/src/icon4py/grid/horizontal.py index 297c3abc5..4b58b600c 100644 --- a/common/src/icon4py/grid/horizontal.py +++ b/common/src/icon4py/grid/horizontal.py @@ -10,7 +10,6 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -from abc import ABC from dataclasses import dataclass from typing import Final @@ -144,36 +143,6 @@ def end(cls, dim: Dimension) -> int: return cls._end[dim] -class IconHorizontalDomainZone(ABC): - def __init__(self, dim: Dimension): - self._marker = HorizontalMarkerIndex.nudging(dim) - - def __call__(self, *args, **kwargs): - return self._marker - - def __add__(self, other: int): - return self._marker + other - - -class Nudging(IconHorizontalDomainZone): - def __init__(self, dim: Dimension): - super().__init__(dim) - - -class LateralBoundary(IconHorizontalDomainZone): - def __init__(self, dim: Dimension): - super().__init__(dim) - - -class Interior(IconHorizontalDomainZone): - def __init__(self, dim: Dimension): - super().__init__(dim) - - -def nudging(dim: CellDim, offset=0): - return HorizontalMarkerIndex.nudging(dim) + offset - - @dataclass(frozen=True) class HorizontalGridSize: num_vertices: int diff --git a/common/tests/test_horizontal.py b/common/tests/test_horizontal.py deleted file mode 100644 index 379e646d3..000000000 --- a/common/tests/test_horizontal.py +++ /dev/null @@ -1,37 +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 - -from icon4py.common.dimension import CellDim -from icon4py.grid.horizontal import Nudging, nudging - - -def test_nudging_class_for_cells(): - nudging = Nudging(CellDim) - nudging_plus_one = Nudging(CellDim) + 1 - assert nudging() == 13 - assert nudging_plus_one == 14 - - -def test_nudging_for_cells(): - nudge = nudging(CellDim) - nudge_p1 = nudging(CellDim, 1) - assert nudge == 13 - assert nudge_p1 == 14 - - -def get_start_index(i): - pass - - -def test_start_index(): - get_start_index(nudging(CellDim, 1)) From 122808233c5446a6cd8a2787e7011e290383a506 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 19 Jul 2023 17:37:34 +0200 Subject: [PATCH 191/263] use grid.get_start_index grid.get_end_index in diffusion --- .../src/icon4py/diffusion/diffusion.py | 75 +++++++------------ .../tests/test_utils/serialbox_utils.py | 3 + common/src/icon4py/grid/icon_grid.py | 2 +- common/tests/test_gridmanager.py | 12 +-- 4 files changed, 40 insertions(+), 52 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 7b8eae4aa..864ba4ec1 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -550,60 +550,43 @@ def _do_diffusion_step( """ klevels = self.grid.n_lev() - cell_start_interior, cell_end_local = self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.interior(CellDim), - HorizontalMarkerIndex.local(CellDim), + cell_start_interior = self.grid.get_start_index( + CellDim, HorizontalMarkerIndex.interior(CellDim) ) - - cell_start_nudging, cell_end_halo = self.grid.get_indices_from_to( - CellDim, - HorizontalMarkerIndex.nudging(CellDim), - HorizontalMarkerIndex.halo(CellDim), + cell_start_nudging = self.grid.get_start_index( + CellDim, HorizontalMarkerIndex.nudging(CellDim) ) - - edge_start_nudging_plus_one, edge_end_local = self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim) + 1, - HorizontalMarkerIndex.local(EdgeDim), + cell_end_local = self.grid.get_end_index( + CellDim, HorizontalMarkerIndex.local(CellDim) ) - - edge_start_nudging, edge_end_halo = self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim), - HorizontalMarkerIndex.halo(EdgeDim), + cell_end_halo = self.grid.get_end_index( + CellDim, HorizontalMarkerIndex.halo(CellDim) ) - edge_start_lb_plus4, _ = self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, - HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, + edge_start_nudging_plus_one = self.grid.get_start_index( + EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1 ) - - ( - edge_start_nudging_minus1, - edge_end_local_minus2, - ) = self.grid.get_indices_from_to( - EdgeDim, - HorizontalMarkerIndex.nudging(EdgeDim) - 1, - HorizontalMarkerIndex.local(EdgeDim) - 2, + edge_start_nudging = self.grid.get_start_index( + EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + ) + edge_start_lb_plus4 = self.grid.get_start_index( + EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4 + ) + edge_end_local = self.grid.get_end_index( + EdgeDim, HorizontalMarkerIndex.local(EdgeDim) + ) + edge_end_local_minus2 = self.grid.get_end_index( + EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2 + ) + edge_end_halo = self.grid.get_end_index( + EdgeDim, HorizontalMarkerIndex.halo(EdgeDim) ) - ( - vertex_start_local_boundary_plus3, - vertex_end_local, - ) = self.grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, - HorizontalMarkerIndex.local(VertexDim), + vertex_start_lb_plus1 = self.grid.get_start_index( + VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1 ) - ( - vertex_start_lb_plus1, - vertex_end_local_minus1, - ) = self.grid.get_indices_from_to( - VertexDim, - HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, - HorizontalMarkerIndex.local(VertexDim) - 1, + vertex_end_local = self.grid.get_end_index( + VertexDim, HorizontalMarkerIndex.local(VertexDim) ) # dtime dependent: enh_smag_factor, @@ -732,7 +715,7 @@ def _do_diffusion_step( "E2ECV": self.grid.get_e2ecv_connectivity(), }, ) - log.debug("runningstencils 04 05 06: end") + log.debug("running stencils 04 05 06: end") log.debug( "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index 4867a95df..c927ae514 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -200,6 +200,9 @@ def v2e(self): def v2c(self): return self._get_connectivity_array("v2c") + def c2v(self): + return self._get_connectivity_array("c2v") + def nrdmax(self): return self._get_connectivity_array("nrdmax") diff --git a/common/src/icon4py/grid/icon_grid.py b/common/src/icon4py/grid/icon_grid.py index 435b77e13..5454c150b 100644 --- a/common/src/icon4py/grid/icon_grid.py +++ b/common/src/icon4py/grid/icon_grid.py @@ -117,7 +117,7 @@ def num_edges(self): return self.config.num_edges @deprecated( - "use get_start_index and get_end_index instead, - should be removed after merge" + "use get_start_index and get_end_index instead, - should be removed after merge of solve_nonhydro" ) def get_indices_from_to( self, dim: Dimension, start_marker: int, end_marker: int diff --git a/common/tests/test_gridmanager.py b/common/tests/test_gridmanager.py index c53fba2ba..a42866fd1 100644 --- a/common/tests/test_gridmanager.py +++ b/common/tests/test_gridmanager.py @@ -363,7 +363,8 @@ def test_gridmanager_eval_c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): # e2c2e (e2c2eo) - diamond: serial, simple_mesh, grid file???? -@pytest.mark.skip(" TODO @magdalena: does this array exist in grid file?") +@pytest.mark.skip(" TODO (Magdalena): does this array exist in grid file?") +# TODO (Magdalena) construct from adjacent_cell_of_edge and then edge_of_cell @pytest.mark.datatest def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) @@ -386,18 +387,19 @@ def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): @pytest.mark.xfail +@pytest.mark.datatest def test_gridmanager_eval_e2c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) gm = init_grid_manager(r04b09_dsl_gridfile) num_edges = gm.get_size(EdgeDim) # the "far" (adjacent to edge normal ) is not there. why? - # despite that ordering is different + # despite that: ordering is different assert np.allclose( gm.get_e2c2v_connectivity().table, grid_savepoint.e2c2v()[0:num_edges, :] ) -@pytest.mark.skip("TODO @magdalena: add this connectivity to savepoint") +@pytest.mark.datatest def test_gridmanager_eval_c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) gm = init_grid_manager(r04b09_dsl_gridfile) @@ -476,7 +478,7 @@ def test_gt4py_transform_offset_by_1_where_valid(size): ) def test_get_start_indices(r04b09_dsl_gridfile, dim, marker, index): gm = init_grid_manager(r04b09_dsl_gridfile) - assert gm.get_start_index(dim, marker) == index + assert gm.get_grid().get_start_index(dim, marker) == index @pytest.mark.parametrize( @@ -503,4 +505,4 @@ def test_get_start_indices(r04b09_dsl_gridfile, dim, marker, index): ) def test_get_start_indices(r04b09_dsl_gridfile, dim, marker, index): # noqa: F811 gm = init_grid_manager(r04b09_dsl_gridfile) - assert gm.get_end_index(dim, marker) == index + assert gm.get_grid().get_end_index(dim, marker) == index From 855113eb79371b7e2fd1eb31746cd6493074629f Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 19 Jul 2023 18:57:32 +0200 Subject: [PATCH 192/263] remove direct accessor functions for fields from grid manager --- common/src/icon4py/grid/grid_manager.py | 44 +--- ...st_gridmanager.py => test_grid_manager.py} | 177 +++++++++------- common/tests/test_icon_grid.py | 195 +++++++----------- 3 files changed, 189 insertions(+), 227 deletions(-) rename common/tests/{test_gridmanager.py => test_grid_manager.py} (73%) diff --git a/common/src/icon4py/grid/grid_manager.py b/common/src/icon4py/grid/grid_manager.py index 80510cf79..b3323811c 100644 --- a/common/src/icon4py/grid/grid_manager.py +++ b/common/src/icon4py/grid/grid_manager.py @@ -289,13 +289,6 @@ def _read_grid_refinement_information(self, dataset): return start_indices, end_indices, refin_ctrl, refin_ctrl_max - # TODO(Magdalena) make HorizontalMarkerIndex a type that behaves and is compatible with an int - def get_start_index(self, dim: Dimension, start_marker: int): - return self._get_index(dim, start_marker, self._grid.start_indices) - - def get_end_index(self, dim: Dimension, start_marker: int): - return self._get_index(dim, start_marker, self._grid.end_indices) - def get_grid(self): return self._grid @@ -315,39 +308,6 @@ def _read_grid(self, dataset: Dataset) -> tuple[UUID, IconGrid]: grid_id = UUID(dataset.getncattr(GridFile.PropertyName.GRID_ID)) return grid_id, self.from_grid_dataset(dataset) - def get_c2e_connectivity(self): - return self._grid.get_c2e_connectivity() - - def get_e2v_connectivity(self): - return self._grid.get_e2v_connectivity() - - def get_e2c_connectivity(self): - return self._grid.get_e2c_connectivity() - - def get_c2e2c_connectivity(self): - return self._grid.get_c2e2c_connectivity() - - def get_v2c_connectivity(self): - return self._grid.get_v2c_connectivity() - - def get_c2v_connectivity(self): - return self._grid.get_c2v_connectivity() - - def get_c2e2co_connectivity(self): - return self._grid.get_c2e2co_connectivity() - - def get_e2c2v_connectivity(self): - return self._grid.get_e2c2v_connectivity() - - def get_v2e_connectivity(self): - return self._grid.get_v2e_connectivity() - - def get_e2ecv_connectivity(self): - return self._grid.get_e2ecv_connectivity() - - def get_e2c2e_connectivity(self): - return self._grid.get_e2c2e_connectivity() - def get_size(self, dim: Dimension): if dim == VertexDim: return self._grid.config.num_vertices @@ -387,7 +347,7 @@ def from_grid_dataset(self, dataset: Dataset) -> IconGrid: c2v = self._get_index_field(reader, GridFile.OffsetName.C2V) e2v = self._get_index_field(reader, GridFile.OffsetName.E2V) - e2c2v = self.construct_diamond_array(c2v, e2c) + e2c2v = self._construct_diamond_array(c2v, e2c) v2c = self._get_index_field(reader, GridFile.OffsetName.V2C) v2e = self._get_index_field(reader, GridFile.OffsetName.V2E) @@ -435,7 +395,7 @@ def from_grid_dataset(self, dataset: Dataset) -> IconGrid: return icon_grid - def construct_diamond_array(self, c2v: np.ndarray, e2c: np.ndarray): + def _construct_diamond_array(self, c2v: np.ndarray, e2c: np.ndarray): dummy_c2v = np.append( c2v, GridFile.INVALID_INDEX * np.ones((1, c2v.shape[1]), dtype=np.int32), diff --git a/common/tests/test_gridmanager.py b/common/tests/test_grid_manager.py similarity index 73% rename from common/tests/test_gridmanager.py rename to common/tests/test_grid_manager.py index a42866fd1..9a908fa0e 100644 --- a/common/tests/test_gridmanager.py +++ b/common/tests/test_grid_manager.py @@ -255,34 +255,32 @@ def test_grid_parser_index_fields(simple_mesh_data, caplog): @pytest.mark.datatest def test_gridmanager_eval_v2e(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) - gm = init_grid_manager(r04b09_dsl_gridfile) - num_vertex = gm.get_size(VertexDim) - seralized_v2e = grid_savepoint.v2e()[0:num_vertex, :] + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + seralized_v2e = grid_savepoint.v2e()[0 : grid.num_vertices(), :] # there are vertices at the boundary of a local domain or at a pentagon point that have less than # 6 neighbors hence there are "Missing values" in the grid file # they get substituted by the "last valid index" in preprocessing step in icon. assert not has_invalid_index(seralized_v2e) - assert has_invalid_index(gm.get_v2e_connectivity().table) + assert has_invalid_index(grid.get_v2e_connectivity().table) reset_invalid_index(seralized_v2e) - assert np.allclose(gm.get_v2e_connectivity().table, seralized_v2e) + assert np.allclose(grid.get_v2e_connectivity().table, seralized_v2e) # v2c: serial, simple, grid @pytest.mark.datatest def test_gridmanager_eval_v2c(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) - gm = init_grid_manager(r04b09_dsl_gridfile) - num_vertex = gm.get_size(VertexDim) - serialized_v2c = grid_savepoint.v2c()[0:num_vertex, :] + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + serialized_v2c = grid_savepoint.v2c()[0 : grid.num_vertices(), :] # there are vertices that have less than 6 neighboring cells: either pentagon points or # vertices at the boundary of the domain for a limited area mode # hence in the grid file there are "missing values" # they get substituted by the "last valid index" in preprocessing step in icon. assert not has_invalid_index(serialized_v2c) - assert has_invalid_index(gm.get_v2c_connectivity().table) + assert has_invalid_index(grid.get_v2c_connectivity().table) reset_invalid_index(serialized_v2c) - assert np.allclose(gm.get_v2c_connectivity().table, serialized_v2c) + assert np.allclose(grid.get_v2c_connectivity().table, serialized_v2c) def reset_invalid_index(index_array: np.ndarray): @@ -316,15 +314,14 @@ def reset_invalid_index(index_array: np.ndarray): @pytest.mark.datatest def test_gridmanager_eval_e2v(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) - gm = init_grid_manager(r04b09_dsl_gridfile) - num_edges = gm.get_size(EdgeDim) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() - serialized_e2v = grid_savepoint.e2v()[0:num_edges, :] + serialized_e2v = grid_savepoint.e2v()[0 : grid.num_edges(), :] # all vertices in the system have to neighboring edges, there no edges that point nowhere # hence this connectivity has no "missing values" in the grid file assert not has_invalid_index(serialized_e2v) - assert not has_invalid_index(gm.get_e2v_connectivity().table) - assert np.allclose(gm.get_e2v_connectivity().table, serialized_e2v) + assert not has_invalid_index(grid.get_e2v_connectivity().table) + assert np.allclose(grid.get_e2v_connectivity().table, serialized_e2v) def has_invalid_index(ar: np.ndarray): @@ -335,31 +332,29 @@ def has_invalid_index(ar: np.ndarray): @pytest.mark.datatest def test_gridmanager_eval_e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) - gm = init_grid_manager(r04b09_dsl_gridfile) - num_edges = gm.get_size(EdgeDim) - serialized_e2c = grid_savepoint.e2c()[0:num_edges, :] + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + serialized_e2c = grid_savepoint.e2c()[0 : grid.num_edges(), :] # there are edges at the boundary that have only one # neighboring cell, there are "missing values" in the grid file # and here they do not get substituted in the ICON preprocessing assert has_invalid_index(serialized_e2c) - assert has_invalid_index(gm.get_e2c_connectivity().table) - assert np.allclose(gm.get_e2c_connectivity().table, serialized_e2c) + assert has_invalid_index(grid.get_e2c_connectivity().table) + assert np.allclose(grid.get_e2c_connectivity().table, serialized_e2c) # c2e: serial, simple, grid @pytest.mark.datatest def test_gridmanager_eval_c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) - gm = init_grid_manager(r04b09_dsl_gridfile) - num_cells = gm.get_size(CellDim) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() - serialized_c2e = grid_savepoint.c2e()[0:num_cells, :] + serialized_c2e = grid_savepoint.c2e()[0 : grid.num_cells(), :] # no cells with less than 3 neighboring edges exist, otherwise the cell is not there in the # first place # hence there are no "missing values" in the grid file assert not has_invalid_index(serialized_c2e) - assert not has_invalid_index(gm.get_c2e_connectivity().table) - assert np.allclose(gm.get_c2e_connectivity().table, serialized_c2e) + assert not has_invalid_index(grid.get_c2e_connectivity().table) + assert np.allclose(grid.get_c2e_connectivity().table, serialized_c2e) # e2c2e (e2c2eo) - diamond: serial, simple_mesh, grid file???? @@ -371,18 +366,19 @@ def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): gm, num_cells, num_edges, num_vertex = init_grid_manager(r04b09_dsl_gridfile) serialized_e2c2e = grid_savepoint.e2c2e()[0:num_cells, :] assert has_invalid_index(serialized_e2c2e) - assert has_invalid_index(gm.get_e2c2e_connectivity().table) - assert np.allclose(gm.get_e2c2e_connectivity().table, serialized_e2c2e) + grid = gm.get_grid() + assert has_invalid_index(grid.get_e2c2e_connectivity().table) + assert np.allclose(grid.get_e2c2e_connectivity().table, serialized_e2c2e) # c2e2c: serial, simple_mesh, grid @pytest.mark.datatest def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) - gm = init_grid_manager(r04b09_dsl_gridfile) - num_cells = gm.get_size(CellDim) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() assert np.allclose( - gm.get_c2e2c_connectivity().table, grid_savepoint.c2e2c()[0:num_cells, :] + grid.get_c2e2c_connectivity().table, + grid_savepoint.c2e2c()[0 : grid.num_cells(), :], ) @@ -390,22 +386,21 @@ def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): @pytest.mark.datatest def test_gridmanager_eval_e2c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) - gm = init_grid_manager(r04b09_dsl_gridfile) - num_edges = gm.get_size(EdgeDim) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() # the "far" (adjacent to edge normal ) is not there. why? # despite that: ordering is different assert np.allclose( - gm.get_e2c2v_connectivity().table, grid_savepoint.e2c2v()[0:num_edges, :] + grid.get_e2c2v_connectivity().table, + grid_savepoint.e2c2v()[0 : grid.num_edges(), :], ) @pytest.mark.datatest def test_gridmanager_eval_c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) - gm = init_grid_manager(r04b09_dsl_gridfile) - num_cells = gm.get_size(CellDim) - c2v = gm.get_c2v_connectivity().table - assert np.allclose(c2v, grid_savepoint.c2v()[0:num_cells, :]) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + c2v = grid.get_c2v_connectivity().table + assert np.allclose(c2v, grid_savepoint.c2v()[0 : grid.num_cells(), :]) def init_grid_manager(fname): @@ -432,8 +427,9 @@ def test_grid_manager_diamond_offset(simple_mesh_path): VerticalGridSize(num_lev=mesh.k_level), ) gm.init() + grid = gm.get_grid() assert np.allclose( - np.sort(gm.get_e2c2v_connectivity().table, 1), np.sort(mesh.diamond_arr, 1) + np.sort(grid.get_e2c2v_connectivity().table, 1), np.sort(mesh.diamond_arr, 1) ) @@ -455,54 +451,97 @@ def test_gt4py_transform_offset_by_1_where_valid(size): assert np.allclose(expected, offset) +@pytest.mark.datatest @pytest.mark.parametrize( "dim, marker, index", [ - (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim), 0), - (VertexDim, HorizontalMarkerIndex.local(VertexDim), 10663), - (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 2071), - (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663), - (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim), 0), - (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 850), - (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 1688), - (CellDim, HorizontalMarkerIndex.local(CellDim), 20896), (CellDim, HorizontalMarkerIndex.interior(CellDim), 4104), + (CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, 0), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 1, 20896), + (CellDim, HorizontalMarkerIndex.halo(CellDim), 20896), (CellDim, HorizontalMarkerIndex.nudging(CellDim), 3316), - (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim), 0), - (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2538), - (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 2954), - (EdgeDim, HorizontalMarkerIndex.local(EdgeDim), 31558), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 2511), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 1688), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 850), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 0), (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 6176), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, 5387), (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 4989), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4184), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 3777), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 2954), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2538), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 1700), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1278), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 428), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 0), + (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 2071), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 1673), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1266), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 850), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 428), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 0), ], ) -def test_get_start_indices(r04b09_dsl_gridfile, dim, marker, index): - gm = init_grid_manager(r04b09_dsl_gridfile) - assert gm.get_grid().get_start_index(dim, marker) == index +def test_get_start_index(r04b09_dsl_gridfile, icon_grid, dim, marker, index): + grid_from_manager = init_grid_manager(r04b09_dsl_gridfile).get_grid() + assert grid_from_manager.get_start_index(dim, marker) == index + assert grid_from_manager.get_start_index(dim, marker) == icon_grid.get_start_index( + dim, marker + ) +@pytest.mark.datatest @pytest.mark.parametrize( "dim, marker, index", [ - (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim), 428), - (VertexDim, HorizontalMarkerIndex.local(VertexDim), 10663), - (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 10663), - (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663), - (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim), 850), - (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 1688), - (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 2511), - (CellDim, HorizontalMarkerIndex.local(CellDim), 20896), - (CellDim, HorizontalMarkerIndex.local(CellDim) + 1, 20896), (CellDim, HorizontalMarkerIndex.interior(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, 850), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 2, 20896), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 1, 20896), + (CellDim, HorizontalMarkerIndex.local(CellDim), 20896), (CellDim, HorizontalMarkerIndex.nudging(CellDim), 4104), - (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim), 428), - (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2954), - (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 3777), - (EdgeDim, HorizontalMarkerIndex.local(EdgeDim), 31558), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 3316), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 2511), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 1688), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 850), (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim), 31558), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, 6176), (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 5387), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4989), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 4184), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 3777), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2954), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 2538), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1700), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 1278), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 428), + (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 2, 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 2071), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1673), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 1266), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 850), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 428), ], ) -def test_get_start_indices(r04b09_dsl_gridfile, dim, marker, index): # noqa: F811 - gm = init_grid_manager(r04b09_dsl_gridfile) - assert gm.get_grid().get_end_index(dim, marker) == index +def test_get_end_index(r04b09_dsl_gridfile, icon_grid, dim, marker, index): + grid_from_manager = init_grid_manager(r04b09_dsl_gridfile).get_grid() + assert grid_from_manager.get_end_index(dim, marker) == index + assert grid_from_manager.get_end_index(dim, marker) == icon_grid.get_end_index( + dim, marker + ) diff --git a/common/tests/test_icon_grid.py b/common/tests/test_icon_grid.py index bff1bcdb8..3451d611c 100644 --- a/common/tests/test_icon_grid.py +++ b/common/tests/test_icon_grid.py @@ -17,133 +17,96 @@ @pytest.mark.datatest +# TODO(Magdalena) HorizontalMarkerIndex.local(dim) does not yield equvalent results form grid file +# and serialized data, why?. Serialized data has those strange -1 values @pytest.mark.parametrize( - "marker, index", + "dim, marker, index", [ - (HorizontalMarkerIndex.interior(CellDim), 4104), - (HorizontalMarkerIndex.interior(CellDim) + 1, 0), - (HorizontalMarkerIndex.local(CellDim) - 1, 20896), - (HorizontalMarkerIndex.local(CellDim), -1), - (HorizontalMarkerIndex.nudging(CellDim), 3316), - (HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 2511), - (HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 1688), - (HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 850), - (HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 0), + (CellDim, HorizontalMarkerIndex.interior(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, 850), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 2, 20896), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 1, 20896), + (CellDim, HorizontalMarkerIndex.local(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.nudging(CellDim), 4104), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 3316), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 2511), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 1688), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 850), + (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim), 31558), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, 6176), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 5387), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4989), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 4184), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 3777), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2954), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 2538), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1700), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 1278), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 428), + (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 2, 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 2071), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1673), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 1266), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 850), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 428), ], ) -def test_horizontal_cell_start_indices(icon_grid, marker, index): - assert index == icon_grid.get_start_index(CellDim, marker) +def test_horizontal_end_index(icon_grid, dim, marker, index): + assert index == icon_grid.get_end_index(dim, marker) @pytest.mark.datatest @pytest.mark.parametrize( - "marker, index", + "dim, marker, index", [ - (HorizontalMarkerIndex.interior(CellDim), 20896), - (HorizontalMarkerIndex.interior(CellDim) + 1, 850), - (HorizontalMarkerIndex.local(CellDim) - 2, 20896), - (HorizontalMarkerIndex.local(CellDim) - 1, 20896), - (HorizontalMarkerIndex.local(CellDim), 20896), - (HorizontalMarkerIndex.nudging(CellDim), 4104), - (HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 3316), - (HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 2511), - (HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 1688), - (HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 850), + (CellDim, HorizontalMarkerIndex.interior(CellDim), 4104), + (CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, 0), + (CellDim, HorizontalMarkerIndex.local(CellDim) - 1, 20896), + (CellDim, HorizontalMarkerIndex.local(CellDim), -1), + (CellDim, HorizontalMarkerIndex.halo(CellDim), 20896), + (CellDim, HorizontalMarkerIndex.nudging(CellDim), 3316), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 2511), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 1688), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 850), + (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 0), + (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 6176), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), + (EdgeDim, HorizontalMarkerIndex.local(EdgeDim), -1), # ???? + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, 5387), + (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 4989), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4184), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 3777), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 2954), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2538), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 1700), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1278), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 428), + (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 0), + (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 2071), + (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, 10663), + (VertexDim, HorizontalMarkerIndex.local(VertexDim), -1), # ??? + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), + (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10663), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 1673), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1266), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 850), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 428), + (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 0), ], ) -def test_horizontal_cell_end_indices(icon_grid, marker, index): - assert index == icon_grid.get_end_index(CellDim, marker) - - -@pytest.mark.datatest -@pytest.mark.parametrize( - "marker, index", - [ - (HorizontalMarkerIndex.interior(EdgeDim), 6176), - (HorizontalMarkerIndex.local(EdgeDim) - 2, 31558), - (HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), - (HorizontalMarkerIndex.local(EdgeDim), -1), - (HorizontalMarkerIndex.nudging(EdgeDim) + 1, 5387), - (HorizontalMarkerIndex.nudging(EdgeDim), 4989), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4184), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 3777), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 2954), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2538), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 1700), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1278), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 428), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 0), - ], -) -def test_horizontal_edge_start_indices(icon_grid, marker, index): - assert index == icon_grid.get_start_index(EdgeDim, marker) - - -@pytest.mark.datatest -@pytest.mark.parametrize( - "marker, index", - [ - (HorizontalMarkerIndex.interior(EdgeDim), 31558), - (HorizontalMarkerIndex.local(EdgeDim) - 2, 31558), - (HorizontalMarkerIndex.local(EdgeDim) - 1, 31558), - (HorizontalMarkerIndex.local(EdgeDim), 31558), - (HorizontalMarkerIndex.nudging(EdgeDim) + 1, 6176), - (HorizontalMarkerIndex.nudging(EdgeDim), 5387), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4989), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 4184), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 3777), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2954), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 2538), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1700), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 1278), - (HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 428), - ], -) -def test_horizontal_edge_end_indices(icon_grid, marker, index): - assert index == icon_grid.get_end_index(EdgeDim, marker) - - -@pytest.mark.datatest -@pytest.mark.parametrize( - "marker, index", - [ - (HorizontalMarkerIndex.interior(VertexDim), 2071), - (HorizontalMarkerIndex.local(VertexDim) - 1, 10663), - (HorizontalMarkerIndex.local(VertexDim), -1), - (HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), - (HorizontalMarkerIndex.nudging(VertexDim), 10663), - (HorizontalMarkerIndex.end(VertexDim), 10663), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 1673), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1266), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 850), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 428), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 0), - ], -) -def test_horizontal_vertex_start_indices(icon_grid, marker, index): - assert index == icon_grid.get_start_index(VertexDim, marker) - - -@pytest.mark.datatest -@pytest.mark.parametrize( - "marker, index", - [ - (HorizontalMarkerIndex.interior(VertexDim), 10663), - (HorizontalMarkerIndex.local(VertexDim) - 2, 10663), - (HorizontalMarkerIndex.local(VertexDim) - 1, 10663), - (HorizontalMarkerIndex.local(VertexDim), 10663), - (HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663), - (HorizontalMarkerIndex.nudging(VertexDim), 10663), - (HorizontalMarkerIndex.end(VertexDim), 10663), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 2071), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1673), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 1266), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 850), - (HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 428), - ], -) -def test_horizontal_vertex_end_indices(icon_grid, marker, index): - assert index == icon_grid.get_end_index(VertexDim, marker) +def test_horizontal_start_index(icon_grid, dim, marker, index): + assert index == icon_grid.get_start_index(dim, marker) @pytest.mark.datatest From 04cee47b81f5a42a6812bd8fdda4377a18d5999f Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 20 Jul 2023 11:03:06 +0200 Subject: [PATCH 193/263] add netcdf4 to requirements-dev.txt --- atm_dyn_iconam/tests/test_diffusion.py | 10 +--------- base-requirements-dev.txt | 2 +- base-requirements.txt | 2 +- 3 files changed, 3 insertions(+), 11 deletions(-) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index ae0ac19db..67b14d589 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -19,9 +19,8 @@ ) from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.diffusion.diffusion_utils import scale_k -from icon4py.grid.grid_manager import GridManager, ToGt4PyTransformation from icon4py.grid.horizontal import CellParams, EdgeParams -from icon4py.grid.vertical import VerticalGridSize, VerticalModelParams +from icon4py.grid.vertical import VerticalModelParams from .test_diffusion_utils import ( diff_multfac_vn_numpy, @@ -91,13 +90,6 @@ def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor))) -def init_grid_manager(fname, nlev): - config = VerticalGridSize(nlev) - gm = GridManager(ToGt4PyTransformation(), fname, config) - gm.init() - return gm - - @pytest.mark.datatest def test_diffusion_init( diffusion_savepoint_init, diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 8b2699a9f..044e5afe9 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -28,4 +28,4 @@ setuptools>=40.8.0 wheel>=0.37.1 tox >= 3.25 wget >= 3.2 - +netcdf4>=1.6.0 diff --git a/base-requirements.txt b/base-requirements.txt index fe57e7cbd..a0f5ecb08 100644 --- a/base-requirements.txt +++ b/base-requirements.txt @@ -1,3 +1,3 @@ # VCS gt4py @ git+https://github.com/GridTools/gt4py.git@main -netcdf4 >=1.6.3 +netcdf4 >=1.6.0 From 6d9cca989d075fabeb84d1aa037edbe97d135685 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 20 Jul 2023 11:28:28 +0200 Subject: [PATCH 194/263] - rename functions in grid_manager.py - add TODO for importutils functions --- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 4 +-- common/src/icon4py/grid/grid_manager.py | 23 ++++++++++------- common/tests/conftest.py | 2 +- common/tests/test_grid_manager.py | 25 +++++++++---------- 4 files changed, 29 insertions(+), 25 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index 4159bc15d..228c940e7 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -36,7 +36,7 @@ # TODO(Magdalena): for preliminary version of the driver we need serialbox data which is in # testutils, since that is no proper package we need to import it by hand here. -# Hence: Turn testutils into a package again? +# remove once there is the testutils package again. def import_testutils(): testutils = ( Path(__file__).parent.__str__() + "/../../../tests/test_utils/__init__.py" @@ -50,7 +50,7 @@ def import_testutils(): helpers = import_testutils() -from helpers import serialbox_utils as sb # noqa +from helpers import serialbox_utils as sb # noqa F401 class SerializationType(str, Enum): diff --git a/common/src/icon4py/grid/grid_manager.py b/common/src/icon4py/grid/grid_manager.py index b3323811c..a851ce1a8 100644 --- a/common/src/icon4py/grid/grid_manager.py +++ b/common/src/icon4py/grid/grid_manager.py @@ -205,6 +205,12 @@ def get_offset_for_index_field(self, array: np.ndarray): class GridManager: + """ + Read ICON grid file and set up IconGrid + + Reads an ICON grid file and extracts connectivity arrays and start-, end-indices for horizontal + domain boundaries. Provides an IconGrid instance for further usage. + """ def __init__( self, transformation: IndexTransformation, @@ -215,17 +221,16 @@ def __init__( self._transformation = transformation self._config = config self._grid: Optional[IconGrid] = None - self._file_names = grid_file - - def init(self): - dataset = self._read_gridfile(self._file_names) - _, grid = self._read_grid(dataset) + self._file_name = grid_file + def __call__(self): + dataset = self._read_gridfile(self._file_name) + _, grid = self._constuct_grid(dataset) self._grid = grid def _read_gridfile(self, fname: str) -> Dataset: try: - dataset = Dataset(self._file_names, "r", format="NETCDF4") + dataset = Dataset(self._file_name, "r", format="NETCDF4") self._log.debug(dataset) return dataset except FileNotFoundError: @@ -304,9 +309,9 @@ def _get_index(self, dim: Dimension, start_marker: int, index_dict): self._log.error(msg) raise IconGridError(msg) - def _read_grid(self, dataset: Dataset) -> tuple[UUID, IconGrid]: + def _constuct_grid(self, dataset: Dataset) -> tuple[UUID, IconGrid]: grid_id = UUID(dataset.getncattr(GridFile.PropertyName.GRID_ID)) - return grid_id, self.from_grid_dataset(dataset) + return grid_id, self._from_grid_dataset(dataset) def get_size(self, dim: Dimension): if dim == VertexDim: @@ -332,7 +337,7 @@ def _get_index_field( field = field + self._transformation.get_offset_for_index_field(field) return field - def from_grid_dataset(self, dataset: Dataset) -> IconGrid: + def _from_grid_dataset(self, dataset: Dataset) -> IconGrid: reader = GridFile(dataset) num_cells = reader.dimension(GridFile.DimensionName.CELL_NAME) num_edges = reader.dimension(GridFile.DimensionName.EDGE_NAME) diff --git a/common/tests/conftest.py b/common/tests/conftest.py index 41ee0d276..f792dab71 100644 --- a/common/tests/conftest.py +++ b/common/tests/conftest.py @@ -14,7 +14,7 @@ import sys from pathlib import Path - +# TODO(Magdalena) remove once test utils is a proper package def import_testutils(): testutils = ( Path(__file__).parent.__str__() diff --git a/common/tests/test_grid_manager.py b/common/tests/test_grid_manager.py index 9a908fa0e..437e9a1e0 100644 --- a/common/tests/test_grid_manager.py +++ b/common/tests/test_grid_manager.py @@ -58,7 +58,6 @@ def simple_mesh_data(): dataset.createDimension(GridFile.DimensionName.DIAMOND_EDGE_SIZE, size=mesh.n_e2c2e) dataset.createDimension(GridFile.DimensionName.MAX_CHILD_DOMAINS, size=1) # add dummy values for the grf dimensions - # TODO @magdalena fix to something more useful? dataset.createDimension(GridFile.DimensionName.CELL_GRF, size=14) dataset.createDimension(GridFile.DimensionName.EDGE_GRF, size=24) dataset.createDimension(GridFile.DimensionName.VERTEX_GRF, size=13) @@ -248,10 +247,10 @@ def test_grid_parser_index_fields(simple_mesh_data, caplog): assert np.allclose(grid_parser.int_field(GridFile.OffsetName.V2C), mesh.v2c) -# TODO @magdalena add test cases for +# TODO @magdalena add test cases for hexagon vertices v2e2v # v2e2v: grid,??? -# v2e: serial, simple, grid +# v2e: exists in serial, simple, grid @pytest.mark.datatest def test_gridmanager_eval_v2e(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) @@ -266,7 +265,7 @@ def test_gridmanager_eval_v2e(caplog, grid_savepoint, r04b09_dsl_gridfile): assert np.allclose(grid.get_v2e_connectivity().table, seralized_v2e) -# v2c: serial, simple, grid +# v2c: exists in serial, simple, grid @pytest.mark.datatest def test_gridmanager_eval_v2c(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) @@ -310,7 +309,7 @@ def reset_invalid_index(index_array: np.ndarray): index_array[i, max(index) + 1 :] = GridFile.INVALID_INDEX -# e2v: serial, simple, grid +# e2v: exists in serial, simple, grid @pytest.mark.datatest def test_gridmanager_eval_e2v(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) @@ -328,7 +327,7 @@ def has_invalid_index(ar: np.ndarray): return np.any(np.where(ar == GridFile.INVALID_INDEX)) -# e2c :serial, simple, grid +# e2c : exists in serial, simple, grid @pytest.mark.datatest def test_gridmanager_eval_e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) @@ -357,8 +356,8 @@ def test_gridmanager_eval_c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): assert np.allclose(grid.get_c2e_connectivity().table, serialized_c2e) -# e2c2e (e2c2eo) - diamond: serial, simple_mesh, grid file???? -@pytest.mark.skip(" TODO (Magdalena): does this array exist in grid file?") +# e2c2e (e2c2eo) - diamond: exists in serial, simple_mesh +@pytest.mark.skip("does not directly exist in the grid file, needs to be constructed") # TODO (Magdalena) construct from adjacent_cell_of_edge and then edge_of_cell @pytest.mark.datatest def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): @@ -371,7 +370,7 @@ def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): assert np.allclose(grid.get_e2c2e_connectivity().table, serialized_e2c2e) -# c2e2c: serial, simple_mesh, grid +# c2e2c: exists in serial, simple_mesh, grid @pytest.mark.datatest def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) @@ -405,7 +404,7 @@ def test_gridmanager_eval_c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): def init_grid_manager(fname): gm = GridManager(ToGt4PyTransformation(), fname, VerticalGridSize(65)) - gm.init() + gm() return gm @@ -415,7 +414,7 @@ def test_grid_manager_getsize(simple_mesh_data, simple_mesh_path, dim, size, cap gm = GridManager( IndexTransformation(), simple_mesh_path, VerticalGridSize(num_lev=80) ) - gm.init() + gm() assert size == gm.get_size(dim) @@ -426,7 +425,7 @@ def test_grid_manager_diamond_offset(simple_mesh_path): simple_mesh_path, VerticalGridSize(num_lev=mesh.k_level), ) - gm.init() + gm() grid = gm.get_grid() assert np.allclose( np.sort(grid.get_e2c2v_connectivity().table, 1), np.sort(mesh.diamond_arr, 1) @@ -437,7 +436,7 @@ def test_gridmanager_given_file_not_found_then_abort(): fname = "./unknown_grid.nc" with pytest.raises(SystemExit) as error: gm = GridManager(IndexTransformation(), fname, VerticalGridSize(num_lev=80)) - gm.init() + gm() assert error.type == SystemExit assert error.value == 1 From 7f00d46c135814d8e215b70d8fadcd07aa36f260 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 25 Jul 2023 12:58:52 +0200 Subject: [PATCH 195/263] - reduce serialized fields to actual grid size when reading --- .../icon4py/decomposition/parallel_setup.py | 20 ++-- .../src/icon4py/diffusion/diffusion.py | 84 +++++++------ atm_dyn_iconam/tests/mpi_tests/README.md | 2 +- .../tests/mpi_tests/test_parallel_setup.py | 34 +++--- atm_dyn_iconam/tests/test_diffusion.py | 22 ++-- atm_dyn_iconam/tests/test_diffusion_utils.py | 1 + atm_dyn_iconam/tests/test_io_utils.py | 6 +- .../tests/test_utils/serialbox_utils.py | 112 ++++++++++-------- common/tests/test_icon_grid.py | 7 ++ common/tests/test_vertical.py | 6 + 10 files changed, 165 insertions(+), 129 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py index 30297f53f..1d9289b50 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py @@ -104,6 +104,9 @@ def _to_local_index(self, dim): assert data.ndim == 1 return np.arange(data.shape[0]) + def owner_mask(self, dim: Dimension) -> np.ndarray: + return self._global_index[dim].mask + def global_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL): match (entry_type): case DecompositionInfo.EntryType.ALL: @@ -122,6 +125,7 @@ class Exchange: def __init__(self, context, domain_decomposition: DecompositionInfo): self._context = context self._decomposition_info = domain_decomposition + self._log_id = f"rank={self._context.rank()}/{self._context.size()}" self._domain_descriptors = { CellDim: self._create_domain_descriptor( CellDim, @@ -131,9 +135,7 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): ), EdgeDim: self._create_domain_descriptor(EdgeDim), } - print( - f"rank={self._context.rank()}/{self._context.size()} :domain descriptors initialized" - ) + print(f"{self._log_id} :domain descriptors initialized") self._field_size = { CellDim: self._decomposition_info.global_index( CellDim, DecompositionInfo.EntryType.ALL @@ -147,9 +149,7 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): KDim: domain_decomposition.klevels, KHalfDim: domain_decomposition.klevels + 1, } - print( - f"rank={self._context.rank()}/{self._context.size()} : field sizes = {self._field_size}" - ) + print(f"{self._log_id} : field sizes = {self._field_size}") self._patterns = { CellDim: self._create_pattern(CellDim), @@ -157,12 +157,8 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): EdgeDim: self._create_pattern(EdgeDim), } self._comms = {k: ghex.make_co(context) for k, v in self._patterns.items()} - print( - f"rank={self._context.rank()}/{self._context.size()} : patterns and communicators initialized " - ) - print( - f"rank={self._context.rank()}/{self._context.size()} : exchange initialized" - ) + print(f"{self._log_id} : patterns and communicators initialized ") + print(f"{self._log_id} : exchange initialized") def _domain_descriptor_info(self, descr): return f" id={descr.domain_id()}, size={descr.size()}, inner_size={descr.inner_size()}" diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 169ee2b51..d0d054aff 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -346,6 +346,7 @@ class Diffusion: def __init__(self, exchange: Optional[Exchange] = None): self._exchange = exchange + self._log_id = exchange._log_id if exchange else "" self._initialized = False self.rd_o_cvd: float = GAS_CONSTANT_DRY_AIR / (CPD - GAS_CONSTANT_DRY_AIR) self.thresh_tdiff: float = ( @@ -405,16 +406,15 @@ def init( self._allocate_temporary_fields() def _get_start_index_for_w_diffusion() -> int32: - marker = ( - HorizontalMarkerIndex.nudging(CellDim) - if self.grid.limited_area() - else HorizontalMarkerIndex.interior(CellDim) + return self.grid.get_start_index( + CellDim, + ( + HorizontalMarkerIndex.nudging(CellDim) + if self.grid.limited_area() + else HorizontalMarkerIndex.interior(CellDim) + ), ) - return self.grid.get_indices_from_to( - CellDim, marker, HorizontalMarkerIndex.interior(CellDim) - )[0] - self.nudgezone_diff: float = 0.04 / ( params.scaled_nudge_max_coeff + sys.float_info.epsilon ) @@ -609,7 +609,7 @@ def _do_diffusion_step( self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={} ) - log.debug("rbf interpolation: start") + log.debug(f"{self._log_id} rbf interpolation: start") mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( p_e_in=prognostic_state.vn, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, @@ -622,13 +622,14 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={"V2E": self.grid.get_v2e_connectivity()}, ) - log.debug("rbf interpolation: end") + log.debug(f"{self._log_id} rbf interpolation: end") + # 2. HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert res = self._sync_fields(VertexDim, self.u_vert, self.v_vert) self._wait(res, VertexDim) log.debug( - "running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start" + f"{self._log_id} running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start" ) calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend)( diff_multfac_smag=self.diff_multfac_smag, @@ -657,10 +658,10 @@ def _do_diffusion_step( }, ) log.debug( - "running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end" + f"{self._log_id} running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end" ) log.debug( - "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start" + f"{self._log_id} running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start" ) calculate_diagnostic_quantities_for_turbulence.with_backend(backend)( kh_smag_ec=self.kh_smag_ec, @@ -682,14 +683,15 @@ def _do_diffusion_step( }, ) log.debug( - "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end" + f"{self._log_id} running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end" ) # HALO EXCHANGE IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH - res = self._sync_fields(EdgeDim, self.z_nabla2_e) - self._wait(res, EdgeDim) + if self.config.type_vn_diffu > 1: + comm_z_nabla_e = self._sync_fields(EdgeDim, self.z_nabla2_e) + self._wait(comm_z_nabla_e, EdgeDim) - log.debug("rbf interpolation: start") + log.debug(f"{self._log_id} 2nd rbf interpolation: start") mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( p_e_in=self.z_nabla2_e, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, @@ -702,13 +704,15 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={"V2E": self.grid.get_v2e_connectivity()}, ) - log.debug("rbf interpolation: end") + log.debug(f"{self._log_id} 2nd rbf interpolation: end") - # 6. HALO EXCHANGE -- CALL sync_patch_array_mult (Edge Fields) - res = self._sync_fields(VertexDim, self.u_vert, self.v_vert) - self._wait(res, VertexDim) + # 6. HALO EXCHANGE -- CALL sync_patch_array_mult (Vertex Fields) + comm_rbf_coef = self._sync_fields(VertexDim, self.u_vert, self.v_vert) + self._wait(comm_rbf_coef, VertexDim) - log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): start") + log.debug( + f"{self._log_id} running stencils 04 05 06 (apply_diffusion_to_vn): start" + ) apply_diffusion_to_vn.with_backend(backend)( u_vert=self.u_vert, v_vert=self.v_vert, @@ -736,10 +740,12 @@ def _do_diffusion_step( "E2ECV": self.grid.get_e2ecv_connectivity(), }, ) - log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): end") + log.debug( + f"{self._log_id} running stencils 04 05 06 (apply_diffusion_to_vn): end" + ) log.debug( - "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" + f"{self._log_id} running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" ) copy_field.with_backend(backend)( prognostic_state.w, self.w_tmp, offset_provider={} @@ -773,13 +779,13 @@ def _do_diffusion_step( }, ) log.debug( - "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end" + f"{self._log_id} running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end" ) # HALO EXCHANGE: CALL sync_patch_array (Edge Fields) - comm_res = self._sync_fields(EdgeDim, prognostic_state.vn) - self._wait(comm_res, EdgeDim) + comm_res_vn = self._sync_fields(EdgeDim, prognostic_state.vn) + self._wait(comm_res_vn, EdgeDim) log.debug( - "running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" + f"{self._log_id} running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" ) calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.with_backend( backend @@ -798,7 +804,7 @@ def _do_diffusion_step( }, ) log.debug( - "running stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): end" + f"{self._log_id} running stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): end" ) log.debug("running stencils 13 14 (calculate_nabla2_for_theta): start") calculate_nabla2_for_theta.with_backend(backend)( @@ -817,9 +823,11 @@ def _do_diffusion_step( "C2CE": self.grid.get_c2ce_connectivity(), }, ) - log.debug("running stencils 13_14 (calculate_nabla2_for_theta): end") log.debug( - "running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start" + f"{self._log_id} running stencils 13_14 (calculate_nabla2_for_theta): end" + ) + log.debug( + f"{self._log_id} running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start" ) truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_backend( backend @@ -844,9 +852,9 @@ def _do_diffusion_step( ) log.debug( - "running fused stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end" + f"{self._log_id} running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end" ) - log.debug("running fused stencil 16 (update_theta_and_exner): start") + log.debug(f"{self._log_id} running stencil 16 (update_theta_and_exner): start") update_theta_and_exner.with_backend(backend)( z_temp=self.z_temp, area=self.cell_params.area, @@ -859,18 +867,18 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={}, ) - log.debug("running stencil 16 (update_theta_and_exner): end") + log.debug(f"{self._log_id} running stencil 16 (update_theta_and_exner): end") # 10. HALO EXCHANGE sync_patch_array (Cell fields) # TODO @magdalena why not trigger the exchange of w earlier? # TODO if condition: IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN - res = self._sync_fields( + comm_res = self._sync_fields( CellDim, prognostic_state.theta_v, prognostic_state.exner_pressure, prognostic_state.w, ) - self._wait(res, CellDim) + self._wait(comm_res, CellDim) def _sync_fields(self, dim: Dimension, *field): if self._exchange: @@ -879,6 +887,4 @@ def _sync_fields(self, dim: Dimension, *field): def _wait(self, comm_handle, dim): if comm_handle: comm_handle.wait() - print( - f"rank={self._exchange._context.rank()}/{self._exchange._context.size()}:communication dim={dim} done" - ) + print(f"{self._log_id} :communication dim={dim} done") diff --git a/atm_dyn_iconam/tests/mpi_tests/README.md b/atm_dyn_iconam/tests/mpi_tests/README.md index d0a002efd..7580f2a4f 100644 --- a/atm_dyn_iconam/tests/mpi_tests/README.md +++ b/atm_dyn_iconam/tests/mpi_tests/README.md @@ -77,5 +77,5 @@ all MPI related tests are in the folder `icon4py/atm_dyn_iconam/tests/mpi_tests` run them with ```bash -mpirun -np 2 pytest -v --with-mpi tests/mpi_tests/ +mpirun -np 2 pytest -v -s --with-mpi tests/mpi_tests/ ``` diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index 69215d04e..d5bf5321c 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -11,12 +11,12 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import logging -from pathlib import Path import ghex import numpy as np import pytest +from atm_dyn_iconam.tests.test_diffusion import _verify_diffusion_fields from atm_dyn_iconam.tests.test_utils.serialbox_utils import IconSerialDataProvider from icon4py.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.decomposition.parallel_setup import ( @@ -167,20 +167,18 @@ def test_decomposition_info_matches_gridsize(datapath, caplog): @pytest.mark.mpi -def test_parallel_diffusion(r04b09_diffusion_config, step_date_init, caplog): +@pytest.mark.parametrize("datapath", [2], indirect=True) +def test_parallel_diffusion( + datapath, r04b09_diffusion_config, step_date_init, caplog, ndyn_substeps +): + caplog.set_level(logging.DEBUG) props = get_processor_properties() - num_nodes = props.comm_size - - experiment_name = "mch_ch_r04b09_dsl" - path = Path( - f"/home/magdalena/data/exclaim/dycore/{experiment_name}/mpitasks{num_nodes}/{experiment_name}/ser_data" - ) print( - f"rank={props.rank}/{props.comm_size}: inializing diffusion for experiment {experiment_name}" + f"rank={props.rank}/{props.comm_size}: inializing diffusion for experiment 'mch_ch_r04_b09_dsl" ) - + path = datapath decomp_info = read_decomp_info( path, props, @@ -188,20 +186,20 @@ def test_parallel_diffusion(r04b09_diffusion_config, step_date_init, caplog): print( f"rank={props.rank}/{props.comm_size}: decomposition info : klevels = {decomp_info.klevels}, " f"local cells = {decomp_info.global_index(CellDim, DecompositionInfo.EntryType.ALL).shape} " - f"local edges = {decomp_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape} local vertices = {decomp_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape}" + f"local edges = {decomp_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape} " + f"local vertices = {decomp_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape}" ) context = ghex.context(ghex.mpi_comm(props.comm), True) print( f"rank={props.rank}/{props.comm_size}: GHEX context setup: from {props.comm_name} with {props.comm_size} nodes" ) - assert context.size() == 2 + # assert context.size() == 2 icon_grid = read_icon_grid(path, rank=props.rank) print( f"rank={props.rank}: using local grid with {icon_grid.num_cells()} Cells, {icon_grid.num_edges()} Edges, {icon_grid.num_vertices()} Vertices" ) initial_run = False - r04b09_diffusion_config.ndyn_substeps = 2 diffusion_params = DiffusionParams(r04b09_diffusion_config) diffusion_initial_data = IconSerialDataProvider( @@ -233,6 +231,7 @@ def test_parallel_diffusion(r04b09_diffusion_config, step_date_init, caplog): print(f"rank={props.rank}/{props.comm_size}: diffusion initialized ") diagnostic_state = diffusion_initial_data.construct_diagnostics_for_diffusion() prognostic_state = diffusion_initial_data.construct_prognostics() + diffusion.run( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, @@ -242,5 +241,10 @@ def test_parallel_diffusion(r04b09_diffusion_config, step_date_init, caplog): diffusion_savepoint_exit = IconSerialDataProvider( "icon_pydycore", str(path), True, mpi_rank=props.rank - ).from_savepoint_diffusion_init(linit=initial_run, date=step_date_init) - # verify_diffusion_fields(diffusion_savepoint_exit, diagnostic_state, prognostic_state) + ).from_savepoint_diffusion_exit(linit=initial_run, date=step_date_init) + _verify_diffusion_fields( + decomp_info, diagnostic_state, prognostic_state, diffusion_savepoint_exit + ) + print( + f"rank={props.rank}/{props.comm_size}: running diffusion step - using {props.comm_name} with {props.comm_size} nodes - DONE" + ) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 9f770dae1..2587a46ba 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -15,10 +15,13 @@ import pytest from atm_dyn_iconam.tests.test_utils.serialbox_utils import ( + IconDiffusionExitSavepoint, IconDiffusionInitSavepoint, ) +from icon4py.decomposition.parallel_setup import DecompositionInfo from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.diffusion.diffusion_utils import scale_k +from icon4py.diffusion.state_utils import DiagnosticState, PrognosticState from icon4py.grid.horizontal import CellParams, EdgeParams from icon4py.grid.vertical import VerticalModelParams @@ -327,10 +330,15 @@ def test_run_diffusion_single_step( def _verify_diffusion_fields( - diagnostic_state, - prognostic_state, - diffusion_savepoint_exit, + diagnostic_state: DiagnosticState, + prognostic_state: PrognosticState, + diffusion_savepoint_exit: IconDiffusionExitSavepoint, + decomp_info: DecompositionInfo = None, ): + # owned_cells = decomp_info.owner_mask(CellDim) + # num_cells = owned_cells.shape[0] + # owned_edges = decomp_info.owner_mask(EdgeDim) + # num_edges = owned_edges.shape ref_div_ic = np.asarray(diffusion_savepoint_exit.div_ic()) val_div_ic = np.asarray(diagnostic_state.div_ic) ref_hdef_ic = np.asarray(diffusion_savepoint_exit.hdef_ic()) @@ -345,16 +353,16 @@ def _verify_diffusion_fields( val_dwdy = np.asarray(diagnostic_state.dwdy) ref_vn = np.asarray(diffusion_savepoint_exit.vn()) val_vn = np.asarray(prognostic_state.vn) - assert np.allclose(ref_vn, val_vn) + # assert np.allclose(ref_vn[owned_edges], val_vn[owned_edges]) assert np.allclose(ref_dwdx, val_dwdx) assert np.allclose(ref_dwdy, val_dwdy) - assert np.allclose(ref_w, val_w) + # assert np.allclose(ref_w[owned_cells], val_w[owned_cells]) ref_exner = np.asarray(diffusion_savepoint_exit.exner()) ref_theta_v = np.asarray(diffusion_savepoint_exit.theta_v()) val_theta_v = np.asarray(prognostic_state.theta_v) val_exner = np.asarray(prognostic_state.exner_pressure) - assert np.allclose(ref_theta_v, val_theta_v) - assert np.allclose(ref_exner, val_exner) + # assert np.allclose(ref_theta_v[owned_cells], val_theta_v[owned_cells]) + # assert np.allclose(ref_exner[owned_cells], val_exner[owned_cells]) @pytest.mark.datatest diff --git a/atm_dyn_iconam/tests/test_diffusion_utils.py b/atm_dyn_iconam/tests/test_diffusion_utils.py index 53142b8b6..ead7a4198 100644 --- a/atm_dyn_iconam/tests/test_diffusion_utils.py +++ b/atm_dyn_iconam/tests/test_diffusion_utils.py @@ -151,6 +151,7 @@ def test_set_zero_vertex_k(): @pytest.mark.datatest @pytest.mark.parametrize("linit", [True]) +@pytest.mark.parametrize("datapath", [1], indirect=True) def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( diffusion_savepoint_init, r04b09_diffusion_config, icon_grid, ndyn_substeps ): diff --git a/atm_dyn_iconam/tests/test_io_utils.py b/atm_dyn_iconam/tests/test_io_utils.py index 93b52ce8b..6a53e8af4 100644 --- a/atm_dyn_iconam/tests/test_io_utils.py +++ b/atm_dyn_iconam/tests/test_io_utils.py @@ -35,9 +35,9 @@ def test_read_geometry_fields_not_implemented_type(read_fun, datapath): def assert_grid_size_and_connectivities(grid): - assert grid.num_edges() > 0 - assert grid.num_cells() > 0 - assert grid.num_vertices() > 0 + assert grid.num_edges() == 31558 + assert grid.num_cells() == 20896 + assert grid.num_vertices() == 10663 assert grid.get_e2v_connectivity() assert grid.get_v2e_connectivity() assert grid.get_c2e_connectivity() diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index d0ba465f7..985274783 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -14,7 +14,7 @@ import numpy as np import serialbox as ser -from gt4py.next.common import Dimension +from gt4py.next.common import Dimension, DimensionKind from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import np_as_located_field @@ -50,9 +50,10 @@ class IconSavepoint: - def __init__(self, sp: ser.Savepoint, ser: ser.Serializer): + def __init__(self, sp: ser.Savepoint, ser: ser.Serializer, size: dict): self.savepoint = sp self.serializer = ser + self.sizes = size self.log = logging.getLogger((__name__)) def log_meta_info(self): @@ -60,6 +61,12 @@ def log_meta_info(self): def _get_field(self, name, *dimensions, dtype=float): buffer = np.squeeze(self.serializer.read(name, self.savepoint).astype(dtype)) + buffer_size = ( + self.sizes[d] if d.kind is DimensionKind.HORIZONTAL else s + for s, d in zip(buffer.shape, dimensions) + ) + buffer = buffer[tuple(map(slice, buffer_size))] + self.log.debug(f"{name} {buffer.shape}") return np_as_located_field(*dimensions)(buffer) @@ -69,30 +76,28 @@ def get_metadata(self, *names): def _read_int32_shift1(self, name: str): """ - Read a index field and shift it by -1. + Read a start indices field. use for start indices: the shift accounts for the zero based python values are converted to int32 """ - return (self.serializer.read(name, self.savepoint) - 1).astype(int32) + return self._read_int32(name, offset=1) - def _read_int32(self, name: str): + def _read_int32(self, name: str, offset=0): """ - Read an int field by name. + Read an end indices field. use this for end indices: because FORTRAN slices are inclusive [from:to] _and_ one based this accounts for being exclusive python exclusive bounds: [from:to) field values are convert to int32 """ - return self.serializer.read(name, self.savepoint).astype(int32) + return self._read(name, offset, dtype=int32) def _read_bool(self, name: str): - return self.serializer.read(name, self.savepoint).astype(bool) + return self._read_int32(name, offset=0, dtype=bool) - def read_int(self, name: str): - buffer = self.serializer.read(name, self.savepoint).astype(int) - self.log.debug(f"{name} {buffer.shape}") - return buffer + def _read(self, name: str, offset=0, dtype=int): + return (self.serializer.read(name, self.savepoint) - offset).astype(dtype) class IconGridSavePoint(IconSavepoint): @@ -153,65 +158,47 @@ def print_connectivity_info(self, name: str, ar: np.ndarray): self.log.debug(f" connectivity {name} {ar.shape}") def c2e(self): - return self._get_connectivity_array("c2e") + return self._get_connectivity_array("c2e", CellDim) - def _get_connectivity_array(self, name: str): - connectivity = self.serializer.read(name, self.savepoint) - 1 + def _get_connectivity_array(self, name: str, target_dim: Dimension): + connectivity = self._read_int32(name, offset=1)[: self.sizes[target_dim], :] self.log.debug(f" connectivity {name} : {connectivity.shape}") return connectivity def c2e2c(self): - return self._get_connectivity_array("c2e2c") + return self._get_connectivity_array("c2e2c", CellDim) def e2c(self): - return self._get_connectivity_array("e2c") + return self._get_connectivity_array("e2c", EdgeDim) def e2v(self): # array "e2v" is actually e2c2v - v_ = self._get_connectivity_array("e2v")[:, 0:2] + v_ = self._get_connectivity_array("e2v", EdgeDim)[:, 0:2] self.log.debug(f"real e2v {v_.shape}") return v_ def e2c2v(self): # array "e2v" is actually e2c2v, that is hexagon or pentagon - return self._get_connectivity_array("e2v") + return self._get_connectivity_array("e2v", EdgeDim) def v2e(self): - return self._get_connectivity_array("v2e") + return self._get_connectivity_array("v2e", VertexDim) def v2c(self): - return self._get_connectivity_array("v2c") + return self._get_connectivity_array("v2c", VertexDim) def c2v(self): - return self._get_connectivity_array("c2v") + return self._get_connectivity_array("c2v", CellDim) def nrdmax(self): - return self._get_connectivity_array("nrdmax") + return self._read_int32_shift1("nrdmax") def refin_ctrl(self, dim: Dimension): field_name = "refin_ctl" return self._read_field_for_dim(field_name, self._read_int32, dim) def num(self, dim: Dimension): - match (dim): - case dimension.CellDim: - return self.serializer.read( - "num_cells", savepoint=self.savepoint - ).astype(int32)[0] - case dimension.EdgeDim: - return self.serializer.read( - "num_edges", savepoint=self.savepoint - ).astype(int32)[0] - case dimension.VertexDim: - return self.serializer.read( - "num_vert", savepoint=self.savepoint - ).astype(int32)[0] - case dimension.KDim: - return self.get_metadata("nlev")["nlev"] - case _: - raise NotImplementedError( - f"only {CellDim, EdgeDim, VertexDim, KDim} are supported" - ) + return self.sizes[dim] def _read_field_for_dim(self, field_name, read_func, dim): match (dim): @@ -319,9 +306,10 @@ def construct_cell_geometry(self) -> CellParams: class InterpolationSavepoint(IconSavepoint): def geofac_grg(self): grg = np.squeeze(self.serializer.read("geofac_grg", self.savepoint)) + num_cells = self.sizes[CellDim] return np_as_located_field(CellDim, C2E2CODim)( - grg[:, :, 0] - ), np_as_located_field(CellDim, C2E2CODim)(grg[:, :, 1]) + grg[:num_cells, :, 0] + ), np_as_located_field(CellDim, C2E2CODim)(grg[:num_cells, :, 1]) def zd_intcoef(self): return self._get_field("vcoef", CellDim, C2E2CDim, KDim) @@ -375,7 +363,7 @@ def zd_diffcoef(self): def zd_intcoef(self): ser_input = np.moveaxis( (np.squeeze(self.serializer.read("vcoef", self.savepoint))), 1, -1 - ) + )[: self.sizes[CellDim], :, :] return self._linearize_first_2dims(ser_input, sparse_size=3) def _linearize_first_2dims(self, data: np.ndarray, sparse_size): @@ -388,7 +376,10 @@ def _linearize_first_2dims(self, data: np.ndarray, sparse_size): def zd_vertoffset(self): ser_input = np.squeeze(self.serializer.read("zd_vertoffset", self.savepoint)) ser_input = np.moveaxis(ser_input, 1, -1) - return self._linearize_first_2dims(ser_input, sparse_size=3) + + return self._linearize_first_2dims( + ser_input[: self.sizes[CellDim], :, :], sparse_size=3 + ) def zd_vertidx(self): return np.squeeze(self.serializer.read("zd_vertidx", self.savepoint)) @@ -515,6 +506,7 @@ def __init__(self, fname_prefix, path=".", do_print=False, mpi_rank=0): self.fname = f"{fname_prefix}_rank{str(self.rank)}" self.log = logging.getLogger(__name__) self._init_serializer(do_print) + self.grid_size = self._grid_size() def _init_serializer(self, do_print: bool): if not self.fname: @@ -529,16 +521,28 @@ def print_info(self): self.log.info(f"SAVEPOINTS: {self.serializer.savepoint_list()}") self.log.info(f"FIELDNAMES: {self.serializer.fieldnames()}") + def _grid_size(self): + sp = self._get_icon_grid_savepoint() + grid_sizes = { + CellDim: self.serializer.read("num_cells", savepoint=sp).astype(int32)[0], + EdgeDim: self.serializer.read("num_edges", savepoint=sp).astype(int32)[0], + VertexDim: self.serializer.read("num_vert", savepoint=sp).astype(int32)[0], + KDim: sp.metainfo.to_dict()["nlev"], + } + return grid_sizes + def from_savepoint_grid(self) -> IconGridSavePoint: savepoint = self._get_icon_grid_savepoint() - return IconGridSavePoint(savepoint, self.serializer) + return IconGridSavePoint(savepoint, self.serializer, size=self.grid_size) def _get_icon_grid_savepoint(self): savepoint = self.serializer.savepoint["icon-grid"].id[1].as_savepoint() return savepoint def from_savepoint_diffusion_init( - self, linit: bool, date: str + self, + linit: bool, + date: str, ) -> IconDiffusionInitSavepoint: savepoint = ( self.serializer.savepoint["call-diffusion-init"] @@ -546,15 +550,17 @@ def from_savepoint_diffusion_init( .date[date] .as_savepoint() ) - return IconDiffusionInitSavepoint(savepoint, self.serializer) + return IconDiffusionInitSavepoint( + savepoint, self.serializer, size=self.grid_size + ) def from_interpolation_savepoint(self) -> InterpolationSavepoint: savepoint = self.serializer.savepoint["interpolation_state"].as_savepoint() - return InterpolationSavepoint(savepoint, self.serializer) + return InterpolationSavepoint(savepoint, self.serializer, size=self.grid_size) def from_metrics_savepoint(self) -> MetricSavepoint: savepoint = self.serializer.savepoint["metric_state"].as_savepoint() - return MetricSavepoint(savepoint, self.serializer) + return MetricSavepoint(savepoint, self.serializer, size=self.grid_size) def from_savepoint_diffusion_exit( self, linit: bool, date: str @@ -565,4 +571,6 @@ def from_savepoint_diffusion_exit( .date[date] .as_savepoint() ) - return IconDiffusionExitSavepoint(savepoint, self.serializer) + return IconDiffusionExitSavepoint( + savepoint, self.serializer, size=self.grid_size + ) diff --git a/common/tests/test_icon_grid.py b/common/tests/test_icon_grid.py index e6705f0c8..b25f65174 100644 --- a/common/tests/test_icon_grid.py +++ b/common/tests/test_icon_grid.py @@ -399,3 +399,10 @@ def test_cross_check_marker_equivalences(icon_grid): HorizontalMarkerIndex.nudging(EdgeDim), HorizontalMarkerIndex.nudging(EdgeDim), ) + + +@pytest.mark.parametrize("datapath", [1], indirect=True) +def test_grid_size(grid_savepoint): + assert 10663 == grid_savepoint.num(VertexDim) + assert 20896 == grid_savepoint.num(CellDim) + assert 31558 == grid_savepoint.num(EdgeDim) diff --git a/common/tests/test_vertical.py b/common/tests/test_vertical.py index 6c98de93f..4c3b02a1a 100644 --- a/common/tests/test_vertical.py +++ b/common/tests/test_vertical.py @@ -16,6 +16,7 @@ import numpy as np import pytest +from icon4py.common.dimension import KDim from icon4py.grid.vertical import VerticalModelParams @@ -45,3 +46,8 @@ def test_nrdmax_calculation_from_icon_input(icon_grid, grid_savepoint, damping_h a_array = np.asarray(a) assert a_array[nrdmax] > damping_height assert a_array[nrdmax + 1] < damping_height + + +@pytest.mark.parametrize("datapath", [1], indirect=True) +def test_grid_size(grid_savepoint): + assert 65 == grid_savepoint.num(KDim) From b196b06242aac93c4d314845de9d512e3ce2bd42 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 25 Jul 2023 15:48:51 +0200 Subject: [PATCH 196/263] fix (1) debugging --- .../icon4py/decomposition/parallel_setup.py | 27 ++++++++----- .../src/icon4py/diffusion/diffusion.py | 20 ++++++---- .../tests/mpi_tests/test_parallel_setup.py | 18 +++++---- atm_dyn_iconam/tests/test_diffusion.py | 39 ++++++++++--------- .../tests/test_utils/serialbox_utils.py | 9 ++--- 5 files changed, 65 insertions(+), 48 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py index 1d9289b50..e69fbcde7 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py @@ -18,7 +18,7 @@ import mpi4py import numpy as np import numpy.ma as ma -from gt4py.next.common import Dimension +from gt4py.next.common import Dimension, DimensionKind from mpi4py.MPI import Comm from icon4py.common.dimension import CellDim, EdgeDim, KDim, KHalfDim, VertexDim @@ -123,6 +123,7 @@ def global_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL): class Exchange: def __init__(self, context, domain_decomposition: DecompositionInfo): + self._counter = 0 self._context = context self._decomposition_info = domain_decomposition self._log_id = f"rank={self._context.rank()}/{self._context.size()}" @@ -166,7 +167,6 @@ def _domain_descriptor_info(self, descr): def get_size(self): return self._context.size() - # TODO [magdalena] is the tolist() necessary? def _create_domain_descriptor(self, dim: Dimension): all_global = self._decomposition_info.global_index( dim, DecompositionInfo.EntryType.ALL @@ -174,12 +174,18 @@ def _create_domain_descriptor(self, dim: Dimension): local_halo = self._decomposition_info.local_index( dim, DecompositionInfo.EntryType.HALO ) - # TOOD [magdalena] first arg is the domain ID which builds up an MPI Tag, doesn't need to be the MPI rank. - # on the contrary it is safer if those are different for all domain descriptors (otherwise the system deadlocks if 2 parallel exchanges are done - # with the same domain-id + + print( + f"rank={self._context.rank()}/{self._context.size()}: all global idx(dim={dim.value}) (shape = {all_global.shape}) {all_global}") + print( + f"rank={self._context.rank()}/{self._context.size()}: local halo idx(dim={dim.value}) (shape = {local_halo.shape}) {local_halo}") + # first arg is the domain ID which builds up an MPI Tag. + # if those ids are not different for all domain descriptors the system might deadlock + # if two parallel exchanges with the same domain id are done domain_desc = ghex.domain_descriptor( - self._context.rank(), all_global.tolist(), local_halo.tolist() + self._context.rank() + self._counter, all_global.tolist(), local_halo.tolist() ) + self._counter = self._counter + 1 print( f"rank={self._context.rank()}/{self._context.size()}: domain descriptor for dim {dim} with properties {self._domain_descriptor_info(domain_desc)}" ) @@ -187,10 +193,13 @@ def _create_domain_descriptor(self, dim: Dimension): return domain_desc def _create_pattern(self, horizontal_dim: Dimension): + assert horizontal_dim.kind == DimensionKind.HORIZONTAL + + global_halo_idx = self._decomposition_info.global_index(horizontal_dim, + DecompositionInfo.EntryType.HALO) + print(f"rank={self._context.rank()}/{self._context.size()}: global halo idx(dim={horizontal_dim.value}) (shape = {global_halo_idx.shape}) {global_halo_idx}") halo_generator = ghex.halo_generator_with_gids( - self._decomposition_info.global_index( - horizontal_dim, DecompositionInfo.EntryType.HALO - ) + global_halo_idx ) print( f"rank={self._context.rank()}/{self._context.size()}: halo generator for dim={horizontal_dim} created" diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index d0d054aff..fd1b46575 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -13,6 +13,7 @@ import functools import logging import math +import pdb import sys from collections import namedtuple from dataclasses import InitVar, dataclass, field @@ -625,8 +626,10 @@ def _do_diffusion_step( log.debug(f"{self._log_id} rbf interpolation: end") # 2. HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert + print(f"{self._log_id} communication rbf extrapolation of vn - start") res = self._sync_fields(VertexDim, self.u_vert, self.v_vert) - self._wait(res, VertexDim) + self._wait(res) + print(f"{self._log_id} communication rbf extrapolation of vn - end") log.debug( f"{self._log_id} running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start" @@ -687,9 +690,10 @@ def _do_diffusion_step( ) # HALO EXCHANGE IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH + if self.config.type_vn_diffu > 1: comm_z_nabla_e = self._sync_fields(EdgeDim, self.z_nabla2_e) - self._wait(comm_z_nabla_e, EdgeDim) + self._wait(comm_z_nabla_e) log.debug(f"{self._log_id} 2nd rbf interpolation: start") mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( @@ -707,8 +711,10 @@ def _do_diffusion_step( log.debug(f"{self._log_id} 2nd rbf interpolation: end") # 6. HALO EXCHANGE -- CALL sync_patch_array_mult (Vertex Fields) + print(f"{self._log_id} communication rbf extrapolation of z_nable2_e - start") comm_rbf_coef = self._sync_fields(VertexDim, self.u_vert, self.v_vert) - self._wait(comm_rbf_coef, VertexDim) + self._wait(comm_rbf_coef) + print(f"{self._log_id} communication rbf extrapolation of z_nable2_e - start") log.debug( f"{self._log_id} running stencils 04 05 06 (apply_diffusion_to_vn): start" @@ -783,7 +789,7 @@ def _do_diffusion_step( ) # HALO EXCHANGE: CALL sync_patch_array (Edge Fields) comm_res_vn = self._sync_fields(EdgeDim, prognostic_state.vn) - self._wait(comm_res_vn, EdgeDim) + self._wait(comm_res_vn) log.debug( f"{self._log_id} running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" ) @@ -878,13 +884,13 @@ def _do_diffusion_step( prognostic_state.exner_pressure, prognostic_state.w, ) - self._wait(comm_res, CellDim) + self._wait(comm_res) def _sync_fields(self, dim: Dimension, *field): if self._exchange: return self._exchange.exchange(dim, *field) - def _wait(self, comm_handle, dim): + def _wait(self, comm_handle): if comm_handle: comm_handle.wait() - print(f"{self._log_id} :communication dim={dim} done") + print(f"{self._log_id} :communication ={comm_handle} done") diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index d5bf5321c..14f900753 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -13,6 +13,7 @@ import logging import ghex +import mpi4py.MPI import numpy as np import pytest @@ -22,7 +23,7 @@ from icon4py.decomposition.parallel_setup import ( DecompositionInfo, Exchange, - get_processor_properties, + get_processor_properties, finalize_mpi, ) from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.driver.io_utils import ( @@ -166,13 +167,13 @@ def test_decomposition_info_matches_gridsize(datapath, caplog): ) -@pytest.mark.mpi +#@pytest.mark.mpi @pytest.mark.parametrize("datapath", [2], indirect=True) def test_parallel_diffusion( - datapath, r04b09_diffusion_config, step_date_init, caplog, ndyn_substeps + datapath, r04b09_diffusion_config, step_date_init, ndyn_substeps,caplog ): - caplog.set_level(logging.DEBUG) + caplog.set_level(logging.WARN) props = get_processor_properties() print( @@ -199,12 +200,11 @@ def test_parallel_diffusion( print( f"rank={props.rank}: using local grid with {icon_grid.num_cells()} Cells, {icon_grid.num_edges()} Edges, {icon_grid.num_vertices()} Vertices" ) - initial_run = False diffusion_params = DiffusionParams(r04b09_diffusion_config) diffusion_initial_data = IconSerialDataProvider( "icon_pydycore", str(path), True, mpi_rank=props.rank - ).from_savepoint_diffusion_init(linit=initial_run, date=step_date_init) + ).from_savepoint_diffusion_init(linit=False, date=step_date_init) (edge_geometry, cell_geometry, vertical_geometry) = read_geometry_fields( path, rank=props.rank ) @@ -241,10 +241,12 @@ def test_parallel_diffusion( diffusion_savepoint_exit = IconSerialDataProvider( "icon_pydycore", str(path), True, mpi_rank=props.rank - ).from_savepoint_diffusion_exit(linit=initial_run, date=step_date_init) + ).from_savepoint_diffusion_exit(linit=False, date=step_date_init) _verify_diffusion_fields( - decomp_info, diagnostic_state, prognostic_state, diffusion_savepoint_exit + diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, diffusion_savepoint=diffusion_savepoint_exit, decomp_info=decomp_info ) print( f"rank={props.rank}/{props.comm_size}: running diffusion step - using {props.comm_name} with {props.comm_size} nodes - DONE" ) + + diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 2587a46ba..97a63a58f 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -10,6 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import pdb import numpy as np import pytest @@ -18,6 +19,7 @@ IconDiffusionExitSavepoint, IconDiffusionInitSavepoint, ) +from icon4py.common.dimension import EdgeDim, CellDim from icon4py.decomposition.parallel_setup import DecompositionInfo from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.diffusion.diffusion_utils import scale_k @@ -332,37 +334,36 @@ def test_run_diffusion_single_step( def _verify_diffusion_fields( diagnostic_state: DiagnosticState, prognostic_state: PrognosticState, - diffusion_savepoint_exit: IconDiffusionExitSavepoint, - decomp_info: DecompositionInfo = None, + diffusion_savepoint: IconDiffusionExitSavepoint, + decomp_info: DecompositionInfo, ): - # owned_cells = decomp_info.owner_mask(CellDim) - # num_cells = owned_cells.shape[0] - # owned_edges = decomp_info.owner_mask(EdgeDim) - # num_edges = owned_edges.shape - ref_div_ic = np.asarray(diffusion_savepoint_exit.div_ic()) + owned_cells = decomp_info.owner_mask(CellDim) + owned_edges = decomp_info.owner_mask(EdgeDim) + ref_div_ic = np.asarray(diffusion_savepoint.div_ic()) val_div_ic = np.asarray(diagnostic_state.div_ic) - ref_hdef_ic = np.asarray(diffusion_savepoint_exit.hdef_ic()) + ref_hdef_ic = np.asarray(diffusion_savepoint.hdef_ic()) val_hdef_ic = np.asarray(diagnostic_state.hdef_ic) assert np.allclose(ref_div_ic, val_div_ic) assert np.allclose(ref_hdef_ic, val_hdef_ic) - ref_w = np.asarray(diffusion_savepoint_exit.w()) + ref_w = np.asarray(diffusion_savepoint.w()) val_w = np.asarray(prognostic_state.w) - ref_dwdx = np.asarray(diffusion_savepoint_exit.dwdx()) + ref_dwdx = np.asarray(diffusion_savepoint.dwdx()) val_dwdx = np.asarray(diagnostic_state.dwdx) - ref_dwdy = np.asarray(diffusion_savepoint_exit.dwdy()) + ref_dwdy = np.asarray(diffusion_savepoint.dwdy()) val_dwdy = np.asarray(diagnostic_state.dwdy) - ref_vn = np.asarray(diffusion_savepoint_exit.vn()) + ref_vn = np.asarray(diffusion_savepoint.vn()) val_vn = np.asarray(prognostic_state.vn) - # assert np.allclose(ref_vn[owned_edges], val_vn[owned_edges]) assert np.allclose(ref_dwdx, val_dwdx) assert np.allclose(ref_dwdy, val_dwdy) - # assert np.allclose(ref_w[owned_cells], val_w[owned_cells]) - ref_exner = np.asarray(diffusion_savepoint_exit.exner()) - ref_theta_v = np.asarray(diffusion_savepoint_exit.theta_v()) + # pdb.set_trace() + assert np.allclose(ref_vn[owned_edges, :], val_vn[owned_edges, :]) + assert np.allclose(ref_w[owned_cells, :], val_w[owned_cells, :]) + ref_exner = np.asarray(diffusion_savepoint.exner()) + ref_theta_v = np.asarray(diffusion_savepoint.theta_v()) val_theta_v = np.asarray(prognostic_state.theta_v) val_exner = np.asarray(prognostic_state.exner_pressure) - # assert np.allclose(ref_theta_v[owned_cells], val_theta_v[owned_cells]) - # assert np.allclose(ref_exner[owned_cells], val_exner[owned_cells]) + assert np.allclose(ref_theta_v[owned_cells,:], val_theta_v[owned_cells, :]) + assert np.allclose(ref_exner[owned_cells,:], val_exner[owned_cells, :]) @pytest.mark.datatest @@ -416,5 +417,5 @@ def test_run_diffusion_initial_step( _verify_diffusion_fields( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, - diffusion_savepoint_exit=diffusion_savepoint_exit, + diffusion_savepoint=diffusion_savepoint_exit, ) diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index 985274783..a0adb57d9 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -94,7 +94,7 @@ def _read_int32(self, name: str, offset=0): return self._read(name, offset, dtype=int32) def _read_bool(self, name: str): - return self._read_int32(name, offset=0, dtype=bool) + return self._read(name, offset=0, dtype=bool) def _read(self, name: str, offset=0, dtype=int): return (self.serializer.read(name, self.savepoint) - offset).astype(dtype) @@ -235,10 +235,9 @@ def construct_decomposition_info(self): ) def _get_decomp_fields(self, dim: Dimension): - index = self.global_index(dim) - number = self.num(dim) - mask = self.owner_mask(dim)[0:number] - return dim, index, mask + global_index = self.global_index(dim) + mask = self.owner_mask(dim)[0:self.num(dim)] + return dim, global_index, mask def construct_icon_grid(self) -> IconGrid: From a7a4a34f1ac3684fa84d7be0eebb4245f83687dc Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 26 Jul 2023 14:43:55 +0200 Subject: [PATCH 197/263] fix (2) work around for ghex seg faults: keep field reference passed to field_descriptor alive. --- .../icon4py/decomposition/parallel_setup.py | 106 +++++++++++------- .../src/icon4py/diffusion/diffusion.py | 55 +++++---- .../tests/mpi_tests/test_parallel_setup.py | 23 +++- atm_dyn_iconam/tests/test_diffusion.py | 13 ++- 4 files changed, 130 insertions(+), 67 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py index e69fbcde7..97b32408f 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py @@ -14,14 +14,14 @@ from enum import Enum from typing import Union -import ghex.unstructured as ghex import mpi4py import numpy as np import numpy.ma as ma +from ghex import unstructured as ghex from gt4py.next.common import Dimension, DimensionKind from mpi4py.MPI import Comm -from icon4py.common.dimension import CellDim, EdgeDim, KDim, KHalfDim, VertexDim +from icon4py.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.decomposition.decomposed import ProcessProperties from icon4py.diffusion.diffusion_utils import builder @@ -64,6 +64,24 @@ def finalize_mpi(): log.info("finalizing MPI") MPI.Finalize() +class DomainDescriptorIdGenerator(): + _counter = 0 + _roundtrips = 0 + + def __init__(self, context): + self._comm_size = context.size() + self._roundtrips = context.rank() + self._base = self._roundtrips * self._comm_size + + def __call__(self): + id = self._base + self._counter + if self._counter + 1 >= self._comm_size: + self._roundtrips = self._roundtrips + self._comm_size + self._base = self._roundtrips * self._comm_size + self._counter = 0 + else: + self._counter = self._counter + 1 + return id class DecompositionInfo: class EntryType(int, Enum): @@ -123,10 +141,10 @@ def global_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL): class Exchange: def __init__(self, context, domain_decomposition: DecompositionInfo): - self._counter = 0 self._context = context + self._domain_id_gen = DomainDescriptorIdGenerator(context) self._decomposition_info = domain_decomposition - self._log_id = f"rank={self._context.rank()}/{self._context.size()}" + self._log_id = f"rank={self._context.rank()}/{self._context.size()}>>>" self._domain_descriptors = { CellDim: self._create_domain_descriptor( CellDim, @@ -136,37 +154,27 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): ), EdgeDim: self._create_domain_descriptor(EdgeDim), } - print(f"{self._log_id} :domain descriptors initialized") - self._field_size = { - CellDim: self._decomposition_info.global_index( - CellDim, DecompositionInfo.EntryType.ALL - ).shape[0], - EdgeDim: self._decomposition_info.global_index( - EdgeDim, DecompositionInfo.EntryType.ALL - ).shape[0], - VertexDim: self._decomposition_info.global_index( - VertexDim, DecompositionInfo.EntryType.ALL - ).shape[0], - KDim: domain_decomposition.klevels, - KHalfDim: domain_decomposition.klevels + 1, - } - print(f"{self._log_id} : field sizes = {self._field_size}") + print(f"{self._log_id}: domain descriptors initialized") self._patterns = { CellDim: self._create_pattern(CellDim), VertexDim: self._create_pattern(VertexDim), EdgeDim: self._create_pattern(EdgeDim), } - self._comms = {k: ghex.make_co(context) for k, v in self._patterns.items()} - print(f"{self._log_id} : patterns and communicators initialized ") - print(f"{self._log_id} : exchange initialized") + print(f"{self._log_id}: patterns initialized ") + self._comm = ghex.make_co(context) + print(f"{self._log_id}: exchange initialized") def _domain_descriptor_info(self, descr): - return f" id={descr.domain_id()}, size={descr.size()}, inner_size={descr.inner_size()}" + return f" id={descr.domain_id()}, size={descr.size()}, inner_size={descr.inner_size()} (halo size = {descr.size() - descr.inner_size()})" def get_size(self): return self._context.size() + + def my_rank(self): + return self._context.rank() + def _create_domain_descriptor(self, dim: Dimension): all_global = self._decomposition_info.global_index( dim, DecompositionInfo.EntryType.ALL @@ -175,21 +183,20 @@ def _create_domain_descriptor(self, dim: Dimension): dim, DecompositionInfo.EntryType.HALO ) - print( - f"rank={self._context.rank()}/{self._context.size()}: all global idx(dim={dim.value}) (shape = {all_global.shape}) {all_global}") - print( - f"rank={self._context.rank()}/{self._context.size()}: local halo idx(dim={dim.value}) (shape = {local_halo.shape}) {local_halo}") + # print( + # f"rank={self._context.rank()}/{self._context.size()}: all global idx(dim={dim.value}) (shape = {all_global.shape}) {all_global}") + # print( + # f"rank={self._context.rank()}/{self._context.size()}: local halo idx(dim={dim.value}) (shape = {local_halo.shape}) {local_halo}") + # first arg is the domain ID which builds up an MPI Tag. # if those ids are not different for all domain descriptors the system might deadlock # if two parallel exchanges with the same domain id are done domain_desc = ghex.domain_descriptor( - self._context.rank() + self._counter, all_global.tolist(), local_halo.tolist() + self._domain_id_gen(), all_global.tolist(), local_halo.tolist() ) - self._counter = self._counter + 1 print( - f"rank={self._context.rank()}/{self._context.size()}: domain descriptor for dim {dim} with properties {self._domain_descriptor_info(domain_desc)}" + f"{self._log_id}: domain descriptor for dim {dim} with properties {self._domain_descriptor_info(domain_desc)}" ) - return domain_desc def _create_pattern(self, horizontal_dim: Dimension): @@ -197,34 +204,49 @@ def _create_pattern(self, horizontal_dim: Dimension): global_halo_idx = self._decomposition_info.global_index(horizontal_dim, DecompositionInfo.EntryType.HALO) - print(f"rank={self._context.rank()}/{self._context.size()}: global halo idx(dim={horizontal_dim.value}) (shape = {global_halo_idx.shape}) {global_halo_idx}") + #print(f"rank={self._context.rank()}/{self._context.size()}: global halo idx(dim={horizontal_dim.value}) (shape = {global_halo_idx.shape}) {global_halo_idx}") halo_generator = ghex.halo_generator_with_gids( global_halo_idx ) print( - f"rank={self._context.rank()}/{self._context.size()}: halo generator for dim={horizontal_dim} created" + f"{self._log_id}: halo generator for dim={horizontal_dim} created" ) pattern = ghex.make_pattern( self._context, halo_generator, [self._domain_descriptors[horizontal_dim]] ) print( - f"rank={self._context.rank()}/{self._context.size()}: pattern for dim={horizontal_dim} and {self._domain_descriptors[horizontal_dim]} created" + f"{self._log_id}: pattern for dim={horizontal_dim} and {self._domain_descriptors[horizontal_dim]} created" ) return pattern + def prepare_field(self, dim:Dimension, field): + assert dim in [CellDim, EdgeDim, VertexDim] + pattern = self._patterns[dim] + assert pattern is not None + domain_descriptor = self._domain_descriptors[dim] + assert domain_descriptor is not None + descriptor = ghex.field_descriptor(domain_descriptor, np.asarray(field)) + return field, pattern(descriptor) + + + + + def exchange(self, dim: Dimension, *fields): assert dim in [CellDim, EdgeDim, VertexDim] - horizontal_size = self._field_size[dim] pattern = self._patterns[dim] assert pattern is not None - fields = [np.asarray(f)[:horizontal_size, :] for f in fields] - shapes = list(map(lambda f: f.shape, fields)) - print( - f"rank = {self._context.rank()}/{self._context.size()}: communicating fields of dim = {dim} : shapes = {shapes}" - ) + fields = [np.asarray(f) for f in fields] + for f in fields: + print( + f"{self._log_id}: communicating field of dim = {dim} : shape = {f.shape}" + ) + domain_descriptor = self._domain_descriptors[dim] + print(f"{self._log_id}: applying pattern to field_descriptor of field f={f.shape} with domain descriptor {self._domain_descriptor_info(domain_descriptor)} ") patterns_of_field = [ - pattern(ghex.field_descriptor(self._domain_descriptors[dim], f)) + pattern(ghex.field_descriptor(domain_descriptor, f)) for f in fields ] - return self._comms[dim].exchange(patterns_of_field) + + return self._comm.exchange(patterns_of_field) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index fd1b46575..eb3959eb4 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -87,7 +87,7 @@ cached_backend = run_gtfn_cached compiled_backend = run_gtfn imperative_backend = run_gtfn_imperative -backend = compiled_backend # +backend = run_gtfn_cached # class DiffusionType(int, Enum): @@ -347,7 +347,7 @@ class Diffusion: def __init__(self, exchange: Optional[Exchange] = None): self._exchange = exchange - self._log_id = exchange._log_id if exchange else "" + self._log_id = f"rank={exchange.my_rank()}/{exchange.get_size()}" if exchange else "" self._initialized = False self.rd_o_cvd: float = GAS_CONSTANT_DRY_AIR / (CPD - GAS_CONSTANT_DRY_AIR) self.thresh_tdiff: float = ( @@ -610,7 +610,7 @@ def _do_diffusion_step( self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={} ) - log.debug(f"{self._log_id} rbf interpolation: start") + log.debug(f"{self._log_id} rbf interpolation 1: start") mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( p_e_in=prognostic_state.vn, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, @@ -623,12 +623,15 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={"V2E": self.grid.get_v2e_connectivity()}, ) - log.debug(f"{self._log_id} rbf interpolation: end") + log.debug(f"{self._log_id} rbf interpolation 1: end") # 2. HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert print(f"{self._log_id} communication rbf extrapolation of vn - start") - res = self._sync_fields(VertexDim, self.u_vert, self.v_vert) - self._wait(res) + #res = self._sync_fields(VertexDim, self.u_vert, self.v_vert) + vf, pattern_v = self._exchange.prepare_field(VertexDim, self.v_vert) + uf, pattern_u = self._exchange.prepare_field(VertexDim, self.u_vert) + h = self._exchange._comm.exchange([pattern_u, pattern_v]) + h.wait() print(f"{self._log_id} communication rbf extrapolation of vn - end") log.debug( @@ -691,9 +694,9 @@ def _do_diffusion_step( # HALO EXCHANGE IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH - if self.config.type_vn_diffu > 1: - comm_z_nabla_e = self._sync_fields(EdgeDim, self.z_nabla2_e) - self._wait(comm_z_nabla_e) + # if self.config.type_vn_diffu > 1: + # comm_z_nabla_e = self._sync_fields(EdgeDim, self.z_nabla2_e) + # self._wait(comm_z_nabla_e) log.debug(f"{self._log_id} 2nd rbf interpolation: start") mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( @@ -712,9 +715,11 @@ def _do_diffusion_step( # 6. HALO EXCHANGE -- CALL sync_patch_array_mult (Vertex Fields) print(f"{self._log_id} communication rbf extrapolation of z_nable2_e - start") - comm_rbf_coef = self._sync_fields(VertexDim, self.u_vert, self.v_vert) - self._wait(comm_rbf_coef) - print(f"{self._log_id} communication rbf extrapolation of z_nable2_e - start") + vf, pattern_v = self._exchange.prepare_field(VertexDim, self.v_vert) + uf, pattern_u = self._exchange.prepare_field(VertexDim, self.u_vert) + h = self._exchange._comm.exchange([pattern_u, pattern_v]) + h.wait() + print(f"{self._log_id} communication rbf extrapolation of z_nable2_e - end") log.debug( f"{self._log_id} running stencils 04 05 06 (apply_diffusion_to_vn): start" @@ -788,8 +793,12 @@ def _do_diffusion_step( f"{self._log_id} running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end" ) # HALO EXCHANGE: CALL sync_patch_array (Edge Fields) - comm_res_vn = self._sync_fields(EdgeDim, prognostic_state.vn) - self._wait(comm_res_vn) + #comm_res_vn = self._sync_fields(EdgeDim, prognostic_state.vn) + #self._wait(comm_res_vn) + print(f"{self._log_id}: communication of vn - start") + vn_f, pattern_vn = self._exchange.prepare_field(EdgeDim, prognostic_state.vn) + h_vn = self._exchange._comm.exchange([pattern_vn]) + log.debug( f"{self._log_id} running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" ) @@ -877,14 +886,18 @@ def _do_diffusion_step( # 10. HALO EXCHANGE sync_patch_array (Cell fields) # TODO @magdalena why not trigger the exchange of w earlier? # TODO if condition: IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN + h_vn.wait() + print(f"{self._log_id}: communication of vn - end") + # comm_c = self._sync_fields( + # CellDim, + # prognostic_state.theta_v, + # prognostic_state.exner_pressure, + # prognostic_state.w, + # ) + # self._wait(comm_c) + + theta_f, pattern_theta = self._exchange.prepare_field(CellDim, prognostic_state.theta_v) - comm_res = self._sync_fields( - CellDim, - prognostic_state.theta_v, - prognostic_state.exner_pressure, - prognostic_state.w, - ) - self._wait(comm_res) def _sync_fields(self, dim: Dimension, *field): if self._exchange: diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index 14f900753..26244a0e2 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -23,7 +23,7 @@ from icon4py.decomposition.parallel_setup import ( DecompositionInfo, Exchange, - get_processor_properties, finalize_mpi, + get_processor_properties, finalize_mpi, DomainDescriptorIdGenerator, ) from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.driver.io_utils import ( @@ -137,6 +137,24 @@ def test_processor_properties_from_comm_world(mpi, caplog): assert props.comm_name == mpi.COMM_WORLD.Get_name() + +@pytest.mark.mpi +def test_domain_descriptor_id(): + props = get_processor_properties() + size = props.comm_size + context = ghex.context(ghex.mpi_comm(props.comm), True) + id_gen = DomainDescriptorIdGenerator(context) + id1 = id_gen() + assert id1 == props.comm_size * props.rank + assert id1 < props.comm_size * (props.rank + 1) + for i in range(1, 2*size): + id = id_gen() + assert id == id1 + i + assert id >= 4 * props.rank + assert id < 4 * ( + props.rank + 1) + + @pytest.mark.mpi @pytest.mark.parametrize("datapath", [2], indirect=True) def test_decomposition_info_matches_gridsize(datapath, caplog): @@ -167,7 +185,7 @@ def test_decomposition_info_matches_gridsize(datapath, caplog): ) -#@pytest.mark.mpi +@pytest.mark.mpi @pytest.mark.parametrize("datapath", [2], indirect=True) def test_parallel_diffusion( datapath, r04b09_diffusion_config, step_date_init, ndyn_substeps,caplog @@ -248,5 +266,6 @@ def test_parallel_diffusion( print( f"rank={props.rank}/{props.comm_size}: running diffusion step - using {props.comm_name} with {props.comm_size} nodes - DONE" ) + del context diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 97a63a58f..1a1891357 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -351,11 +351,13 @@ def _verify_diffusion_fields( val_dwdx = np.asarray(diagnostic_state.dwdx) ref_dwdy = np.asarray(diffusion_savepoint.dwdy()) val_dwdy = np.asarray(diagnostic_state.dwdy) - ref_vn = np.asarray(diffusion_savepoint.vn()) - val_vn = np.asarray(prognostic_state.vn) assert np.allclose(ref_dwdx, val_dwdx) assert np.allclose(ref_dwdy, val_dwdy) + # TODO (Magdalena) verification on 2 nodes fine until here.. # pdb.set_trace() + ref_vn = np.asarray(diffusion_savepoint.vn()) + val_vn = np.asarray(prognostic_state.vn) + print_diff_into(ref_vn, val_vn) assert np.allclose(ref_vn[owned_edges, :], val_vn[owned_edges, :]) assert np.allclose(ref_w[owned_cells, :], val_w[owned_cells, :]) ref_exner = np.asarray(diffusion_savepoint.exner()) @@ -366,6 +368,13 @@ def _verify_diffusion_fields( assert np.allclose(ref_exner[owned_cells,:], val_exner[owned_cells, :]) +def print_diff_into(ref_vn, val_vn): + diff_vn = np.abs(ref_vn - val_vn) + print(f"max diff vn {np.max(diff_vn)}") + d = 0.00001 + print(f"number of diffs > {d}: {np.count_nonzero(diff_vn > d)}") + + @pytest.mark.datatest @pytest.mark.parametrize("linit", [True]) @pytest.mark.parametrize("datapath", [1], indirect=True) From d46b2c9efd88053872ba6dbd8ba49cc085ea43c0 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 26 Jul 2023 14:51:44 +0200 Subject: [PATCH 198/263] fix (3) remove old communication calls. --- .../src/icon4py/diffusion/diffusion.py | 31 ++++++------------- 1 file changed, 10 insertions(+), 21 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index eb3959eb4..172724e30 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -694,9 +694,10 @@ def _do_diffusion_step( # HALO EXCHANGE IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH - # if self.config.type_vn_diffu > 1: - # comm_z_nabla_e = self._sync_fields(EdgeDim, self.z_nabla2_e) - # self._wait(comm_z_nabla_e) + if self.config.type_vn_diffu > 1: + z_nablae_f, pattern = self._exchange.prepare_field(EdgeDim, self.z_nabla2_e) + h_z = self._exchange._comm.exchange(pattern) + h_z.wait() log.debug(f"{self._log_id} 2nd rbf interpolation: start") mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( @@ -793,8 +794,6 @@ def _do_diffusion_step( f"{self._log_id} running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end" ) # HALO EXCHANGE: CALL sync_patch_array (Edge Fields) - #comm_res_vn = self._sync_fields(EdgeDim, prognostic_state.vn) - #self._wait(comm_res_vn) print(f"{self._log_id}: communication of vn - start") vn_f, pattern_vn = self._exchange.prepare_field(EdgeDim, prognostic_state.vn) h_vn = self._exchange._comm.exchange([pattern_vn]) @@ -888,22 +887,12 @@ def _do_diffusion_step( # TODO if condition: IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN h_vn.wait() print(f"{self._log_id}: communication of vn - end") - # comm_c = self._sync_fields( - # CellDim, - # prognostic_state.theta_v, - # prognostic_state.exner_pressure, - # prognostic_state.w, - # ) - # self._wait(comm_c) - + print(f"{self._log_id}: communication of prognostic cell fields: theta, w, exner - start") theta_f, pattern_theta = self._exchange.prepare_field(CellDim, prognostic_state.theta_v) + exner_f, pattern_exner = self._exchange.prepare_field(CellDim, prognostic_state.exner_pressure) + w_f, pattern_w = self._exchange.prepare_field(CellDim, prognostic_state.w) + h = self._exchange._comm.exchange([pattern_theta, pattern_exner, pattern_w]) + h.wait() + print(f"{self._log_id}: communication of prognostic cell fields: theta, w, exner - end") - def _sync_fields(self, dim: Dimension, *field): - if self._exchange: - return self._exchange.exchange(dim, *field) - - def _wait(self, comm_handle): - if comm_handle: - comm_handle.wait() - print(f"{self._log_id} :communication ={comm_handle} done") From ce3f771e4ae3a31017cc0395c1da35cb12bb97dd Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 26 Jul 2023 16:07:30 +0200 Subject: [PATCH 199/263] verify parallel diffusion: initial and regular step --- .../src/icon4py/diffusion/diffusion.py | 38 +++++++++++++++---- .../tests/mpi_tests/test_parallel_setup.py | 28 ++++++++------ atm_dyn_iconam/tests/test_diffusion.py | 32 ++++++++-------- 3 files changed, 63 insertions(+), 35 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 172724e30..f2916d3e3 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -522,6 +522,7 @@ def initial_run( smag_limit, 0.0, ) + self._sync_cell_fields(prognostic_state) def run( self, @@ -543,6 +544,25 @@ def run( smag_limit=self.smag_limit, smag_offset=self.smag_offset, ) + if not self.config.lhdiff_rcf : + self._sync_cell_fields(prognostic_state) + + def _sync_cell_fields(self, prognostic_state): + """ + Communicate theta_v, exner and w. + + communication only done in original code if the following condition applies: + IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN + """ + print( + f"{self._log_id}: communication of prognostic cell fields: theta, w, exner - start") + theta_f, pattern_theta = self._exchange.prepare_field(CellDim, prognostic_state.theta_v) + exner_f, pattern_exner = self._exchange.prepare_field(CellDim, + prognostic_state.exner_pressure) + w_f, pattern_w = self._exchange.prepare_field(CellDim, prognostic_state.w) + handle_cell_comm = self._exchange._comm.exchange([pattern_theta, pattern_exner, pattern_w]) + handle_cell_comm.wait() + print(f"{self._log_id}: communication of prognostic cell fields: theta, w, exner - end") def _do_diffusion_step( self, @@ -755,6 +775,9 @@ def _do_diffusion_step( log.debug( f"{self._log_id} running stencils 04 05 06 (apply_diffusion_to_vn): end" ) + print(f"{self._log_id}: communication of vn - start") + vn_f, pattern_vn = self._exchange.prepare_field(EdgeDim, prognostic_state.vn) + handle_edge_comm = self._exchange._comm.exchange([pattern_vn]) log.debug( f"{self._log_id} running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" @@ -793,10 +816,8 @@ def _do_diffusion_step( log.debug( f"{self._log_id} running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end" ) - # HALO EXCHANGE: CALL sync_patch_array (Edge Fields) - print(f"{self._log_id}: communication of vn - start") - vn_f, pattern_vn = self._exchange.prepare_field(EdgeDim, prognostic_state.vn) - h_vn = self._exchange._comm.exchange([pattern_vn]) + # HALO EXCHANGE: CALL sync_patch_array (Edge Fields) (vn) (original location) + log.debug( f"{self._log_id} running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" @@ -882,17 +903,18 @@ def _do_diffusion_step( offset_provider={}, ) log.debug(f"{self._log_id} running stencil 16 (update_theta_and_exner): end") + handle_edge_comm.wait() # need to do this here, since we currently only use 1 communication object. + print(f"{self._log_id}: communication of vn - end") # 10. HALO EXCHANGE sync_patch_array (Cell fields) # TODO @magdalena why not trigger the exchange of w earlier? + # TODO if condition: IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN - h_vn.wait() - print(f"{self._log_id}: communication of vn - end") print(f"{self._log_id}: communication of prognostic cell fields: theta, w, exner - start") theta_f, pattern_theta = self._exchange.prepare_field(CellDim, prognostic_state.theta_v) exner_f, pattern_exner = self._exchange.prepare_field(CellDim, prognostic_state.exner_pressure) w_f, pattern_w = self._exchange.prepare_field(CellDim, prognostic_state.w) - h = self._exchange._comm.exchange([pattern_theta, pattern_exner, pattern_w]) - h.wait() + #handle_cell_comm = self._exchange._comm.exchange([pattern_theta, pattern_exner, pattern_w]) + #handle_cell_comm.wait() print(f"{self._log_id}: communication of prognostic cell fields: theta, w, exner - end") diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index 26244a0e2..ff004387a 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -149,7 +149,7 @@ def test_domain_descriptor_id(): assert id1 < props.comm_size * (props.rank + 1) for i in range(1, 2*size): id = id_gen() - assert id == id1 + i + assert id > id1 assert id >= 4 * props.rank assert id < 4 * ( props.rank + 1) @@ -185,10 +185,12 @@ def test_decomposition_info_matches_gridsize(datapath, caplog): ) -@pytest.mark.mpi +#@pytest.mark.mpi @pytest.mark.parametrize("datapath", [2], indirect=True) +@pytest.mark.parametrize("ndyn_substeps", [5]) +@pytest.mark.parametrize("linit", [True, False]) def test_parallel_diffusion( - datapath, r04b09_diffusion_config, step_date_init, ndyn_substeps,caplog + datapath, r04b09_diffusion_config, step_date_init, linit, caplog, ndyn_substeps ): caplog.set_level(logging.WARN) @@ -222,7 +224,7 @@ def test_parallel_diffusion( diffusion_initial_data = IconSerialDataProvider( "icon_pydycore", str(path), True, mpi_rank=props.rank - ).from_savepoint_diffusion_init(linit=False, date=step_date_init) + ).from_savepoint_diffusion_init(linit=linit, date=step_date_init) (edge_geometry, cell_geometry, vertical_geometry) = read_geometry_fields( path, rank=props.rank ) @@ -249,19 +251,21 @@ def test_parallel_diffusion( print(f"rank={props.rank}/{props.comm_size}: diffusion initialized ") diagnostic_state = diffusion_initial_data.construct_diagnostics_for_diffusion() prognostic_state = diffusion_initial_data.construct_prognostics() - - diffusion.run( - diagnostic_state=diagnostic_state, - prognostic_state=prognostic_state, - dtime=dtime, - ) + if linit: + diffusion.initial_run(diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime) + else: + diffusion.run( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + ) print(f"rank={props.rank}/{props.comm_size}: diffusion run ") diffusion_savepoint_exit = IconSerialDataProvider( "icon_pydycore", str(path), True, mpi_rank=props.rank - ).from_savepoint_diffusion_exit(linit=False, date=step_date_init) + ).from_savepoint_diffusion_exit(linit=linit, date=step_date_init) _verify_diffusion_fields( - diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, diffusion_savepoint=diffusion_savepoint_exit, decomp_info=decomp_info + diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, diffusion_savepoint=diffusion_savepoint_exit ) print( f"rank={props.rank}/{props.comm_size}: running diffusion step - using {props.comm_name} with {props.comm_size} nodes - DONE" diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 1a1891357..cc89ce1dd 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -335,10 +335,7 @@ def _verify_diffusion_fields( diagnostic_state: DiagnosticState, prognostic_state: PrognosticState, diffusion_savepoint: IconDiffusionExitSavepoint, - decomp_info: DecompositionInfo, ): - owned_cells = decomp_info.owner_mask(CellDim) - owned_edges = decomp_info.owner_mask(EdgeDim) ref_div_ic = np.asarray(diffusion_savepoint.div_ic()) val_div_ic = np.asarray(diagnostic_state.div_ic) ref_hdef_ic = np.asarray(diffusion_savepoint.hdef_ic()) @@ -353,26 +350,31 @@ def _verify_diffusion_fields( val_dwdy = np.asarray(diagnostic_state.dwdy) assert np.allclose(ref_dwdx, val_dwdx) assert np.allclose(ref_dwdy, val_dwdy) - # TODO (Magdalena) verification on 2 nodes fine until here.. - # pdb.set_trace() + ref_vn = np.asarray(diffusion_savepoint.vn()) val_vn = np.asarray(prognostic_state.vn) - print_diff_into(ref_vn, val_vn) - assert np.allclose(ref_vn[owned_edges, :], val_vn[owned_edges, :]) - assert np.allclose(ref_w[owned_cells, :], val_w[owned_cells, :]) + assert np.allclose(ref_vn, val_vn) + assert np.allclose(ref_w, val_w) ref_exner = np.asarray(diffusion_savepoint.exner()) ref_theta_v = np.asarray(diffusion_savepoint.theta_v()) val_theta_v = np.asarray(prognostic_state.theta_v) val_exner = np.asarray(prognostic_state.exner_pressure) - assert np.allclose(ref_theta_v[owned_cells,:], val_theta_v[owned_cells, :]) - assert np.allclose(ref_exner[owned_cells,:], val_exner[owned_cells, :]) + assert np.allclose(ref_theta_v, val_theta_v) + assert np.allclose(ref_exner, val_exner) + #print(f" cell halo indices: shape = {halo_cells.shape}{halo_cells}") + #print_diff_info(ref_exner[halo_cells, :], val_exner[halo_cells, :]) + #print_diff_info(ref_w[halo_cells, :], val_w[halo_cells, :]) + #print_diff_info(ref_theta_v[halo_cells, :], val_theta_v[halo_cells, :]) + + -def print_diff_into(ref_vn, val_vn): - diff_vn = np.abs(ref_vn - val_vn) - print(f"max diff vn {np.max(diff_vn)}") - d = 0.00001 - print(f"number of diffs > {d}: {np.count_nonzero(diff_vn > d)}") +def print_diff_info(ref, value): + difference = np.abs(ref - value) + print("differences ----- ") + print(f" max diff vn {np.max(difference)}") + d = 0.000001 + print(f"number of diffs > {d}: {np.count_nonzero(difference > d)}") @pytest.mark.datatest From 2c6998fdd11cf4ea77c4e9e79cdf194b3e765ca9 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 26 Jul 2023 17:54:54 +0200 Subject: [PATCH 200/263] add test for domain_descriptor id --- .../icon4py/decomposition/parallel_setup.py | 5 +++- .../src/icon4py/diffusion/diffusion.py | 24 ++++--------------- .../tests/mpi_tests/test_parallel_setup.py | 19 ++++++++++----- 3 files changed, 22 insertions(+), 26 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py index 97b32408f..e61af245d 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py @@ -229,7 +229,8 @@ def prepare_field(self, dim:Dimension, field): return field, pattern(descriptor) - + def sync_patterns(self, patterns:list): + return self._comm.exchange(patterns) def exchange(self, dim: Dimension, *fields): @@ -250,3 +251,5 @@ def exchange(self, dim: Dimension, *fields): ] return self._comm.exchange(patterns_of_field) + + diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index f2916d3e3..3b0c42687 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -560,7 +560,7 @@ def _sync_cell_fields(self, prognostic_state): exner_f, pattern_exner = self._exchange.prepare_field(CellDim, prognostic_state.exner_pressure) w_f, pattern_w = self._exchange.prepare_field(CellDim, prognostic_state.w) - handle_cell_comm = self._exchange._comm.exchange([pattern_theta, pattern_exner, pattern_w]) + handle_cell_comm = self._exchange.sync_patterns([pattern_theta, pattern_exner, pattern_w]) handle_cell_comm.wait() print(f"{self._log_id}: communication of prognostic cell fields: theta, w, exner - end") @@ -650,7 +650,7 @@ def _do_diffusion_step( #res = self._sync_fields(VertexDim, self.u_vert, self.v_vert) vf, pattern_v = self._exchange.prepare_field(VertexDim, self.v_vert) uf, pattern_u = self._exchange.prepare_field(VertexDim, self.u_vert) - h = self._exchange._comm.exchange([pattern_u, pattern_v]) + h = self._exchange.sync_patterns([pattern_u, pattern_v]) h.wait() print(f"{self._log_id} communication rbf extrapolation of vn - end") @@ -716,7 +716,7 @@ def _do_diffusion_step( if self.config.type_vn_diffu > 1: z_nablae_f, pattern = self._exchange.prepare_field(EdgeDim, self.z_nabla2_e) - h_z = self._exchange._comm.exchange(pattern) + h_z = self._exchange.sync_patterns(pattern) h_z.wait() log.debug(f"{self._log_id} 2nd rbf interpolation: start") @@ -738,7 +738,7 @@ def _do_diffusion_step( print(f"{self._log_id} communication rbf extrapolation of z_nable2_e - start") vf, pattern_v = self._exchange.prepare_field(VertexDim, self.v_vert) uf, pattern_u = self._exchange.prepare_field(VertexDim, self.u_vert) - h = self._exchange._comm.exchange([pattern_u, pattern_v]) + h = self._exchange.sync_patterns([pattern_u, pattern_v]) h.wait() print(f"{self._log_id} communication rbf extrapolation of z_nable2_e - end") @@ -777,7 +777,7 @@ def _do_diffusion_step( ) print(f"{self._log_id}: communication of vn - start") vn_f, pattern_vn = self._exchange.prepare_field(EdgeDim, prognostic_state.vn) - handle_edge_comm = self._exchange._comm.exchange([pattern_vn]) + handle_edge_comm = self._exchange.sync_patterns([pattern_vn]) log.debug( f"{self._log_id} running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" @@ -816,8 +816,6 @@ def _do_diffusion_step( log.debug( f"{self._log_id} running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end" ) - # HALO EXCHANGE: CALL sync_patch_array (Edge Fields) (vn) (original location) - log.debug( f"{self._log_id} running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" @@ -905,16 +903,4 @@ def _do_diffusion_step( log.debug(f"{self._log_id} running stencil 16 (update_theta_and_exner): end") handle_edge_comm.wait() # need to do this here, since we currently only use 1 communication object. print(f"{self._log_id}: communication of vn - end") - # 10. HALO EXCHANGE sync_patch_array (Cell fields) - # TODO @magdalena why not trigger the exchange of w earlier? - - # TODO if condition: IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN - print(f"{self._log_id}: communication of prognostic cell fields: theta, w, exner - start") - theta_f, pattern_theta = self._exchange.prepare_field(CellDim, prognostic_state.theta_v) - exner_f, pattern_exner = self._exchange.prepare_field(CellDim, prognostic_state.exner_pressure) - w_f, pattern_w = self._exchange.prepare_field(CellDim, prognostic_state.w) - #handle_cell_comm = self._exchange._comm.exchange([pattern_theta, pattern_exner, pattern_w]) - #handle_cell_comm.wait() - print(f"{self._log_id}: communication of prognostic cell fields: theta, w, exner - end") - diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index ff004387a..0856ca501 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -139,7 +139,8 @@ def test_processor_properties_from_comm_world(mpi, caplog): @pytest.mark.mpi -def test_domain_descriptor_id(): +@pytest.mark.parametrize("num", [1, 2, 3 ]) +def test_domain_descriptor_id_are_globally_unique(num): props = get_processor_properties() size = props.comm_size context = ghex.context(ghex.mpi_comm(props.comm), True) @@ -147,12 +148,18 @@ def test_domain_descriptor_id(): id1 = id_gen() assert id1 == props.comm_size * props.rank assert id1 < props.comm_size * (props.rank + 1) - for i in range(1, 2*size): + ids = [] + ids.append(id1) + for i in range(1, num * size): id = id_gen() assert id > id1 - assert id >= 4 * props.rank - assert id < 4 * ( - props.rank + 1) + ids.append(id) + all_ids = props.comm.gather(ids, root=0) + if props.rank == 0: + all_ids = np.asarray(all_ids).flatten() + assert len(all_ids) == size * size * num + assert len(all_ids) == len(set(all_ids)) + @pytest.mark.mpi @@ -185,7 +192,7 @@ def test_decomposition_info_matches_gridsize(datapath, caplog): ) -#@pytest.mark.mpi +@pytest.mark.mpi @pytest.mark.parametrize("datapath", [2], indirect=True) @pytest.mark.parametrize("ndyn_substeps", [5]) @pytest.mark.parametrize("linit", [True, False]) From 3406090e86b1d699cc4ca28c9ced857a4fff77e0 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 27 Jul 2023 12:37:59 +0200 Subject: [PATCH 201/263] combine exchange and field_descriptors --- .../icon4py/decomposition/parallel_setup.py | 35 +++++-------------- .../src/icon4py/diffusion/diffusion.py | 26 ++++++-------- 2 files changed, 18 insertions(+), 43 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py index e61af245d..3581b86f2 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import logging +from dataclasses import dataclass from enum import Enum from typing import Union @@ -20,6 +21,7 @@ from ghex import unstructured as ghex from gt4py.next.common import Dimension, DimensionKind from mpi4py.MPI import Comm +from typing_extensions import deprecated from icon4py.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.decomposition.decomposed import ProcessProperties @@ -219,37 +221,16 @@ def _create_pattern(self, horizontal_dim: Dimension): ) return pattern - def prepare_field(self, dim:Dimension, field): - assert dim in [CellDim, EdgeDim, VertexDim] - pattern = self._patterns[dim] - assert pattern is not None - domain_descriptor = self._domain_descriptors[dim] - assert domain_descriptor is not None - descriptor = ghex.field_descriptor(domain_descriptor, np.asarray(field)) - return field, pattern(descriptor) - - def sync_patterns(self, patterns:list): - return self._comm.exchange(patterns) - - - def exchange(self, dim: Dimension, *fields): + def exchange(self, dim:Dimension, *fields:tuple): assert dim in [CellDim, EdgeDim, VertexDim] pattern = self._patterns[dim] assert pattern is not None - fields = [np.asarray(f) for f in fields] - for f in fields: - print( - f"{self._log_id}: communicating field of dim = {dim} : shape = {f.shape}" - ) domain_descriptor = self._domain_descriptors[dim] - print(f"{self._log_id}: applying pattern to field_descriptor of field f={f.shape} with domain descriptor {self._domain_descriptor_info(domain_descriptor)} ") - - patterns_of_field = [ - pattern(ghex.field_descriptor(domain_descriptor, f)) - for f in fields - ] - - return self._comm.exchange(patterns_of_field) + assert domain_descriptor is not None + applied_patterns = [pattern(ghex.field_descriptor(domain_descriptor, np.asarray(f))) for f in + fields] + handle = self._comm.exchange(applied_patterns) + return handle, applied_patterns diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 3b0c42687..ba2f07119 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -556,11 +556,12 @@ def _sync_cell_fields(self, prognostic_state): """ print( f"{self._log_id}: communication of prognostic cell fields: theta, w, exner - start") - theta_f, pattern_theta = self._exchange.prepare_field(CellDim, prognostic_state.theta_v) - exner_f, pattern_exner = self._exchange.prepare_field(CellDim, - prognostic_state.exner_pressure) - w_f, pattern_w = self._exchange.prepare_field(CellDim, prognostic_state.w) - handle_cell_comm = self._exchange.sync_patterns([pattern_theta, pattern_exner, pattern_w]) + #w_f, pattern_w = self._exchange.prepare_field(CellDim, prognostic_state.w) + #theta_f, pattern_theta = self._exchange.prepare_field(CellDim, prognostic_state.theta_v) + #exner_f, pattern_exner = self._exchange.prepare_field(CellDim,prognostic_state.exner_pressure) + #patterns_cell = [pattern_theta, pattern_exner, pattern_w] + #handle_cell_comm = self._exchange.sync_patterns(patterns_cell) + handle_cell_comm, field_refs = self._exchange.exchange(CellDim, prognostic_state.w, prognostic_state.theta_v, prognostic_state.exner_pressure) handle_cell_comm.wait() print(f"{self._log_id}: communication of prognostic cell fields: theta, w, exner - end") @@ -647,10 +648,7 @@ def _do_diffusion_step( # 2. HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert print(f"{self._log_id} communication rbf extrapolation of vn - start") - #res = self._sync_fields(VertexDim, self.u_vert, self.v_vert) - vf, pattern_v = self._exchange.prepare_field(VertexDim, self.v_vert) - uf, pattern_u = self._exchange.prepare_field(VertexDim, self.u_vert) - h = self._exchange.sync_patterns([pattern_u, pattern_v]) + h, rbf_patterns = self._exchange.exchange(VertexDim, self.u_vert, self.v_vert) h.wait() print(f"{self._log_id} communication rbf extrapolation of vn - end") @@ -715,8 +713,7 @@ def _do_diffusion_step( # HALO EXCHANGE IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH if self.config.type_vn_diffu > 1: - z_nablae_f, pattern = self._exchange.prepare_field(EdgeDim, self.z_nabla2_e) - h_z = self._exchange.sync_patterns(pattern) + h_z, p = self._exchange.exchange(EdgeDim, self.z_nabla2_e) h_z.wait() log.debug(f"{self._log_id} 2nd rbf interpolation: start") @@ -736,9 +733,7 @@ def _do_diffusion_step( # 6. HALO EXCHANGE -- CALL sync_patch_array_mult (Vertex Fields) print(f"{self._log_id} communication rbf extrapolation of z_nable2_e - start") - vf, pattern_v = self._exchange.prepare_field(VertexDim, self.v_vert) - uf, pattern_u = self._exchange.prepare_field(VertexDim, self.u_vert) - h = self._exchange.sync_patterns([pattern_u, pattern_v]) + h, rbf_patterns = self._exchange.exchange(VertexDim, self.u_vert, self.v_vert) h.wait() print(f"{self._log_id} communication rbf extrapolation of z_nable2_e - end") @@ -776,8 +771,7 @@ def _do_diffusion_step( f"{self._log_id} running stencils 04 05 06 (apply_diffusion_to_vn): end" ) print(f"{self._log_id}: communication of vn - start") - vn_f, pattern_vn = self._exchange.prepare_field(EdgeDim, prognostic_state.vn) - handle_edge_comm = self._exchange.sync_patterns([pattern_vn]) + handle_edge_comm, vn_f_ref = self._exchange.exchange(EdgeDim, prognostic_state.vn) log.debug( f"{self._log_id} running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" From d1487cff1bfdfc09dc879e76bde7b7333c23e415 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 27 Jul 2023 14:08:08 +0200 Subject: [PATCH 202/263] make serial code run again --- .../src/icon4py/decomposition/decomposed.py | 229 ++++++++++++++++++ .../icon4py/decomposition/parallel_setup.py | 179 -------------- .../src/icon4py/diffusion/diffusion.py | 41 ++-- atm_dyn_iconam/src/icon4py/driver/README.md | 2 +- .../src/icon4py/driver/dycore_driver.py | 8 - atm_dyn_iconam/src/icon4py/driver/io_utils.py | 3 +- atm_dyn_iconam/tests/mpi_tests/README.md | 2 +- ...t_parallel_setup.py => test_decomposed.py} | 36 +-- atm_dyn_iconam/tests/test_diffusion.py | 9 - .../tests/test_utils/serialbox_utils.py | 4 +- base-requirements-dev.txt | 1 - base-requirements.txt | 1 + 12 files changed, 276 insertions(+), 239 deletions(-) rename atm_dyn_iconam/tests/mpi_tests/{test_parallel_setup.py => test_decomposed.py} (92%) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index 95fa174b6..2b8acee74 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -11,7 +11,19 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from dataclasses import dataclass +from enum import Enum +from typing import Protocol + +import mpi4py import mpi4py.MPI +import numpy as np +import numpy.ma as ma +from ghex import unstructured as ghex +from gt4py.next import Dimension + +from icon4py.common.dimension import CellDim, DimensionKind, EdgeDim, VertexDim +from icon4py.diffusion.diffusion_utils import builder class ProcessProperties: @@ -40,3 +52,220 @@ def comm(self): @classmethod def from_mpi_comm(cls, comm: mpi4py.MPI.Comm): return ProcessProperties(comm) + + +class DomainDescriptorIdGenerator: + _counter = 0 + _roundtrips = 0 + + def __init__(self, context): + self._comm_size = context.size() + self._roundtrips = context.rank() + self._base = self._roundtrips * self._comm_size + + def __call__(self): + next_id = self._base + self._counter + if self._counter + 1 >= self._comm_size: + self._roundtrips = self._roundtrips + self._comm_size + self._base = self._roundtrips * self._comm_size + self._counter = 0 + else: + self._counter = self._counter + 1 + return next_id + + +class DecompositionInfo: + class EntryType(int, Enum): + ALL = (0,) + OWNED = (1,) + HALO = 2 + + @builder + def with_dimension( + self, dim: Dimension, global_index: np.ndarray, owner_mask: np.ndarray + ): + masked_global_index = ma.array(global_index, mask=owner_mask) + self._global_index[dim] = masked_global_index + + def __init__(self, klevels: int): + self._global_index = {} + self._klevels = klevels + + @property + def klevels(self): + return self._klevels + + def local_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL): + match (entry_type): + case DecompositionInfo.EntryType.ALL: + return self._to_local_index(dim) + case DecompositionInfo.EntryType.HALO: + index = self._to_local_index(dim) + mask = self._global_index[dim].mask + return index[~mask] + case DecompositionInfo.EntryType.OWNED: + index = self._to_local_index(dim) + mask = self._global_index[dim].mask + return index[mask] + + def _to_local_index(self, dim): + data = ma.getdata(self._global_index[dim], subok=False) + assert data.ndim == 1 + return np.arange(data.shape[0]) + + def owner_mask(self, dim: Dimension) -> np.ndarray: + return self._global_index[dim].mask + + def global_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL): + match (entry_type): + case DecompositionInfo.EntryType.ALL: + return ma.getdata(self._global_index[dim], subok=False) + case DecompositionInfo.EntryType.OWNED: + global_index = self._global_index[dim] + return ma.getdata(global_index[global_index.mask]) + case DecompositionInfo.EntryType.HALO: + global_index = self._global_index[dim] + return ma.getdata(global_index[~global_index.mask]) + case _: + raise NotImplementedError() + + +class ExchangeResult(Protocol): + def wait(self): + ... + + def is_ready(self) -> bool: + ... + + +class ExchangeRuntime(Protocol): + def exchange(self, dim: Dimension, *fields: tuple) -> ExchangeResult: + ... + + def get_size(self): + ... + + def my_rank(self): + ... + + def wait(self): + pass + + def is_ready(self) -> bool: + return True + + +@dataclass +class SingleNode: + def exchange(self, dim: Dimension, *fields: tuple) -> ExchangeResult: + return SingleNodeResult() + + def my_rank(self): + return 0 + + def get_size(self): + return 1 + + +class SingleNodeResult: + def wait(self): + pass + + def is_ready(self) -> bool: + return True + + +class MultiNode: + def __init__(self, context, domain_decomposition: DecompositionInfo): + self._context = context + self._domain_id_gen = DomainDescriptorIdGenerator(context) + self._decomposition_info = domain_decomposition + self._log_id = f"rank={self._context.rank()}/{self._context.size()}>>>" + self._domain_descriptors = { + CellDim: self._create_domain_descriptor( + CellDim, + ), + VertexDim: self._create_domain_descriptor( + VertexDim, + ), + EdgeDim: self._create_domain_descriptor(EdgeDim), + } + print(f"{self._log_id}: domain descriptors initialized") + + self._patterns = { + CellDim: self._create_pattern(CellDim), + VertexDim: self._create_pattern(VertexDim), + EdgeDim: self._create_pattern(EdgeDim), + } + print(f"{self._log_id}: patterns initialized ") + self._comm = ghex.make_co(context) + print(f"{self._log_id}: exchange initialized") + + def _domain_descriptor_info(self, descr): + return f" id={descr.domain_id()}, size={descr.size()}, inner_size={descr.inner_size()} (halo size = {descr.size() - descr.inner_size()})" + + def get_size(self): + return self._context.size() + + def my_rank(self): + return self._context.rank() + + def _create_domain_descriptor(self, dim: Dimension): + all_global = self._decomposition_info.global_index( + dim, DecompositionInfo.EntryType.ALL + ) + local_halo = self._decomposition_info.local_index( + dim, DecompositionInfo.EntryType.HALO + ) + # first arg is the domain ID which builds up an MPI Tag. + # if those ids are not different for all domain descriptors the system might deadlock + # if two parallel exchanges with the same domain id are done + domain_desc = ghex.domain_descriptor( + self._domain_id_gen(), all_global.tolist(), local_halo.tolist() + ) + print( + f"{self._log_id}: domain descriptor for dim {dim} with properties {self._domain_descriptor_info(domain_desc)}" + ) + return domain_desc + + def _create_pattern(self, horizontal_dim: Dimension): + assert horizontal_dim.kind == DimensionKind.HORIZONTAL + + global_halo_idx = self._decomposition_info.global_index( + horizontal_dim, DecompositionInfo.EntryType.HALO + ) + halo_generator = ghex.halo_generator_with_gids(global_halo_idx) + print(f"{self._log_id}: halo generator for dim={horizontal_dim} created") + pattern = ghex.make_pattern( + self._context, halo_generator, [self._domain_descriptors[horizontal_dim]] + ) + print( + f"{self._log_id}: pattern for dim={horizontal_dim} and {self._domain_descriptors[horizontal_dim]} created" + ) + return pattern + + def exchange(self, dim: Dimension, *fields: tuple): + assert dim in [CellDim, EdgeDim, VertexDim] + pattern = self._patterns[dim] + assert pattern is not None + domain_descriptor = self._domain_descriptors[dim] + assert domain_descriptor is not None + applied_patterns = [ + pattern(ghex.field_descriptor(domain_descriptor, np.asarray(f))) + for f in fields + ] + handle = self._comm.exchange(applied_patterns) + return MultiNodeResult(handle, applied_patterns) + + +@dataclass +class MultiNodeResult: + handle: ... + pattern_refs: ... + + def wait(self): + self.handle.wait() + del self.pattern_refs + + def is_ready(self) -> bool: + return self.handle.is_ready() diff --git a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py index 3581b86f2..4e28cb8eb 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py @@ -11,25 +11,15 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import logging -from dataclasses import dataclass -from enum import Enum from typing import Union import mpi4py -import numpy as np -import numpy.ma as ma -from ghex import unstructured as ghex -from gt4py.next.common import Dimension, DimensionKind from mpi4py.MPI import Comm -from typing_extensions import deprecated -from icon4py.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.decomposition.decomposed import ProcessProperties -from icon4py.diffusion.diffusion_utils import builder mpi4py.rc.initialize = False -# mpi4py.rc.finalize = False CommId = Union[int, Comm, None] log = logging.getLogger(__name__) @@ -65,172 +55,3 @@ def finalize_mpi(): if not MPI.Is_finalized(): log.info("finalizing MPI") MPI.Finalize() - -class DomainDescriptorIdGenerator(): - _counter = 0 - _roundtrips = 0 - - def __init__(self, context): - self._comm_size = context.size() - self._roundtrips = context.rank() - self._base = self._roundtrips * self._comm_size - - def __call__(self): - id = self._base + self._counter - if self._counter + 1 >= self._comm_size: - self._roundtrips = self._roundtrips + self._comm_size - self._base = self._roundtrips * self._comm_size - self._counter = 0 - else: - self._counter = self._counter + 1 - return id - -class DecompositionInfo: - class EntryType(int, Enum): - ALL = (0,) - OWNED = (1,) - HALO = 2 - - @builder - def with_dimension( - self, dim: Dimension, global_index: np.ndarray, owner_mask: np.ndarray - ): - masked_global_index = ma.array(global_index, mask=owner_mask) - self._global_index[dim] = masked_global_index - - def __init__(self, klevels: int): - self._global_index = {} - self._klevels = klevels - - @property - def klevels(self): - return self._klevels - - def local_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL): - match (entry_type): - case DecompositionInfo.EntryType.ALL: - return self._to_local_index(dim) - case DecompositionInfo.EntryType.HALO: - index = self._to_local_index(dim) - mask = self._global_index[dim].mask - return index[~mask] - case DecompositionInfo.EntryType.OWNED: - index = self._to_local_index(dim) - mask = self._global_index[dim].mask - return index[mask] - - def _to_local_index(self, dim): - data = ma.getdata(self._global_index[dim], subok=False) - assert data.ndim == 1 - return np.arange(data.shape[0]) - - def owner_mask(self, dim: Dimension) -> np.ndarray: - return self._global_index[dim].mask - - def global_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL): - match (entry_type): - case DecompositionInfo.EntryType.ALL: - return ma.getdata(self._global_index[dim], subok=False) - case DecompositionInfo.EntryType.OWNED: - global_index = self._global_index[dim] - return ma.getdata(global_index[global_index.mask]) - case DecompositionInfo.EntryType.HALO: - global_index = self._global_index[dim] - return ma.getdata(global_index[~global_index.mask]) - case _: - raise NotImplementedError() - - -class Exchange: - def __init__(self, context, domain_decomposition: DecompositionInfo): - self._context = context - self._domain_id_gen = DomainDescriptorIdGenerator(context) - self._decomposition_info = domain_decomposition - self._log_id = f"rank={self._context.rank()}/{self._context.size()}>>>" - self._domain_descriptors = { - CellDim: self._create_domain_descriptor( - CellDim, - ), - VertexDim: self._create_domain_descriptor( - VertexDim, - ), - EdgeDim: self._create_domain_descriptor(EdgeDim), - } - print(f"{self._log_id}: domain descriptors initialized") - - self._patterns = { - CellDim: self._create_pattern(CellDim), - VertexDim: self._create_pattern(VertexDim), - EdgeDim: self._create_pattern(EdgeDim), - } - print(f"{self._log_id}: patterns initialized ") - self._comm = ghex.make_co(context) - print(f"{self._log_id}: exchange initialized") - - def _domain_descriptor_info(self, descr): - return f" id={descr.domain_id()}, size={descr.size()}, inner_size={descr.inner_size()} (halo size = {descr.size() - descr.inner_size()})" - - def get_size(self): - return self._context.size() - - - def my_rank(self): - return self._context.rank() - - def _create_domain_descriptor(self, dim: Dimension): - all_global = self._decomposition_info.global_index( - dim, DecompositionInfo.EntryType.ALL - ) - local_halo = self._decomposition_info.local_index( - dim, DecompositionInfo.EntryType.HALO - ) - - # print( - # f"rank={self._context.rank()}/{self._context.size()}: all global idx(dim={dim.value}) (shape = {all_global.shape}) {all_global}") - # print( - # f"rank={self._context.rank()}/{self._context.size()}: local halo idx(dim={dim.value}) (shape = {local_halo.shape}) {local_halo}") - - # first arg is the domain ID which builds up an MPI Tag. - # if those ids are not different for all domain descriptors the system might deadlock - # if two parallel exchanges with the same domain id are done - domain_desc = ghex.domain_descriptor( - self._domain_id_gen(), all_global.tolist(), local_halo.tolist() - ) - print( - f"{self._log_id}: domain descriptor for dim {dim} with properties {self._domain_descriptor_info(domain_desc)}" - ) - return domain_desc - - def _create_pattern(self, horizontal_dim: Dimension): - assert horizontal_dim.kind == DimensionKind.HORIZONTAL - - global_halo_idx = self._decomposition_info.global_index(horizontal_dim, - DecompositionInfo.EntryType.HALO) - #print(f"rank={self._context.rank()}/{self._context.size()}: global halo idx(dim={horizontal_dim.value}) (shape = {global_halo_idx.shape}) {global_halo_idx}") - halo_generator = ghex.halo_generator_with_gids( - global_halo_idx - ) - print( - f"{self._log_id}: halo generator for dim={horizontal_dim} created" - ) - pattern = ghex.make_pattern( - self._context, halo_generator, [self._domain_descriptors[horizontal_dim]] - ) - print( - f"{self._log_id}: pattern for dim={horizontal_dim} and {self._domain_descriptors[horizontal_dim]} created" - ) - return pattern - - - def exchange(self, dim:Dimension, *fields:tuple): - assert dim in [CellDim, EdgeDim, VertexDim] - pattern = self._patterns[dim] - assert pattern is not None - domain_descriptor = self._domain_descriptors[dim] - assert domain_descriptor is not None - applied_patterns = [pattern(ghex.field_descriptor(domain_descriptor, np.asarray(f))) for f in - fields] - handle = self._comm.exchange(applied_patterns) - return handle, applied_patterns - - diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index ba2f07119..c6c1d0b69 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -13,7 +13,6 @@ import functools import logging import math -import pdb import sys from collections import namedtuple from dataclasses import InitVar, dataclass, field @@ -59,7 +58,7 @@ GAS_CONSTANT_DRY_AIR, ) from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim -from icon4py.decomposition.parallel_setup import Exchange +from icon4py.decomposition.decomposed import ExchangeRuntime, SingleNode from icon4py.diffusion.diffusion_utils import ( copy_field, init_diffusion_local_fields_for_regular_timestep, @@ -345,9 +344,11 @@ def diffusion_type_5_smagorinski_factor(config: DiffusionConfig): class Diffusion: """Class that configures diffusion and does one diffusion step.""" - def __init__(self, exchange: Optional[Exchange] = None): + def __init__(self, exchange: ExchangeRuntime = SingleNode()): self._exchange = exchange - self._log_id = f"rank={exchange.my_rank()}/{exchange.get_size()}" if exchange else "" + self._log_id = ( + f"rank={exchange.my_rank()}/{exchange.get_size()}" if exchange else "" + ) self._initialized = False self.rd_o_cvd: float = GAS_CONSTANT_DRY_AIR / (CPD - GAS_CONSTANT_DRY_AIR) self.thresh_tdiff: float = ( @@ -544,7 +545,7 @@ def run( smag_limit=self.smag_limit, smag_offset=self.smag_offset, ) - if not self.config.lhdiff_rcf : + if not self.config.lhdiff_rcf: self._sync_cell_fields(prognostic_state) def _sync_cell_fields(self, prognostic_state): @@ -555,15 +556,18 @@ def _sync_cell_fields(self, prognostic_state): IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN """ print( - f"{self._log_id}: communication of prognostic cell fields: theta, w, exner - start") - #w_f, pattern_w = self._exchange.prepare_field(CellDim, prognostic_state.w) - #theta_f, pattern_theta = self._exchange.prepare_field(CellDim, prognostic_state.theta_v) - #exner_f, pattern_exner = self._exchange.prepare_field(CellDim,prognostic_state.exner_pressure) - #patterns_cell = [pattern_theta, pattern_exner, pattern_w] - #handle_cell_comm = self._exchange.sync_patterns(patterns_cell) - handle_cell_comm, field_refs = self._exchange.exchange(CellDim, prognostic_state.w, prognostic_state.theta_v, prognostic_state.exner_pressure) + f"{self._log_id}: communication of prognostic cell fields: theta, w, exner - start" + ) + handle_cell_comm = self._exchange.exchange( + CellDim, + prognostic_state.w, + prognostic_state.theta_v, + prognostic_state.exner_pressure, + ) handle_cell_comm.wait() - print(f"{self._log_id}: communication of prognostic cell fields: theta, w, exner - end") + print( + f"{self._log_id}: communication of prognostic cell fields: theta, w, exner - end" + ) def _do_diffusion_step( self, @@ -648,7 +652,7 @@ def _do_diffusion_step( # 2. HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert print(f"{self._log_id} communication rbf extrapolation of vn - start") - h, rbf_patterns = self._exchange.exchange(VertexDim, self.u_vert, self.v_vert) + h = self._exchange.exchange(VertexDim, self.u_vert, self.v_vert) h.wait() print(f"{self._log_id} communication rbf extrapolation of vn - end") @@ -713,7 +717,7 @@ def _do_diffusion_step( # HALO EXCHANGE IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH if self.config.type_vn_diffu > 1: - h_z, p = self._exchange.exchange(EdgeDim, self.z_nabla2_e) + h_z = self._exchange.exchange(EdgeDim, self.z_nabla2_e) h_z.wait() log.debug(f"{self._log_id} 2nd rbf interpolation: start") @@ -733,7 +737,7 @@ def _do_diffusion_step( # 6. HALO EXCHANGE -- CALL sync_patch_array_mult (Vertex Fields) print(f"{self._log_id} communication rbf extrapolation of z_nable2_e - start") - h, rbf_patterns = self._exchange.exchange(VertexDim, self.u_vert, self.v_vert) + h = self._exchange.exchange(VertexDim, self.u_vert, self.v_vert) h.wait() print(f"{self._log_id} communication rbf extrapolation of z_nable2_e - end") @@ -771,7 +775,7 @@ def _do_diffusion_step( f"{self._log_id} running stencils 04 05 06 (apply_diffusion_to_vn): end" ) print(f"{self._log_id}: communication of vn - start") - handle_edge_comm, vn_f_ref = self._exchange.exchange(EdgeDim, prognostic_state.vn) + handle_edge_comm = self._exchange.exchange(EdgeDim, prognostic_state.vn) log.debug( f"{self._log_id} running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" @@ -895,6 +899,5 @@ def _do_diffusion_step( offset_provider={}, ) log.debug(f"{self._log_id} running stencil 16 (update_theta_and_exner): end") - handle_edge_comm.wait() # need to do this here, since we currently only use 1 communication object. + handle_edge_comm.wait() # need to do this here, since we currently only use 1 communication object. print(f"{self._log_id}: communication of vn - end") - diff --git a/atm_dyn_iconam/src/icon4py/driver/README.md b/atm_dyn_iconam/src/icon4py/driver/README.md index a4839c882..c6b80a936 100644 --- a/atm_dyn_iconam/src/icon4py/driver/README.md +++ b/atm_dyn_iconam/src/icon4py/driver/README.md @@ -10,7 +10,7 @@ The code is meant to be changed and enlarged as we port new parts of the model. ``` cd atm_dyn_iconam/src/icon4py -python driver/dycore_driver.py ../../tests/ser_icondata/mch_ch_r04b09_dsl/ser_data --n_steps=2 --run_path=/home/magdalena/temp/icon +python driver/dycore_driver.py ../../../testdata/ser_icondata/mpitask1/mch_ch_r04b09_dsl/ser_data --n_steps=2 --run_path=/home/magdalena/temp/icon ``` #### remarks diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index 6843fed37..48da0299f 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -20,7 +20,6 @@ from devtools import Timer from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn -from icon4py.decomposition.parallel_setup import get_processor_properties from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.diffusion.diffusion_utils import copy_diagnostic_and_prognostics from icon4py.diffusion.state_utils import DiagnosticState, PrognosticState @@ -29,7 +28,6 @@ SIMULATION_START_DATE, configure_logging, import_testutils, - read_decomp_info, read_geometry_fields, read_icon_grid, read_initial_state, @@ -179,12 +177,6 @@ def initialize(n_time_steps, file_path: Path): log.info(f"reading configuration: experiment {experiment_name}") config = read_config(experiment_name, n_time_steps=n_time_steps) - parallel_props = get_processor_properties() - decomp_info = read_decomp_info( - f"/home/magdalena/data/exclaim/dycore/{experiment_name}/mpitasks2/{experiment_name}/ser_data", - parallel_props, - ) - log.info("initializing the grid") icon_grid = read_icon_grid(file_path) log.info("reading input fields") diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index beb46fee8..0969d5313 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -17,8 +17,7 @@ from enum import Enum from pathlib import Path -from icon4py.decomposition.decomposed import ProcessProperties -from icon4py.decomposition.parallel_setup import DecompositionInfo +from icon4py.decomposition.decomposed import DecompositionInfo, ProcessProperties from icon4py.diffusion.state_utils import ( DiagnosticState, InterpolationState, diff --git a/atm_dyn_iconam/tests/mpi_tests/README.md b/atm_dyn_iconam/tests/mpi_tests/README.md index 7580f2a4f..1d495fa95 100644 --- a/atm_dyn_iconam/tests/mpi_tests/README.md +++ b/atm_dyn_iconam/tests/mpi_tests/README.md @@ -5,7 +5,7 @@ The parallelized code uses [GHEX](https://github.com/ghex-org/GHEX) with MPI for halo exchanges. GHEX has a CMake build but no setup script for pip, so it needs to be installed manually: 1. You need a running MPI installation in the system. -2. clone GHEX +2. You need to have boost (headers) installed in the system 3clone GHEX ```bash cd {icon4py}/_external_src diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py similarity index 92% rename from atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py rename to atm_dyn_iconam/tests/mpi_tests/test_decomposed.py index 0856ca501..d69f2baf0 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py @@ -13,18 +13,18 @@ import logging import ghex -import mpi4py.MPI import numpy as np import pytest from atm_dyn_iconam.tests.test_diffusion import _verify_diffusion_fields from atm_dyn_iconam.tests.test_utils.serialbox_utils import IconSerialDataProvider from icon4py.common.dimension import CellDim, EdgeDim, VertexDim -from icon4py.decomposition.parallel_setup import ( +from icon4py.decomposition.decomposed import ( DecompositionInfo, - Exchange, - get_processor_properties, finalize_mpi, DomainDescriptorIdGenerator, + DomainDescriptorIdGenerator, + MultiNode, ) +from icon4py.decomposition.parallel_setup import get_processor_properties from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.driver.io_utils import ( SerializationType, @@ -137,9 +137,8 @@ def test_processor_properties_from_comm_world(mpi, caplog): assert props.comm_name == mpi.COMM_WORLD.Get_name() - @pytest.mark.mpi -@pytest.mark.parametrize("num", [1, 2, 3 ]) +@pytest.mark.parametrize("num", [1, 2, 3]) def test_domain_descriptor_id_are_globally_unique(num): props = get_processor_properties() size = props.comm_size @@ -150,10 +149,10 @@ def test_domain_descriptor_id_are_globally_unique(num): assert id1 < props.comm_size * (props.rank + 1) ids = [] ids.append(id1) - for i in range(1, num * size): - id = id_gen() - assert id > id1 - ids.append(id) + for _ in range(1, num * size): + next_id = id_gen() + assert next_id > id1 + ids.append(next_id) all_ids = props.comm.gather(ids, root=0) if props.rank == 0: all_ids = np.asarray(all_ids).flatten() @@ -161,7 +160,6 @@ def test_domain_descriptor_id_are_globally_unique(num): assert len(all_ids) == len(set(all_ids)) - @pytest.mark.mpi @pytest.mark.parametrize("datapath", [2], indirect=True) def test_decomposition_info_matches_gridsize(datapath, caplog): @@ -221,7 +219,7 @@ def test_parallel_diffusion( print( f"rank={props.rank}/{props.comm_size}: GHEX context setup: from {props.comm_name} with {props.comm_size} nodes" ) - # assert context.size() == 2 + assert context.size() == 2 icon_grid = read_icon_grid(path, rank=props.rank) print( @@ -241,7 +239,7 @@ def test_parallel_diffusion( print( f"rank={props.rank}/{props.comm_size}: setup: using {props.comm_name} with {props.comm_size} nodes" ) - exchange = Exchange(context, decomp_info) + exchange = MultiNode(context, decomp_info) diffusion = Diffusion(exchange) @@ -259,7 +257,11 @@ def test_parallel_diffusion( diagnostic_state = diffusion_initial_data.construct_diagnostics_for_diffusion() prognostic_state = diffusion_initial_data.construct_prognostics() if linit: - diffusion.initial_run(diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime) + diffusion.initial_run( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + ) else: diffusion.run( diagnostic_state=diagnostic_state, @@ -272,11 +274,11 @@ def test_parallel_diffusion( "icon_pydycore", str(path), True, mpi_rank=props.rank ).from_savepoint_diffusion_exit(linit=linit, date=step_date_init) _verify_diffusion_fields( - diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, diffusion_savepoint=diffusion_savepoint_exit + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + diffusion_savepoint=diffusion_savepoint_exit, ) print( f"rank={props.rank}/{props.comm_size}: running diffusion step - using {props.comm_name} with {props.comm_size} nodes - DONE" ) del context - - diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index cc89ce1dd..3e583c515 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -10,7 +10,6 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -import pdb import numpy as np import pytest @@ -19,8 +18,6 @@ IconDiffusionExitSavepoint, IconDiffusionInitSavepoint, ) -from icon4py.common.dimension import EdgeDim, CellDim -from icon4py.decomposition.parallel_setup import DecompositionInfo from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.diffusion.diffusion_utils import scale_k from icon4py.diffusion.state_utils import DiagnosticState, PrognosticState @@ -361,12 +358,6 @@ def _verify_diffusion_fields( val_exner = np.asarray(prognostic_state.exner_pressure) assert np.allclose(ref_theta_v, val_theta_v) assert np.allclose(ref_exner, val_exner) - #print(f" cell halo indices: shape = {halo_cells.shape}{halo_cells}") - #print_diff_info(ref_exner[halo_cells, :], val_exner[halo_cells, :]) - #print_diff_info(ref_w[halo_cells, :], val_w[halo_cells, :]) - #print_diff_info(ref_theta_v[halo_cells, :], val_theta_v[halo_cells, :]) - - def print_diff_info(ref, value): diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index a0adb57d9..8152fbea4 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -35,7 +35,7 @@ V2EDim, VertexDim, ) -from icon4py.decomposition.parallel_setup import DecompositionInfo +from icon4py.decomposition.decomposed import DecompositionInfo from icon4py.diffusion.diffusion import VectorTuple from icon4py.diffusion.state_utils import ( DiagnosticState, @@ -236,7 +236,7 @@ def construct_decomposition_info(self): def _get_decomp_fields(self, dim: Dimension): global_index = self.global_index(dim) - mask = self.owner_mask(dim)[0:self.num(dim)] + mask = self.owner_mask(dim)[0 : self.num(dim)] return dim, global_index, mask def construct_icon_grid(self) -> IconGrid: diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 3a5b086e6..bd590d63e 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -31,4 +31,3 @@ wget >= 3.2 netcdf4>=1.6.0 mpi4py<=3.1.4 pytest-mpi>=0.6 -netcdf4>=1.6.0 diff --git a/base-requirements.txt b/base-requirements.txt index a0f5ecb08..06f0a1341 100644 --- a/base-requirements.txt +++ b/base-requirements.txt @@ -1,3 +1,4 @@ # VCS gt4py @ git+https://github.com/GridTools/gt4py.git@main netcdf4 >=1.6.0 +mpi4py<=3.1.4 From 08b8e10fea71a608ce3928da58ed44fd6b67a0fd Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 27 Jul 2023 16:57:50 +0200 Subject: [PATCH 203/263] clean ups - convert print statements to logging - add logging filter to config to capture rank info --- .../src/icon4py/decomposition/decomposed.py | 50 ++++---------- .../icon4py/decomposition/parallel_setup.py | 40 +++++++++++- .../src/icon4py/diffusion/diffusion.py | 65 +++++++++---------- .../src/icon4py/driver/dycore_driver.py | 9 +-- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 10 +-- .../tests/mpi_tests/test_decomposed.py | 5 +- atm_dyn_iconam/tests/test_diffusion.py | 8 --- 7 files changed, 96 insertions(+), 91 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index 2b8acee74..2d7fbb301 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -10,13 +10,11 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - +import logging from dataclasses import dataclass from enum import Enum from typing import Protocol -import mpi4py -import mpi4py.MPI import numpy as np import numpy.ma as ma from ghex import unstructured as ghex @@ -26,32 +24,12 @@ from icon4py.diffusion.diffusion_utils import builder -class ProcessProperties: - def __init__(self, comm: mpi4py.MPI.Comm): - self._communicator_name: str = comm.Get_name() - self._rank: int = comm.Get_rank() - self._comm_size = comm.Get_size() - self._comm = comm - @property - def rank(self): - return self._rank - @property - def comm_name(self): - return self._communicator_name +log = logging.getLogger(__name__) - @property - def comm_size(self): - return self._comm_size - @property - def comm(self): - return self._comm - @classmethod - def from_mpi_comm(cls, comm: mpi4py.MPI.Comm): - return ProcessProperties(comm) class DomainDescriptorIdGenerator: @@ -180,7 +158,6 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): self._context = context self._domain_id_gen = DomainDescriptorIdGenerator(context) self._decomposition_info = domain_decomposition - self._log_id = f"rank={self._context.rank()}/{self._context.size()}>>>" self._domain_descriptors = { CellDim: self._create_domain_descriptor( CellDim, @@ -190,19 +167,19 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): ), EdgeDim: self._create_domain_descriptor(EdgeDim), } - print(f"{self._log_id}: domain descriptors initialized") + log.info(f"domain descriptors for dimensions {self._domain_descriptors.keys()} initialized") self._patterns = { CellDim: self._create_pattern(CellDim), VertexDim: self._create_pattern(VertexDim), EdgeDim: self._create_pattern(EdgeDim), } - print(f"{self._log_id}: patterns initialized ") + log.info(f"patterns for dimensions {self._patterns.keys()} initialized ") self._comm = ghex.make_co(context) - print(f"{self._log_id}: exchange initialized") + log.info(f"communication object initialized") def _domain_descriptor_info(self, descr): - return f" id={descr.domain_id()}, size={descr.size()}, inner_size={descr.inner_size()} (halo size = {descr.size() - descr.inner_size()})" + return f" domain_descriptor=[id='{descr.domain_id()}', size='{descr.size()}', inner_size='{descr.inner_size()}' (halo size='{descr.size() - descr.inner_size()}')" def get_size(self): return self._context.size() @@ -223,8 +200,8 @@ def _create_domain_descriptor(self, dim: Dimension): domain_desc = ghex.domain_descriptor( self._domain_id_gen(), all_global.tolist(), local_halo.tolist() ) - print( - f"{self._log_id}: domain descriptor for dim {dim} with properties {self._domain_descriptor_info(domain_desc)}" + log.debug( + f"domain descriptor for dim='{dim.value}' with properties {self._domain_descriptor_info(domain_desc)} created" ) return domain_desc @@ -235,26 +212,27 @@ def _create_pattern(self, horizontal_dim: Dimension): horizontal_dim, DecompositionInfo.EntryType.HALO ) halo_generator = ghex.halo_generator_with_gids(global_halo_idx) - print(f"{self._log_id}: halo generator for dim={horizontal_dim} created") + log.debug(f"halo generator for dim='{horizontal_dim.value}' created") pattern = ghex.make_pattern( self._context, halo_generator, [self._domain_descriptors[horizontal_dim]] ) - print( - f"{self._log_id}: pattern for dim={horizontal_dim} and {self._domain_descriptors[horizontal_dim]} created" + log.debug( + f"pattern for dim='{horizontal_dim.value}' and {self._domain_descriptor_info(self._domain_descriptors[horizontal_dim])} created" ) return pattern def exchange(self, dim: Dimension, *fields: tuple): assert dim in [CellDim, EdgeDim, VertexDim] pattern = self._patterns[dim] - assert pattern is not None + assert pattern is not None, f"pattern for {dim.value} not found" domain_descriptor = self._domain_descriptors[dim] - assert domain_descriptor is not None + assert domain_descriptor is not None, f"domain descriptor for {dim.value} not found" applied_patterns = [ pattern(ghex.field_descriptor(domain_descriptor, np.asarray(f))) for f in fields ] handle = self._comm.exchange(applied_patterns) + log.info(f"exchange for {len(fields)} fields of dimension ='{dim.value}' initiated.") return MultiNodeResult(handle, applied_patterns) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py index 4e28cb8eb..080690335 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py @@ -16,7 +16,6 @@ import mpi4py from mpi4py.MPI import Comm -from icon4py.decomposition.decomposed import ProcessProperties mpi4py.rc.initialize = False @@ -55,3 +54,42 @@ def finalize_mpi(): if not MPI.Is_finalized(): log.info("finalizing MPI") MPI.Finalize() + + +class ProcessProperties: + def __init__(self, comm: mpi4py.MPI.Comm): + self._communicator_name: str = comm.Get_name() + self._rank: int = comm.Get_rank() + self._comm_size = comm.Get_size() + self._comm = comm + + @property + def rank(self): + return self._rank + + @property + def comm_name(self): + return self._communicator_name + + @property + def comm_size(self): + return self._comm_size + + @property + def comm(self): + return self._comm + + @classmethod + def from_mpi_comm(cls, comm: mpi4py.MPI.Comm): + return ProcessProperties(comm) + +class ParallelLogger(logging.Filter): + def __init__(self, processProperties: ProcessProperties = None): + super().__init__() + self._rank_info = "" + if processProperties and processProperties.comm_size > 1: + self._rank_info = f"rank={processProperties.rank}/{processProperties.comm_size} [{processProperties.comm_name}] >>>" + + def filter(self, record: logging.LogRecord) -> bool: + record.rank = self._rank_info + return True diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index c6c1d0b69..8b23c60c7 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -346,14 +346,11 @@ class Diffusion: def __init__(self, exchange: ExchangeRuntime = SingleNode()): self._exchange = exchange - self._log_id = ( - f"rank={exchange.my_rank()}/{exchange.get_size()}" if exchange else "" - ) self._initialized = False self.rd_o_cvd: float = GAS_CONSTANT_DRY_AIR / (CPD - GAS_CONSTANT_DRY_AIR) self.thresh_tdiff: float = ( -5.0 - ) # threshold temperature deviation from neighboring grid points hat activates extra diffusion against runaway cooling + ) #: threshold temperature deviation from neighboring grid points hat activates extra diffusion against runaway cooling self.grid: Optional[IconGrid] = None self.config: Optional[DiffusionConfig] = None self.params: Optional[DiffusionParams] = None @@ -555,8 +552,8 @@ def _sync_cell_fields(self, prognostic_state): communication only done in original code if the following condition applies: IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN """ - print( - f"{self._log_id}: communication of prognostic cell fields: theta, w, exner - start" + log.debug( + f"communication of prognostic cell fields: theta, w, exner - start" ) handle_cell_comm = self._exchange.exchange( CellDim, @@ -565,8 +562,8 @@ def _sync_cell_fields(self, prognostic_state): prognostic_state.exner_pressure, ) handle_cell_comm.wait() - print( - f"{self._log_id}: communication of prognostic cell fields: theta, w, exner - end" + log.debug( + f"communication of prognostic cell fields: theta, w, exner - done" ) def _do_diffusion_step( @@ -635,7 +632,7 @@ def _do_diffusion_step( self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={} ) - log.debug(f"{self._log_id} rbf interpolation 1: start") + log.debug(f"rbf interpolation 1: start") mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( p_e_in=prognostic_state.vn, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, @@ -648,16 +645,16 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={"V2E": self.grid.get_v2e_connectivity()}, ) - log.debug(f"{self._log_id} rbf interpolation 1: end") + log.debug(f"rbf interpolation 1: end") # 2. HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert - print(f"{self._log_id} communication rbf extrapolation of vn - start") + log.debug(f"communication rbf extrapolation of vn - start") h = self._exchange.exchange(VertexDim, self.u_vert, self.v_vert) h.wait() - print(f"{self._log_id} communication rbf extrapolation of vn - end") + log.debug(f"communication rbf extrapolation of vn - end") log.debug( - f"{self._log_id} running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start" + f"running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start" ) calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend)( diff_multfac_smag=self.diff_multfac_smag, @@ -686,10 +683,10 @@ def _do_diffusion_step( }, ) log.debug( - f"{self._log_id} running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end" + f"running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end" ) log.debug( - f"{self._log_id} running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start" + f"running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start" ) calculate_diagnostic_quantities_for_turbulence.with_backend(backend)( kh_smag_ec=self.kh_smag_ec, @@ -711,16 +708,18 @@ def _do_diffusion_step( }, ) log.debug( - f"{self._log_id} running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end" + f"running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end" ) # HALO EXCHANGE IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH if self.config.type_vn_diffu > 1: + log.debug(f"communication rbf extrapolation of z_nable2_e - start") h_z = self._exchange.exchange(EdgeDim, self.z_nabla2_e) h_z.wait() + log.debug(f"communication rbf extrapolation of z_nable2_e - end") - log.debug(f"{self._log_id} 2nd rbf interpolation: start") + log.debug(f"2nd rbf interpolation: start") mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( p_e_in=self.z_nabla2_e, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, @@ -733,16 +732,16 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={"V2E": self.grid.get_v2e_connectivity()}, ) - log.debug(f"{self._log_id} 2nd rbf interpolation: end") + log.debug(f"2nd rbf interpolation: end") # 6. HALO EXCHANGE -- CALL sync_patch_array_mult (Vertex Fields) - print(f"{self._log_id} communication rbf extrapolation of z_nable2_e - start") + log.debug(f"communication rbf extrapolation of z_nable2_e - start") h = self._exchange.exchange(VertexDim, self.u_vert, self.v_vert) h.wait() - print(f"{self._log_id} communication rbf extrapolation of z_nable2_e - end") + log.debug(f"communication rbf extrapolation of z_nable2_e - end") log.debug( - f"{self._log_id} running stencils 04 05 06 (apply_diffusion_to_vn): start" + f"running stencils 04 05 06 (apply_diffusion_to_vn): start" ) apply_diffusion_to_vn.with_backend(backend)( u_vert=self.u_vert, @@ -772,13 +771,13 @@ def _do_diffusion_step( }, ) log.debug( - f"{self._log_id} running stencils 04 05 06 (apply_diffusion_to_vn): end" + f"running stencils 04 05 06 (apply_diffusion_to_vn): end" ) - print(f"{self._log_id}: communication of vn - start") + log.debug(f"communication of prognistic.vn : start") handle_edge_comm = self._exchange.exchange(EdgeDim, prognostic_state.vn) log.debug( - f"{self._log_id} running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" + f"running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" ) copy_field.with_backend(backend)( prognostic_state.w, self.w_tmp, offset_provider={} @@ -812,11 +811,11 @@ def _do_diffusion_step( }, ) log.debug( - f"{self._log_id} running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end" + f"running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end" ) log.debug( - f"{self._log_id} running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" + f"running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" ) calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.with_backend( backend @@ -835,7 +834,7 @@ def _do_diffusion_step( }, ) log.debug( - f"{self._log_id} running stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): end" + f"running stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): end" ) log.debug("running stencils 13 14 (calculate_nabla2_for_theta): start") calculate_nabla2_for_theta.with_backend(backend)( @@ -855,10 +854,10 @@ def _do_diffusion_step( }, ) log.debug( - f"{self._log_id} running stencils 13_14 (calculate_nabla2_for_theta): end" + f"running stencils 13_14 (calculate_nabla2_for_theta): end" ) log.debug( - f"{self._log_id} running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start" + f"running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start" ) truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_backend( backend @@ -883,9 +882,9 @@ def _do_diffusion_step( ) log.debug( - f"{self._log_id} running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end" + f"running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end" ) - log.debug(f"{self._log_id} running stencil 16 (update_theta_and_exner): start") + log.debug(f"running stencil 16 (update_theta_and_exner): start") update_theta_and_exner.with_backend(backend)( z_temp=self.z_temp, area=self.cell_params.area, @@ -898,6 +897,6 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={}, ) - log.debug(f"{self._log_id} running stencil 16 (update_theta_and_exner): end") + log.debug(f"running stencil 16 (update_theta_and_exner): end") handle_edge_comm.wait() # need to do this here, since we currently only use 1 communication object. - print(f"{self._log_id}: communication of vn - end") + log.debug(f"communication of prognogistic.vn - end") diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index 48da0299f..a279eb130 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -20,6 +20,7 @@ from devtools import Timer from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn +from icon4py.decomposition.parallel_setup import get_processor_properties from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.diffusion.diffusion_utils import copy_diagnostic_and_prognostics from icon4py.diffusion.state_utils import DiagnosticState, PrognosticState @@ -38,11 +39,6 @@ helpers = import_testutils() from helpers import serialbox_utils as sb_utils # noqa - -helpers = import_testutils() -from helpers import serialbox_utils as sb_utils # noqa - - log = logging.getLogger(__name__) @@ -241,7 +237,8 @@ def run( """ start_time = datetime.now().astimezone(pytz.UTC) - configure_logging(run_path, start_time) + parallel_props = get_processor_properties() + configure_logging(run_path, start_time, parallel_props) log.info(f"Starting ICON dycore run: {datetime.isoformat(start_time)}") log.info(f"input args: input_path={input_path}, n_time_steps={n_steps}") timeloop, diagnostic_state, prognostic_state = initialize(n_steps, Path(input_path)) diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index 0969d5313..13e9afbcd 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -17,7 +17,8 @@ from enum import Enum from pathlib import Path -from icon4py.decomposition.decomposed import DecompositionInfo, ProcessProperties +from icon4py.decomposition.decomposed import DecompositionInfo +from icon4py.decomposition.parallel_setup import ProcessProperties, ParallelLogger from icon4py.diffusion.state_utils import ( DiagnosticState, InterpolationState, @@ -192,7 +193,7 @@ def read_static_fields( raise NotImplementedError(SB_ONLY_MSG) -def configure_logging(run_path: str, start_time) -> None: +def configure_logging(run_path: str, start_time, processor_procs = None) -> None: """ Configure logging. @@ -218,9 +219,10 @@ def configure_logging(run_path: str, start_time) -> None: filename=logfile, ) console_handler = logging.StreamHandler() + console_handler.addFilter(ParallelLogger(processor_procs)) formatter = logging.Formatter( - "%(asctime)s %(filename)-20s : %(funcName)-20s: %(levelname)-8s %(message)s" + "%(rank)-20s %(asctime)s %(filename)-20s : %(funcName)-20s: %(levelname)-8s %(message)s" ) console_handler.setFormatter(formatter) - console_handler.setLevel(logging.INFO) + console_handler.setLevel(logging.DEBUG) logging.getLogger("").addHandler(console_handler) diff --git a/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py b/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py index d69f2baf0..72bf19f44 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py @@ -195,10 +195,9 @@ def test_decomposition_info_matches_gridsize(datapath, caplog): @pytest.mark.parametrize("ndyn_substeps", [5]) @pytest.mark.parametrize("linit", [True, False]) def test_parallel_diffusion( - datapath, r04b09_diffusion_config, step_date_init, linit, caplog, ndyn_substeps + datapath, r04b09_diffusion_config, step_date_init, linit, ndyn_substeps ): - caplog.set_level(logging.WARN) props = get_processor_properties() print( @@ -219,7 +218,7 @@ def test_parallel_diffusion( print( f"rank={props.rank}/{props.comm_size}: GHEX context setup: from {props.comm_name} with {props.comm_size} nodes" ) - assert context.size() == 2 + assert context.size() == 2, "unexpected context sise" icon_grid = read_icon_grid(path, rank=props.rank) print( diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 3e583c515..fee747293 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -360,14 +360,6 @@ def _verify_diffusion_fields( assert np.allclose(ref_exner, val_exner) -def print_diff_info(ref, value): - difference = np.abs(ref - value) - print("differences ----- ") - print(f" max diff vn {np.max(difference)}") - d = 0.000001 - print(f"number of diffs > {d}: {np.count_nonzero(difference > d)}") - - @pytest.mark.datatest @pytest.mark.parametrize("linit", [True]) @pytest.mark.parametrize("datapath", [1], indirect=True) From 9f2530b25dd575c69107cf3d429ce5171c40db78 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 2 Aug 2023 13:31:38 +0200 Subject: [PATCH 204/263] clean up tests make tests run with 1, 2, 4 mpi nodes make tests run without mpi --- .../src/icon4py/decomposition/decomposed.py | 19 +-- .../icon4py/decomposition/parallel_setup.py | 28 ++-- .../src/icon4py/diffusion/diffusion.py | 20 +-- .../src/icon4py/driver/dycore_driver.py | 1 + atm_dyn_iconam/src/icon4py/driver/io_utils.py | 5 +- atm_dyn_iconam/tests/conftest.py | 14 +- atm_dyn_iconam/tests/mpi_tests/README.md | 18 ++- atm_dyn_iconam/tests/mpi_tests/common.py | 22 +++ .../tests/mpi_tests/test_decomposed.py | 148 ++---------------- .../mpi_tests/test_parallel_diffusion.py | 121 ++++++++++++++ .../tests/mpi_tests/test_parallel_setup.py | 50 ++++++ atm_dyn_iconam/tests/test_diffusion.py | 72 ++++----- atm_dyn_iconam/tests/test_diffusion_utils.py | 1 - atm_dyn_iconam/tests/test_utils/fixtures.py | 4 +- .../tests/test_utils/serialbox_utils.py | 31 ++-- common/tests/test_grid_manager.py | 34 ++-- 16 files changed, 327 insertions(+), 261 deletions(-) create mode 100644 atm_dyn_iconam/tests/mpi_tests/common.py create mode 100644 atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py create mode 100644 atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index 2d7fbb301..cb6049c0b 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -24,14 +24,9 @@ from icon4py.diffusion.diffusion_utils import builder - - log = logging.getLogger(__name__) - - - class DomainDescriptorIdGenerator: _counter = 0 _roundtrips = 0 @@ -167,7 +162,9 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): ), EdgeDim: self._create_domain_descriptor(EdgeDim), } - log.info(f"domain descriptors for dimensions {self._domain_descriptors.keys()} initialized") + log.info( + f"domain descriptors for dimensions {self._domain_descriptors.keys()} initialized" + ) self._patterns = { CellDim: self._create_pattern(CellDim), @@ -176,7 +173,7 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): } log.info(f"patterns for dimensions {self._patterns.keys()} initialized ") self._comm = ghex.make_co(context) - log.info(f"communication object initialized") + log.info("communication object initialized") def _domain_descriptor_info(self, descr): return f" domain_descriptor=[id='{descr.domain_id()}', size='{descr.size()}', inner_size='{descr.inner_size()}' (halo size='{descr.size() - descr.inner_size()}')" @@ -226,13 +223,17 @@ def exchange(self, dim: Dimension, *fields: tuple): pattern = self._patterns[dim] assert pattern is not None, f"pattern for {dim.value} not found" domain_descriptor = self._domain_descriptors[dim] - assert domain_descriptor is not None, f"domain descriptor for {dim.value} not found" + assert ( + domain_descriptor is not None + ), f"domain descriptor for {dim.value} not found" applied_patterns = [ pattern(ghex.field_descriptor(domain_descriptor, np.asarray(f))) for f in fields ] handle = self._comm.exchange(applied_patterns) - log.info(f"exchange for {len(fields)} fields of dimension ='{dim.value}' initiated.") + log.info( + f"exchange for {len(fields)} fields of dimension ='{dim.value}' initiated." + ) return MultiNodeResult(handle, applied_patterns) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py index 080690335..7e00c541e 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py @@ -11,22 +11,19 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import logging -from typing import Union +from typing import Optional, Union import mpi4py from mpi4py.MPI import Comm - mpi4py.rc.initialize = False CommId = Union[int, Comm, None] log = logging.getLogger(__name__) -def get_processor_properties(comm_id: CommId = None): - init_mpi() - +def get_processor_properties(with_mpi=False, comm_id: CommId = None): def _get_current_comm_or_comm_world(comm_id: CommId) -> Comm: if isinstance(comm_id, int): comm = Comm.f2py(comm_id) @@ -36,8 +33,12 @@ def _get_current_comm_or_comm_world(comm_id: CommId) -> Comm: comm = mpi4py.MPI.COMM_WORLD return comm - current_comm = _get_current_comm_or_comm_world(comm_id) - return ProcessProperties.from_mpi_comm(current_comm) + if with_mpi: + init_mpi() + current_comm = _get_current_comm_or_comm_world(comm_id) + return ProcessProperties.from_mpi_comm(current_comm) + else: + return ProcessProperties.from_single_node() def init_mpi(): @@ -57,10 +58,10 @@ def finalize_mpi(): class ProcessProperties: - def __init__(self, comm: mpi4py.MPI.Comm): - self._communicator_name: str = comm.Get_name() - self._rank: int = comm.Get_rank() - self._comm_size = comm.Get_size() + def __init__(self, comm: Optional[mpi4py.MPI.Comm]): + self._communicator_name: str = comm.Get_name() if comm else "" + self._rank: int = comm.Get_rank() if comm else 0 + self._comm_size = comm.Get_size() if comm else 1 self._comm = comm @property @@ -83,6 +84,11 @@ def comm(self): def from_mpi_comm(cls, comm: mpi4py.MPI.Comm): return ProcessProperties(comm) + @classmethod + def from_single_node(cls): + return ProcessProperties(None) + + class ParallelLogger(logging.Filter): def __init__(self, processProperties: ProcessProperties = None): super().__init__() diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 8b23c60c7..3aa99b62d 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -552,9 +552,7 @@ def _sync_cell_fields(self, prognostic_state): communication only done in original code if the following condition applies: IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN """ - log.debug( - f"communication of prognostic cell fields: theta, w, exner - start" - ) + log.debug(f"communication of prognostic cell fields: theta, w, exner - start") handle_cell_comm = self._exchange.exchange( CellDim, prognostic_state.w, @@ -562,9 +560,7 @@ def _sync_cell_fields(self, prognostic_state): prognostic_state.exner_pressure, ) handle_cell_comm.wait() - log.debug( - f"communication of prognostic cell fields: theta, w, exner - done" - ) + log.debug(f"communication of prognostic cell fields: theta, w, exner - done") def _do_diffusion_step( self, @@ -740,9 +736,7 @@ def _do_diffusion_step( h.wait() log.debug(f"communication rbf extrapolation of z_nable2_e - end") - log.debug( - f"running stencils 04 05 06 (apply_diffusion_to_vn): start" - ) + log.debug(f"running stencils 04 05 06 (apply_diffusion_to_vn): start") apply_diffusion_to_vn.with_backend(backend)( u_vert=self.u_vert, v_vert=self.v_vert, @@ -770,9 +764,7 @@ def _do_diffusion_step( "E2ECV": self.grid.get_e2ecv_connectivity(), }, ) - log.debug( - f"running stencils 04 05 06 (apply_diffusion_to_vn): end" - ) + log.debug(f"running stencils 04 05 06 (apply_diffusion_to_vn): end") log.debug(f"communication of prognistic.vn : start") handle_edge_comm = self._exchange.exchange(EdgeDim, prognostic_state.vn) @@ -853,9 +845,7 @@ def _do_diffusion_step( "C2CE": self.grid.get_c2ce_connectivity(), }, ) - log.debug( - f"running stencils 13_14 (calculate_nabla2_for_theta): end" - ) + log.debug(f"running stencils 13_14 (calculate_nabla2_for_theta): end") log.debug( f"running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start" ) diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index a279eb130..4fe0699f1 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -39,6 +39,7 @@ helpers = import_testutils() from helpers import serialbox_utils as sb_utils # noqa + log = logging.getLogger(__name__) diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index 13e9afbcd..e3ae2d338 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -18,7 +18,7 @@ from pathlib import Path from icon4py.decomposition.decomposed import DecompositionInfo -from icon4py.decomposition.parallel_setup import ProcessProperties, ParallelLogger +from icon4py.decomposition.parallel_setup import ParallelLogger, ProcessProperties from icon4py.diffusion.state_utils import ( DiagnosticState, InterpolationState, @@ -136,7 +136,6 @@ def read_geometry_fields( raise NotImplementedError(SB_ONLY_MSG) -# /home/magdalena/data/exclaim/dycore/mch_ch_r04b09_dsl/node2/mch_ch_r04b09_dsl/icon_grid def read_decomp_info( path: Path, procs_props: ProcessProperties, @@ -193,7 +192,7 @@ def read_static_fields( raise NotImplementedError(SB_ONLY_MSG) -def configure_logging(run_path: str, start_time, processor_procs = None) -> None: +def configure_logging(run_path: str, start_time, processor_procs=None) -> None: """ Configure logging. diff --git a/atm_dyn_iconam/tests/conftest.py b/atm_dyn_iconam/tests/conftest.py index eb6d1b956..f76c20b45 100644 --- a/atm_dyn_iconam/tests/conftest.py +++ b/atm_dyn_iconam/tests/conftest.py @@ -11,12 +11,12 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import mpi4py + import pytest from gt4py.next.program_processors.runners.roundtrip import executor from atm_dyn_iconam.tests.test_utils.simple_mesh import SimpleMesh -from icon4py.diffusion.diffusion import DiffusionConfig +from icon4py.diffusion.diffusion import DiffusionConfig, DiffusionType from .test_utils.fixtures import ( # noqa F401 damping_height, @@ -30,14 +30,6 @@ ) -@pytest.fixture -def mpi(): - from mpi4py import MPI - - mpi4py.rc.initialize = False - return MPI - - @pytest.fixture def ndyn_substeps(): """ @@ -125,7 +117,7 @@ def r04b09_diffusion_config(ndyn_substeps) -> DiffusionConfig: from the default. """ return DiffusionConfig( - diffusion_type=5, + diffusion_type=DiffusionType.SMAGORINSKY_4TH_ORDER, hdiff_w=True, hdiff_vn=True, type_t_diffu=2, diff --git a/atm_dyn_iconam/tests/mpi_tests/README.md b/atm_dyn_iconam/tests/mpi_tests/README.md index 1d495fa95..55cec7e57 100644 --- a/atm_dyn_iconam/tests/mpi_tests/README.md +++ b/atm_dyn_iconam/tests/mpi_tests/README.md @@ -4,8 +4,20 @@ The parallelized code uses [GHEX](https://github.com/ghex-org/GHEX) with MPI for halo exchanges. GHEX has a CMake build but no setup script for pip, so it needs to be installed manually: -1. You need a running MPI installation in the system. -2. You need to have boost (headers) installed in the system 3clone GHEX +1. You need a running MPI installation in the system. On linux (apt base system) do + +```bash +sudo apt-get install libopenmpi-dev +``` + +on MacOs + +```bash +brew install mpich +``` + +2. You need to have boost (headers) installed in the system +3. clone GHEX: ```bash cd {icon4py}/_external_src @@ -63,7 +75,7 @@ make test ## will now run the python tests #### use GHEX bindings from icon4py -simply create a sym link the the installation above: +simply create a sym link the installation above: ``` cd {icon4py}/.venv/lib/python3.10/site-packages diff --git a/atm_dyn_iconam/tests/mpi_tests/common.py b/atm_dyn_iconam/tests/mpi_tests/common.py new file mode 100644 index 000000000..68de58e6c --- /dev/null +++ b/atm_dyn_iconam/tests/mpi_tests/common.py @@ -0,0 +1,22 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from pathlib import Path + +from icon4py.decomposition.parallel_setup import get_processor_properties + + +base_path = Path(__file__).parent.parent.parent.parent.joinpath("testdata/ser_icondata") + +props = get_processor_properties(with_mpi=True) +path = base_path.joinpath(f"mpitask{props.comm_size}/mch_ch_r04b09_dsl/ser_data") diff --git a/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py b/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py index 72bf19f44..0af3bd2fd 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py @@ -10,28 +10,21 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -import logging import ghex import numpy as np import pytest -from atm_dyn_iconam.tests.test_diffusion import _verify_diffusion_fields -from atm_dyn_iconam.tests.test_utils.serialbox_utils import IconSerialDataProvider +from atm_dyn_iconam.tests.mpi_tests.common import path, props from icon4py.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.decomposition.decomposed import ( DecompositionInfo, DomainDescriptorIdGenerator, - MultiNode, ) -from icon4py.decomposition.parallel_setup import get_processor_properties -from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.driver.io_utils import ( SerializationType, read_decomp_info, - read_geometry_fields, read_icon_grid, - read_static_fields, ) @@ -45,12 +38,9 @@ """ -props = get_processor_properties() - -@pytest.mark.skipif( - props.comm_size > 2, reason="input files only available for 1 or 2 nodes" -) +@pytest.mark.mpi +@pytest.mark.skipif(props.comm_size != 2, reason="runs on 2 nodes only") @pytest.mark.parametrize( ("dim, owned, total"), ( @@ -59,11 +49,9 @@ (VertexDim, (5373, 5290), (5455, 5456)), ), ) -@pytest.mark.parametrize("datapath", [2], indirect=True) -def test_decomposition_info_masked(mpi, datapath, dim, owned, total, caplog): - props = get_processor_properties() +def test_decomposition_info_masked(dim, owned, total, caplog): my_rank = props.rank - decomposition_info = read_decomp_info(datapath, props, SerializationType.SB) + decomposition_info = read_decomp_info(path, props, SerializationType.SB) all_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.ALL) my_total = total[my_rank] my_owned = owned[my_rank] @@ -81,10 +69,7 @@ def test_decomposition_info_masked(mpi, datapath, dim, owned, total, caplog): _assert_index_partitioning(all_indices, halo_indices, owned_indices) -@pytest.mark.skipif( - props.comm_size > 2, reason="input files only available for 1 or 2 nodes" -) -@pytest.mark.parametrize("datapath", [2], indirect=True) +@pytest.mark.skipif(props.comm_size != 2, reason="runs on 2 nodes only") @pytest.mark.parametrize( ("dim, owned, total"), ( @@ -93,11 +78,10 @@ def test_decomposition_info_masked(mpi, datapath, dim, owned, total, caplog): (VertexDim, (5373, 5290), (5455, 5456)), ), ) -def test_decomposition_info_local_index(mpi, datapath, dim, owned, total, caplog): - props = get_processor_properties() +def test_decomposition_info_local_index(dim, owned, total, caplog): my_rank = props.rank - decomposition_info = read_decomp_info(datapath, props, SerializationType.SB) + decomposition_info = read_decomp_info(path, props, SerializationType.SB) all_indices = decomposition_info.local_index(dim, DecompositionInfo.EntryType.ALL) my_total = total[my_rank] my_owned = owned[my_rank] @@ -128,19 +112,10 @@ def _assert_index_partitioning(all_indices, halo_indices, owned_indices): assert set(halos_list) | set(owned_list) == set(all_list) -@pytest.mark.mpi -def test_processor_properties_from_comm_world(mpi, caplog): - caplog.set_level(logging.DEBUG) - props = get_processor_properties() - - assert props.rank < mpi.COMM_WORLD.Get_size() - assert props.comm_name == mpi.COMM_WORLD.Get_name() - - @pytest.mark.mpi @pytest.mark.parametrize("num", [1, 2, 3]) def test_domain_descriptor_id_are_globally_unique(num): - props = get_processor_properties() + size = props.comm_size context = ghex.context(ghex.mpi_comm(props.comm), True) id_gen = DomainDescriptorIdGenerator(context) @@ -161,15 +136,17 @@ def test_domain_descriptor_id_are_globally_unique(num): @pytest.mark.mpi -@pytest.mark.parametrize("datapath", [2], indirect=True) -def test_decomposition_info_matches_gridsize(datapath, caplog): - props = get_processor_properties() +@pytest.mark.skipif( + props.comm_size not in (1, 2, 4), + reason="input files only available for 1 or 2 nodes", +) +def test_decomposition_info_matches_gridsize(caplog): decomposition_info = read_decomp_info( - datapath, + path, props, SerializationType.SB, ) - icon_grid = read_icon_grid(datapath, props.rank) + icon_grid = read_icon_grid(path, props.rank) assert ( decomposition_info.global_index( dim=CellDim, entry_type=DecompositionInfo.EntryType.ALL @@ -188,96 +165,3 @@ def test_decomposition_info_matches_gridsize(datapath, caplog): ] == icon_grid.num_edges() ) - - -@pytest.mark.mpi -@pytest.mark.parametrize("datapath", [2], indirect=True) -@pytest.mark.parametrize("ndyn_substeps", [5]) -@pytest.mark.parametrize("linit", [True, False]) -def test_parallel_diffusion( - datapath, r04b09_diffusion_config, step_date_init, linit, ndyn_substeps -): - - props = get_processor_properties() - - print( - f"rank={props.rank}/{props.comm_size}: inializing diffusion for experiment 'mch_ch_r04_b09_dsl" - ) - path = datapath - decomp_info = read_decomp_info( - path, - props, - ) - print( - f"rank={props.rank}/{props.comm_size}: decomposition info : klevels = {decomp_info.klevels}, " - f"local cells = {decomp_info.global_index(CellDim, DecompositionInfo.EntryType.ALL).shape} " - f"local edges = {decomp_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape} " - f"local vertices = {decomp_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape}" - ) - context = ghex.context(ghex.mpi_comm(props.comm), True) - print( - f"rank={props.rank}/{props.comm_size}: GHEX context setup: from {props.comm_name} with {props.comm_size} nodes" - ) - assert context.size() == 2, "unexpected context sise" - - icon_grid = read_icon_grid(path, rank=props.rank) - print( - f"rank={props.rank}: using local grid with {icon_grid.num_cells()} Cells, {icon_grid.num_edges()} Edges, {icon_grid.num_vertices()} Vertices" - ) - diffusion_params = DiffusionParams(r04b09_diffusion_config) - - diffusion_initial_data = IconSerialDataProvider( - "icon_pydycore", str(path), True, mpi_rank=props.rank - ).from_savepoint_diffusion_init(linit=linit, date=step_date_init) - (edge_geometry, cell_geometry, vertical_geometry) = read_geometry_fields( - path, rank=props.rank - ) - (metric_state, interpolation_state) = read_static_fields(path, rank=props.rank) - - dtime = diffusion_initial_data.get_metadata("dtime").get("dtime") - print( - f"rank={props.rank}/{props.comm_size}: setup: using {props.comm_name} with {props.comm_size} nodes" - ) - exchange = MultiNode(context, decomp_info) - - diffusion = Diffusion(exchange) - - diffusion.init( - grid=icon_grid, - config=r04b09_diffusion_config, - params=diffusion_params, - vertical_params=vertical_geometry, - metric_state=metric_state, - interpolation_state=interpolation_state, - edge_params=edge_geometry, - cell_params=cell_geometry, - ) - print(f"rank={props.rank}/{props.comm_size}: diffusion initialized ") - diagnostic_state = diffusion_initial_data.construct_diagnostics_for_diffusion() - prognostic_state = diffusion_initial_data.construct_prognostics() - if linit: - diffusion.initial_run( - diagnostic_state=diagnostic_state, - prognostic_state=prognostic_state, - dtime=dtime, - ) - else: - diffusion.run( - diagnostic_state=diagnostic_state, - prognostic_state=prognostic_state, - dtime=dtime, - ) - print(f"rank={props.rank}/{props.comm_size}: diffusion run ") - - diffusion_savepoint_exit = IconSerialDataProvider( - "icon_pydycore", str(path), True, mpi_rank=props.rank - ).from_savepoint_diffusion_exit(linit=linit, date=step_date_init) - _verify_diffusion_fields( - diagnostic_state=diagnostic_state, - prognostic_state=prognostic_state, - diffusion_savepoint=diffusion_savepoint_exit, - ) - print( - f"rank={props.rank}/{props.comm_size}: running diffusion step - using {props.comm_name} with {props.comm_size} nodes - DONE" - ) - del context diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py new file mode 100644 index 000000000..e9fd2eda6 --- /dev/null +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py @@ -0,0 +1,121 @@ +# 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 ghex +import pytest + +from atm_dyn_iconam.tests.mpi_tests.common import path, props +from atm_dyn_iconam.tests.test_diffusion import _verify_diffusion_fields +from atm_dyn_iconam.tests.test_utils.serialbox_utils import IconSerialDataProvider +from icon4py.common.dimension import CellDim, EdgeDim, VertexDim +from icon4py.decomposition.decomposed import DecompositionInfo, MultiNode +from icon4py.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.driver.io_utils import ( + read_decomp_info, + read_geometry_fields, + read_icon_grid, + read_static_fields, +) + + +@pytest.mark.mpi +@pytest.mark.parametrize("ndyn_substeps", [2]) +@pytest.mark.parametrize("linit", [True, False]) +@pytest.mark.skipif( + props.comm_size not in (1, 2, 4), + reason="input files only available for 1 or 2 nodes", +) +def test_parallel_diffusion( + r04b09_diffusion_config, step_date_init, linit, ndyn_substeps +): + + print( + f"rank={props.rank}/{props.comm_size}: inializing diffusion for experiment 'mch_ch_r04_b09_dsl" + ) + decomp_info = read_decomp_info( + path, + props, + ) + print( + f"rank={props.rank}/{props.comm_size}: decomposition info : klevels = {decomp_info.klevels}, " + f"local cells = {decomp_info.global_index(CellDim, DecompositionInfo.EntryType.ALL).shape} " + f"local edges = {decomp_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape} " + f"local vertices = {decomp_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape}" + ) + context = ghex.context(ghex.mpi_comm(props.comm), True) + print( + f"rank={props.rank}/{props.comm_size}: GHEX context setup: from {props.comm_name} with {props.comm_size} nodes" + ) + + icon_grid = read_icon_grid(path, rank=props.rank) + print( + f"rank={props.rank}: using local grid with {icon_grid.num_cells()} Cells, {icon_grid.num_edges()} Edges, {icon_grid.num_vertices()} Vertices" + ) + diffusion_params = DiffusionParams(r04b09_diffusion_config) + + diffusion_initial_data = IconSerialDataProvider( + "icon_pydycore", str(path), True, mpi_rank=props.rank + ).from_savepoint_diffusion_init(linit=linit, date=step_date_init) + (edge_geometry, cell_geometry, vertical_geometry) = read_geometry_fields( + path, rank=props.rank + ) + (metric_state, interpolation_state) = read_static_fields(path, rank=props.rank) + + dtime = diffusion_initial_data.get_metadata("dtime").get("dtime") + print( + f"rank={props.rank}/{props.comm_size}: setup: using {props.comm_name} with {props.comm_size} nodes" + ) + exchange = MultiNode(context, decomp_info) + + diffusion = Diffusion(exchange) + + diffusion.init( + grid=icon_grid, + config=r04b09_diffusion_config, + params=diffusion_params, + vertical_params=vertical_geometry, + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry, + ) + print(f"rank={props.rank}/{props.comm_size}: diffusion initialized ") + diagnostic_state = diffusion_initial_data.construct_diagnostics_for_diffusion() + prognostic_state = diffusion_initial_data.construct_prognostics() + if linit: + diffusion.initial_run( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + ) + else: + diffusion.run( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + ) + print(f"rank={props.rank}/{props.comm_size}: diffusion run ") + + diffusion_savepoint_exit = IconSerialDataProvider( + "icon_pydycore", str(path), True, mpi_rank=props.rank + ).from_savepoint_diffusion_exit(linit=linit, date=step_date_init) + _verify_diffusion_fields( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + diffusion_savepoint=diffusion_savepoint_exit, + ) + print( + f"rank={props.rank}/{props.comm_size}: running diffusion step - using {props.comm_name} with {props.comm_size} nodes - DONE" + ) + del context diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py new file mode 100644 index 000000000..cfdfc49ec --- /dev/null +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -0,0 +1,50 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + + +import pytest +from mpi4py import MPI + +from icon4py.decomposition.parallel_setup import ( + get_processor_properties, + init_mpi, +) + + +@pytest.mark.mpi +def test_parallel_properties_from_comm_world(): + props = get_processor_properties(with_mpi=True) + assert props.rank < props.comm_size + assert props.comm_name == "MPI_COMM_WORLD" + + +@pytest.mark.mpi(min_size=2) +def test_parallel_properties_from_mpi_comm(): + init_mpi() + world = MPI.COMM_WORLD + group = world.Get_group() + pair = group.Incl([0, 1]) + comm = world.Create(pair) + if comm != MPI.COMM_NULL: + comm.Set_name("my_comm") + props = get_processor_properties(with_mpi=True, comm_id=comm) + assert props.rank < props.comm_size + assert props.comm_size == 2 + assert props.comm_name == "my_comm" + + +def test_single_node_properties(): + props = get_processor_properties() + assert props.comm_size == 1 + assert props.rank == 0 + assert props.comm_name == "" diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index fee747293..0b5234275 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -93,7 +93,6 @@ def test_smagorinski_factor_diffusion_type_5(r04b09_diffusion_config): @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_diffusion_init( diffusion_savepoint_init, interpolation_savepoint, @@ -108,9 +107,8 @@ def test_diffusion_init( additional_parameters = DiffusionParams(config) vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) - meta = diffusion_savepoint_init.get_metadata("nlev", "linit", "date") + meta = diffusion_savepoint_init.get_metadata("linit", "date") - assert meta["nlev"] == 65 assert meta["linit"] is False assert meta["date"] == step_date_init @@ -193,7 +191,6 @@ def _verify_init_values_against_savepoint( @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_verify_diffusion_init_against_first_regular_savepoint( diffusion_savepoint_init, interpolation_savepoint, @@ -230,7 +227,6 @@ def test_verify_diffusion_init_against_first_regular_savepoint( @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) @pytest.mark.parametrize("step_date_init", ["2021-06-20T12:00:50.000"]) def test_verify_diffusion_init_against_other_regular_savepoint( r04b09_diffusion_config, @@ -267,8 +263,39 @@ def test_verify_diffusion_init_against_other_regular_savepoint( _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) +def _verify_diffusion_fields( + diagnostic_state: DiagnosticState, + prognostic_state: PrognosticState, + diffusion_savepoint: IconDiffusionExitSavepoint, +): + ref_div_ic = np.asarray(diffusion_savepoint.div_ic()) + val_div_ic = np.asarray(diagnostic_state.div_ic) + ref_hdef_ic = np.asarray(diffusion_savepoint.hdef_ic()) + val_hdef_ic = np.asarray(diagnostic_state.hdef_ic) + assert np.allclose(ref_div_ic, val_div_ic) + assert np.allclose(ref_hdef_ic, val_hdef_ic) + ref_w = np.asarray(diffusion_savepoint.w()) + val_w = np.asarray(prognostic_state.w) + ref_dwdx = np.asarray(diffusion_savepoint.dwdx()) + val_dwdx = np.asarray(diagnostic_state.dwdx) + ref_dwdy = np.asarray(diffusion_savepoint.dwdy()) + val_dwdy = np.asarray(diagnostic_state.dwdy) + assert np.allclose(ref_dwdx, val_dwdx) + assert np.allclose(ref_dwdy, val_dwdy) + + ref_vn = np.asarray(diffusion_savepoint.vn()) + val_vn = np.asarray(prognostic_state.vn) + assert np.allclose(ref_vn, val_vn) + assert np.allclose(ref_w, val_w) + ref_exner = np.asarray(diffusion_savepoint.exner()) + ref_theta_v = np.asarray(diffusion_savepoint.theta_v()) + val_theta_v = np.asarray(prognostic_state.theta_v) + val_exner = np.asarray(prognostic_state.exner_pressure) + assert np.allclose(ref_theta_v, val_theta_v) + assert np.allclose(ref_exner, val_exner) + + @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) @pytest.mark.parametrize( "step_date_init, step_date_exit", [ @@ -328,41 +355,8 @@ def test_run_diffusion_single_step( ) -def _verify_diffusion_fields( - diagnostic_state: DiagnosticState, - prognostic_state: PrognosticState, - diffusion_savepoint: IconDiffusionExitSavepoint, -): - ref_div_ic = np.asarray(diffusion_savepoint.div_ic()) - val_div_ic = np.asarray(diagnostic_state.div_ic) - ref_hdef_ic = np.asarray(diffusion_savepoint.hdef_ic()) - val_hdef_ic = np.asarray(diagnostic_state.hdef_ic) - assert np.allclose(ref_div_ic, val_div_ic) - assert np.allclose(ref_hdef_ic, val_hdef_ic) - ref_w = np.asarray(diffusion_savepoint.w()) - val_w = np.asarray(prognostic_state.w) - ref_dwdx = np.asarray(diffusion_savepoint.dwdx()) - val_dwdx = np.asarray(diagnostic_state.dwdx) - ref_dwdy = np.asarray(diffusion_savepoint.dwdy()) - val_dwdy = np.asarray(diagnostic_state.dwdy) - assert np.allclose(ref_dwdx, val_dwdx) - assert np.allclose(ref_dwdy, val_dwdy) - - ref_vn = np.asarray(diffusion_savepoint.vn()) - val_vn = np.asarray(prognostic_state.vn) - assert np.allclose(ref_vn, val_vn) - assert np.allclose(ref_w, val_w) - ref_exner = np.asarray(diffusion_savepoint.exner()) - ref_theta_v = np.asarray(diffusion_savepoint.theta_v()) - val_theta_v = np.asarray(prognostic_state.theta_v) - val_exner = np.asarray(prognostic_state.exner_pressure) - assert np.allclose(ref_theta_v, val_theta_v) - assert np.allclose(ref_exner, val_exner) - - @pytest.mark.datatest @pytest.mark.parametrize("linit", [True]) -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_run_diffusion_initial_step( diffusion_savepoint_init, diffusion_savepoint_exit, diff --git a/atm_dyn_iconam/tests/test_diffusion_utils.py b/atm_dyn_iconam/tests/test_diffusion_utils.py index ead7a4198..53142b8b6 100644 --- a/atm_dyn_iconam/tests/test_diffusion_utils.py +++ b/atm_dyn_iconam/tests/test_diffusion_utils.py @@ -151,7 +151,6 @@ def test_set_zero_vertex_k(): @pytest.mark.datatest @pytest.mark.parametrize("linit", [True]) -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( diffusion_savepoint_init, r04b09_diffusion_config, icon_grid, ndyn_substeps ): diff --git a/atm_dyn_iconam/tests/test_utils/fixtures.py b/atm_dyn_iconam/tests/test_utils/fixtures.py index 0fce7da2d..4230001df 100644 --- a/atm_dyn_iconam/tests/test_utils/fixtures.py +++ b/atm_dyn_iconam/tests/test_utils/fixtures.py @@ -41,8 +41,8 @@ @pytest.fixture -def datapath(setup_icon_data, request): - local_path = f"mpitask{request.param}/mch_ch_r04b09_dsl/ser_data" +def datapath(setup_icon_data): + local_path = "mpitask1/mch_ch_r04b09_dsl/ser_data" return data_path.joinpath(local_path) diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index 8152fbea4..eea4f0ff9 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -360,25 +360,32 @@ def zd_diffcoef(self): return self._get_field("zd_diffcoef", CellDim, KDim) def zd_intcoef(self): - ser_input = np.moveaxis( - (np.squeeze(self.serializer.read("vcoef", self.savepoint))), 1, -1 - )[: self.sizes[CellDim], :, :] - return self._linearize_first_2dims(ser_input, sparse_size=3) + return self._read_and_reorder_sparse_field("vcoef", CellDim) - def _linearize_first_2dims(self, data: np.ndarray, sparse_size): + def _read_and_reorder_sparse_field( + self, name: str, horizontal_dim: Dimension, sparse_size=3 + ): + ser_input = np.squeeze(self.serializer.read(name, self.savepoint))[ + : self.sizes[horizontal_dim], :, : + ] + if ser_input.shape[1] != sparse_size: + ser_input = np.moveaxis((ser_input), 1, -1) + + return self._linearize_first_2dims( + ser_input, sparse_size=sparse_size, target_dims=(CECDim, KDim) + ) + + def _linearize_first_2dims( + self, data: np.ndarray, sparse_size: int, target_dims: tuple[Dimension, ...] + ): old_shape = data.shape assert old_shape[1] == sparse_size - return np_as_located_field(CECDim, KDim)( + return np_as_located_field(*target_dims)( data.reshape(old_shape[0] * old_shape[1], old_shape[2]) ) def zd_vertoffset(self): - ser_input = np.squeeze(self.serializer.read("zd_vertoffset", self.savepoint)) - ser_input = np.moveaxis(ser_input, 1, -1) - - return self._linearize_first_2dims( - ser_input[: self.sizes[CellDim], :, :], sparse_size=3 - ) + return self._read_and_reorder_sparse_field("zd_vertoffset", CellDim) def zd_vertidx(self): return np.squeeze(self.serializer.read("zd_vertidx", self.savepoint)) diff --git a/common/tests/test_grid_manager.py b/common/tests/test_grid_manager.py index 451d22ed0..12b164d0d 100644 --- a/common/tests/test_grid_manager.py +++ b/common/tests/test_grid_manager.py @@ -221,7 +221,6 @@ def test_gridparser_dimension(simple_mesh_data): @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_gridfile_vertex_cell_edge_dimensions(grid_savepoint, r04b09_dsl_gridfile): data = Dataset(r04b09_dsl_gridfile, "r") grid_file = GridFile(data) @@ -254,7 +253,6 @@ def test_grid_parser_index_fields(simple_mesh_data, caplog): # v2e: exists in serial, simple, grid @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_gridmanager_eval_v2e(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() @@ -270,7 +268,6 @@ def test_gridmanager_eval_v2e(caplog, grid_savepoint, r04b09_dsl_gridfile): # v2c: exists in serial, simple, grid @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_gridmanager_eval_v2c(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() @@ -315,7 +312,6 @@ def reset_invalid_index(index_array: np.ndarray): # e2v: exists in serial, simple, grid @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_gridmanager_eval_e2v(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() @@ -334,7 +330,6 @@ def has_invalid_index(ar: np.ndarray): # e2c : exists in serial, simple, grid @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_gridmanager_eval_e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() @@ -349,7 +344,6 @@ def test_gridmanager_eval_e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): # c2e: serial, simple, grid @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_gridmanager_eval_c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() @@ -364,8 +358,18 @@ def test_gridmanager_eval_c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): # e2c2e (e2c2eo) - diamond: exists in serial, simple_mesh +@pytest.mark.datatest +def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + assert np.allclose( + grid.get_c2e2c_connectivity().table, + grid_savepoint.c2e2c()[0 : grid.num_cells(), :], + ) + + +# c2e2c: exists in serial, simple_mesh, grid @pytest.mark.skip("does not directly exist in the grid file, needs to be constructed") -@pytest.mark.parametrize("datapath", [1], indirect=True) # TODO (Magdalena) construct from adjacent_cell_of_edge and then edge_of_cell @pytest.mark.datatest def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): @@ -378,20 +382,7 @@ def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): assert np.allclose(grid.get_e2c2e_connectivity().table, serialized_e2c2e) -# c2e2c: exists in serial, simple_mesh, grid -@pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) -def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): - caplog.set_level(logging.DEBUG) - grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() - assert np.allclose( - grid.get_c2e2c_connectivity().table, - grid_savepoint.c2e2c()[0 : grid.num_cells(), :], - ) - - @pytest.mark.xfail -@pytest.mark.parametrize("datapath", [1], indirect=True) @pytest.mark.datatest def test_gridmanager_eval_e2c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) @@ -405,7 +396,6 @@ def test_gridmanager_eval_e2c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_gridmanager_eval_c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() @@ -462,7 +452,6 @@ def test_gt4py_transform_offset_by_1_where_valid(size): @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) @pytest.mark.parametrize( "dim, marker, index", [ @@ -509,7 +498,6 @@ def test_get_start_index(r04b09_dsl_gridfile, icon_grid, dim, marker, index): @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) @pytest.mark.parametrize( "dim, marker, index", [ From 8186de50bd30229c095e5309d42b3538dcb4318f Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 2 Aug 2023 14:35:19 +0200 Subject: [PATCH 205/263] clean up tests make tests run with 1, 2, 4 mpi nodes make tests run without mpi --- atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py index 7e00c541e..d0665af11 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py @@ -96,6 +96,8 @@ def __init__(self, processProperties: ProcessProperties = None): if processProperties and processProperties.comm_size > 1: self._rank_info = f"rank={processProperties.rank}/{processProperties.comm_size} [{processProperties.comm_name}] >>>" - def filter(self, record: logging.LogRecord) -> bool: + def filter( + self, record: logging.LogRecord + ) -> bool: record.rank = self._rank_info return True From fcf4519189a3196219ac049530ec910cbec616bc Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 2 Aug 2023 17:23:03 +0200 Subject: [PATCH 206/263] fix dummy driver to run with MPI --- .../src/icon4py/decomposition/decomposed.py | 29 ++++++++++--- .../icon4py/decomposition/parallel_setup.py | 6 +-- atm_dyn_iconam/src/icon4py/driver/README.md | 23 +++++++++-- .../src/icon4py/driver/dycore_driver.py | 41 ++++++++++++------- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 31 +++++--------- .../tests/mpi_tests/test_decomposed.py | 22 ++++++++++ .../mpi_tests/test_parallel_diffusion.py | 8 ++-- atm_dyn_iconam/tests/test_utils/fixtures.py | 3 +- 8 files changed, 108 insertions(+), 55 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index cb6049c0b..5ffc82568 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -15,12 +15,14 @@ from enum import Enum from typing import Protocol +import ghex +import ghex.unstructured as unstructured import numpy as np import numpy.ma as ma -from ghex import unstructured as ghex from gt4py.next import Dimension from icon4py.common.dimension import CellDim, DimensionKind, EdgeDim, VertexDim +from icon4py.decomposition.parallel_setup import ProcessProperties from icon4py.diffusion.diffusion_utils import builder @@ -128,6 +130,21 @@ def is_ready(self) -> bool: return True +def create_exchange( + props: ProcessProperties, decomp_info: DecompositionInfo +) -> ExchangeRuntime: + """ + Create an Exchange depending on the runtime size. + + Depending on the number of processor a SingleNode version is returned or a GHEX context created and a Multinode returned. + """ + if props.comm_size > 1: + context = ghex.context(ghex.mpi_comm(props.comm), True) + return MultiNode(context, decomp_info) + else: + return SingleNode() + + @dataclass class SingleNode: def exchange(self, dim: Dimension, *fields: tuple) -> ExchangeResult: @@ -172,7 +189,7 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): EdgeDim: self._create_pattern(EdgeDim), } log.info(f"patterns for dimensions {self._patterns.keys()} initialized ") - self._comm = ghex.make_co(context) + self._comm = unstructured.make_co(context) log.info("communication object initialized") def _domain_descriptor_info(self, descr): @@ -194,7 +211,7 @@ def _create_domain_descriptor(self, dim: Dimension): # first arg is the domain ID which builds up an MPI Tag. # if those ids are not different for all domain descriptors the system might deadlock # if two parallel exchanges with the same domain id are done - domain_desc = ghex.domain_descriptor( + domain_desc = unstructured.domain_descriptor( self._domain_id_gen(), all_global.tolist(), local_halo.tolist() ) log.debug( @@ -208,9 +225,9 @@ def _create_pattern(self, horizontal_dim: Dimension): global_halo_idx = self._decomposition_info.global_index( horizontal_dim, DecompositionInfo.EntryType.HALO ) - halo_generator = ghex.halo_generator_with_gids(global_halo_idx) + halo_generator = unstructured.halo_generator_with_gids(global_halo_idx) log.debug(f"halo generator for dim='{horizontal_dim.value}' created") - pattern = ghex.make_pattern( + pattern = unstructured.make_pattern( self._context, halo_generator, [self._domain_descriptors[horizontal_dim]] ) log.debug( @@ -227,7 +244,7 @@ def exchange(self, dim: Dimension, *fields: tuple): domain_descriptor is not None ), f"domain descriptor for {dim.value} not found" applied_patterns = [ - pattern(ghex.field_descriptor(domain_descriptor, np.asarray(f))) + pattern(unstructured.field_descriptor(domain_descriptor, np.asarray(f))) for f in fields ] handle = self._comm.exchange(applied_patterns) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py index d0665af11..1dd9828b4 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py @@ -94,10 +94,8 @@ def __init__(self, processProperties: ProcessProperties = None): super().__init__() self._rank_info = "" if processProperties and processProperties.comm_size > 1: - self._rank_info = f"rank={processProperties.rank}/{processProperties.comm_size} [{processProperties.comm_name}] >>>" + self._rank_info = f"rank={processProperties.rank}/{processProperties.comm_size} [{processProperties.comm_name}] " - def filter( - self, record: logging.LogRecord - ) -> bool: + def filter(self, record: logging.LogRecord) -> bool: record.rank = self._rank_info return True diff --git a/atm_dyn_iconam/src/icon4py/driver/README.md b/atm_dyn_iconam/src/icon4py/driver/README.md index c6b80a936..069362ae3 100644 --- a/atm_dyn_iconam/src/icon4py/driver/README.md +++ b/atm_dyn_iconam/src/icon4py/driver/README.md @@ -8,13 +8,28 @@ The code is meant to be changed and enlarged as we port new parts of the model. ## usage -``` +```bash cd atm_dyn_iconam/src/icon4py python driver/dycore_driver.py ../../../testdata/ser_icondata/mpitask1/mch_ch_r04b09_dsl/ser_data --n_steps=2 --run_path=/home/magdalena/temp/icon ``` +or if running in parallel + +```bash +cd atm_dyn_iconam/src/icon4py +mpirun -np 2 python driver/dycore_driver.py ../../../testdata/ser_icondata/mpitask2/mch_ch_r04b09_dsl/ser_data --mpi=True --n_steps=2 --run_path=/home/magdalena/temp/icon + +``` + #### remarks -- first (required) arg is the folder where the serialized input data is stored. The same input data s for the unit tests is used. It can be obtained from [polybox](https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download) -- the serialized data used contains only 5 timesteps so `--n_steps > 2` will throw an exception. -- the code logs to file and to console. Debug logging is only going to file. The log directory can be changed with the --run_path option. +- First (required) arg is the folder where the serialized input data is stored. The same input data s for the unit tests is used. It can be obtained from [polybox](https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download) +- The driver heavily relies on serialized input data. The paths in the example are where the data is put when downloaded via the unit tests. + - data for a serial (single node) run can be downloaded from https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p. +- parallel runs are possible if corresponding data is provided, which is currently available for test with 2 or 4 MPI processes: + + - 2 processes: https://polybox.ethz.ch/index.php/s/NUQjmJcMEoQxFiK + - 4 processes: https://polybox.ethz.ch/index.php/s/QC7xt7xLT5xeVN5 + +- The serialized data used contains only 5 timesteps so `--n_steps > 2` will throw an exception. +- The code logs to file and to console. Debug logging is only going to file. The log directory can be changed with the --run_path option. diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index 4fe0699f1..2875afe03 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -20,7 +20,11 @@ from devtools import Timer from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn -from icon4py.decomposition.parallel_setup import get_processor_properties +from icon4py.decomposition.decomposed import create_exchange +from icon4py.decomposition.parallel_setup import ( + ProcessProperties, + get_processor_properties, +) from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.diffusion.diffusion_utils import copy_diagnostic_and_prognostics from icon4py.diffusion.state_utils import DiagnosticState, PrognosticState @@ -29,6 +33,7 @@ SIMULATION_START_DATE, configure_logging, import_testutils, + read_decomp_info, read_geometry_fields, read_icon_grid, read_initial_state, @@ -153,7 +158,7 @@ def __call__( timer.summary(True) -def initialize(n_time_steps, file_path: Path): +def initialize(n_time_steps, file_path: Path, props: ProcessProperties): """ Inititalize the driver run. @@ -174,15 +179,20 @@ def initialize(n_time_steps, file_path: Path): log.info(f"reading configuration: experiment {experiment_name}") config = read_config(experiment_name, n_time_steps=n_time_steps) - log.info("initializing the grid") - icon_grid = read_icon_grid(file_path) - log.info("reading input fields") - (edge_geometry, cell_geometry, vertical_geometry) = read_geometry_fields(file_path) + decomp_info = read_decomp_info(file_path, props) + + log.info(f"initializing the grid from '{file_path}'") + icon_grid = read_icon_grid(file_path, rank=props.rank) + log.info(f"reading input fields from '{file_path}'") + (edge_geometry, cell_geometry, vertical_geometry) = read_geometry_fields( + file_path, rank=props.rank + ) (metric_state, interpolation_state) = read_static_fields(file_path) log.info("initializing diffusion") diffusion_params = DiffusionParams(config.diffusion_config) - diffusion = Diffusion() + exchange = create_exchange(props, decomp_info) + diffusion = Diffusion(exchange) diffusion.init( icon_grid, config.diffusion_config, @@ -194,7 +204,9 @@ def initialize(n_time_steps, file_path: Path): cell_geometry, ) - data_provider, diagnostic_state, prognostic_state = read_initial_state(file_path) + data_provider, diagnostic_state, prognostic_state = read_initial_state( + file_path, rank=props.rank + ) atmo_non_hydro = DummyAtmoNonHydro(data_provider) atmo_non_hydro.init(config=config.dycore_config) @@ -213,11 +225,8 @@ def initialize(n_time_steps, file_path: Path): @click.option( "--n_steps", default=5, help="number of time steps to run, max 5 is supported" ) -def run( - input_path, - run_path, - n_steps, -): +@click.option("--mpi", default=False, help="whether or not you are running with mpi") +def run(input_path, run_path, n_steps, mpi): """ Run the driver. @@ -238,11 +247,13 @@ def run( """ start_time = datetime.now().astimezone(pytz.UTC) - parallel_props = get_processor_properties() + parallel_props = get_processor_properties(with_mpi=mpi) configure_logging(run_path, start_time, parallel_props) log.info(f"Starting ICON dycore run: {datetime.isoformat(start_time)}") log.info(f"input args: input_path={input_path}, n_time_steps={n_steps}") - timeloop, diagnostic_state, prognostic_state = initialize(n_steps, Path(input_path)) + timeloop, diagnostic_state, prognostic_state = initialize( + n_steps, Path(input_path), parallel_props + ) log.info("dycore configuring: DONE") log.info("timeloop: START") diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index e3ae2d338..6f41ee775 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -99,7 +99,7 @@ def read_initial_state( """ data_provider = sb.IconSerialDataProvider( - "icon_pydycore", str(gridfile_path), False + "icon_pydycore", str(gridfile_path), False, mpi_rank=rank ) init_savepoint = data_provider.from_savepoint_diffusion_init( linit=True, date=SIMULATION_START_DATE @@ -110,7 +110,7 @@ def read_initial_state( def read_geometry_fields( - path: Path, ser_type: SerializationType = SerializationType.SB, rank=0 + path: Path, rank=0, ser_type: SerializationType = SerializationType.SB ) -> tuple[EdgeParams, CellParams, VerticalModelParams]: """ Read fields containing grid properties. @@ -150,28 +150,15 @@ def read_decomp_info( raise NotImplementedError(SB_ONLY_MSG) -def read_grid( - path: Path, - procs_props: ProcessProperties, - ser_type=SerializationType.SB, -) -> IconGrid: - if ser_type == SerializationType.SB: - sp = sb.IconSerialDataProvider( - "icon_grid", str(path.absolute()), True, procs_props.rank - ) - return sp.from_savepoint_grid().construct_icon_grid() - else: - raise NotImplementedError(SB_ONLY_MSG) - - def read_static_fields( - path: Path, ser_type: SerializationType = SerializationType.SB, rank=0 + path: Path, rank=0, ser_type: SerializationType = SerializationType.SB ) -> tuple[MetricState, InterpolationState]: """ Read fields for metric and interpolation state. Args: path: path to the serialized input data + rank: mpi rank, defaults to 0 for serial run ser_type: (optional) defaults to SB=serialbox, type of input data to be read Returns: @@ -192,7 +179,9 @@ def read_static_fields( raise NotImplementedError(SB_ONLY_MSG) -def configure_logging(run_path: str, start_time, processor_procs=None) -> None: +def configure_logging( + run_path: str, start_time, processor_procs: ProcessProperties = None +) -> None: """ Configure logging. @@ -219,9 +208,11 @@ def configure_logging(run_path: str, start_time, processor_procs=None) -> None: ) console_handler = logging.StreamHandler() console_handler.addFilter(ParallelLogger(processor_procs)) - formatter = logging.Formatter( - "%(rank)-20s %(asctime)s %(filename)-20s : %(funcName)-20s: %(levelname)-8s %(message)s" + + log_format = ( + "{rank} {asctime} - {filename}: {funcName:<20}: {levelname:<7} {message}" ) + formatter = logging.Formatter(fmt=log_format, style="{", defaults={"rank": None}) console_handler.setFormatter(formatter) console_handler.setLevel(logging.DEBUG) logging.getLogger("").addHandler(console_handler) diff --git a/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py b/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py index 0af3bd2fd..62d9c6584 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py @@ -20,7 +20,11 @@ from icon4py.decomposition.decomposed import ( DecompositionInfo, DomainDescriptorIdGenerator, + MultiNode, + SingleNode, + create_exchange, ) +from icon4py.decomposition.parallel_setup import ProcessProperties from icon4py.driver.io_utils import ( SerializationType, read_decomp_info, @@ -165,3 +169,21 @@ def test_decomposition_info_matches_gridsize(caplog): ] == icon_grid.num_edges() ) + + +@pytest.mark.mpi +def test_create_multinode(): + decomp_info = read_decomp_info(path, props) + exchange = create_exchange(props, decomp_info) + if props.comm_size > 1: + assert isinstance(exchange, MultiNode) + else: + assert isinstance(exchange, SingleNode) + + +def test_create_single_node_without_mpi(): + props = ProcessProperties.from_single_node() + decomp_info = read_decomp_info(path, props) + exchange = create_exchange(props, decomp_info) + + assert isinstance(exchange, SingleNode) diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py index e9fd2eda6..54d9840a5 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py @@ -12,14 +12,13 @@ # SPDX-License-Identifier: GPL-3.0-or-later -import ghex import pytest from atm_dyn_iconam.tests.mpi_tests.common import path, props from atm_dyn_iconam.tests.test_diffusion import _verify_diffusion_fields from atm_dyn_iconam.tests.test_utils.serialbox_utils import IconSerialDataProvider from icon4py.common.dimension import CellDim, EdgeDim, VertexDim -from icon4py.decomposition.decomposed import DecompositionInfo, MultiNode +from icon4py.decomposition.decomposed import DecompositionInfo, create_exchange from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.driver.io_utils import ( read_decomp_info, @@ -53,7 +52,6 @@ def test_parallel_diffusion( f"local edges = {decomp_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape} " f"local vertices = {decomp_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape}" ) - context = ghex.context(ghex.mpi_comm(props.comm), True) print( f"rank={props.rank}/{props.comm_size}: GHEX context setup: from {props.comm_name} with {props.comm_size} nodes" ) @@ -76,7 +74,7 @@ def test_parallel_diffusion( print( f"rank={props.rank}/{props.comm_size}: setup: using {props.comm_name} with {props.comm_size} nodes" ) - exchange = MultiNode(context, decomp_info) + exchange = create_exchange(props, decomp_info) diffusion = Diffusion(exchange) @@ -118,4 +116,4 @@ def test_parallel_diffusion( print( f"rank={props.rank}/{props.comm_size}: running diffusion step - using {props.comm_name} with {props.comm_size} nodes - DONE" ) - del context + diff --git a/atm_dyn_iconam/tests/test_utils/fixtures.py b/atm_dyn_iconam/tests/test_utils/fixtures.py index 4230001df..035928348 100644 --- a/atm_dyn_iconam/tests/test_utils/fixtures.py +++ b/atm_dyn_iconam/tests/test_utils/fixtures.py @@ -23,7 +23,8 @@ "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download" ) r02b04_global_grid_uri = "https://polybox.ethz.ch/index.php/s/0EM8O8U53GKGsst/download" -data_uri = "https://polybox.ethz.ch/index.php/s/LcAbscZqnsx4WCf/download" +# data_uri = "https://polybox.ethz.ch/index.php/s/LcAbscZqnsx4WCf/download" +data_uri = "https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p/download" base_path = Path(__file__).parent.parent.parent.parent.joinpath("testdata") data_path = base_path.joinpath("ser_icondata") From b919cc2f35309a6ec1f0c8734e9ec286b5fdb124 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 2 Aug 2023 23:11:03 +0200 Subject: [PATCH 207/263] add download_data fixture for parallel tests --- .../src/icon4py/decomposition/decomposed.py | 2 +- atm_dyn_iconam/tests/mpi_tests/common.py | 27 ++++++++++++++++--- .../tests/mpi_tests/test_decomposed.py | 10 +++---- .../mpi_tests/test_parallel_diffusion.py | 4 +-- .../tests/mpi_tests/test_parallel_setup.py | 1 + atm_dyn_iconam/tests/test_utils/fixtures.py | 2 +- 6 files changed, 34 insertions(+), 12 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index 5ffc82568..b25abff28 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -33,7 +33,7 @@ class DomainDescriptorIdGenerator: _counter = 0 _roundtrips = 0 - def __init__(self, context): + def __init__(self, context:ProcessProperties): self._comm_size = context.size() self._roundtrips = context.rank() self._base = self._roundtrips * self._comm_size diff --git a/atm_dyn_iconam/tests/mpi_tests/common.py b/atm_dyn_iconam/tests/mpi_tests/common.py index 68de58e6c..f5c609c7a 100644 --- a/atm_dyn_iconam/tests/mpi_tests/common.py +++ b/atm_dyn_iconam/tests/mpi_tests/common.py @@ -13,10 +13,31 @@ from pathlib import Path -from icon4py.decomposition.parallel_setup import get_processor_properties +import pytest +from atm_dyn_iconam.tests.test_utils.data_handling import download_and_extract +from icon4py.decomposition.parallel_setup import get_processor_properties -base_path = Path(__file__).parent.parent.parent.parent.joinpath("testdata/ser_icondata") props = get_processor_properties(with_mpi=True) -path = base_path.joinpath(f"mpitask{props.comm_size}/mch_ch_r04b09_dsl/ser_data") +base_path = Path(__file__).parent.parent.parent.parent.joinpath(f"testdata/ser_icondata/mpitask{props.comm_size}/") +path = base_path.joinpath("mch_ch_r04b09_dsl/ser_data") +data_uris = {1: "https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p/download" +, 2: "https://polybox.ethz.ch/index.php/s/NUQjmJcMEoQxFiK/download", 4: "https://polybox.ethz.ch/index.php/s/QC7xt7xLT5xeVN5/download"} + +@pytest.fixture(scope="session") +def download_data(): + """ + Get the binary ICON data from a remote server. + + Session scoped fixture which is a prerequisite of all the other fixtures in this file. + """ + try: + uri = data_uris[props.comm_size] + + data_file = path.joinpath(f"mch_ch_r04b09_dsl_mpitask{props.comm_size}.tar.gz").name + if props.rank == 0: + download_and_extract(uri, base_path, data_file) + props.comm.barrier() + except KeyError: + assert False, f"no data for communicator of size {props.comm_size} exists, use 1, 2 or 4" diff --git a/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py b/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py index 62d9c6584..9232e864d 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py @@ -15,7 +15,7 @@ import numpy as np import pytest -from atm_dyn_iconam.tests.mpi_tests.common import path, props +from atm_dyn_iconam.tests.mpi_tests.common import path, props, download_data from icon4py.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.decomposition.decomposed import ( DecompositionInfo, @@ -53,7 +53,7 @@ (VertexDim, (5373, 5290), (5455, 5456)), ), ) -def test_decomposition_info_masked(dim, owned, total, caplog): +def test_decomposition_info_masked(dim, owned, total, caplog, download_data): my_rank = props.rank decomposition_info = read_decomp_info(path, props, SerializationType.SB) all_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.ALL) @@ -82,7 +82,7 @@ def test_decomposition_info_masked(dim, owned, total, caplog): (VertexDim, (5373, 5290), (5455, 5456)), ), ) -def test_decomposition_info_local_index(dim, owned, total, caplog): +def test_decomposition_info_local_index(dim, owned, total, caplog, download_data): my_rank = props.rank decomposition_info = read_decomp_info(path, props, SerializationType.SB) @@ -144,7 +144,7 @@ def test_domain_descriptor_id_are_globally_unique(num): props.comm_size not in (1, 2, 4), reason="input files only available for 1 or 2 nodes", ) -def test_decomposition_info_matches_gridsize(caplog): +def test_decomposition_info_matches_gridsize(caplog, download_data): decomposition_info = read_decomp_info( path, props, @@ -172,7 +172,7 @@ def test_decomposition_info_matches_gridsize(caplog): @pytest.mark.mpi -def test_create_multinode(): +def test_create_multinode(download_data): decomp_info = read_decomp_info(path, props) exchange = create_exchange(props, decomp_info) if props.comm_size > 1: diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py index 54d9840a5..b2f4997ae 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py @@ -14,7 +14,7 @@ import pytest -from atm_dyn_iconam.tests.mpi_tests.common import path, props +from atm_dyn_iconam.tests.mpi_tests.common import path, props, download_data from atm_dyn_iconam.tests.test_diffusion import _verify_diffusion_fields from atm_dyn_iconam.tests.test_utils.serialbox_utils import IconSerialDataProvider from icon4py.common.dimension import CellDim, EdgeDim, VertexDim @@ -36,7 +36,7 @@ reason="input files only available for 1 or 2 nodes", ) def test_parallel_diffusion( - r04b09_diffusion_config, step_date_init, linit, ndyn_substeps + r04b09_diffusion_config, step_date_init, linit, ndyn_substeps, download_data ): print( diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index cfdfc49ec..602c81c19 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -48,3 +48,4 @@ def test_single_node_properties(): assert props.comm_size == 1 assert props.rank == 0 assert props.comm_name == "" + diff --git a/atm_dyn_iconam/tests/test_utils/fixtures.py b/atm_dyn_iconam/tests/test_utils/fixtures.py index 035928348..135aa3403 100644 --- a/atm_dyn_iconam/tests/test_utils/fixtures.py +++ b/atm_dyn_iconam/tests/test_utils/fixtures.py @@ -23,7 +23,7 @@ "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download" ) r02b04_global_grid_uri = "https://polybox.ethz.ch/index.php/s/0EM8O8U53GKGsst/download" -# data_uri = "https://polybox.ethz.ch/index.php/s/LcAbscZqnsx4WCf/download" +data_uri = "https://polybox.ethz.ch/index.php/s/LcAbscZqnsx4WCf/download" data_uri = "https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p/download" base_path = Path(__file__).parent.parent.parent.parent.joinpath("testdata") From 920578307ed6e3b2bfefccc4fa18e5a5e80e7600 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 3 Aug 2023 09:49:37 +0200 Subject: [PATCH 208/263] move ghex context creation inside MultiNodeRuntime --- .../src/icon4py/decomposition/decomposed.py | 19 ++++----- atm_dyn_iconam/tests/mpi_tests/common.py | 22 +++++++--- .../tests/mpi_tests/test_decomposed.py | 41 +++++++++++-------- .../mpi_tests/test_parallel_diffusion.py | 25 +++++++---- .../tests/mpi_tests/test_parallel_setup.py | 1 - 5 files changed, 64 insertions(+), 44 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index b25abff28..2a893a76e 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -33,9 +33,9 @@ class DomainDescriptorIdGenerator: _counter = 0 _roundtrips = 0 - def __init__(self, context:ProcessProperties): - self._comm_size = context.size() - self._roundtrips = context.rank() + def __init__(self, parallel_props: ProcessProperties): + self._comm_size = parallel_props.comm_size + self._roundtrips = parallel_props.rank self._base = self._roundtrips * self._comm_size def __call__(self): @@ -139,8 +139,7 @@ def create_exchange( Depending on the number of processor a SingleNode version is returned or a GHEX context created and a Multinode returned. """ if props.comm_size > 1: - context = ghex.context(ghex.mpi_comm(props.comm), True) - return MultiNode(context, decomp_info) + return GHexMultiNode(props, decomp_info) else: return SingleNode() @@ -165,10 +164,10 @@ def is_ready(self) -> bool: return True -class MultiNode: - def __init__(self, context, domain_decomposition: DecompositionInfo): - self._context = context - self._domain_id_gen = DomainDescriptorIdGenerator(context) +class GHexMultiNode: + def __init__(self, props, domain_decomposition: DecompositionInfo): + self._context = ghex.context(ghex.mpi_comm(props.comm), True) + self._domain_id_gen = DomainDescriptorIdGenerator(props) self._decomposition_info = domain_decomposition self._domain_descriptors = { CellDim: self._create_domain_descriptor( @@ -189,7 +188,7 @@ def __init__(self, context, domain_decomposition: DecompositionInfo): EdgeDim: self._create_pattern(EdgeDim), } log.info(f"patterns for dimensions {self._patterns.keys()} initialized ") - self._comm = unstructured.make_co(context) + self._comm = unstructured.make_co(self._context) log.info("communication object initialized") def _domain_descriptor_info(self, descr): diff --git a/atm_dyn_iconam/tests/mpi_tests/common.py b/atm_dyn_iconam/tests/mpi_tests/common.py index f5c609c7a..a44f481eb 100644 --- a/atm_dyn_iconam/tests/mpi_tests/common.py +++ b/atm_dyn_iconam/tests/mpi_tests/common.py @@ -20,10 +20,16 @@ props = get_processor_properties(with_mpi=True) -base_path = Path(__file__).parent.parent.parent.parent.joinpath(f"testdata/ser_icondata/mpitask{props.comm_size}/") -path = base_path.joinpath("mch_ch_r04b09_dsl/ser_data") -data_uris = {1: "https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p/download" -, 2: "https://polybox.ethz.ch/index.php/s/NUQjmJcMEoQxFiK/download", 4: "https://polybox.ethz.ch/index.php/s/QC7xt7xLT5xeVN5/download"} +base_path = Path(__file__).parent.parent.parent.parent.joinpath( + f"testdata/ser_icondata/mpitask{props.comm_size}/" +) +data_path = base_path.joinpath("mch_ch_r04b09_dsl/ser_data") +data_uris = { + 1: "https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p/download", + 2: "https://polybox.ethz.ch/index.php/s/NUQjmJcMEoQxFiK/download", + 4: "https://polybox.ethz.ch/index.php/s/QC7xt7xLT5xeVN5/download", +} + @pytest.fixture(scope="session") def download_data(): @@ -35,9 +41,13 @@ def download_data(): try: uri = data_uris[props.comm_size] - data_file = path.joinpath(f"mch_ch_r04b09_dsl_mpitask{props.comm_size}.tar.gz").name + data_file = data_path.joinpath( + f"mch_ch_r04b09_dsl_mpitask{props.comm_size}.tar.gz" + ).name if props.rank == 0: download_and_extract(uri, base_path, data_file) props.comm.barrier() except KeyError: - assert False, f"no data for communicator of size {props.comm_size} exists, use 1, 2 or 4" + raise AssertionError( + f"no data for communicator of size {props.comm_size} exists, use 1, 2 or 4" + ) diff --git a/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py b/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py index 9232e864d..b3f23f9c0 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py @@ -11,16 +11,19 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import ghex import numpy as np import pytest -from atm_dyn_iconam.tests.mpi_tests.common import path, props, download_data +from atm_dyn_iconam.tests.mpi_tests.common import ( # noqa F401 + data_path, + download_data, + props, +) from icon4py.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.decomposition.decomposed import ( DecompositionInfo, DomainDescriptorIdGenerator, - MultiNode, + GHexMultiNode, SingleNode, create_exchange, ) @@ -53,9 +56,11 @@ (VertexDim, (5373, 5290), (5455, 5456)), ), ) -def test_decomposition_info_masked(dim, owned, total, caplog, download_data): +def test_decomposition_info_masked( + dim, owned, total, caplog, download_data # noqa F811 +): my_rank = props.rank - decomposition_info = read_decomp_info(path, props, SerializationType.SB) + decomposition_info = read_decomp_info(data_path, props, SerializationType.SB) all_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.ALL) my_total = total[my_rank] my_owned = owned[my_rank] @@ -82,10 +87,11 @@ def test_decomposition_info_masked(dim, owned, total, caplog, download_data): (VertexDim, (5373, 5290), (5455, 5456)), ), ) -def test_decomposition_info_local_index(dim, owned, total, caplog, download_data): - +def test_decomposition_info_local_index( + dim, owned, total, caplog, download_data # noqa F811 +): my_rank = props.rank - decomposition_info = read_decomp_info(path, props, SerializationType.SB) + decomposition_info = read_decomp_info(data_path, props, SerializationType.SB) all_indices = decomposition_info.local_index(dim, DecompositionInfo.EntryType.ALL) my_total = total[my_rank] my_owned = owned[my_rank] @@ -121,8 +127,7 @@ def _assert_index_partitioning(all_indices, halo_indices, owned_indices): def test_domain_descriptor_id_are_globally_unique(num): size = props.comm_size - context = ghex.context(ghex.mpi_comm(props.comm), True) - id_gen = DomainDescriptorIdGenerator(context) + id_gen = DomainDescriptorIdGenerator(parallel_props=props) id1 = id_gen() assert id1 == props.comm_size * props.rank assert id1 < props.comm_size * (props.rank + 1) @@ -144,13 +149,13 @@ def test_domain_descriptor_id_are_globally_unique(num): props.comm_size not in (1, 2, 4), reason="input files only available for 1 or 2 nodes", ) -def test_decomposition_info_matches_gridsize(caplog, download_data): +def test_decomposition_info_matches_gridsize(caplog, download_data): # noqa F811 decomposition_info = read_decomp_info( - path, + data_path, props, SerializationType.SB, ) - icon_grid = read_icon_grid(path, props.rank) + icon_grid = read_icon_grid(data_path, props.rank) assert ( decomposition_info.global_index( dim=CellDim, entry_type=DecompositionInfo.EntryType.ALL @@ -172,18 +177,18 @@ def test_decomposition_info_matches_gridsize(caplog, download_data): @pytest.mark.mpi -def test_create_multinode(download_data): - decomp_info = read_decomp_info(path, props) +def test_create_multi_pytenode_runtime_with_mpi(download_data): # noqa F811 + decomp_info = read_decomp_info(data_path, props) exchange = create_exchange(props, decomp_info) if props.comm_size > 1: - assert isinstance(exchange, MultiNode) + assert isinstance(exchange, GHexMultiNode) else: assert isinstance(exchange, SingleNode) -def test_create_single_node_without_mpi(): +def test_create_single_node_runtime_without_mpi(): props = ProcessProperties.from_single_node() - decomp_info = read_decomp_info(path, props) + decomp_info = read_decomp_info(data_path, props) exchange = create_exchange(props, decomp_info) assert isinstance(exchange, SingleNode) diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py index b2f4997ae..4468ccd0c 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py @@ -14,7 +14,11 @@ import pytest -from atm_dyn_iconam.tests.mpi_tests.common import path, props, download_data +from atm_dyn_iconam.tests.mpi_tests.common import ( # noqa F401 + data_path, + download_data, + props, +) from atm_dyn_iconam.tests.test_diffusion import _verify_diffusion_fields from atm_dyn_iconam.tests.test_utils.serialbox_utils import IconSerialDataProvider from icon4py.common.dimension import CellDim, EdgeDim, VertexDim @@ -36,14 +40,18 @@ reason="input files only available for 1 or 2 nodes", ) def test_parallel_diffusion( - r04b09_diffusion_config, step_date_init, linit, ndyn_substeps, download_data + r04b09_diffusion_config, + step_date_init, + linit, + ndyn_substeps, + download_data, # noqa: F811 ): print( f"rank={props.rank}/{props.comm_size}: inializing diffusion for experiment 'mch_ch_r04_b09_dsl" ) decomp_info = read_decomp_info( - path, + data_path, props, ) print( @@ -56,19 +64,19 @@ def test_parallel_diffusion( f"rank={props.rank}/{props.comm_size}: GHEX context setup: from {props.comm_name} with {props.comm_size} nodes" ) - icon_grid = read_icon_grid(path, rank=props.rank) + icon_grid = read_icon_grid(data_path, rank=props.rank) print( f"rank={props.rank}: using local grid with {icon_grid.num_cells()} Cells, {icon_grid.num_edges()} Edges, {icon_grid.num_vertices()} Vertices" ) diffusion_params = DiffusionParams(r04b09_diffusion_config) diffusion_initial_data = IconSerialDataProvider( - "icon_pydycore", str(path), True, mpi_rank=props.rank + "icon_pydycore", str(data_path), True, mpi_rank=props.rank ).from_savepoint_diffusion_init(linit=linit, date=step_date_init) (edge_geometry, cell_geometry, vertical_geometry) = read_geometry_fields( - path, rank=props.rank + data_path, rank=props.rank ) - (metric_state, interpolation_state) = read_static_fields(path, rank=props.rank) + (metric_state, interpolation_state) = read_static_fields(data_path, rank=props.rank) dtime = diffusion_initial_data.get_metadata("dtime").get("dtime") print( @@ -106,7 +114,7 @@ def test_parallel_diffusion( print(f"rank={props.rank}/{props.comm_size}: diffusion run ") diffusion_savepoint_exit = IconSerialDataProvider( - "icon_pydycore", str(path), True, mpi_rank=props.rank + "icon_pydycore", str(data_path), True, mpi_rank=props.rank ).from_savepoint_diffusion_exit(linit=linit, date=step_date_init) _verify_diffusion_fields( diagnostic_state=diagnostic_state, @@ -116,4 +124,3 @@ def test_parallel_diffusion( print( f"rank={props.rank}/{props.comm_size}: running diffusion step - using {props.comm_name} with {props.comm_size} nodes - DONE" ) - diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py index 602c81c19..cfdfc49ec 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_setup.py @@ -48,4 +48,3 @@ def test_single_node_properties(): assert props.comm_size == 1 assert props.rank == 0 assert props.comm_name == "" - From e8e877ee1a90036533f73c8b030a5841b8556211 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 3 Aug 2023 11:31:06 +0200 Subject: [PATCH 209/263] remove duplicate URIs --- atm_dyn_iconam/tests/mpi_tests/common.py | 15 ++++----------- atm_dyn_iconam/tests/test_utils/fixtures.py | 13 +++++++++---- 2 files changed, 13 insertions(+), 15 deletions(-) diff --git a/atm_dyn_iconam/tests/mpi_tests/common.py b/atm_dyn_iconam/tests/mpi_tests/common.py index a44f481eb..d17432fb9 100644 --- a/atm_dyn_iconam/tests/mpi_tests/common.py +++ b/atm_dyn_iconam/tests/mpi_tests/common.py @@ -11,24 +11,17 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from pathlib import Path import pytest from atm_dyn_iconam.tests.test_utils.data_handling import download_and_extract +from atm_dyn_iconam.tests.test_utils.fixtures import base_path, data_uris from icon4py.decomposition.parallel_setup import get_processor_properties props = get_processor_properties(with_mpi=True) -base_path = Path(__file__).parent.parent.parent.parent.joinpath( - f"testdata/ser_icondata/mpitask{props.comm_size}/" -) -data_path = base_path.joinpath("mch_ch_r04b09_dsl/ser_data") -data_uris = { - 1: "https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p/download", - 2: "https://polybox.ethz.ch/index.php/s/NUQjmJcMEoQxFiK/download", - 4: "https://polybox.ethz.ch/index.php/s/QC7xt7xLT5xeVN5/download", -} +path = base_path.joinpath("mpitask{props.comm_size}") +data_path = path.joinpath("mch_ch_r04b09_dsl/ser_data") @pytest.fixture(scope="session") @@ -45,7 +38,7 @@ def download_data(): f"mch_ch_r04b09_dsl_mpitask{props.comm_size}.tar.gz" ).name if props.rank == 0: - download_and_extract(uri, base_path, data_file) + download_and_extract(uri, path, data_file) props.comm.barrier() except KeyError: raise AssertionError( diff --git a/atm_dyn_iconam/tests/test_utils/fixtures.py b/atm_dyn_iconam/tests/test_utils/fixtures.py index 135aa3403..6488e9b4f 100644 --- a/atm_dyn_iconam/tests/test_utils/fixtures.py +++ b/atm_dyn_iconam/tests/test_utils/fixtures.py @@ -19,13 +19,18 @@ from .serialbox_utils import IconSerialDataProvider +base_path = Path(__file__).parent.parent.parent.parent.joinpath("testdata") + +data_uris = { + 1: "https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p/download", + 2: "https://polybox.ethz.ch/index.php/s/NUQjmJcMEoQxFiK/download", + 4: "https://polybox.ethz.ch/index.php/s/QC7xt7xLT5xeVN5/download", +} mch_ch_r04b09_dsl_grid_uri = ( "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download" ) r02b04_global_grid_uri = "https://polybox.ethz.ch/index.php/s/0EM8O8U53GKGsst/download" -data_uri = "https://polybox.ethz.ch/index.php/s/LcAbscZqnsx4WCf/download" -data_uri = "https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p/download" -base_path = Path(__file__).parent.parent.parent.parent.joinpath("testdata") + data_path = base_path.joinpath("ser_icondata") data_file = data_path.joinpath("mch_ch_r04b09_dsl.tar.gz").name @@ -54,7 +59,7 @@ def setup_icon_data(): Session scoped fixture which is a prerequisite of all the other fixtures in this file. """ - download_and_extract(data_uri, data_path, data_file) + download_and_extract(data_uris[1], data_path, data_file) @pytest.fixture From 20be71d745ff96b28ee50742c5d63b26a883510c Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 3 Aug 2023 11:55:35 +0200 Subject: [PATCH 210/263] fixes: - remove empty file - remove parametrized decorator for datapath --- .../icon4py/diffusion/interpolation_state.py | 12 ---- atm_dyn_iconam/tests/test_diffusion.py | 64 +++++++++---------- atm_dyn_iconam/tests/test_io_utils.py | 4 -- atm_dyn_iconam/tests/test_state_utils.py | 1 - common/tests/test_grid_manager.py | 22 +++---- common/tests/test_icon_grid.py | 7 -- common/tests/test_vertical.py | 2 - 7 files changed, 43 insertions(+), 69 deletions(-) delete mode 100644 atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py diff --git a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py b/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py deleted file mode 100644 index 15dfdb009..000000000 --- a/atm_dyn_iconam/src/icon4py/diffusion/interpolation_state.py +++ /dev/null @@ -1,12 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 0b5234275..98ab9d0f9 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -263,38 +263,6 @@ def test_verify_diffusion_init_against_other_regular_savepoint( _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion) -def _verify_diffusion_fields( - diagnostic_state: DiagnosticState, - prognostic_state: PrognosticState, - diffusion_savepoint: IconDiffusionExitSavepoint, -): - ref_div_ic = np.asarray(diffusion_savepoint.div_ic()) - val_div_ic = np.asarray(diagnostic_state.div_ic) - ref_hdef_ic = np.asarray(diffusion_savepoint.hdef_ic()) - val_hdef_ic = np.asarray(diagnostic_state.hdef_ic) - assert np.allclose(ref_div_ic, val_div_ic) - assert np.allclose(ref_hdef_ic, val_hdef_ic) - ref_w = np.asarray(diffusion_savepoint.w()) - val_w = np.asarray(prognostic_state.w) - ref_dwdx = np.asarray(diffusion_savepoint.dwdx()) - val_dwdx = np.asarray(diagnostic_state.dwdx) - ref_dwdy = np.asarray(diffusion_savepoint.dwdy()) - val_dwdy = np.asarray(diagnostic_state.dwdy) - assert np.allclose(ref_dwdx, val_dwdx) - assert np.allclose(ref_dwdy, val_dwdy) - - ref_vn = np.asarray(diffusion_savepoint.vn()) - val_vn = np.asarray(prognostic_state.vn) - assert np.allclose(ref_vn, val_vn) - assert np.allclose(ref_w, val_w) - ref_exner = np.asarray(diffusion_savepoint.exner()) - ref_theta_v = np.asarray(diffusion_savepoint.theta_v()) - val_theta_v = np.asarray(prognostic_state.theta_v) - val_exner = np.asarray(prognostic_state.exner_pressure) - assert np.allclose(ref_theta_v, val_theta_v) - assert np.allclose(ref_exner, val_exner) - - @pytest.mark.datatest @pytest.mark.parametrize( "step_date_init, step_date_exit", @@ -355,6 +323,38 @@ def test_run_diffusion_single_step( ) +def _verify_diffusion_fields( + diagnostic_state: DiagnosticState, + prognostic_state: PrognosticState, + diffusion_savepoint: IconDiffusionExitSavepoint, +): + ref_div_ic = np.asarray(diffusion_savepoint.div_ic()) + val_div_ic = np.asarray(diagnostic_state.div_ic) + ref_hdef_ic = np.asarray(diffusion_savepoint.hdef_ic()) + val_hdef_ic = np.asarray(diagnostic_state.hdef_ic) + assert np.allclose(ref_div_ic, val_div_ic) + assert np.allclose(ref_hdef_ic, val_hdef_ic) + ref_w = np.asarray(diffusion_savepoint.w()) + val_w = np.asarray(prognostic_state.w) + ref_dwdx = np.asarray(diffusion_savepoint.dwdx()) + val_dwdx = np.asarray(diagnostic_state.dwdx) + ref_dwdy = np.asarray(diffusion_savepoint.dwdy()) + val_dwdy = np.asarray(diagnostic_state.dwdy) + assert np.allclose(ref_dwdx, val_dwdx) + assert np.allclose(ref_dwdy, val_dwdy) + + ref_vn = np.asarray(diffusion_savepoint.vn()) + val_vn = np.asarray(prognostic_state.vn) + assert np.allclose(ref_vn, val_vn) + assert np.allclose(ref_w, val_w) + ref_exner = np.asarray(diffusion_savepoint.exner()) + ref_theta_v = np.asarray(diffusion_savepoint.theta_v()) + val_theta_v = np.asarray(prognostic_state.theta_v) + val_exner = np.asarray(prognostic_state.exner_pressure) + assert np.allclose(ref_theta_v, val_theta_v) + assert np.allclose(ref_exner, val_exner) + + @pytest.mark.datatest @pytest.mark.parametrize("linit", [True]) def test_run_diffusion_initial_step( diff --git a/atm_dyn_iconam/tests/test_io_utils.py b/atm_dyn_iconam/tests/test_io_utils.py index 6a53e8af4..2f92b6617 100644 --- a/atm_dyn_iconam/tests/test_io_utils.py +++ b/atm_dyn_iconam/tests/test_io_utils.py @@ -25,7 +25,6 @@ @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) @pytest.mark.parametrize( "read_fun", (read_geometry_fields, read_static_fields, read_icon_grid) ) @@ -48,14 +47,12 @@ def assert_grid_size_and_connectivities(grid): @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_read_icon_grid_for_type_sb(datapath): grid = read_icon_grid(datapath, ser_type=SerializationType.SB) assert_grid_size_and_connectivities(grid) @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_read_static_fields_for_type_sb(datapath): metric_state, interpolation_state = read_static_fields( datapath, ser_type=SerializationType.SB @@ -65,7 +62,6 @@ def test_read_static_fields_for_type_sb(datapath): @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_read_geometry_fields_for_type_sb(datapath): edge_geometry, cell_geometry, vertical_geometry = read_geometry_fields( datapath, ser_type=SerializationType.SB diff --git a/atm_dyn_iconam/tests/test_state_utils.py b/atm_dyn_iconam/tests/test_state_utils.py index dcfc7c3ba..243ab01e2 100644 --- a/atm_dyn_iconam/tests/test_state_utils.py +++ b/atm_dyn_iconam/tests/test_state_utils.py @@ -16,7 +16,6 @@ @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_verify_geofac_n2s_field_manipulation(interpolation_savepoint, icon_grid): geofac_n2s = np.asarray(interpolation_savepoint.geofac_n2s()) int_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() diff --git a/common/tests/test_grid_manager.py b/common/tests/test_grid_manager.py index 12b164d0d..d70e01420 100644 --- a/common/tests/test_grid_manager.py +++ b/common/tests/test_grid_manager.py @@ -359,19 +359,8 @@ def test_gridmanager_eval_c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): # e2c2e (e2c2eo) - diamond: exists in serial, simple_mesh @pytest.mark.datatest -def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): - caplog.set_level(logging.DEBUG) - grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() - assert np.allclose( - grid.get_c2e2c_connectivity().table, - grid_savepoint.c2e2c()[0 : grid.num_cells(), :], - ) - - -# c2e2c: exists in serial, simple_mesh, grid @pytest.mark.skip("does not directly exist in the grid file, needs to be constructed") # TODO (Magdalena) construct from adjacent_cell_of_edge and then edge_of_cell -@pytest.mark.datatest def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): caplog.set_level(logging.DEBUG) gm, num_cells, num_edges, num_vertex = init_grid_manager(r04b09_dsl_gridfile) @@ -382,6 +371,17 @@ def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile): assert np.allclose(grid.get_e2c2e_connectivity().table, serialized_e2c2e) +# c2e2c: exists in serial, simple_mesh, grid +@pytest.mark.datatest +def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile): + caplog.set_level(logging.DEBUG) + grid = init_grid_manager(r04b09_dsl_gridfile).get_grid() + assert np.allclose( + grid.get_c2e2c_connectivity().table, + grid_savepoint.c2e2c()[0 : grid.num_cells(), :], + ) + + @pytest.mark.xfail @pytest.mark.datatest def test_gridmanager_eval_e2c2v(caplog, grid_savepoint, r04b09_dsl_gridfile): diff --git a/common/tests/test_icon_grid.py b/common/tests/test_icon_grid.py index b25f65174..3a7db7bbc 100644 --- a/common/tests/test_icon_grid.py +++ b/common/tests/test_icon_grid.py @@ -17,7 +17,6 @@ @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) # TODO(Magdalena) HorizontalMarkerIndex.local(dim) does not yield equvalent results form grid file # and serialized data, why?. Serialized data has those strange -1 values @pytest.mark.parametrize( @@ -66,7 +65,6 @@ def test_horizontal_end_index(icon_grid, dim, marker, index): @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) @pytest.mark.parametrize( "dim, marker, index", [ @@ -112,7 +110,6 @@ def test_horizontal_start_index(icon_grid, dim, marker, index): @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) @pytest.mark.parametrize( "start_marker, end_marker, expected_bounds", [ @@ -189,7 +186,6 @@ def test_horizontal_cell_markers(icon_grid, start_marker, end_marker, expected_b @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) @pytest.mark.parametrize( "start_marker, end_marker, expected_bounds", [ @@ -286,7 +282,6 @@ def test_horizontal_edge_markers(icon_grid, start_marker, end_marker, expected_b @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) @pytest.mark.parametrize( "start_marker, end_marker, expected_bounds", [ @@ -357,7 +352,6 @@ def test_horizontal_vertex_markers( @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_cross_check_marker_equivalences(icon_grid): """Check actual equivalences of calculated markers.""" # TODO(Magdalena): This should go away once we refactor these markers in a good way, such that no calculation need to be done with them anymore. @@ -401,7 +395,6 @@ def test_cross_check_marker_equivalences(icon_grid): ) -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_grid_size(grid_savepoint): assert 10663 == grid_savepoint.num(VertexDim) assert 20896 == grid_savepoint.num(CellDim) diff --git a/common/tests/test_vertical.py b/common/tests/test_vertical.py index 4c3b02a1a..9579927f6 100644 --- a/common/tests/test_vertical.py +++ b/common/tests/test_vertical.py @@ -35,7 +35,6 @@ def test_nrdmax_calculation(max_h, damping, delta): @pytest.mark.datatest -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_nrdmax_calculation_from_icon_input(icon_grid, grid_savepoint, damping_height): a = grid_savepoint.vct_a() nrdmax = grid_savepoint.nrdmax() @@ -48,6 +47,5 @@ def test_nrdmax_calculation_from_icon_input(icon_grid, grid_savepoint, damping_h assert a_array[nrdmax + 1] < damping_height -@pytest.mark.parametrize("datapath", [1], indirect=True) def test_grid_size(grid_savepoint): assert 65 == grid_savepoint.num(KDim) From 66e49c3b6e3fcc8e88c1aa672c25076a3bae65b4 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 3 Aug 2023 12:12:59 +0200 Subject: [PATCH 211/263] remove unnecessary f-strings --- .../src/icon4py/diffusion/diffusion.py | 58 +++++++++---------- atm_dyn_iconam/tests/mpi_tests/common.py | 2 +- atm_dyn_iconam/tests/test_utils/fixtures.py | 2 +- 3 files changed, 31 insertions(+), 31 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 3aa99b62d..2522c870a 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -552,7 +552,7 @@ def _sync_cell_fields(self, prognostic_state): communication only done in original code if the following condition applies: IF ( .NOT. lhdiff_rcf .OR. linit .OR. (iforcing /= inwp .AND. iforcing /= iaes) ) THEN """ - log.debug(f"communication of prognostic cell fields: theta, w, exner - start") + log.debug("communication of prognostic cell fields: theta, w, exner - start") handle_cell_comm = self._exchange.exchange( CellDim, prognostic_state.w, @@ -560,7 +560,7 @@ def _sync_cell_fields(self, prognostic_state): prognostic_state.exner_pressure, ) handle_cell_comm.wait() - log.debug(f"communication of prognostic cell fields: theta, w, exner - done") + log.debug("communication of prognostic cell fields: theta, w, exner - done") def _do_diffusion_step( self, @@ -628,7 +628,7 @@ def _do_diffusion_step( self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={} ) - log.debug(f"rbf interpolation 1: start") + log.debug("rbf interpolation 1: start") mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( p_e_in=prognostic_state.vn, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, @@ -641,16 +641,16 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={"V2E": self.grid.get_v2e_connectivity()}, ) - log.debug(f"rbf interpolation 1: end") + log.debug("rbf interpolation 1: end") # 2. HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert - log.debug(f"communication rbf extrapolation of vn - start") + log.debug("communication rbf extrapolation of vn - start") h = self._exchange.exchange(VertexDim, self.u_vert, self.v_vert) h.wait() - log.debug(f"communication rbf extrapolation of vn - end") + log.debug("communication rbf extrapolation of vn - end") log.debug( - f"running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start" + "running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start" ) calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend)( diff_multfac_smag=self.diff_multfac_smag, @@ -679,10 +679,10 @@ def _do_diffusion_step( }, ) log.debug( - f"running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end" + "running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end" ) log.debug( - f"running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start" + "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start" ) calculate_diagnostic_quantities_for_turbulence.with_backend(backend)( kh_smag_ec=self.kh_smag_ec, @@ -704,18 +704,18 @@ def _do_diffusion_step( }, ) log.debug( - f"running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end" + "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end" ) # HALO EXCHANGE IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH if self.config.type_vn_diffu > 1: - log.debug(f"communication rbf extrapolation of z_nable2_e - start") + log.debug("communication rbf extrapolation of z_nable2_e - start") h_z = self._exchange.exchange(EdgeDim, self.z_nabla2_e) h_z.wait() - log.debug(f"communication rbf extrapolation of z_nable2_e - end") + log.debug("communication rbf extrapolation of z_nable2_e - end") - log.debug(f"2nd rbf interpolation: start") + log.debug("2nd rbf interpolation: start") mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( p_e_in=self.z_nabla2_e, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, @@ -728,15 +728,15 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={"V2E": self.grid.get_v2e_connectivity()}, ) - log.debug(f"2nd rbf interpolation: end") + log.debug("2nd rbf interpolation: end") # 6. HALO EXCHANGE -- CALL sync_patch_array_mult (Vertex Fields) - log.debug(f"communication rbf extrapolation of z_nable2_e - start") + log.debug("communication rbf extrapolation of z_nable2_e - start") h = self._exchange.exchange(VertexDim, self.u_vert, self.v_vert) h.wait() - log.debug(f"communication rbf extrapolation of z_nable2_e - end") + log.debug("communication rbf extrapolation of z_nable2_e - end") - log.debug(f"running stencils 04 05 06 (apply_diffusion_to_vn): start") + log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): start") apply_diffusion_to_vn.with_backend(backend)( u_vert=self.u_vert, v_vert=self.v_vert, @@ -764,12 +764,12 @@ def _do_diffusion_step( "E2ECV": self.grid.get_e2ecv_connectivity(), }, ) - log.debug(f"running stencils 04 05 06 (apply_diffusion_to_vn): end") - log.debug(f"communication of prognistic.vn : start") + log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): end") + log.debug("communication of prognistic.vn : start") handle_edge_comm = self._exchange.exchange(EdgeDim, prognostic_state.vn) log.debug( - f"running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" + "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" ) copy_field.with_backend(backend)( prognostic_state.w, self.w_tmp, offset_provider={} @@ -803,11 +803,11 @@ def _do_diffusion_step( }, ) log.debug( - f"running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end" + "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end" ) log.debug( - f"running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" + "running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" ) calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.with_backend( backend @@ -826,7 +826,7 @@ def _do_diffusion_step( }, ) log.debug( - f"running stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): end" + "running stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): end" ) log.debug("running stencils 13 14 (calculate_nabla2_for_theta): start") calculate_nabla2_for_theta.with_backend(backend)( @@ -845,9 +845,9 @@ def _do_diffusion_step( "C2CE": self.grid.get_c2ce_connectivity(), }, ) - log.debug(f"running stencils 13_14 (calculate_nabla2_for_theta): end") + log.debug("running stencils 13_14 (calculate_nabla2_for_theta): end") log.debug( - f"running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start" + "running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start" ) truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_backend( backend @@ -872,9 +872,9 @@ def _do_diffusion_step( ) log.debug( - f"running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end" + "running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end" ) - log.debug(f"running stencil 16 (update_theta_and_exner): start") + log.debug("running stencil 16 (update_theta_and_exner): start") update_theta_and_exner.with_backend(backend)( z_temp=self.z_temp, area=self.cell_params.area, @@ -887,6 +887,6 @@ def _do_diffusion_step( vertical_end=klevels, offset_provider={}, ) - log.debug(f"running stencil 16 (update_theta_and_exner): end") + log.debug("running stencil 16 (update_theta_and_exner): end") handle_edge_comm.wait() # need to do this here, since we currently only use 1 communication object. - log.debug(f"communication of prognogistic.vn - end") + log.debug("communication of prognogistic.vn - end") diff --git a/atm_dyn_iconam/tests/mpi_tests/common.py b/atm_dyn_iconam/tests/mpi_tests/common.py index d17432fb9..5feb16cc1 100644 --- a/atm_dyn_iconam/tests/mpi_tests/common.py +++ b/atm_dyn_iconam/tests/mpi_tests/common.py @@ -20,7 +20,7 @@ props = get_processor_properties(with_mpi=True) -path = base_path.joinpath("mpitask{props.comm_size}") +path = base_path.joinpath(f"mpitask{props.comm_size}") data_path = path.joinpath("mch_ch_r04b09_dsl/ser_data") diff --git a/atm_dyn_iconam/tests/test_utils/fixtures.py b/atm_dyn_iconam/tests/test_utils/fixtures.py index 6488e9b4f..5b821c5eb 100644 --- a/atm_dyn_iconam/tests/test_utils/fixtures.py +++ b/atm_dyn_iconam/tests/test_utils/fixtures.py @@ -55,7 +55,7 @@ def datapath(setup_icon_data): @pytest.fixture(scope="session") def setup_icon_data(): """ - Get the binary ICON data from a remote server. + Get the binary ICON data for single node run from a remote server. Session scoped fixture which is a prerequisite of all the other fixtures in this file. """ From a1b9490f15899c7120f8e78135e90e4ff879717f Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 8 Aug 2023 09:33:00 +0200 Subject: [PATCH 212/263] (fix) import correct path from fixtures to common.py --- atm_dyn_iconam/tests/mpi_tests/common.py | 4 ++-- atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/atm_dyn_iconam/tests/mpi_tests/common.py b/atm_dyn_iconam/tests/mpi_tests/common.py index 5feb16cc1..6d2ab265c 100644 --- a/atm_dyn_iconam/tests/mpi_tests/common.py +++ b/atm_dyn_iconam/tests/mpi_tests/common.py @@ -15,12 +15,12 @@ import pytest from atm_dyn_iconam.tests.test_utils.data_handling import download_and_extract -from atm_dyn_iconam.tests.test_utils.fixtures import base_path, data_uris +from atm_dyn_iconam.tests.test_utils.fixtures import data_path, data_uris from icon4py.decomposition.parallel_setup import get_processor_properties props = get_processor_properties(with_mpi=True) -path = base_path.joinpath(f"mpitask{props.comm_size}") +path = data_path.joinpath(f"mpitask{props.comm_size}") data_path = path.joinpath("mch_ch_r04b09_dsl/ser_data") diff --git a/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py b/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py index 4468ccd0c..6382c260d 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_parallel_diffusion.py @@ -66,7 +66,7 @@ def test_parallel_diffusion( icon_grid = read_icon_grid(data_path, rank=props.rank) print( - f"rank={props.rank}: using local grid with {icon_grid.num_cells()} Cells, {icon_grid.num_edges()} Edges, {icon_grid.num_vertices()} Vertices" + f"rank={props.rank}/{props.comm_size}: using local grid with {icon_grid.num_cells()} Cells, {icon_grid.num_edges()} Edges, {icon_grid.num_vertices()} Vertices" ) diffusion_params = DiffusionParams(r04b09_diffusion_config) From 84a99e00b9455ccf03e29896f5ceb2f192991a95 Mon Sep 17 00:00:00 2001 From: Andreas Jocksch Date: Tue, 8 Aug 2023 13:13:44 +0200 Subject: [PATCH 213/263] calculation of c_lin_e (#254) calculation of c_lin_e in numpy --------- Co-authored-by: Magdalena Luz --- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 2 +- .../field_management/interpolation_fields.py | 51 ++++++++++++++ atm_dyn_iconam/tests/test_diffusion.py | 34 +++++----- .../tests/test_interpolation_fields.py | 52 ++++++++++++++ atm_dyn_iconam/tests/test_utils/fixtures.py | 2 +- .../tests/test_utils/serialbox_utils.py | 68 ++++++++++++------- common/src/icon4py/grid/grid_manager.py | 3 +- common/tests/conftest.py | 1 + 8 files changed, 169 insertions(+), 44 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/field_management/interpolation_fields.py create mode 100644 atm_dyn_iconam/tests/test_interpolation_fields.py diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index 228c940e7..e5d5ee484 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -50,7 +50,7 @@ def import_testutils(): helpers = import_testutils() -from helpers import serialbox_utils as sb # noqa F401 +from helpers import serialbox_utils as sb # noqa F401 class SerializationType(str, Enum): diff --git a/atm_dyn_iconam/src/icon4py/field_management/interpolation_fields.py b/atm_dyn_iconam/src/icon4py/field_management/interpolation_fields.py new file mode 100644 index 000000000..cd7fa4ab7 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/field_management/interpolation_fields.py @@ -0,0 +1,51 @@ +# 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 + + +def compute_c_lin_e( + edge_cell_length: np.array, + inv_dual_edge_length: np.array, + owner_mask: np.array, + second_boundary_layer_start_index: np.int32, +) -> np.array: + """ + Compute E2C average inverse distance. + + Args: + edge_cell_length: numpy array, representing a Field[[EdgeDim, E2CDim], float] + inv_dual_edge_length: inverse dual edge length, numpy array representing a Field[[EdgeDim], float] + owner_mask: numpy array, representing a Field[[EdgeDim], bool]boolean field, True for all edges owned by this compute node + second_boundary_layer_start_index: start index of the 2nd boundary line: c_lin_e is not calculated for the first boundary layer + + Returns: c_lin_e: numpy array representing Field[[EdgeDim, E2CDim], float] + + """ + c_lin_e_ = edge_cell_length[:, 1] * inv_dual_edge_length + c_lin_e = np.transpose(np.vstack((c_lin_e_, (1.0 - 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) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 67b14d589..98ab9d0f9 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -15,10 +15,12 @@ import pytest from atm_dyn_iconam.tests.test_utils.serialbox_utils import ( + IconDiffusionExitSavepoint, IconDiffusionInitSavepoint, ) from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.diffusion.diffusion_utils import scale_k +from icon4py.diffusion.state_utils import DiagnosticState, PrognosticState from icon4py.grid.horizontal import CellParams, EdgeParams from icon4py.grid.vertical import VerticalModelParams @@ -105,9 +107,8 @@ def test_diffusion_init( additional_parameters = DiffusionParams(config) vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) - meta = diffusion_savepoint_init.get_metadata("nlev", "linit", "date") + meta = diffusion_savepoint_init.get_metadata("linit", "date") - assert meta["nlev"] == 65 assert meta["linit"] is False assert meta["date"] == step_date_init @@ -323,30 +324,31 @@ def test_run_diffusion_single_step( def _verify_diffusion_fields( - diagnostic_state, - prognostic_state, - diffusion_savepoint_exit, + diagnostic_state: DiagnosticState, + prognostic_state: PrognosticState, + diffusion_savepoint: IconDiffusionExitSavepoint, ): - ref_div_ic = np.asarray(diffusion_savepoint_exit.div_ic()) + ref_div_ic = np.asarray(diffusion_savepoint.div_ic()) val_div_ic = np.asarray(diagnostic_state.div_ic) - ref_hdef_ic = np.asarray(diffusion_savepoint_exit.hdef_ic()) + ref_hdef_ic = np.asarray(diffusion_savepoint.hdef_ic()) val_hdef_ic = np.asarray(diagnostic_state.hdef_ic) assert np.allclose(ref_div_ic, val_div_ic) assert np.allclose(ref_hdef_ic, val_hdef_ic) - ref_w = np.asarray(diffusion_savepoint_exit.w()) + ref_w = np.asarray(diffusion_savepoint.w()) val_w = np.asarray(prognostic_state.w) - ref_dwdx = np.asarray(diffusion_savepoint_exit.dwdx()) + ref_dwdx = np.asarray(diffusion_savepoint.dwdx()) val_dwdx = np.asarray(diagnostic_state.dwdx) - ref_dwdy = np.asarray(diffusion_savepoint_exit.dwdy()) + ref_dwdy = np.asarray(diffusion_savepoint.dwdy()) val_dwdy = np.asarray(diagnostic_state.dwdy) - ref_vn = np.asarray(diffusion_savepoint_exit.vn()) - val_vn = np.asarray(prognostic_state.vn) - assert np.allclose(ref_vn, val_vn) assert np.allclose(ref_dwdx, val_dwdx) assert np.allclose(ref_dwdy, val_dwdy) + + ref_vn = np.asarray(diffusion_savepoint.vn()) + val_vn = np.asarray(prognostic_state.vn) + assert np.allclose(ref_vn, val_vn) assert np.allclose(ref_w, val_w) - ref_exner = np.asarray(diffusion_savepoint_exit.exner()) - ref_theta_v = np.asarray(diffusion_savepoint_exit.theta_v()) + ref_exner = np.asarray(diffusion_savepoint.exner()) + ref_theta_v = np.asarray(diffusion_savepoint.theta_v()) val_theta_v = np.asarray(prognostic_state.theta_v) val_exner = np.asarray(prognostic_state.exner_pressure) assert np.allclose(ref_theta_v, val_theta_v) @@ -403,5 +405,5 @@ def test_run_diffusion_initial_step( _verify_diffusion_fields( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, - diffusion_savepoint_exit=diffusion_savepoint_exit, + diffusion_savepoint=diffusion_savepoint_exit, ) diff --git a/atm_dyn_iconam/tests/test_interpolation_fields.py b/atm_dyn_iconam/tests/test_interpolation_fields.py new file mode 100644 index 000000000..2c5a3d140 --- /dev/null +++ b/atm_dyn_iconam/tests/test_interpolation_fields.py @@ -0,0 +1,52 @@ +# 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 icon4py.common.dimension import EdgeDim +from icon4py.field_management.interpolation_fields import compute_c_lin_e +from icon4py.grid.horizontal import HorizontalMarkerIndex + +@pytest.mark.datatest +def test_compute_c_lin_e( + grid_savepoint, interpolation_savepoint, icon_grid +): + inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() + edge_cell_length = grid_savepoint.edge_cell_length() + owner_mask = grid_savepoint.e_owner_mask() + c_lin_e_ref = interpolation_savepoint.c_lin_e() + lateral_boundary = icon_grid.get_start_index( + EdgeDim, + HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, + ) + c_lin_e = compute_c_lin_e( + np.asarray(edge_cell_length), + np.asarray(inv_dual_edge_length), + np.asarray(owner_mask), + lateral_boundary, + ) + + assert np.allclose(c_lin_e, c_lin_e_ref) diff --git a/atm_dyn_iconam/tests/test_utils/fixtures.py b/atm_dyn_iconam/tests/test_utils/fixtures.py index 278e0aba4..ce304fdfa 100644 --- a/atm_dyn_iconam/tests/test_utils/fixtures.py +++ b/atm_dyn_iconam/tests/test_utils/fixtures.py @@ -20,7 +20,7 @@ mch_ch_r04b09_dsl_grid_uri = ( - "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download" + "https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p/download" ) r02b04_global_grid_uri = "https://polybox.ethz.ch/index.php/s/0EM8O8U53GKGsst/download" data_uri = "https://polybox.ethz.ch/index.php/s/LcAbscZqnsx4WCf/download" diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index c927ae514..481bb08b2 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -68,27 +68,28 @@ def get_metadata(self, *names): def _read_int32_shift1(self, name: str): """ - Read a index field and shift it by -1. + Read a start indices field. - use for start indeces: the shift accounts for the zero based python + use for start indices: the shift accounts for the zero based python values are converted to int32 """ - return (self.serializer.read(name, self.savepoint) - 1).astype(int32) + return self._read_int32(name, offset=1) - def _read_int32(self, name: str): + def _read_int32(self, name: str, offset=0): """ - Read a int field by name. + Read an end indices field. use this for end indices: because FORTRAN slices are inclusive [from:to] _and_ one based this accounts for being exclusive python exclusive bounds: [from:to) field values are convert to int32 """ - return self.serializer.read(name, self.savepoint).astype(int32) + return self._read(name, offset, dtype=int32) - def read_int(self, name: str): - buffer = self.serializer.read(name, self.savepoint).astype(int) - self.log.debug(f"{name} {buffer.shape}") - return buffer + def _read_bool(self, name: str): + return self._read(name, offset=0, dtype=bool) + + def _read(self, name: str, offset=0, dtype=int): + return (self.serializer.read(name, self.savepoint) - offset).astype(dtype) class IconGridSavePoint(IconSavepoint): @@ -125,6 +126,9 @@ def edge_areas(self): def inv_dual_edge_length(self): return self._get_field("inv_dual_edge_length", EdgeDim) + def edge_cell_length(self): + return self._get_field("edge_cell_length", EdgeDim, E2CDim) + def cells_start_index(self): return self._read_int32_shift1("c_start_index") @@ -145,6 +149,12 @@ def edge_end_index(self): # one off accounts for being exclusive [from:to) return self.serializer.read("e_end_index", self.savepoint) + def c_owner_mask(self): + return self._get_field("c_owner_mask", CellDim, dtype=bool) + + def e_owner_mask(self): + return self._get_field("e_owner_mask", CellDim, dtype=bool) + def print_connectivity_info(self, name: str, ar: np.ndarray): self.log.debug(f" connectivity {name} {ar.shape}") @@ -207,22 +217,21 @@ def nrdmax(self): return self._get_connectivity_array("nrdmax") def construct_icon_grid(self) -> IconGrid: - sp_meta = self.get_metadata( - "nproma", "nlev", "num_vert", "num_cells", "num_edges" - ) + cell_starts = self.cells_start_index() cell_ends = self.cells_end_index() vertex_starts = self.vertex_start_index() vertex_ends = self.vertex_end_index() edge_starts = self.edge_start_index() edge_ends = self.edge_end_index() + nproma = self.get_metadata("nproma")["nproma"] config = GridConfig( horizontal_config=HorizontalGridSize( - num_vertices=sp_meta["nproma"], # or rather "num_vert" - num_cells=sp_meta["nproma"], # or rather "num_cells" - num_edges=sp_meta["nproma"], # or rather "num_edges" + num_vertices=nproma, # or rather "num_vert" + num_cells=nproma, # or rather "num_cells" + num_edges=nproma, # or rather "num_edges" ), - vertical_config=VerticalGridSize(num_lev=sp_meta["nlev"]), + vertical_config=VerticalGridSize(num_lev=self.num(KDim)), ) c2e2c = self.c2e2c() c2e2c0 = np.column_stack(((np.asarray(range(c2e2c.shape[0]))), c2e2c)) @@ -312,6 +321,9 @@ def construct_interpolation_state_for_diffusion(self) -> InterpolationState: nudgecoeff_e=self.nudgecoeff_e(), ) + def c_lin_e(self): + return self._get_field("c_lin_e", EdgeDim, E2CDim) + class MetricSavepoint(IconSavepoint): def construct_metric_state(self) -> MetricState: @@ -328,22 +340,28 @@ def zd_diffcoef(self): return self._get_field("zd_diffcoef", CellDim, KDim) def zd_intcoef(self): - ser_input = np.moveaxis( - (np.squeeze(self.serializer.read("vcoef", self.savepoint))), 1, -1 + return self._read_and_reorder_sparse_field("vcoef") + + def _read_and_reorder_sparse_field(self, name: str, sparse_size=3): + ser_input = np.squeeze(self.serializer.read(name, self.savepoint))[:, :, :] + if ser_input.shape[1] != sparse_size: + ser_input = np.moveaxis((ser_input), 1, -1) + + return self._linearize_first_2dims( + ser_input, sparse_size=sparse_size, target_dims=(CECDim, KDim) ) - return self._linearize_first_2dims(ser_input, sparse_size=3) - def _linearize_first_2dims(self, data: np.ndarray, sparse_size): + def _linearize_first_2dims( + self, data: np.ndarray, sparse_size: int, target_dims: tuple[Dimension, ...] + ): old_shape = data.shape assert old_shape[1] == sparse_size - return np_as_located_field(CECDim, KDim)( + return np_as_located_field(*target_dims)( data.reshape(old_shape[0] * old_shape[1], old_shape[2]) ) def zd_vertoffset(self): - ser_input = np.squeeze(self.serializer.read("zd_vertoffset", self.savepoint)) - ser_input = np.moveaxis(ser_input, 1, -1) - return self._linearize_first_2dims(ser_input, sparse_size=3) + return self._read_and_reorder_sparse_field("zd_vertoffset") def zd_vertidx(self): return np.squeeze(self.serializer.read("zd_vertidx", self.savepoint)) diff --git a/common/src/icon4py/grid/grid_manager.py b/common/src/icon4py/grid/grid_manager.py index a851ce1a8..f958b5a44 100644 --- a/common/src/icon4py/grid/grid_manager.py +++ b/common/src/icon4py/grid/grid_manager.py @@ -206,11 +206,12 @@ def get_offset_for_index_field(self, array: np.ndarray): class GridManager: """ - Read ICON grid file and set up IconGrid + Read ICON grid file and set up IconGrid. Reads an ICON grid file and extracts connectivity arrays and start-, end-indices for horizontal domain boundaries. Provides an IconGrid instance for further usage. """ + def __init__( self, transformation: IndexTransformation, diff --git a/common/tests/conftest.py b/common/tests/conftest.py index f792dab71..dafefd09d 100644 --- a/common/tests/conftest.py +++ b/common/tests/conftest.py @@ -14,6 +14,7 @@ import sys from pathlib import Path + # TODO(Magdalena) remove once test utils is a proper package def import_testutils(): testutils = ( From e544e0476ee9dcabff81283792f141246e8aafaf Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 8 Aug 2023 13:24:50 +0200 Subject: [PATCH 214/263] - add testdata and generated netcdf fields to git ignore. - pre-commit --- .gitignore | 2 ++ atm_dyn_iconam/tests/test_interpolation_fields.py | 5 ++--- common/tests/test_vertical.py | 2 +- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/.gitignore b/.gitignore index 56ff4f8b6..e7620d5cf 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,8 @@ _local _external_src _reports tmp +testdata +simple_mesh*.nc ### GT4Py #### .gt_cache/ diff --git a/atm_dyn_iconam/tests/test_interpolation_fields.py b/atm_dyn_iconam/tests/test_interpolation_fields.py index 2c5a3d140..bd08abd1e 100644 --- a/atm_dyn_iconam/tests/test_interpolation_fields.py +++ b/atm_dyn_iconam/tests/test_interpolation_fields.py @@ -30,10 +30,9 @@ from icon4py.field_management.interpolation_fields import compute_c_lin_e from icon4py.grid.horizontal import HorizontalMarkerIndex + @pytest.mark.datatest -def test_compute_c_lin_e( - grid_savepoint, interpolation_savepoint, icon_grid -): +def test_compute_c_lin_e(grid_savepoint, interpolation_savepoint, icon_grid): inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() edge_cell_length = grid_savepoint.edge_cell_length() owner_mask = grid_savepoint.e_owner_mask() diff --git a/common/tests/test_vertical.py b/common/tests/test_vertical.py index da509a6f5..7b1947db5 100644 --- a/common/tests/test_vertical.py +++ b/common/tests/test_vertical.py @@ -34,7 +34,7 @@ def test_nrdmax_calculation(max_h, damping, delta): @pytest.mark.datatest -def test_nrdmax_calculation_from_icon_input(icon_grid, grid_savepoint, damping_height): +def test_nrdmax_calculation_from_icon_input(grid_savepoint, damping_height): a = grid_savepoint.vct_a() nrdmax = grid_savepoint.nrdmax() vertical_params = VerticalModelParams( From b26db2ba747c1dec96d78bdc19a326d51d79e774 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 8 Aug 2023 15:29:46 +0200 Subject: [PATCH 215/263] state classes used in diffusion: rename files and classes --- .../src/icon4py/diffusion/diffusion.py | 26 +++++++++---------- .../{state_utils.py => diffusion_states.py} | 6 ++--- .../src/icon4py/driver/dycore_driver.py | 8 +++--- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 14 +++++----- atm_dyn_iconam/tests/test_diffusion.py | 14 +++++----- ...tate_utils.py => test_diffusion_states.py} | 0 .../tests/test_utils/serialbox_utils.py | 20 +++++++------- .../py2f/wrappers/diffusion_wrapper.py | 14 +++++----- 8 files changed, 51 insertions(+), 51 deletions(-) rename atm_dyn_iconam/src/icon4py/diffusion/{state_utils.py => diffusion_states.py} (97%) rename atm_dyn_iconam/tests/{test_state_utils.py => test_diffusion_states.py} (100%) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index d754555f4..1d7370c77 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -65,10 +65,10 @@ setup_fields_for_initial_step, zero_field, ) -from icon4py.diffusion.state_utils import ( - DiagnosticState, - InterpolationState, - MetricState, +from icon4py.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, PrognosticState, ) from icon4py.grid.horizontal import CellParams, EdgeParams, HorizontalMarkerIndex @@ -352,8 +352,8 @@ def __init__(self): self.config: Optional[DiffusionConfig] = None self.params: Optional[DiffusionParams] = None self.vertical_params: Optional[VerticalModelParams] = None - self.interpolation_state: InterpolationState = None - self.metric_state: MetricState = None + self.interpolation_state: DiffusionInterpolationState = None + self.metric_state: DiffusionMetricState = None self.diff_multfac_w: Optional[float] = None self.diff_multfac_n2w: Field[[KDim], float] = None self.smag_offset: Optional[float] = None @@ -370,8 +370,8 @@ def init( config: DiffusionConfig, params: DiffusionParams, vertical_params: VerticalModelParams, - metric_state: MetricState, - interpolation_state: InterpolationState, + metric_state: DiffusionMetricState, + interpolation_state: DiffusionInterpolationState, edge_params: EdgeParams, cell_params: CellParams, ): @@ -394,8 +394,8 @@ def init( self.params: DiffusionParams = params self.grid = grid self.vertical_params = vertical_params - self.metric_state: MetricState = metric_state - self.interpolation_state: InterpolationState = interpolation_state + self.metric_state: DiffusionMetricState = metric_state + self.interpolation_state: DiffusionInterpolationState = interpolation_state self.edge_params = edge_params self.cell_params = cell_params @@ -485,7 +485,7 @@ def _index_field(dim: Dimension, size=None): def initial_run( self, - diagnostic_state: DiagnosticState, + diagnostic_state: DiffusionDiagnosticState, prognostic_state: PrognosticState, dtime: float, ): @@ -521,7 +521,7 @@ def initial_run( def run( self, - diagnostic_state: DiagnosticState, + diagnostic_state: DiffusionDiagnosticState, prognostic_state: PrognosticState, dtime: float, ): @@ -542,7 +542,7 @@ def run( def _do_diffusion_step( self, - diagnostic_state: DiagnosticState, + diagnostic_state: DiffusionDiagnosticState, prognostic_state: PrognosticState, dtime: float, diff_multfac_vn: Field[[KDim], float], diff --git a/atm_dyn_iconam/src/icon4py/diffusion/state_utils.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_states.py similarity index 97% rename from atm_dyn_iconam/src/icon4py/diffusion/state_utils.py rename to atm_dyn_iconam/src/icon4py/diffusion/diffusion_states.py index dc1bc9944..f0b0dd2af 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/state_utils.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_states.py @@ -31,7 +31,7 @@ @dataclass(frozen=True) -class DiagnosticState: +class DiffusionDiagnosticState: """Represents the diagnostic fields needed in diffusion.""" # fields for 3D elements in turbdiff @@ -51,7 +51,7 @@ class DiagnosticState: @dataclass(frozen=True) -class MetricState: +class DiffusionMetricState: """Represents the metric state fields needed in diffusion.""" theta_ref_mc: Field[[CellDim, KDim], float] @@ -65,7 +65,7 @@ class MetricState: @dataclass(frozen=True) -class InterpolationState: +class DiffusionInterpolationState: """Represents the ICON interpolation state needed in diffusion.""" e_bln_c_s: Field[ diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index aa6703bfc..dd389125c 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -22,7 +22,7 @@ from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.diffusion.diffusion_utils import copy_diagnostic_and_prognostics -from icon4py.diffusion.state_utils import DiagnosticState, PrognosticState +from icon4py.diffusion.diffusion_states import DiffusionDiagnosticState, PrognosticState from icon4py.driver.icon_configuration import IconRunConfig, read_config from icon4py.driver.io_utils import ( SIMULATION_START_DATE, @@ -62,7 +62,7 @@ def _dynamics_timestep(self, dtime): def do_dynamics_substepping( self, dtime, - diagnostic_state: DiagnosticState, + diagnostic_state: DiffusionDiagnosticState, prognostic_state: PrognosticState, ): for _ in range(self.config.n_substeps): @@ -113,7 +113,7 @@ def _full_name(self, func: Callable): def _timestep( self, - diagnostic_state: DiagnosticState, + diagnostic_state: DiffusionDiagnosticState, prognostic_state: PrognosticState, ): @@ -128,7 +128,7 @@ def _timestep( def __call__( self, - diagnostic_state: DiagnosticState, + diagnostic_state: DiffusionDiagnosticState, prognostic_state: PrognosticState, ): log.info( diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index e5d5ee484..57d062bad 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -17,10 +17,10 @@ from enum import Enum from pathlib import Path -from icon4py.diffusion.state_utils import ( - DiagnosticState, - InterpolationState, - MetricState, +from icon4py.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, PrognosticState, ) from icon4py.grid.horizontal import CellParams, EdgeParams @@ -82,7 +82,7 @@ def read_icon_grid( def read_initial_state( gridfile_path: Path, -) -> tuple[sb.IconSerialDataProvider, DiagnosticState, PrognosticState]: +) -> tuple[sb.IconSerialDataProvider, DiffusionDiagnosticState, PrognosticState]: """ Read prognostic and diagnostic state from serialized data. @@ -134,7 +134,7 @@ def read_geometry_fields( def read_static_fields( path: Path, ser_type: SerializationType = SerializationType.SB -) -> tuple[MetricState, InterpolationState]: +) -> tuple[DiffusionMetricState, DiffusionInterpolationState]: """ Read fields for metric and interpolation state. @@ -154,7 +154,7 @@ def read_static_fields( interpolation_state = ( dataprovider.from_interpolation_savepoint().construct_interpolation_state_for_diffusion() ) - metric_state = dataprovider.from_metrics_savepoint().construct_metric_state() + metric_state = dataprovider.from_metrics_savepoint().construct_metric_state_for_diffusion() return metric_state, interpolation_state else: raise NotImplementedError(SB_ONLY_MSG) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 98ab9d0f9..4f2179334 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -20,7 +20,7 @@ ) from icon4py.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.diffusion.diffusion_utils import scale_k -from icon4py.diffusion.state_utils import DiagnosticState, PrognosticState +from icon4py.diffusion.diffusion_states import DiffusionDiagnosticState, PrognosticState from icon4py.grid.horizontal import CellParams, EdgeParams from icon4py.grid.vertical import VerticalModelParams @@ -115,7 +115,7 @@ def test_diffusion_init( interpolation_state = ( interpolation_savepoint.construct_interpolation_state_for_diffusion() ) - metric_state = metrics_savepoint.construct_metric_state() + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() edge_params = grid_savepoint.construct_edge_geometry() cell_params = grid_savepoint.construct_cell_geometry() @@ -209,7 +209,7 @@ def test_verify_diffusion_init_against_first_regular_savepoint( interpolation_state = ( interpolation_savepoint.construct_interpolation_state_for_diffusion() ) - metric_state = metrics_savepoint.construct_metric_state() + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() diffusion = Diffusion() diffusion.init( @@ -244,7 +244,7 @@ def test_verify_diffusion_init_against_other_regular_savepoint( interpolation_state = ( interpolation_savepoint.construct_interpolation_state_for_diffusion() ) - metric_state = metrics_savepoint.construct_metric_state() + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() edge_params = grid_savepoint.construct_edge_geometry() cell_params = grid_savepoint.construct_cell_geometry() @@ -288,7 +288,7 @@ def test_run_diffusion_single_step( interpolation_state = ( interpolation_savepoint.construct_interpolation_state_for_diffusion() ) - metric_state = metrics_savepoint.construct_metric_state() + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() diagnostic_state = diffusion_savepoint_init.construct_diagnostics_for_diffusion() prognostic_state = diffusion_savepoint_init.construct_prognostics() vct_a = grid_savepoint.vct_a() @@ -324,7 +324,7 @@ def test_run_diffusion_single_step( def _verify_diffusion_fields( - diagnostic_state: DiagnosticState, + diagnostic_state: DiffusionDiagnosticState, prognostic_state: PrognosticState, diffusion_savepoint: IconDiffusionExitSavepoint, ): @@ -373,7 +373,7 @@ def test_run_diffusion_initial_step( interpolation_state = ( interpolation_savepoint.construct_interpolation_state_for_diffusion() ) - metric_state = metrics_savepoint.construct_metric_state() + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() diagnostic_state = diffusion_savepoint_init.construct_diagnostics_for_diffusion() prognostic_state = diffusion_savepoint_init.construct_prognostics() vct_a = grid_savepoint.vct_a() diff --git a/atm_dyn_iconam/tests/test_state_utils.py b/atm_dyn_iconam/tests/test_diffusion_states.py similarity index 100% rename from atm_dyn_iconam/tests/test_state_utils.py rename to atm_dyn_iconam/tests/test_diffusion_states.py diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index 481bb08b2..0fe596079 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -36,10 +36,10 @@ VertexDim, ) from icon4py.diffusion.diffusion import VectorTuple -from icon4py.diffusion.state_utils import ( - DiagnosticState, - InterpolationState, - MetricState, +from icon4py.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, PrognosticState, ) from icon4py.grid.horizontal import CellParams, EdgeParams, HorizontalGridSize @@ -308,9 +308,9 @@ def rbf_vec_coeff_v2(self): def nudgecoeff_e(self): return self._get_field("nudgecoeff_e", EdgeDim) - def construct_interpolation_state_for_diffusion(self) -> InterpolationState: + def construct_interpolation_state_for_diffusion(self) -> DiffusionInterpolationState: grg = self.geofac_grg() - return InterpolationState( + return DiffusionInterpolationState( e_bln_c_s=as_1D_sparse_field(self.e_bln_c_s(), CEDim), rbf_coeff_1=self.rbf_vec_coeff_v1(), rbf_coeff_2=self.rbf_vec_coeff_v2(), @@ -326,8 +326,8 @@ def c_lin_e(self): class MetricSavepoint(IconSavepoint): - def construct_metric_state(self) -> MetricState: - return MetricState( + def construct_metric_state_for_diffusion(self) -> DiffusionMetricState: + return DiffusionMetricState( mask_hdiff=self.mask_diff(), theta_ref_mc=self.theta_ref_mc(), wgtfac_c=self.wgtfac_c(), @@ -442,8 +442,8 @@ def construct_prognostics(self) -> PrognosticState: theta_v=self.theta_v(), ) - def construct_diagnostics_for_diffusion(self) -> DiagnosticState: - return DiagnosticState( + def construct_diagnostics_for_diffusion(self) -> DiffusionDiagnosticState: + return DiffusionDiagnosticState( hdef_ic=self.hdef_ic(), div_ic=self.div_ic(), dwdx=self.dwdx(), diff --git a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py index 00158a122..db01557f0 100644 --- a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py +++ b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py @@ -30,10 +30,10 @@ DiffusionConfig, DiffusionParams, ) -from icon4py.diffusion.state_utils import ( - DiagnosticState, - InterpolationState, - MetricState, +from icon4py.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, PrognosticState, ) from icon4py.grid.horizontal import CellParams, EdgeParams @@ -99,10 +99,10 @@ def diffusion_init( vertical_params = VerticalModelParams(vct_a=vct_a, rayleigh_damping_height=nrdmax) config: DiffusionConfig = DiffusionConfig(grid, vertical_params) derived_diffusion_params = DiffusionParams(config) - metric_state = MetricState( + metric_state = DiffusionMetricState( theta_ref_mc, wgtfac_c, mask_hdiff, zd_vertidx, zd_diffcoef, zd_intcoef ) - interpolation_state = InterpolationState( + interpolation_state = DiffusionInterpolationState( e_bln_c_s, rbf_coeff_1, rbf_coeff_2, @@ -138,7 +138,7 @@ def diffusion_run( dwdx: Field[[CellDim, KDim], float], dwdy: Field[[CellDim, KDim], float], ): - diagnostic_state = DiagnosticState(hdef_ic, div_ic, dwdx, dwdy) + diagnostic_state = DiffusionDiagnosticState(hdef_ic, div_ic, dwdx, dwdy) prognostic_state = PrognosticState( w=w, vn=vn, From ad12a8dfc13765a4eaf26d421e3936d2ef139412 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 8 Aug 2023 15:40:20 +0200 Subject: [PATCH 216/263] move helper functions for timestepping to dycore_driver.py --- .../src/icon4py/diffusion/diffusion_utils.py | 29 ---------------- .../src/icon4py/driver/dycore_driver.py | 34 +++++++++++++++++-- atm_dyn_iconam/tests/test_dycore_driver.py | 4 +-- 3 files changed, 34 insertions(+), 33 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_utils.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_utils.py index 981aa76c7..ed76dee59 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_utils.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion_utils.py @@ -48,35 +48,6 @@ def _identity_e_k( return field -@program -def copy_diagnostic_and_prognostics( - hdef_ic_new: Field[[CellDim, KDim], float], - hdef_ic: Field[[CellDim, KDim], float], - div_ic_new: Field[[CellDim, KDim], float], - div_ic: Field[[CellDim, KDim], float], - dwdx_new: Field[[CellDim, KDim], float], - dwdx: Field[[CellDim, KDim], float], - dwdy_new: Field[[CellDim, KDim], float], - dwdy: Field[[CellDim, KDim], float], - vn_new: Field[[EdgeDim, KDim], float], - vn: Field[[EdgeDim, KDim], float], - w_new: Field[[CellDim, KDim], float], - w: Field[[CellDim, KDim], float], - exner_new: Field[[CellDim, KDim], float], - exner: Field[[CellDim, KDim], float], - theta_v_new: Field[[CellDim, KDim], float], - theta_v: Field[[CellDim, KDim], float], -): - _identity_c_k(hdef_ic_new, out=hdef_ic) - _identity_c_k(div_ic_new, out=div_ic) - _identity_c_k(dwdx_new, out=dwdx) - _identity_c_k(dwdy_new, out=dwdy) - _identity_e_k(vn_new, out=vn) - _identity_c_k(w_new, out=w) - _identity_c_k(exner_new, out=exner) - _identity_c_k(theta_v_new, out=theta_v) - - @field_operator def _scale_k(field: Field[[KDim], float], factor: float) -> Field[[KDim], float]: return field * factor diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index dd389125c..6eade4082 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -18,11 +18,13 @@ import click import pytz from devtools import Timer +from gt4py.next import program, Field from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn +from icon4py.common.dimension import CellDim, KDim, EdgeDim from icon4py.diffusion.diffusion import Diffusion, DiffusionParams -from icon4py.diffusion.diffusion_utils import copy_diagnostic_and_prognostics from icon4py.diffusion.diffusion_states import DiffusionDiagnosticState, PrognosticState +from icon4py.diffusion.diffusion_utils import _identity_c_k, _identity_e_k from icon4py.driver.icon_configuration import IconRunConfig, read_config from icon4py.driver.io_utils import ( SIMULATION_START_DATE, @@ -41,6 +43,34 @@ log = logging.getLogger(__name__) +# TODO (magdalena) to be removed once there is a proper time stepping +@program +def _copy_diagnostic_and_prognostics( + hdef_ic_new: Field[[CellDim, KDim], float], + hdef_ic: Field[[CellDim, KDim], float], + div_ic_new: Field[[CellDim, KDim], float], + div_ic: Field[[CellDim, KDim], float], + dwdx_new: Field[[CellDim, KDim], float], + dwdx: Field[[CellDim, KDim], float], + dwdy_new: Field[[CellDim, KDim], float], + dwdy: Field[[CellDim, KDim], float], + vn_new: Field[[EdgeDim, KDim], float], + vn: Field[[EdgeDim, KDim], float], + w_new: Field[[CellDim, KDim], float], + w: Field[[CellDim, KDim], float], + exner_new: Field[[CellDim, KDim], float], + exner: Field[[CellDim, KDim], float], + theta_v_new: Field[[CellDim, KDim], float], + theta_v: Field[[CellDim, KDim], float], +): + _identity_c_k(hdef_ic_new, out=hdef_ic) + _identity_c_k(div_ic_new, out=div_ic) + _identity_c_k(dwdx_new, out=dwdx) + _identity_c_k(dwdy_new, out=dwdy) + _identity_e_k(vn_new, out=vn) + _identity_c_k(w_new, out=w) + _identity_c_k(exner_new, out=exner) + _identity_c_k(theta_v_new, out=theta_v) class DummyAtmoNonHydro: def __init__(self, data_provider: sb_utils.IconSerialDataProvider): @@ -72,7 +102,7 @@ def do_dynamics_substepping( ) new_p = sp.construct_prognostics() new_d = sp.construct_diagnostics_for_diffusion() - copy_diagnostic_and_prognostics.with_backend(run_gtfn)( + _copy_diagnostic_and_prognostics.with_backend(run_gtfn)( new_d.hdef_ic, diagnostic_state.hdef_ic, new_d.div_ic, diff --git a/atm_dyn_iconam/tests/test_dycore_driver.py b/atm_dyn_iconam/tests/test_dycore_driver.py index b261dcd6e..0d112c78e 100644 --- a/atm_dyn_iconam/tests/test_dycore_driver.py +++ b/atm_dyn_iconam/tests/test_dycore_driver.py @@ -14,7 +14,7 @@ import numpy as np from icon4py.common.dimension import CellDim, EdgeDim, KDim -from icon4py.diffusion.diffusion_utils import copy_diagnostic_and_prognostics +from icon4py.driver.dycore_driver import _copy_diagnostic_and_prognostics from .test_utils.helpers import random_field, zero_field from .test_utils.simple_mesh import SimpleMesh @@ -39,7 +39,7 @@ def test_copy_diagnostic_and_prognostics(): f8_in = random_field(mesh, CellDim, KDim) f8_out = zero_field(mesh, CellDim, KDim) - copy_diagnostic_and_prognostics( + _copy_diagnostic_and_prognostics( f1_in, f1_out, f2_in, From 2a0b93de28159b4ae2f47e25772a45c767dca2df Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 8 Aug 2023 16:46:49 +0200 Subject: [PATCH 217/263] (fix) wrong dimension when reading field --- atm_dyn_iconam/tests/test_utils/serialbox_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index 0fe596079..f7ef2df19 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -153,7 +153,7 @@ def c_owner_mask(self): return self._get_field("c_owner_mask", CellDim, dtype=bool) def e_owner_mask(self): - return self._get_field("e_owner_mask", CellDim, dtype=bool) + return self._get_field("e_owner_mask", EdgeDim, dtype=bool) def print_connectivity_info(self, name: str, ar: np.ndarray): self.log.debug(f" connectivity {name} {ar.shape}") From ed98b170a6b6a122ee9667ac42a278899e98fecb Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 15 Aug 2023 11:28:29 +0200 Subject: [PATCH 218/263] add pyghex to dependencies --- atm_dyn_iconam/src/icon4py/decomposition/decomposed.py | 3 ++- atm_dyn_iconam/tests/mpi_tests/test_decomposed.py | 2 +- base-requirements-dev.txt | 3 +++ 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index 2a893a76e..f33b9a023 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -165,8 +165,9 @@ def is_ready(self) -> bool: class GHexMultiNode: - def __init__(self, props, domain_decomposition: DecompositionInfo): + def __init__(self, props:ProcessProperties, domain_decomposition: DecompositionInfo): self._context = ghex.context(ghex.mpi_comm(props.comm), True) + #self._context = ghex.context(ghex.mpi_comm(), True) self._domain_id_gen = DomainDescriptorIdGenerator(props) self._decomposition_info = domain_decomposition self._domain_descriptors = { diff --git a/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py b/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py index b3f23f9c0..b06a642f0 100644 --- a/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py +++ b/atm_dyn_iconam/tests/mpi_tests/test_decomposed.py @@ -177,7 +177,7 @@ def test_decomposition_info_matches_gridsize(caplog, download_data): # noqa F81 @pytest.mark.mpi -def test_create_multi_pytenode_runtime_with_mpi(download_data): # noqa F811 +def test_create_multi_node_runtime_with_mpi(download_data): # noqa F811 decomp_info = read_decomp_info(data_path, props) exchange = create_exchange(props, decomp_info) if props.comm_size > 1: diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index bd590d63e..c5a50ae58 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -1,6 +1,7 @@ # VCS -e git+https://github.com/GridTools/gt4py.git@main#egg=gt4py git+https://github.com/GridTools/serialbox#egg=serialbox&subdirectory=src/serialbox-python +git+https://github.com/boeschf/GHEX.git@pipify#subdirectory=bindings/python # PyPI @@ -31,3 +32,5 @@ wget >= 3.2 netcdf4>=1.6.0 mpi4py<=3.1.4 pytest-mpi>=0.6 + + From 5211c82307503316952337da5cb12e4d46cb5b1a Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 15 Aug 2023 11:28:43 +0200 Subject: [PATCH 219/263] convert ParallelProperties to dataclass --- .../icon4py/decomposition/parallel_setup.py | 29 +++++++++---------- 1 file changed, 13 insertions(+), 16 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py index 1dd9828b4..98abbc4e4 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py @@ -10,7 +10,9 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import functools import logging +from dataclasses import dataclass from typing import Optional, Union import mpi4py @@ -56,29 +58,24 @@ def finalize_mpi(): log.info("finalizing MPI") MPI.Finalize() - +@dataclass(frozen=True) class ProcessProperties: - def __init__(self, comm: Optional[mpi4py.MPI.Comm]): - self._communicator_name: str = comm.Get_name() if comm else "" - self._rank: int = comm.Get_rank() if comm else 0 - self._comm_size = comm.Get_size() if comm else 1 - self._comm = comm + comm: Optional[mpi4py.MPI.Comm] = None + - @property + @functools.cached_property def rank(self): - return self._rank + return self.comm.Get_rank() if self.comm else 0 + - @property + @functools.cached_property def comm_name(self): - return self._communicator_name + return self.comm.Get_name() if self.comm else "" - @property + @functools.cached_property def comm_size(self): - return self._comm_size + return self.comm.Get_size() if self.comm else 1 - @property - def comm(self): - return self._comm @classmethod def from_mpi_comm(cls, comm: mpi4py.MPI.Comm): @@ -86,7 +83,7 @@ def from_mpi_comm(cls, comm: mpi4py.MPI.Comm): @classmethod def from_single_node(cls): - return ProcessProperties(None) + return ProcessProperties() class ParallelLogger(logging.Filter): From 6da81ba1cbe0630d53ac8bcbbd36e3f44bc22919 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 15 Aug 2023 13:39:59 +0200 Subject: [PATCH 220/263] (fix) create GHex context without external MPIComm --- atm_dyn_iconam/src/icon4py/decomposition/decomposed.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index f33b9a023..47e82decc 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -166,8 +166,9 @@ def is_ready(self) -> bool: class GHexMultiNode: def __init__(self, props:ProcessProperties, domain_decomposition: DecompositionInfo): - self._context = ghex.context(ghex.mpi_comm(props.comm), True) - #self._context = ghex.context(ghex.mpi_comm(), True) + #TODO (magdalena) fix python build of ghex: cmake find_python_module(mpi4py) does not work + #self._context = ghex.context(ghex.mpi_comm(props.comm), True) + self._context = ghex.context(ghex.mpi_comm(), True) self._domain_id_gen = DomainDescriptorIdGenerator(props) self._decomposition_info = domain_decomposition self._domain_descriptors = { From 9e5101a56ca16d9409f52f4bacb4df9d392c6b3b Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 17 Aug 2023 11:00:44 +0200 Subject: [PATCH 221/263] update README.md for parallel tests. --- atm_dyn_iconam/tests/mpi_tests/README.md | 86 +++++------------------- 1 file changed, 16 insertions(+), 70 deletions(-) diff --git a/atm_dyn_iconam/tests/mpi_tests/README.md b/atm_dyn_iconam/tests/mpi_tests/README.md index 55cec7e57..84b47bd49 100644 --- a/atm_dyn_iconam/tests/mpi_tests/README.md +++ b/atm_dyn_iconam/tests/mpi_tests/README.md @@ -1,10 +1,12 @@ -## running parallel version of diffusion +## Running parallel version of diffusion -### installation +### Installation -The parallelized code uses [GHEX](https://github.com/ghex-org/GHEX) with MPI for halo exchanges. GHEX has a CMake build but no setup script for pip, so it needs to be installed manually: +The parallelized code uses [GHEX](https://github.com/ghex-org/GHEX) with MPI for halo exchanges. +The GHEX python bindings can be installed with pip from VCS. However it has some system +dependencies -1. You need a running MPI installation in the system. On linux (apt base system) do +1.You need a running MPI installation in the system. On linux (apt base system) do ```bash sudo apt-get install libopenmpi-dev @@ -17,77 +19,21 @@ brew install mpich ``` 2. You need to have boost (headers) installed in the system -3. clone GHEX: - -```bash -cd {icon4py}/_external_src -git clone --recursive -b refactoring2 git@github.com:boeschf/GHEX.git -``` - -3. build GHEX - +3. clone and install `icon4py` as described in the [Readme](../../../README.md) install python dependencies +In an existing local `icon4py` clone rerun `pip` install: ``` -cd GHEX -mkdir build -cd build -cmake .. -DCMAKE_BUILD_TYPE=Debug \ --DGHEX_GIT_SUBMODULE=OFF \ --DGHEX_USE_BUNDLED_LIBS=ON \ --DGHEX_USE_BUNDLED_GRIDTOOLS=ON \ --DBUILD_TESTING=OFF \ --DGT_BUILD_TESTING=OFF \ --DGT_INSTALL_EXAMPLES=OFF \ --DGHEX_USE_BUNDLED_OOMPH=ON \ --DGHEX_TRANSPORT_BACKEND=MPI \ --DGHEX_USE_XPMEM=OFF \ --DGHEX_WITH_TESTING=ON \ --DGHEX_BUILD_PYTHON_BINDINGS=ON \ --DGHEX_BUILD_BENCHMARKS=ON \ --DGHEX_BUILD_FORTRAN=ON \ --DGHEX_ENABLE_ATLAS_BINDINGS=OFF \ --DGHEX_ENABLE_PARMETIS_BINDINGS=ON \ --DMETIS_INCLUDE_DIR=/opt/metis/include \ --DMETIS_LIB_DIR=/opt/metis/lib \ --DPARMETIS_INCLUDE_DIR=/opt/parmetis/include \ --DPARMETIS_LIB_DIR=/opt/parmetis/lib \ --DMPIEXEC_PREFLAGS=--oversubscribe \ --DGHEX_USE_GPU=OFF -make -make test ## runs the C++ tests +pip install --src _external_src -r requirements-dev.txt ``` -#### generating python bindings - -```bash -cmake -DGHEX_BUILD_PYTHON_BINDINGS=ON . # turns on python bindings, need pybind11 wo we install it in the -``` - -builds GHEX including python bindings. If the build fails with an error message that `pybind11` was not found, you can either install it int the Python `.venv` in GHEX build folder or set the `pybind11_DIR` variable to some other location. - -```bash -cd pyghex_venv/ -source bin/activate -pip install pybind11 -export pybind11_DIR=./pyghex_venv/lib/python3.10/site-packages/pybind11 -make -make test ## will now run the python tests -``` - -#### use GHEX bindings from icon4py - -simply create a sym link the installation above: - -``` -cd {icon4py}/.venv/lib/python3.10/site-packages -ln -s ../../../../_external_src/GHEX/build/bindings/python/ghex ghex -``` - -### run parallel test - -all MPI related tests are in the folder `icon4py/atm_dyn_iconam/tests/mpi_tests` tests depending on MPI are marked with `@pytest.mark.mpi` are therefore skipped by any serial pytest run. +### Run parallel test +The project includes the `pytest-mpi` utility which augments `pytest` with some MPI +specific utilities. MPI dependent tests are marked with `@pytest.mark.mpi` and are +skipped when not run under mpi (`--with-mpi`) +All MPI related tests are in the folder `icon4py/atm_dyn_iconam/tests/mpi_tests` tests depending on MPI are marked with `@pytest.mark.mpi` are therefore skipped by any serial pytest run. -run them with +Run the tests with ```bash mpirun -np 2 pytest -v -s --with-mpi tests/mpi_tests/ ``` +Test using serialized test data can be run with either 1, 2 or 4 nodes. From 1bf110b5dc85521bd98dc228254f8e56a482a793 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 17 Aug 2023 14:26:07 +0200 Subject: [PATCH 222/263] add pyGHEX to base-requirements.txt remove workaround for ghex/mpi4py build problem in decomposed.py --- atm_dyn_iconam/src/icon4py/decomposition/decomposed.py | 4 +--- base-requirements.txt | 1 + 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index 47e82decc..6593d9425 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -166,9 +166,7 @@ def is_ready(self) -> bool: class GHexMultiNode: def __init__(self, props:ProcessProperties, domain_decomposition: DecompositionInfo): - #TODO (magdalena) fix python build of ghex: cmake find_python_module(mpi4py) does not work - #self._context = ghex.context(ghex.mpi_comm(props.comm), True) - self._context = ghex.context(ghex.mpi_comm(), True) + self._context = ghex.context(ghex.mpi_comm(props.comm), True) self._domain_id_gen = DomainDescriptorIdGenerator(props) self._decomposition_info = domain_decomposition self._domain_descriptors = { diff --git a/base-requirements.txt b/base-requirements.txt index 06f0a1341..d3055912b 100644 --- a/base-requirements.txt +++ b/base-requirements.txt @@ -1,4 +1,5 @@ # VCS gt4py @ git+https://github.com/GridTools/gt4py.git@main +pyghex @ git+https://github.com/boeschf/GHEX.git@pipify#subdirectory=bindings/python netcdf4 >=1.6.0 mpi4py<=3.1.4 From 6f28f8ae608b10f28adf52a2abcd2719dae4c0e0 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 18 Aug 2023 09:19:23 +0200 Subject: [PATCH 223/263] pre-commit fixes --- atm_dyn_iconam/src/icon4py/diffusion/diffusion.py | 12 ++++++------ atm_dyn_iconam/src/icon4py/driver/dycore_driver.py | 11 ++++++++--- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 4 +++- atm_dyn_iconam/tests/test_diffusion.py | 5 ++++- atm_dyn_iconam/tests/test_utils/serialbox_utils.py | 4 +++- 5 files changed, 24 insertions(+), 12 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 1d7370c77..26c425729 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -57,6 +57,12 @@ GAS_CONSTANT_DRY_AIR, ) from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim +from icon4py.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, + PrognosticState, +) from icon4py.diffusion.diffusion_utils import ( copy_field, init_diffusion_local_fields_for_regular_timestep, @@ -65,12 +71,6 @@ setup_fields_for_initial_step, zero_field, ) -from icon4py.diffusion.diffusion_states import ( - DiffusionDiagnosticState, - DiffusionInterpolationState, - DiffusionMetricState, - PrognosticState, -) from icon4py.grid.horizontal import CellParams, EdgeParams, HorizontalMarkerIndex from icon4py.grid.icon_grid import IconGrid from icon4py.grid.vertical import VerticalModelParams diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index 6eade4082..92ca31890 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -18,12 +18,15 @@ import click import pytz from devtools import Timer -from gt4py.next import program, Field +from gt4py.next import Field, program from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn -from icon4py.common.dimension import CellDim, KDim, EdgeDim +from icon4py.common.dimension import CellDim, EdgeDim, KDim from icon4py.diffusion.diffusion import Diffusion, DiffusionParams -from icon4py.diffusion.diffusion_states import DiffusionDiagnosticState, PrognosticState +from icon4py.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + PrognosticState, +) from icon4py.diffusion.diffusion_utils import _identity_c_k, _identity_e_k from icon4py.driver.icon_configuration import IconRunConfig, read_config from icon4py.driver.io_utils import ( @@ -43,6 +46,7 @@ log = logging.getLogger(__name__) + # TODO (magdalena) to be removed once there is a proper time stepping @program def _copy_diagnostic_and_prognostics( @@ -72,6 +76,7 @@ def _copy_diagnostic_and_prognostics( _identity_c_k(exner_new, out=exner) _identity_c_k(theta_v_new, out=theta_v) + class DummyAtmoNonHydro: def __init__(self, data_provider: sb_utils.IconSerialDataProvider): self.config = None diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index 57d062bad..df3dffec6 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -154,7 +154,9 @@ def read_static_fields( interpolation_state = ( dataprovider.from_interpolation_savepoint().construct_interpolation_state_for_diffusion() ) - metric_state = dataprovider.from_metrics_savepoint().construct_metric_state_for_diffusion() + metric_state = ( + dataprovider.from_metrics_savepoint().construct_metric_state_for_diffusion() + ) return metric_state, interpolation_state else: raise NotImplementedError(SB_ONLY_MSG) diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 4f2179334..65871df48 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -19,8 +19,11 @@ IconDiffusionInitSavepoint, ) from icon4py.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + PrognosticState, +) from icon4py.diffusion.diffusion_utils import scale_k -from icon4py.diffusion.diffusion_states import DiffusionDiagnosticState, PrognosticState from icon4py.grid.horizontal import CellParams, EdgeParams from icon4py.grid.vertical import VerticalModelParams diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index f7ef2df19..a93097833 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -308,7 +308,9 @@ def rbf_vec_coeff_v2(self): def nudgecoeff_e(self): return self._get_field("nudgecoeff_e", EdgeDim) - def construct_interpolation_state_for_diffusion(self) -> DiffusionInterpolationState: + def construct_interpolation_state_for_diffusion( + self, + ) -> DiffusionInterpolationState: grg = self.geofac_grg() return DiffusionInterpolationState( e_bln_c_s=as_1D_sparse_field(self.e_bln_c_s(), CEDim), From 07e650b13f9387db9169ff72a09d69e3fc022a93 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 18 Aug 2023 09:48:48 +0200 Subject: [PATCH 224/263] pre-commit fixes --- .../src/icon4py/decomposition/decomposed.py | 4 +++- .../src/icon4py/decomposition/parallel_setup.py | 4 +--- atm_dyn_iconam/src/icon4py/diffusion/diffusion.py | 12 ++++++------ .../src/icon4py/driver/dycore_driver.py | 11 ++++++++--- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 14 ++++++++------ atm_dyn_iconam/tests/mpi_tests/README.md | 15 ++++++--------- atm_dyn_iconam/tests/test_diffusion.py | 5 ++++- .../tests/test_utils/serialbox_utils.py | 4 +++- 8 files changed, 39 insertions(+), 30 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index 6593d9425..d4d6cca55 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -165,7 +165,9 @@ def is_ready(self) -> bool: class GHexMultiNode: - def __init__(self, props:ProcessProperties, domain_decomposition: DecompositionInfo): + def __init__( + self, props: ProcessProperties, domain_decomposition: DecompositionInfo + ): self._context = ghex.context(ghex.mpi_comm(props.comm), True) self._domain_id_gen = DomainDescriptorIdGenerator(props) self._decomposition_info = domain_decomposition diff --git a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py index 98abbc4e4..f848edd5d 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/parallel_setup.py @@ -58,16 +58,15 @@ def finalize_mpi(): log.info("finalizing MPI") MPI.Finalize() + @dataclass(frozen=True) class ProcessProperties: comm: Optional[mpi4py.MPI.Comm] = None - @functools.cached_property def rank(self): return self.comm.Get_rank() if self.comm else 0 - @functools.cached_property def comm_name(self): return self.comm.Get_name() if self.comm else "" @@ -76,7 +75,6 @@ def comm_name(self): def comm_size(self): return self.comm.Get_size() if self.comm else 1 - @classmethod def from_mpi_comm(cls, comm: mpi4py.MPI.Comm): return ProcessProperties(comm) diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py index 50a43b94e..f056a2eea 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py @@ -59,6 +59,12 @@ ) from icon4py.common.dimension import CellDim, EdgeDim, KDim, VertexDim from icon4py.decomposition.decomposed import ExchangeRuntime, SingleNode +from icon4py.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, + PrognosticState, +) from icon4py.diffusion.diffusion_utils import ( copy_field, init_diffusion_local_fields_for_regular_timestep, @@ -67,12 +73,6 @@ setup_fields_for_initial_step, zero_field, ) -from icon4py.diffusion.diffusion_states import ( - DiffusionDiagnosticState, - DiffusionInterpolationState, - DiffusionMetricState, - PrognosticState, -) from icon4py.grid.horizontal import CellParams, EdgeParams, HorizontalMarkerIndex from icon4py.grid.icon_grid import IconGrid from icon4py.grid.vertical import VerticalModelParams diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index ff4df533e..e2d7eff43 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -18,17 +18,20 @@ import click import pytz from devtools import Timer -from gt4py.next import program, Field +from gt4py.next import Field, program from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn +from icon4py.common.dimension import CellDim, EdgeDim, KDim from icon4py.decomposition.decomposed import create_exchange from icon4py.decomposition.parallel_setup import ( ProcessProperties, get_processor_properties, ) -from icon4py.common.dimension import CellDim, KDim, EdgeDim from icon4py.diffusion.diffusion import Diffusion, DiffusionParams -from icon4py.diffusion.diffusion_states import DiffusionDiagnosticState, PrognosticState +from icon4py.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + PrognosticState, +) from icon4py.diffusion.diffusion_utils import _identity_c_k, _identity_e_k from icon4py.driver.icon_configuration import IconRunConfig, read_config from icon4py.driver.io_utils import ( @@ -49,6 +52,7 @@ log = logging.getLogger(__name__) + # TODO (magdalena) to be removed once there is a proper time stepping @program def _copy_diagnostic_and_prognostics( @@ -78,6 +82,7 @@ def _copy_diagnostic_and_prognostics( _identity_c_k(exner_new, out=exner) _identity_c_k(theta_v_new, out=theta_v) + class DummyAtmoNonHydro: def __init__(self, data_provider: sb_utils.IconSerialDataProvider): self.config = None diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index 1dfa22a17..859c07fdc 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -17,14 +17,14 @@ from enum import Enum from pathlib import Path +from icon4py.decomposition.decomposed import DecompositionInfo +from icon4py.decomposition.parallel_setup import ParallelLogger, ProcessProperties from icon4py.diffusion.diffusion_states import ( DiffusionDiagnosticState, - PrognosticState, DiffusionInterpolationState, - DiffusionMetricState) -from icon4py.decomposition.decomposed import DecompositionInfo -from icon4py.decomposition.parallel_setup import ParallelLogger, ProcessProperties - + DiffusionMetricState, + PrognosticState, +) from icon4py.grid.horizontal import CellParams, EdgeParams from icon4py.grid.icon_grid import IconGrid from icon4py.grid.vertical import VerticalModelParams @@ -173,7 +173,9 @@ def read_static_fields( interpolation_state = ( dataprovider.from_interpolation_savepoint().construct_interpolation_state_for_diffusion() ) - metric_state = dataprovider.from_metrics_savepoint().construct_metric_state_for_diffusion() + metric_state = ( + dataprovider.from_metrics_savepoint().construct_metric_state_for_diffusion() + ) return metric_state, interpolation_state else: raise NotImplementedError(SB_ONLY_MSG) diff --git a/atm_dyn_iconam/tests/mpi_tests/README.md b/atm_dyn_iconam/tests/mpi_tests/README.md index 84b47bd49..98b076ee7 100644 --- a/atm_dyn_iconam/tests/mpi_tests/README.md +++ b/atm_dyn_iconam/tests/mpi_tests/README.md @@ -2,9 +2,7 @@ ### Installation -The parallelized code uses [GHEX](https://github.com/ghex-org/GHEX) with MPI for halo exchanges. -The GHEX python bindings can be installed with pip from VCS. However it has some system -dependencies +The parallelized code uses [GHEX](https://github.com/ghex-org/GHEX) with MPI for halo exchanges. The GHEX python bindings can be installed with pip from VCS. However it has some system dependencies 1.You need a running MPI installation in the system. On linux (apt base system) do @@ -19,21 +17,20 @@ brew install mpich ``` 2. You need to have boost (headers) installed in the system -3. clone and install `icon4py` as described in the [Readme](../../../README.md) install python dependencies -In an existing local `icon4py` clone rerun `pip` install: +3. clone and install `icon4py` as described in the [Readme](../../../README.md) install python dependencies In an existing local `icon4py` clone rerun `pip` install: + ``` pip install --src _external_src -r requirements-dev.txt ``` ### Run parallel test -The project includes the `pytest-mpi` utility which augments `pytest` with some MPI -specific utilities. MPI dependent tests are marked with `@pytest.mark.mpi` and are -skipped when not run under mpi (`--with-mpi`) -All MPI related tests are in the folder `icon4py/atm_dyn_iconam/tests/mpi_tests` tests depending on MPI are marked with `@pytest.mark.mpi` are therefore skipped by any serial pytest run. + +The project includes the `pytest-mpi` utility which augments `pytest` with some MPI specific utilities. MPI dependent tests are marked with `@pytest.mark.mpi` and are skipped when not run under mpi (`--with-mpi`) All MPI related tests are in the folder `icon4py/atm_dyn_iconam/tests/mpi_tests` tests depending on MPI are marked with `@pytest.mark.mpi` are therefore skipped by any serial pytest run. Run the tests with ```bash mpirun -np 2 pytest -v -s --with-mpi tests/mpi_tests/ ``` + Test using serialized test data can be run with either 1, 2 or 4 nodes. diff --git a/atm_dyn_iconam/tests/test_diffusion.py b/atm_dyn_iconam/tests/test_diffusion.py index 4f2179334..65871df48 100644 --- a/atm_dyn_iconam/tests/test_diffusion.py +++ b/atm_dyn_iconam/tests/test_diffusion.py @@ -19,8 +19,11 @@ IconDiffusionInitSavepoint, ) from icon4py.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + PrognosticState, +) from icon4py.diffusion.diffusion_utils import scale_k -from icon4py.diffusion.diffusion_states import DiffusionDiagnosticState, PrognosticState from icon4py.grid.horizontal import CellParams, EdgeParams from icon4py.grid.vertical import VerticalModelParams diff --git a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py index 84ad244be..5e4807146 100644 --- a/atm_dyn_iconam/tests/test_utils/serialbox_utils.py +++ b/atm_dyn_iconam/tests/test_utils/serialbox_utils.py @@ -340,7 +340,9 @@ def rbf_vec_coeff_v2(self): def nudgecoeff_e(self): return self._get_field("nudgecoeff_e", EdgeDim) - def construct_interpolation_state_for_diffusion(self) -> DiffusionInterpolationState: + def construct_interpolation_state_for_diffusion( + self, + ) -> DiffusionInterpolationState: grg = self.geofac_grg() return DiffusionInterpolationState( e_bln_c_s=as_1D_sparse_field(self.e_bln_c_s(), CEDim), From 0a801bb1fb83449be558540f3b1d6dd2dd90f426 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 18 Aug 2023 09:51:06 +0200 Subject: [PATCH 225/263] add missing pyty library --- base-requirements-dev.txt | 5 +++-- base-requirements.txt | 3 ++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 044e5afe9..92886481a 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -26,6 +26,7 @@ pytest-factoryboy>=2.0 pytest-xdist[psutil]>=2.2 setuptools>=40.8.0 wheel>=0.37.1 -tox >= 3.25 -wget >= 3.2 +tox>=3.25 +wget>=3.2 netcdf4>=1.6.0 +pytz>=2023.2 diff --git a/base-requirements.txt b/base-requirements.txt index a0f5ecb08..814fde7e6 100644 --- a/base-requirements.txt +++ b/base-requirements.txt @@ -1,3 +1,4 @@ # VCS gt4py @ git+https://github.com/GridTools/gt4py.git@main -netcdf4 >=1.6.0 +netcdf4>=1.6.0 +pytz>=2023.2 From ccaf1bab8fb013d4d2ba18275e6573807de9c7c7 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 18 Aug 2023 11:11:18 +0200 Subject: [PATCH 226/263] (temporarily) instantiate ghex context without passing external communicator --- atm_dyn_iconam/src/icon4py/decomposition/decomposed.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index d4d6cca55..ff05fda88 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -168,7 +168,9 @@ class GHexMultiNode: def __init__( self, props: ProcessProperties, domain_decomposition: DecompositionInfo ): - self._context = ghex.context(ghex.mpi_comm(props.comm), True) + #self._context = ghex.context(ghex.mpi_comm(props.comm), True) + # TODO (magdalena) change back once GHEX build problem is fixed + self._context = ghex.context(ghex.mpi_comm(), True) self._domain_id_gen = DomainDescriptorIdGenerator(props) self._decomposition_info = domain_decomposition self._domain_descriptors = { From 7b880a892f3c370a0d997ad1db1cd61e986858a0 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 18 Aug 2023 13:21:17 +0200 Subject: [PATCH 227/263] remove instantiate ghex context without passing external communicator --- atm_dyn_iconam/src/icon4py/decomposition/decomposed.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index ff05fda88..a3738498d 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -168,8 +168,7 @@ class GHexMultiNode: def __init__( self, props: ProcessProperties, domain_decomposition: DecompositionInfo ): - #self._context = ghex.context(ghex.mpi_comm(props.comm), True) - # TODO (magdalena) change back once GHEX build problem is fixed + self._context = ghex.context(ghex.mpi_comm(props.comm), True) self._context = ghex.context(ghex.mpi_comm(), True) self._domain_id_gen = DomainDescriptorIdGenerator(props) self._decomposition_info = domain_decomposition From cda453f41ad81fc7452a003ebbbfac33a51cb900 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 18 Aug 2023 14:25:21 +0200 Subject: [PATCH 228/263] remove instantiate ghex context without passing external communicator --- atm_dyn_iconam/src/icon4py/decomposition/decomposed.py | 1 - atm_dyn_iconam/tests/mpi_tests/README.md | 5 ++++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py index a3738498d..d4d6cca55 100644 --- a/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py +++ b/atm_dyn_iconam/src/icon4py/decomposition/decomposed.py @@ -169,7 +169,6 @@ def __init__( self, props: ProcessProperties, domain_decomposition: DecompositionInfo ): self._context = ghex.context(ghex.mpi_comm(props.comm), True) - self._context = ghex.context(ghex.mpi_comm(), True) self._domain_id_gen = DomainDescriptorIdGenerator(props) self._decomposition_info = domain_decomposition self._domain_descriptors = { diff --git a/atm_dyn_iconam/tests/mpi_tests/README.md b/atm_dyn_iconam/tests/mpi_tests/README.md index 98b076ee7..0020ffbdb 100644 --- a/atm_dyn_iconam/tests/mpi_tests/README.md +++ b/atm_dyn_iconam/tests/mpi_tests/README.md @@ -33,4 +33,7 @@ Run the tests with mpirun -np 2 pytest -v -s --with-mpi tests/mpi_tests/ ``` -Test using serialized test data can be run with either 1, 2 or 4 nodes. +Test using serialized test data can be run with either 1, 2 or 4 nodes. + +When you run the tests for the first time +the serialized test data will be downloaded from an online storage. That make some time... From 4fa333c79e9a49a9756e8073dbe04fcdec8e0f69 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 18 Aug 2023 18:03:20 +0200 Subject: [PATCH 229/263] moving diffusion code into diffusion namespace package --- .../src/icon4py/driver/dycore_driver.py | 6 +- .../src/icon4py/driver/icon_configuration.py | 2 +- atm_dyn_iconam/src/icon4py/driver/io_utils.py | 2 +- model/atmosphere/diffusion/.bumpversion.cfg | 2 +- model/atmosphere/diffusion/README.md | 4 +- model/atmosphere/diffusion/pyproject.toml | 11 +- .../model/atmosphere}/diffusion/__init__.py | 21 +++ .../model/atmosphere}/diffusion/diffusion.py | 21 +-- .../atmosphere}/diffusion/diffusion_states.py | 0 .../atmosphere}/diffusion/diffusion_utils.py | 6 - .../atmosphere/diffusion/stencils/__init__.py | 0 .../apply_diffusion_to_theta_and_exner.py | 8 +- .../stencils}/apply_diffusion_to_vn.py | 8 +- ...ute_horizontal_gradients_for_turbulance.py | 8 +- .../apply_nabla2_and_nabla4_global_to_vn.py | 0 .../apply_nabla2_and_nabla4_to_vn.py | 0 .../apply_nabla2_to_vn_in_lateral_boundary.py | 0 .../diffusion/stencils}/apply_nabla2_to_w.py | 0 ...pply_nabla2_to_w_in_upper_damping_layer.py | 0 ...te_diagnostic_quantities_for_turbulence.py | 4 +- .../calculate_diagnostics_for_turbulence.py | 0 ..._coefficients_for_grid_point_cold_pools.py | 4 +- ...ate_horizontal_gradients_for_turbulence.py | 0 ...ate_nabla2_and_smag_coefficients_for_vn.py | 0 .../stencils}/calculate_nabla2_for_theta.py | 4 +- .../stencils}/calculate_nabla2_for_w.py | 0 .../stencils}/calculate_nabla2_for_z.py | 0 .../stencils}/calculate_nabla2_of_theta.py | 0 .../diffusion/stencils}/calculate_nabla4.py | 0 ...n_coefficient_for_grid_point_cold_pools.py | 0 ...d_for_grid_point_cold_pools_enhancement.py | 0 ...orary_fields_for_turbulence_diagnostics.py | 0 ...fusion_nabla_of_theta_over_steep_points.py | 0 .../stencils}/update_theta_and_exner.py | 0 model/atmosphere/diffusion/tests/__init__.py | 0 model/atmosphere/diffusion/tests/conftest.py | 66 ++++++++++ ...st_apply_nabla2_and_nabla4_global_to_vn.py | 2 +- .../test_apply_nabla2_and_nabla4_to_vn.py | 6 +- ..._apply_nabla2_to_vn_in_lateral_boundary.py | 2 +- .../tests/test_apply_nabla2_to_w.py | 2 +- ...pply_nabla2_to_w_in_upper_damping_layer.py | 2 +- ...st_calculate_diagnostics_for_turbulence.py | 2 +- ...ate_horizontal_gradients_for_turbulence.py | 2 +- ...ate_nabla2_and_smag_coefficients_for_vn.py | 2 +- .../tests/test_calculate_nabla2_for_w.py | 2 +- .../tests/test_calculate_nabla2_for_z.py | 2 +- .../tests/test_calculate_nabla2_of_theta.py | 2 +- .../tests/test_calculate_nabla4.py | 2 +- .../tests/test_diffusion.py | 10 +- .../tests/test_diffusion_states.py | 0 .../tests/test_diffusion_utils.py | 13 +- ...n_coefficient_for_grid_point_cold_pools.py | 2 +- ...d_for_grid_point_cold_pools_enhancement.py | 2 +- ...orary_fields_for_turbulence_diagnostics.py | 4 +- ...fusion_nabla_of_theta_over_steep_points.py | 2 +- .../tests/test_update_theta_and_exner.py | 2 +- model/atmosphere/dycore/tests/conftest.py | 120 +----------------- .../dycore/tests/test_interpolation_fields.py | 2 +- .../model/common/decomposition/__init__.py | 0 .../model/common/decomposition/decomposed.py | 2 +- .../icon4py/model/common/grid/icon_grid.py | 8 +- .../model/common/test_utils/fixtures.py | 75 +++++++++++ .../common/test_utils/serialbox_utils.py | 6 +- .../common/src/icon4py/model/common/utils.py | 7 + .../py2f/wrappers/diffusion_wrapper.py | 4 +- requirements-dev.txt | 1 + requirements.txt | 1 + tools/tests/liskov/test_external.py | 2 +- 68 files changed, 258 insertions(+), 210 deletions(-) rename {atm_dyn_iconam/src/icon4py => model/atmosphere/diffusion/src/icon4py/model/atmosphere}/diffusion/__init__.py (55%) rename {atm_dyn_iconam/src/icon4py => model/atmosphere/diffusion/src/icon4py/model/atmosphere}/diffusion/diffusion.py (97%) rename {atm_dyn_iconam/src/icon4py => model/atmosphere/diffusion/src/icon4py/model/atmosphere}/diffusion/diffusion_states.py (100%) rename {atm_dyn_iconam/src/icon4py => model/atmosphere/diffusion/src/icon4py/model/atmosphere}/diffusion/diffusion_utils.py (98%) create mode 100644 model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/__init__.py rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/apply_diffusion_to_theta_and_exner.py (89%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/apply_diffusion_to_vn.py (91%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py (90%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/apply_nabla2_and_nabla4_global_to_vn.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/apply_nabla2_and_nabla4_to_vn.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/apply_nabla2_to_vn_in_lateral_boundary.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/apply_nabla2_to_w.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/apply_nabla2_to_w_in_upper_damping_layer.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_diagnostic_quantities_for_turbulence.py (92%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_diagnostics_for_turbulence.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py (90%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_horizontal_gradients_for_turbulence.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_nabla2_and_smag_coefficients_for_vn.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_nabla2_for_theta.py (91%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_nabla2_for_w.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_nabla2_for_z.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_nabla2_of_theta.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/calculate_nabla4.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/enhance_diffusion_coefficient_for_grid_point_cold_pools.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/temporary_field_for_grid_point_cold_pools_enhancement.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/temporary_fields_for_turbulence_diagnostics.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py (100%) rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => diffusion/src/icon4py/model/atmosphere/diffusion/stencils}/update_theta_and_exner.py (100%) create mode 100644 model/atmosphere/diffusion/tests/__init__.py create mode 100644 model/atmosphere/diffusion/tests/conftest.py rename model/atmosphere/{dycore => diffusion}/tests/test_apply_nabla2_and_nabla4_global_to_vn.py (95%) rename model/atmosphere/{dycore => diffusion}/tests/test_apply_nabla2_and_nabla4_to_vn.py (93%) rename model/atmosphere/{dycore => diffusion}/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py (94%) rename model/atmosphere/{dycore => diffusion}/tests/test_apply_nabla2_to_w.py (95%) rename model/atmosphere/{dycore => diffusion}/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py (94%) rename model/atmosphere/{dycore => diffusion}/tests/test_calculate_diagnostics_for_turbulence.py (95%) rename model/atmosphere/{dycore => diffusion}/tests/test_calculate_horizontal_gradients_for_turbulence.py (95%) rename model/atmosphere/{dycore => diffusion}/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py (98%) rename model/atmosphere/{dycore => diffusion}/tests/test_calculate_nabla2_for_w.py (95%) rename model/atmosphere/{dycore => diffusion}/tests/test_calculate_nabla2_for_z.py (95%) rename model/atmosphere/{dycore => diffusion}/tests/test_calculate_nabla2_of_theta.py (95%) rename model/atmosphere/{dycore => diffusion}/tests/test_calculate_nabla4.py (97%) rename model/atmosphere/{dycore => diffusion}/tests/test_diffusion.py (97%) rename model/atmosphere/{dycore => diffusion}/tests/test_diffusion_states.py (100%) rename model/atmosphere/{dycore => diffusion}/tests/test_diffusion_utils.py (94%) rename model/atmosphere/{dycore => diffusion}/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py (92%) rename model/atmosphere/{dycore => diffusion}/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py (94%) rename model/atmosphere/{dycore => diffusion}/tests/test_temporary_fields_for_turbulence_diagnostics.py (95%) rename model/atmosphere/{dycore => diffusion}/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py (97%) rename model/atmosphere/{dycore => diffusion}/tests/test_update_theta_and_exner.py (96%) create mode 100644 model/common/src/icon4py/model/common/decomposition/__init__.py create mode 100644 model/common/src/icon4py/model/common/utils.py diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py index 8b5f20ab0..e52802094 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py @@ -27,12 +27,12 @@ ProcessProperties, get_processor_properties, ) -from icon4py.diffusion.diffusion import Diffusion, DiffusionParams -from icon4py.diffusion.diffusion_states import ( +from model.atmosphere.diffusion.src.icon4py.model.atmosphere.diffusion import Diffusion, DiffusionParams +from model.atmosphere.diffusion.src.icon4py.model.atmosphere.diffusion import ( DiffusionDiagnosticState, PrognosticState, ) -from icon4py.diffusion.diffusion_utils import _identity_c_k, _identity_e_k +from model.atmosphere.diffusion.src.icon4py.model.atmosphere.diffusion import _identity_c_k, _identity_e_k from icon4py.driver.icon_configuration import IconRunConfig, read_config from icon4py.driver.io_utils import ( SIMULATION_START_DATE, diff --git a/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py b/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py index 7e75ebf63..63a42a191 100644 --- a/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py +++ b/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py @@ -14,7 +14,7 @@ from dataclasses import dataclass from typing import Optional -from icon4py.diffusion.diffusion import DiffusionConfig +from model.atmosphere.diffusion.src.icon4py.model.atmosphere.diffusion import DiffusionConfig n_substeps_reduced = 2 diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/atm_dyn_iconam/src/icon4py/driver/io_utils.py index 416d9ff8a..3e5c49494 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/atm_dyn_iconam/src/icon4py/driver/io_utils.py @@ -19,7 +19,7 @@ from icon4py.model.common.decomposition.decomposed import DecompositionInfo from icon4py.model.common.decomposition.parallel_setup import ParallelLogger, ProcessProperties -from icon4py.diffusion.diffusion_states import ( +from model.atmosphere.diffusion.src.icon4py.model.atmosphere.diffusion import ( DiffusionDiagnosticState, DiffusionInterpolationState, DiffusionMetricState, diff --git a/model/atmosphere/diffusion/.bumpversion.cfg b/model/atmosphere/diffusion/.bumpversion.cfg index 1cad57c0b..ef8bea214 100644 --- a/model/atmosphere/diffusion/.bumpversion.cfg +++ b/model/atmosphere/diffusion/.bumpversion.cfg @@ -4,7 +4,7 @@ parse = (?P\d+)\.(?P\d+)(\.(?P\d+))? serialize = {major}.{minor}.{patch} -[bumpversion:file:src/icon4py/model/atmosphere/dycore/__init__.py] +[bumpversion:file:src/icon4py/model/atmosphere/diffusion/__init__.py] parse = \"(?P\d+)\.(?P\d+)(\.(?P\d+))?\" serialize = {major}.{minor}.{patch} diff --git a/model/atmosphere/diffusion/README.md b/model/atmosphere/diffusion/README.md index 714bea17d..51a415319 100644 --- a/model/atmosphere/diffusion/README.md +++ b/model/atmosphere/diffusion/README.md @@ -1,8 +1,8 @@ -# icon4py-atmosphere-dycore +# icon4py-atmosphere-diffusion ## Description -Contains code ported from ICON `src/atm_dyn_iconam`, which is the dynamical core of the ICON model. +Python port of ICON diffusion module. ## Installation instructions diff --git a/model/atmosphere/diffusion/pyproject.toml b/model/atmosphere/diffusion/pyproject.toml index d4aee88f4..6d03d11e1 100644 --- a/model/atmosphere/diffusion/pyproject.toml +++ b/model/atmosphere/diffusion/pyproject.toml @@ -3,8 +3,8 @@ requires = ["setuptools>=61.0", "wheel>=0.40.0"] build-backend = "setuptools.build_meta" [project] -name = "icon4py-atmosphere-dycore" -description = "ICON dynamical core." +name = "icon4py-atmosphere-diffusion" +description = "ICON diffusion." readme = "README.md" requires-python = ">=3.10" license = {file = "LICENSE"} @@ -73,7 +73,7 @@ ignore_errors = true [tool.coverage.run] parallel = true branch = true -source_pkgs = ['dycore'] +source_pkgs = ['diffusion'] [tool.isort] lexicographical = true @@ -110,9 +110,10 @@ show_error_codes = true [tool.pytest.ini_options] testpaths = 'tests' +markers = 'datatest: this test uses binary data' [tool.setuptools.dynamic] -version = {attr = 'icon4py.model.atmosphere.dycore.__init__.__version__'} +version = {attr = 'icon4py.model.atmosphere.diffusion.__init__.__version__'} [tool.setuptools.package-data] -'icon4py.model.atmosphere.dycore' = ['py.typed'] +'icon4py.model.atmosphere.diffusion' = ['py.typed'] diff --git a/atm_dyn_iconam/src/icon4py/diffusion/__init__.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/__init__.py similarity index 55% rename from atm_dyn_iconam/src/icon4py/diffusion/__init__.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/__init__.py index 15dfdb009..5c4011cc5 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/__init__.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/__init__.py @@ -10,3 +10,24 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later + +from typing import Final + +from packaging import version as pkg_version +__all__ = [ + "__author__", + "__copyright__", + "__license__", + "__version__", + "__version_info__", +] + + +__author__: Final = "ETH Zurich and individual contributors" +__copyright__: Final = "Copyright (c) 2014-2022 ETH Zurich" +__license__: Final = "GPL-3.0-or-later" + + +__version__: Final = "0.0.6" +__version_info__: Final = pkg_version.parse(__version__) + diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py similarity index 97% rename from atm_dyn_iconam/src/icon4py/diffusion/diffusion.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py index 66dff8c7c..95ba73c96 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -29,29 +29,30 @@ run_gtfn_imperative, ) -from icon4py.model.atmosphere.dycore.apply_diffusion_to_vn import apply_diffusion_to_vn -from icon4py.model.atmosphere.dycore.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_vn import apply_diffusion_to_vn +from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance, ) -from icon4py.model.atmosphere.dycore.calculate_diagnostic_quantities_for_turbulence import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_diagnostic_quantities_for_turbulence import ( calculate_diagnostic_quantities_for_turbulence, ) -from icon4py.model.atmosphere.dycore.calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools import ( calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools, ) -from icon4py.model.atmosphere.dycore.calculate_nabla2_and_smag_coefficients_for_vn import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_and_smag_coefficients_for_vn import ( calculate_nabla2_and_smag_coefficients_for_vn, ) -from icon4py.model.atmosphere.dycore.calculate_nabla2_for_theta import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_theta import ( calculate_nabla2_for_theta, ) +# TODO (magdalena) should go to common.math from icon4py.model.atmosphere.dycore.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) -from icon4py.model.atmosphere.dycore.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( +from icon4py.model.atmosphere.diffusion.stencils.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( truly_horizontal_diffusion_nabla_of_theta_over_steep_points, ) -from icon4py.model.atmosphere.dycore.update_theta_and_exner import update_theta_and_exner +from icon4py.model.atmosphere.diffusion.stencils.update_theta_and_exner import update_theta_and_exner from icon4py.model.common.constants import ( CPD, DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO, @@ -59,13 +60,13 @@ ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim from icon4py.model.common.decomposition.decomposed import ExchangeRuntime, SingleNode -from icon4py.diffusion.diffusion_states import ( +from icon4py.model.atmosphere.diffusion.diffusion_states import ( DiffusionDiagnosticState, DiffusionInterpolationState, DiffusionMetricState, PrognosticState, ) -from icon4py.diffusion.diffusion_utils import ( +from icon4py.model.atmosphere.diffusion.diffusion_utils import ( copy_field, init_diffusion_local_fields_for_regular_timestep, init_nabla2_factor_in_upper_damping_zone, diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_states.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py similarity index 100% rename from atm_dyn_iconam/src/icon4py/diffusion/diffusion_states.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py diff --git a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_utils.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py similarity index 98% rename from atm_dyn_iconam/src/icon4py/diffusion/diffusion_utils.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py index c390048c5..bec8a87ae 100644 --- a/atm_dyn_iconam/src/icon4py/diffusion/diffusion_utils.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py @@ -244,9 +244,3 @@ def init_nabla2_factor_in_upper_damping_zone( return np_as_located_field(KDim)(buffer) -def builder(func): - def wrapper(self, *args, **kwargs): - func(self, *args, **kwargs) - return self - - return wrapper diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/__init__.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_theta_and_exner.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py similarity index 89% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_theta_and_exner.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py index a0e864646..bceeecf44 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_theta_and_exner.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py @@ -14,16 +14,16 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32 -from icon4py.model.atmosphere.dycore.calculate_nabla2_for_z import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_z import ( _calculate_nabla2_for_z, ) -from icon4py.model.atmosphere.dycore.calculate_nabla2_of_theta import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_of_theta import ( _calculate_nabla2_of_theta, ) -from icon4py.model.atmosphere.dycore.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( +from icon4py.model.atmosphere.diffusion.stencils.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( _truly_horizontal_diffusion_nabla_of_theta_over_steep_points, ) -from icon4py.model.atmosphere.dycore.update_theta_and_exner import ( +from icon4py.model.atmosphere.diffusion.stencils.update_theta_and_exner import ( _update_theta_and_exner, ) from icon4py.model.common.dimension import CECDim, CEDim, CellDim, EdgeDim, KDim diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py similarity index 91% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py index 8da27509f..88b74a92c 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py @@ -14,16 +14,16 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32, where -from icon4py.model.atmosphere.dycore.apply_nabla2_and_nabla4_global_to_vn import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_global_to_vn import ( _apply_nabla2_and_nabla4_global_to_vn, ) -from icon4py.model.atmosphere.dycore.apply_nabla2_and_nabla4_to_vn import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_to_vn import ( _apply_nabla2_and_nabla4_to_vn, ) -from icon4py.model.atmosphere.dycore.apply_nabla2_to_vn_in_lateral_boundary import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_vn_in_lateral_boundary import ( _apply_nabla2_to_vn_in_lateral_boundary, ) -from icon4py.model.atmosphere.dycore.calculate_nabla4 import _calculate_nabla4 +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import _calculate_nabla4 from icon4py.model.common.dimension import ECVDim, EdgeDim, KDim, VertexDim diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py similarity index 90% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py index 2887ed620..dd4f0a264 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py @@ -14,14 +14,14 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where -from icon4py.model.atmosphere.dycore.apply_nabla2_to_w import _apply_nabla2_to_w -from icon4py.model.atmosphere.dycore.apply_nabla2_to_w_in_upper_damping_layer import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w import _apply_nabla2_to_w +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w_in_upper_damping_layer import ( _apply_nabla2_to_w_in_upper_damping_layer, ) -from icon4py.model.atmosphere.dycore.calculate_horizontal_gradients_for_turbulence import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_horizontal_gradients_for_turbulence import ( _calculate_horizontal_gradients_for_turbulence, ) -from icon4py.model.atmosphere.dycore.calculate_nabla2_for_w import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_w import ( _calculate_nabla2_for_w, ) from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_global_to_vn.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_to_vn.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_and_nabla4_to_vn.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_to_vn.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_vn_in_lateral_boundary.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w_in_upper_damping_layer.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_nabla2_to_w_in_upper_damping_layer.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w_in_upper_damping_layer.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostic_quantities_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostic_quantities_for_turbulence.py similarity index 92% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostic_quantities_for_turbulence.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostic_quantities_for_turbulence.py index 8545dc63a..6b497c45a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostic_quantities_for_turbulence.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostic_quantities_for_turbulence.py @@ -14,10 +14,10 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32 -from icon4py.model.atmosphere.dycore.calculate_diagnostics_for_turbulence import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_diagnostics_for_turbulence import ( _calculate_diagnostics_for_turbulence, ) -from icon4py.model.atmosphere.dycore.temporary_fields_for_turbulence_diagnostics import ( +from icon4py.model.atmosphere.diffusion.stencils.temporary_fields_for_turbulence_diagnostics import ( _temporary_fields_for_turbulence_diagnostics, ) from icon4py.model.common.dimension import CEDim, CellDim, EdgeDim, KDim diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostics_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_diagnostics_for_turbulence.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py similarity index 90% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py index 3b43c4e18..6a997cd21 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py @@ -14,10 +14,10 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32 -from icon4py.model.atmosphere.dycore.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( +from icon4py.model.atmosphere.diffusion.stencils.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( _enhance_diffusion_coefficient_for_grid_point_cold_pools, ) -from icon4py.model.atmosphere.dycore.temporary_field_for_grid_point_cold_pools_enhancement import ( +from icon4py.model.atmosphere.diffusion.stencils.temporary_field_for_grid_point_cold_pools_enhancement import ( _temporary_field_for_grid_point_cold_pools_enhancement, ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_horizontal_gradients_for_turbulence.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_horizontal_gradients_for_turbulence.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_horizontal_gradients_for_turbulence.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_and_smag_coefficients_for_vn.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_and_smag_coefficients_for_vn.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_and_smag_coefficients_for_vn.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_theta.py similarity index 91% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_theta.py index d9e16fd99..d72cecd1d 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_theta.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_theta.py @@ -14,10 +14,10 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32 -from icon4py.model.atmosphere.dycore.calculate_nabla2_for_z import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_z import ( _calculate_nabla2_for_z, ) -from icon4py.model.atmosphere.dycore.calculate_nabla2_of_theta import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_of_theta import ( _calculate_nabla2_of_theta, ) from icon4py.model.common.dimension import CEDim, CellDim, EdgeDim, KDim diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_w.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_w.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_w.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_w.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_z.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_z.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_for_z.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_z.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_of_theta.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla2_of_theta.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla4.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/calculate_nabla4.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/enhance_diffusion_coefficient_for_grid_point_cold_pools.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_field_for_grid_point_cold_pools_enhancement.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_field_for_grid_point_cold_pools_enhancement.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_field_for_grid_point_cold_pools_enhancement.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_field_for_grid_point_cold_pools_enhancement.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_fields_for_turbulence_diagnostics.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/temporary_fields_for_turbulence_diagnostics.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/update_theta_and_exner.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_and_exner.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/update_theta_and_exner.py diff --git a/model/atmosphere/diffusion/tests/__init__.py b/model/atmosphere/diffusion/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/model/atmosphere/diffusion/tests/conftest.py b/model/atmosphere/diffusion/tests/conftest.py new file mode 100644 index 000000000..44f095ce8 --- /dev/null +++ b/model/atmosphere/diffusion/tests/conftest.py @@ -0,0 +1,66 @@ +import pytest +from icon4py.model.atmosphere.diffusion.diffusion import DiffusionConfig, DiffusionType + +from icon4py.model.common.test_utils.fixtures import ( # noqa F401 + interpolation_savepoint, + metrics_savepoint, + damping_height, + data_provider, + datapath, + grid_savepoint, + icon_grid, + setup_icon_data, + linit, + ndyn_substeps, + step_date_exit, + step_date_init, + mesh, + backend +) +@pytest.fixture +def r04b09_diffusion_config(ndyn_substeps) -> DiffusionConfig: + """ + Create DiffusionConfig matching MCH_CH_r04b09_dsl. + + Set values to the ones used in the MCH_CH_r04b09_dsl experiment where they differ + from the default. + """ + return DiffusionConfig( + diffusion_type=DiffusionType.SMAGORINSKY_4TH_ORDER, + hdiff_w=True, + hdiff_vn=True, + type_t_diffu=2, + type_vn_diffu=1, + hdiff_efdt_ratio=24.0, + hdiff_w_efdt_ratio=15.0, + smagorinski_scaling_factor=0.025, + zdiffu_t=True, + velocity_boundary_diffusion_denom=150.0, + max_nudging_coeff=0.075, + n_substeps=ndyn_substeps, + ) + + +@pytest.fixture +def diffusion_savepoint_init(data_provider, linit, step_date_init): # noqa F811 + """ + Load data from ICON savepoint at start of diffusion module. + + date of the timestamp to be selected can be set seperately by overriding the 'step_date_init' + fixture, passing 'step_date_init=' + + linit flag can be set by overriding the 'linit' fixture + """ + return data_provider.from_savepoint_diffusion_init(linit=linit, date=step_date_init) + + +@pytest.fixture +def diffusion_savepoint_exit(data_provider, linit, step_date_exit): # noqa F811 + """ + Load data from ICON savepoint at exist of diffusion module. + + date of the timestamp to be selected can be set seperately by overriding the 'step_data' + fixture, passing 'step_data=' + """ + sp = data_provider.from_savepoint_diffusion_exit(linit=linit, date=step_date_exit) + return sp diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/diffusion/tests/test_apply_nabla2_and_nabla4_global_to_vn.py similarity index 95% rename from model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py rename to model/atmosphere/diffusion/tests/test_apply_nabla2_and_nabla4_global_to_vn.py index d9d058adf..02e618fcd 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_global_to_vn.py +++ b/model/atmosphere/diffusion/tests/test_apply_nabla2_and_nabla4_global_to_vn.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.apply_nabla2_and_nabla4_global_to_vn import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_global_to_vn import ( apply_nabla2_and_nabla4_global_to_vn, ) from icon4py.model.common.dimension import EdgeDim, KDim diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_to_vn.py b/model/atmosphere/diffusion/tests/test_apply_nabla2_and_nabla4_to_vn.py similarity index 93% rename from model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_to_vn.py rename to model/atmosphere/diffusion/tests/test_apply_nabla2_and_nabla4_to_vn.py index 6bf222049..a73086bf0 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_and_nabla4_to_vn.py +++ b/model/atmosphere/diffusion/tests/test_apply_nabla2_and_nabla4_to_vn.py @@ -14,9 +14,12 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.apply_nabla2_and_nabla4_to_vn import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_to_vn import ( apply_nabla2_and_nabla4_to_vn, ) +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_global_to_vn import ( + apply_nabla2_and_nabla4_global_to_vn, +) from icon4py.model.common.dimension import EdgeDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field @@ -71,6 +74,7 @@ def reference( class TestApplyNabla2AndNabla4ToVnGlobalMode(StencilTest): PROGRAM = apply_nabla2_and_nabla4_global_to_vn + OUTPUTS = ("vn",) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/diffusion/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py similarity index 94% rename from model/atmosphere/dycore/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py rename to model/atmosphere/diffusion/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py index 792882432..ff44a1f67 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py +++ b/model/atmosphere/diffusion/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.apply_nabla2_to_vn_in_lateral_boundary import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_vn_in_lateral_boundary import ( apply_nabla2_to_vn_in_lateral_boundary, ) from icon4py.model.common.dimension import EdgeDim, KDim diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py b/model/atmosphere/diffusion/tests/test_apply_nabla2_to_w.py similarity index 95% rename from model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py rename to model/atmosphere/diffusion/tests/test_apply_nabla2_to_w.py index c02c61a50..147fd5c59 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w.py +++ b/model/atmosphere/diffusion/tests/test_apply_nabla2_to_w.py @@ -15,7 +15,7 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.apply_nabla2_to_w import apply_nabla2_to_w +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w import apply_nabla2_to_w from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field diff --git a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py b/model/atmosphere/diffusion/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py similarity index 94% rename from model/atmosphere/dycore/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py rename to model/atmosphere/diffusion/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py index 7517e26d7..6608b9008 100644 --- a/model/atmosphere/dycore/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py +++ b/model/atmosphere/diffusion/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py @@ -15,7 +15,7 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.apply_nabla2_to_w_in_upper_damping_layer import ( +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w_in_upper_damping_layer import ( apply_nabla2_to_w_in_upper_damping_layer, ) from icon4py.model.common.dimension import CellDim, KDim diff --git a/model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py b/model/atmosphere/diffusion/tests/test_calculate_diagnostics_for_turbulence.py similarity index 95% rename from model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py rename to model/atmosphere/diffusion/tests/test_calculate_diagnostics_for_turbulence.py index c60af2834..5e4673053 100644 --- a/model/atmosphere/dycore/tests/test_calculate_diagnostics_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/test_calculate_diagnostics_for_turbulence.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.calculate_diagnostics_for_turbulence import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_diagnostics_for_turbulence import ( calculate_diagnostics_for_turbulence, ) from icon4py.model.common.dimension import CellDim, KDim diff --git a/model/atmosphere/dycore/tests/test_calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/test_calculate_horizontal_gradients_for_turbulence.py similarity index 95% rename from model/atmosphere/dycore/tests/test_calculate_horizontal_gradients_for_turbulence.py rename to model/atmosphere/diffusion/tests/test_calculate_horizontal_gradients_for_turbulence.py index c5b0fa1ad..d8f18f442 100644 --- a/model/atmosphere/dycore/tests/test_calculate_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/test_calculate_horizontal_gradients_for_turbulence.py @@ -15,7 +15,7 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.calculate_horizontal_gradients_for_turbulence import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_horizontal_gradients_for_turbulence import ( calculate_horizontal_gradients_for_turbulence, ) from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py similarity index 98% rename from model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py rename to model/atmosphere/diffusion/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index 57ad274ce..dc63e0377 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/diffusion/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -13,7 +13,7 @@ import numpy as np -from icon4py.model.atmosphere.dycore.calculate_nabla2_and_smag_coefficients_for_vn import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_and_smag_coefficients_for_vn import ( calculate_nabla2_and_smag_coefficients_for_vn, ) from icon4py.model.common.dimension import ( diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_for_w.py b/model/atmosphere/diffusion/tests/test_calculate_nabla2_for_w.py similarity index 95% rename from model/atmosphere/dycore/tests/test_calculate_nabla2_for_w.py rename to model/atmosphere/diffusion/tests/test_calculate_nabla2_for_w.py index 695584d1d..7e7e0fa21 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla2_for_w.py +++ b/model/atmosphere/diffusion/tests/test_calculate_nabla2_for_w.py @@ -15,7 +15,7 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.calculate_nabla2_for_w import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_w import ( calculate_nabla2_for_w, ) from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_for_z.py b/model/atmosphere/diffusion/tests/test_calculate_nabla2_for_z.py similarity index 95% rename from model/atmosphere/dycore/tests/test_calculate_nabla2_for_z.py rename to model/atmosphere/diffusion/tests/test_calculate_nabla2_for_z.py index 4b575c47d..a82916ce3 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla2_for_z.py +++ b/model/atmosphere/diffusion/tests/test_calculate_nabla2_for_z.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.calculate_nabla2_for_z import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_z import ( calculate_nabla2_for_z, ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py b/model/atmosphere/diffusion/tests/test_calculate_nabla2_of_theta.py similarity index 95% rename from model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py rename to model/atmosphere/diffusion/tests/test_calculate_nabla2_of_theta.py index 0d6cb63d3..79cbb2ba5 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla2_of_theta.py +++ b/model/atmosphere/diffusion/tests/test_calculate_nabla2_of_theta.py @@ -14,7 +14,7 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.model.atmosphere.dycore.calculate_nabla2_of_theta import ( +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_of_theta import ( calculate_nabla2_of_theta, ) from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim diff --git a/model/atmosphere/dycore/tests/test_calculate_nabla4.py b/model/atmosphere/diffusion/tests/test_calculate_nabla4.py similarity index 97% rename from model/atmosphere/dycore/tests/test_calculate_nabla4.py rename to model/atmosphere/diffusion/tests/test_calculate_nabla4.py index c2a13b8ef..5bc20fe4a 100644 --- a/model/atmosphere/dycore/tests/test_calculate_nabla4.py +++ b/model/atmosphere/diffusion/tests/test_calculate_nabla4.py @@ -14,7 +14,7 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.model.atmosphere.dycore.calculate_nabla4 import calculate_nabla4 +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import calculate_nabla4 from icon4py.model.common.dimension import ( E2C2VDim, ECVDim, diff --git a/model/atmosphere/dycore/tests/test_diffusion.py b/model/atmosphere/diffusion/tests/test_diffusion.py similarity index 97% rename from model/atmosphere/dycore/tests/test_diffusion.py rename to model/atmosphere/diffusion/tests/test_diffusion.py index 43ebd8b4e..ccbaea1b9 100644 --- a/model/atmosphere/dycore/tests/test_diffusion.py +++ b/model/atmosphere/diffusion/tests/test_diffusion.py @@ -14,16 +14,16 @@ import numpy as np import pytest -from atm_dyn_iconam.tests.test_utils.serialbox_utils import ( +from icon4py.model.common.test_utils.serialbox_utils import ( IconDiffusionExitSavepoint, IconDiffusionInitSavepoint, ) -from icon4py.diffusion.diffusion import Diffusion, DiffusionParams -from icon4py.diffusion.diffusion_states import ( +from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.model.atmosphere.diffusion.diffusion_states import ( DiffusionDiagnosticState, PrognosticState, ) -from icon4py.diffusion.diffusion_utils import scale_k +from icon4py.model.atmosphere.diffusion.diffusion_utils import scale_k from icon4py.model.common.grid.horizontal import CellParams, EdgeParams from icon4py.model.common.grid.vertical import VerticalModelParams @@ -181,7 +181,7 @@ def _verify_init_values_against_savepoint( assert savepoint.diff_multfac_w() == diffusion.diff_multfac_w # this is done in diffusion.run(...) because it depends on the dtime - scale_k( + scale_k.with_backend(diffusion.compiled_backend)( diffusion.enh_smag_fac, dtime, diffusion.diff_multfac_smag, offset_provider={} ) assert np.allclose(savepoint.diff_multfac_smag(), diffusion.diff_multfac_smag) diff --git a/model/atmosphere/dycore/tests/test_diffusion_states.py b/model/atmosphere/diffusion/tests/test_diffusion_states.py similarity index 100% rename from model/atmosphere/dycore/tests/test_diffusion_states.py rename to model/atmosphere/diffusion/tests/test_diffusion_states.py diff --git a/model/atmosphere/dycore/tests/test_diffusion_utils.py b/model/atmosphere/diffusion/tests/test_diffusion_utils.py similarity index 94% rename from model/atmosphere/dycore/tests/test_diffusion_utils.py rename to model/atmosphere/diffusion/tests/test_diffusion_utils.py index 53142b8b6..9d36cfddb 100644 --- a/model/atmosphere/dycore/tests/test_diffusion_utils.py +++ b/model/atmosphere/diffusion/tests/test_diffusion_utils.py @@ -14,11 +14,11 @@ import numpy as np import pytest -from atm_dyn_iconam.tests.test_utils.helpers import random_field, zero_field -from atm_dyn_iconam.tests.test_utils.simple_mesh import SimpleMesh -from icon4py.common.dimension import KDim, VertexDim -from icon4py.diffusion.diffusion import DiffusionParams -from icon4py.diffusion.diffusion_utils import ( +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.dimension import KDim, VertexDim +from icon4py.model.atmosphere.diffusion.diffusion import DiffusionParams +from icon4py.model.atmosphere.diffusion.diffusion_utils import ( _en_smag_fac_for_zero_nshift, _setup_runtime_diff_multfac_vn, _setup_smag_limit, @@ -150,9 +150,8 @@ def test_set_zero_vertex_k(): @pytest.mark.datatest -@pytest.mark.parametrize("linit", [True]) def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( - diffusion_savepoint_init, r04b09_diffusion_config, icon_grid, ndyn_substeps + diffusion_savepoint_init, r04b09_diffusion_config, icon_grid ): savepoint = diffusion_savepoint_init config = r04b09_diffusion_config diff --git a/model/atmosphere/dycore/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py similarity index 92% rename from model/atmosphere/dycore/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py rename to model/atmosphere/diffusion/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py index 50e92d4f5..554142226 100644 --- a/model/atmosphere/dycore/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py +++ b/model/atmosphere/diffusion/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( +from icon4py.model.atmosphere.diffusion.stencils.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( enhance_diffusion_coefficient_for_grid_point_cold_pools, ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim diff --git a/model/atmosphere/dycore/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py b/model/atmosphere/diffusion/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py similarity index 94% rename from model/atmosphere/dycore/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py rename to model/atmosphere/diffusion/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py index bae0e5f16..2ab67d0d3 100644 --- a/model/atmosphere/dycore/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py +++ b/model/atmosphere/diffusion/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.temporary_field_for_grid_point_cold_pools_enhancement import ( +from icon4py.model.atmosphere.diffusion.stencils.temporary_field_for_grid_point_cold_pools_enhancement import ( temporary_field_for_grid_point_cold_pools_enhancement, ) from icon4py.model.common.dimension import CellDim, KDim diff --git a/model/atmosphere/dycore/tests/test_temporary_fields_for_turbulence_diagnostics.py b/model/atmosphere/diffusion/tests/test_temporary_fields_for_turbulence_diagnostics.py similarity index 95% rename from model/atmosphere/dycore/tests/test_temporary_fields_for_turbulence_diagnostics.py rename to model/atmosphere/diffusion/tests/test_temporary_fields_for_turbulence_diagnostics.py index d06e3d0fd..521214d70 100644 --- a/model/atmosphere/dycore/tests/test_temporary_fields_for_turbulence_diagnostics.py +++ b/model/atmosphere/diffusion/tests/test_temporary_fields_for_turbulence_diagnostics.py @@ -14,10 +14,10 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.temporary_fields_for_turbulence_diagnostics import ( +from icon4py.model.atmosphere.diffusion.stencils.temporary_fields_for_turbulence_diagnostics import ( temporary_fields_for_turbulence_diagnostics, ) -from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim, CEDim from icon4py.model.common.test_utils.helpers import ( StencilTest, random_field, diff --git a/model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py similarity index 97% rename from model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py rename to model/atmosphere/diffusion/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index 24f9f05c1..3688d917e 100644 --- a/model/atmosphere/dycore/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/diffusion/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.model.atmosphere.dycore.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( +from icon4py.model.atmosphere.diffusion.stencils.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( truly_horizontal_diffusion_nabla_of_theta_over_steep_points, ) from icon4py.model.common.dimension import C2E2CDim, CECDim, CellDim, KDim diff --git a/model/atmosphere/dycore/tests/test_update_theta_and_exner.py b/model/atmosphere/diffusion/tests/test_update_theta_and_exner.py similarity index 96% rename from model/atmosphere/dycore/tests/test_update_theta_and_exner.py rename to model/atmosphere/diffusion/tests/test_update_theta_and_exner.py index a9081c32c..871b32a89 100644 --- a/model/atmosphere/dycore/tests/test_update_theta_and_exner.py +++ b/model/atmosphere/diffusion/tests/test_update_theta_and_exner.py @@ -15,7 +15,7 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.update_theta_and_exner import ( +from icon4py.model.atmosphere.diffusion.stencils.update_theta_and_exner import ( update_theta_and_exner, ) from icon4py.model.common.dimension import CellDim, KDim diff --git a/model/atmosphere/dycore/tests/conftest.py b/model/atmosphere/dycore/tests/conftest.py index 03ca4a3f5..5133429bd 100644 --- a/model/atmosphere/dycore/tests/conftest.py +++ b/model/atmosphere/dycore/tests/conftest.py @@ -11,13 +11,9 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import pytest -from gt4py.next.program_processors.runners.roundtrip import executor -from icon4py.model.common.test_utils.simple_mesh import SimpleMesh -from atm_dyn_iconam.tests.test_utils.simple_mesh import SimpleMesh -from icon4py.diffusion.diffusion import DiffusionConfig, DiffusionType -from .test_utils.fixtures import ( # noqa F401 +from icon4py.model.common.test_utils.fixtures import ( # noqa F401 damping_height, data_provider, datapath, @@ -26,123 +22,11 @@ icon_grid, r04b09_dsl_gridfile, setup_icon_data, + linit,step_date_exit, step_date_init,mesh, backend ) -@pytest.fixture -def ndyn_substeps(): - """ - Return number of dynamical substeps. - Serialized data uses a reduced number (2 instead of the default 5) in order to reduce the amount - of data generated. - """ - return 2 -@pytest.fixture -def linit(): - """ - Set the 'linit' flag for the ICON diffusion data savepoint. - Defaults to False - """ - return False - - -@pytest.fixture -def step_date_init(): - """ - Set the step date for the loaded ICON time stamp at start of module. - - Defaults to 2021-06-20T12:00:10.000' - """ - return "2021-06-20T12:00:10.000" - - -@pytest.fixture -def step_date_exit(): - """ - Set the step date for the loaded ICON time stamp at the end of module. - - Defaults to 2021-06-20T12:00:10.000' - """ - return "2021-06-20T12:00:10.000" - - -@pytest.fixture -def diffusion_savepoint_init(data_provider, linit, step_date_init): # noqa F811 - """ - Load data from ICON savepoint at start of diffusion module. - - date of the timestamp to be selected can be set seperately by overriding the 'step_date_init' - fixture, passing 'step_date_init=' - - linit flag can be set by overriding the 'linit' fixture - """ - return data_provider.from_savepoint_diffusion_init(linit=linit, date=step_date_init) - - -@pytest.fixture -def diffusion_savepoint_exit(data_provider, linit, step_date_exit): # noqa F811 - """ - Load data from ICON savepoint at exist of diffusion module. - - date of the timestamp to be selected can be set seperately by overriding the 'step_data' - fixture, passing 'step_data=' - """ - sp = data_provider.from_savepoint_diffusion_exit(linit=linit, date=step_date_exit) - return sp - - -@pytest.fixture -def interpolation_savepoint(data_provider): # noqa F811 - """Load data from ICON interplation state savepoint.""" - return data_provider.from_interpolation_savepoint() - - -@pytest.fixture -def metrics_savepoint(data_provider): # noqa F811 - """Load data from ICON mestric state savepoint.""" - return data_provider.from_metrics_savepoint() - - -@pytest.fixture -def r04b09_diffusion_config(ndyn_substeps) -> DiffusionConfig: - """ - Create DiffusionConfig matching MCH_CH_r04b09_dsl. - - Set values to the ones used in the MCH_CH_r04b09_dsl experiment where they differ - from the default. - """ - return DiffusionConfig( - diffusion_type=DiffusionType.SMAGORINSKY_4TH_ORDER, - hdiff_w=True, - hdiff_vn=True, - type_t_diffu=2, - type_vn_diffu=1, - hdiff_efdt_ratio=24.0, - hdiff_w_efdt_ratio=15.0, - smagorinski_scaling_factor=0.025, - zdiffu_t=True, - velocity_boundary_diffusion_denom=150.0, - max_nudging_coeff=0.075, - n_substeps=ndyn_substeps, - ) - - -BACKENDS = {"embedded": executor} -MESHES = {"simple_mesh": SimpleMesh()} - - -@pytest.fixture( - ids=MESHES.keys(), - params=MESHES.values(), -) -def mesh(request): - return request.param - - -@pytest.fixture(ids=BACKENDS.keys(), params=BACKENDS.values()) -def backend(request): - return request.param diff --git a/model/atmosphere/dycore/tests/test_interpolation_fields.py b/model/atmosphere/dycore/tests/test_interpolation_fields.py index a0ac48263..37c0b3a31 100644 --- a/model/atmosphere/dycore/tests/test_interpolation_fields.py +++ b/model/atmosphere/dycore/tests/test_interpolation_fields.py @@ -26,7 +26,7 @@ import numpy as np import pytest -from icon4py.common.dimension import EdgeDim +from icon4py.model.common.dimension import EdgeDim from icon4py.field_management.interpolation_fields import compute_c_lin_e from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex diff --git a/model/common/src/icon4py/model/common/decomposition/__init__.py b/model/common/src/icon4py/model/common/decomposition/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/model/common/src/icon4py/model/common/decomposition/decomposed.py b/model/common/src/icon4py/model/common/decomposition/decomposed.py index 2a2749819..d19fdd915 100644 --- a/model/common/src/icon4py/model/common/decomposition/decomposed.py +++ b/model/common/src/icon4py/model/common/decomposition/decomposed.py @@ -23,7 +23,7 @@ from icon4py.model.common.dimension import CellDim, DimensionKind, EdgeDim, VertexDim from icon4py.model.common.decomposition.parallel_setup import ProcessProperties -from icon4py.diffusion.diffusion_utils import builder +from icon4py.model.common.utils import builder log = logging.getLogger(__name__) diff --git a/model/common/src/icon4py/model/common/grid/icon_grid.py b/model/common/src/icon4py/model/common/grid/icon_grid.py index ba860e9bd..4153486cb 100644 --- a/model/common/src/icon4py/model/common/grid/icon_grid.py +++ b/model/common/src/icon4py/model/common/grid/icon_grid.py @@ -30,6 +30,7 @@ ) from icon4py.model.common.grid.horizontal import HorizontalGridSize from icon4py.model.common.grid.vertical import VerticalGridSize +from icon4py.model.common.utils import builder @dataclass( @@ -58,13 +59,6 @@ def num_cells(self): return self.horizontal_config.num_cells -def builder(func): - def wrapper(self, *args, **kwargs): - func(self, *args, **kwargs) - return self - - return wrapper - class IconGrid: def __init__(self): diff --git a/model/common/src/icon4py/model/common/test_utils/fixtures.py b/model/common/src/icon4py/model/common/test_utils/fixtures.py index 8f097eed4..dc3b54d13 100644 --- a/model/common/src/icon4py/model/common/test_utils/fixtures.py +++ b/model/common/src/icon4py/model/common/test_utils/fixtures.py @@ -14,9 +14,13 @@ from pathlib import Path import pytest +from gt4py.next.program_processors.runners.roundtrip import executor from .data_handling import download_and_extract from .serialbox_utils import IconSerialDataProvider +from .simple_mesh import SimpleMesh + + test_utils = Path(__file__).parent model = test_utils.parent.parent @@ -90,6 +94,48 @@ def damping_height(): return 12500 + +@pytest.fixture +def ndyn_substeps(): + """ + Return number of dynamical substeps. + + Serialized data uses a reduced number (2 instead of the default 5) in order to reduce the amount + of data generated. + """ + return 2 + + +@pytest.fixture +def linit(): + """ + Set the 'linit' flag for the ICON diffusion data savepoint. + + Defaults to False + """ + return False + + +@pytest.fixture +def step_date_init(): + """ + Set the step date for the loaded ICON time stamp at start of module. + + Defaults to 2021-06-20T12:00:10.000' + """ + return "2021-06-20T12:00:10.000" + + +@pytest.fixture +def step_date_exit(): + """ + Set the step date for the loaded ICON time stamp at the end of module. + + Defaults to 2021-06-20T12:00:10.000' + """ + return "2021-06-20T12:00:10.000" + + @pytest.fixture(scope="session") def get_grid_files(): """ @@ -105,6 +151,35 @@ def get_grid_files(): ) +@pytest.fixture +def interpolation_savepoint(data_provider): # noqa F811 + """Load data from ICON interplation state savepoint.""" + return data_provider.from_interpolation_savepoint() + + +@pytest.fixture +def metrics_savepoint(data_provider): # noqa F811 + """Load data from ICON mestric state savepoint.""" + return data_provider.from_metrics_savepoint() + + + @pytest.fixture() def r04b09_dsl_gridfile(get_grid_files): return r04b09_dsl_grid_path.joinpath("grid.nc") + +BACKENDS = {"embedded": executor} +MESHES = {"simple_mesh": SimpleMesh()} + + +@pytest.fixture( + ids=MESHES.keys(), + params=MESHES.values(), +) +def mesh(request): + return request.param + + +@pytest.fixture(ids=BACKENDS.keys(), params=BACKENDS.values()) +def backend(request): + return request.param 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 89804e9e9..e642168dc 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 @@ -35,8 +35,8 @@ VertexDim, ) from icon4py.model.common.decomposition.decomposed import DecompositionInfo -from icon4py.diffusion.diffusion import VectorTuple -from icon4py.diffusion.diffusion_states import ( +from icon4py.model.atmosphere.diffusion.diffusion import VectorTuple +from icon4py.model.atmosphere.diffusion.diffusion_states import ( DiffusionDiagnosticState, DiffusionInterpolationState, DiffusionMetricState, @@ -45,7 +45,7 @@ from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalGridSize from icon4py.model.common.grid.icon_grid import GridConfig, IconGrid, VerticalGridSize -from .helpers import as_1D_sparse_field +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field class IconSavepoint: diff --git a/model/common/src/icon4py/model/common/utils.py b/model/common/src/icon4py/model/common/utils.py new file mode 100644 index 000000000..6b24ee54e --- /dev/null +++ b/model/common/src/icon4py/model/common/utils.py @@ -0,0 +1,7 @@ +def builder(func): + """Use as decorator on builder functions.""" + def wrapper(self, *args, **kwargs): + func(self, *args, **kwargs) + return self + + return wrapper diff --git a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py index 59000f08f..54aef39ff 100644 --- a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py +++ b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py @@ -25,12 +25,12 @@ V2EDim, VertexDim, ) -from icon4py.diffusion.diffusion import ( +from model.atmosphere.diffusion.src.icon4py.model.atmosphere.diffusion import ( Diffusion, DiffusionConfig, DiffusionParams, ) -from icon4py.diffusion.diffusion_states import ( +from model.atmosphere.diffusion.src.icon4py.model.atmosphere.diffusion import ( DiffusionDiagnosticState, DiffusionInterpolationState, DiffusionMetricState, diff --git a/requirements-dev.txt b/requirements-dev.txt index 688a17c33..ff9021871 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -3,6 +3,7 @@ # icon4py model -e ./model/atmosphere/dycore -e ./model/common +-e ./model/atmosphere/diffusion # icon4pytools -e ./tools diff --git a/requirements.txt b/requirements.txt index 4b27ce74b..dfcf2e291 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,6 +2,7 @@ # icon4py model ./model/atmosphere/dycore +./model/atmosphere/diffusion ./model/common ./py2f # icon4pytools diff --git a/tools/tests/liskov/test_external.py b/tools/tests/liskov/test_external.py index 739c2ba21..55f2e5f3b 100644 --- a/tools/tests/liskov/test_external.py +++ b/tools/tests/liskov/test_external.py @@ -40,7 +40,7 @@ def test_stencil_collector_invalid_module(): def test_stencil_collector_invalid_member(): - from icon4py.model.atmosphere.dycore import apply_nabla2_to_w + from icon4py.model.atmosphere.diffusion.stencils import apply_nabla2_to_w module_path = Path(apply_nabla2_to_w.__file__) parents = module_path.parents[0] From cb6dffc359ff951bc818a14efcc130cb02468f1b Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 22 Aug 2023 09:29:42 +0200 Subject: [PATCH 230/263] fix icon4py/tools unit tests --- .../atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py | 1 - tools/src/icon4pytools/liskov/external/gt4py.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py index be7b358dc..bf5bb4c3c 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -41,7 +41,6 @@ def mo_intp_rbf_rbf_vec_interpol_vertex( vertical_start: int32, vertical_end: int32, ): - _mo_intp_rbf_rbf_vec_interpol_vertex(p_e_in, ptr_coeff_1, ptr_coeff_2, out=(p_u_out, p_v_out)) _mo_intp_rbf_rbf_vec_interpol_vertex( p_e_in, ptr_coeff_1, diff --git a/tools/src/icon4pytools/liskov/external/gt4py.py b/tools/src/icon4pytools/liskov/external/gt4py.py index 53b815173..b84c27e51 100644 --- a/tools/src/icon4pytools/liskov/external/gt4py.py +++ b/tools/src/icon4pytools/liskov/external/gt4py.py @@ -32,7 +32,7 @@ class UpdateFieldsWithGt4PyStencils(Step): - _STENCIL_PACKAGES = ["atmosphere.dycore"] + _STENCIL_PACKAGES = ["atmosphere.dycore", "atmosphere.diffusion.stencils"] def __init__(self, parsed: IntegrationCodeInterface): self.parsed = parsed From ca9217874db1691e1ddbd924fd311790f0670449 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 22 Aug 2023 13:53:00 +0200 Subject: [PATCH 231/263] move driver code to its own package --- .../mpi_tests/test_parallel_diffusion.py | 11 +- .../diffusion/tests/test_diffusion.py | 2 +- model/atmosphere/dycore/pyproject.toml | 3 +- .../common/tests/mpi_tests/test_decomposed.py | 19 +-- model/driver/.bumpversion.cfg | 10 ++ .../src/icon4py => model}/driver/README.md | 0 model/driver/pyproject.toml | 121 ++++++++++++++++++ model/driver/requirements-dev.txt | 4 + model/driver/requirements.txt | 5 + .../src/icon4py/model}/driver/__init__.py | 19 +++ .../icon4py/model}/driver/dycore_driver.py | 21 ++- .../model}/driver/icon_configuration.py | 3 +- .../src/icon4py/model}/driver/io_utils.py | 31 +---- model/driver/tests/conftest.py | 1 + .../tests/test_dycore_driver.py | 8 +- .../dycore => driver}/tests/test_io_utils.py | 7 +- model/requirements-dev.txt | 2 + model/requirements.txt | 2 + requirements-dev.txt | 4 +- requirements.txt | 3 +- 20 files changed, 203 insertions(+), 73 deletions(-) rename model/{common => atmosphere/diffusion}/tests/mpi_tests/test_parallel_diffusion.py (90%) create mode 100644 model/driver/.bumpversion.cfg rename {atm_dyn_iconam/src/icon4py => model}/driver/README.md (100%) create mode 100644 model/driver/pyproject.toml create mode 100644 model/driver/requirements-dev.txt create mode 100644 model/driver/requirements.txt rename {atm_dyn_iconam/src/icon4py => model/driver/src/icon4py/model}/driver/__init__.py (55%) rename {atm_dyn_iconam/src/icon4py => model/driver/src/icon4py/model}/driver/dycore_driver.py (93%) rename {atm_dyn_iconam/src/icon4py => model/driver/src/icon4py/model}/driver/icon_configuration.py (96%) rename {atm_dyn_iconam/src/icon4py => model/driver/src/icon4py/model}/driver/io_utils.py (89%) create mode 100644 model/driver/tests/conftest.py rename model/{atmosphere/dycore => driver}/tests/test_dycore_driver.py (87%) rename model/{atmosphere/dycore => driver}/tests/test_io_utils.py (96%) diff --git a/model/common/tests/mpi_tests/test_parallel_diffusion.py b/model/atmosphere/diffusion/tests/mpi_tests/test_parallel_diffusion.py similarity index 90% rename from model/common/tests/mpi_tests/test_parallel_diffusion.py rename to model/atmosphere/diffusion/tests/mpi_tests/test_parallel_diffusion.py index 61466aeb4..6a85e23b8 100644 --- a/model/common/tests/mpi_tests/test_parallel_diffusion.py +++ b/model/atmosphere/diffusion/tests/mpi_tests/test_parallel_diffusion.py @@ -14,24 +14,25 @@ import pytest -from .common import ( # noqa F401 +from icon4py.model.common.decomposition.parallel_setup import get_processor_properties +from icon4py.model.common.test_utils.parallel_fixtures import ( # noqa F401 data_path, download_data, props, ) -from icon4py.model.diffusion.test_diffusion import _verify_diffusion_fields +from icon4py.model.atmosphere.diffusion.tests.test_diffusion import _verify_diffusion_fields from icon4py.model.common.test_utils.serialbox_utils import IconSerialDataProvider from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.model.common.decomposition.decomposed import DecompositionInfo, create_exchange -from icon4py.diffusion.diffusion import Diffusion, DiffusionParams -from icon4py.driver.io_utils import ( +from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams +from model.driver.src.icon4py.model.driver.io_utils import ( read_decomp_info, read_geometry_fields, read_icon_grid, read_static_fields, ) - +props = get_processor_properties(with_mpi=True) @pytest.mark.mpi @pytest.mark.parametrize("ndyn_substeps", [2]) @pytest.mark.parametrize("linit", [True, False]) diff --git a/model/atmosphere/diffusion/tests/test_diffusion.py b/model/atmosphere/diffusion/tests/test_diffusion.py index ccbaea1b9..f9a8a3e96 100644 --- a/model/atmosphere/diffusion/tests/test_diffusion.py +++ b/model/atmosphere/diffusion/tests/test_diffusion.py @@ -181,7 +181,7 @@ def _verify_init_values_against_savepoint( assert savepoint.diff_multfac_w() == diffusion.diff_multfac_w # this is done in diffusion.run(...) because it depends on the dtime - scale_k.with_backend(diffusion.compiled_backend)( + scale_k( diffusion.enh_smag_fac, dtime, diffusion.diff_multfac_smag, offset_provider={} ) assert np.allclose(savepoint.diff_multfac_smag(), diffusion.diff_multfac_smag) diff --git a/model/atmosphere/dycore/pyproject.toml b/model/atmosphere/dycore/pyproject.toml index d4aee88f4..af76bc94e 100644 --- a/model/atmosphere/dycore/pyproject.toml +++ b/model/atmosphere/dycore/pyproject.toml @@ -28,7 +28,7 @@ classifiers = [ ] dependencies = [ "gt4py>=1.0.1", - "icon4py_common>=0.0.5", + "icon4py-common>=0.0.5", ] dynamic = ['version'] @@ -110,6 +110,7 @@ show_error_codes = true [tool.pytest.ini_options] testpaths = 'tests' +markers = 'datatest: this test uses binary data' [tool.setuptools.dynamic] version = {attr = 'icon4py.model.atmosphere.dycore.__init__.__version__'} diff --git a/model/common/tests/mpi_tests/test_decomposed.py b/model/common/tests/mpi_tests/test_decomposed.py index 97dcfbf8c..5ab11d4b8 100644 --- a/model/common/tests/mpi_tests/test_decomposed.py +++ b/model/common/tests/mpi_tests/test_decomposed.py @@ -14,7 +14,7 @@ import numpy as np import pytest -from .common import ( +from icon4py.model.common.test_utils.parallel_fixtures import ( data_path, download_data, props, @@ -28,8 +28,7 @@ create_exchange, ) from icon4py.model.common.decomposition.parallel_setup import ProcessProperties -from icon4py.driver.io_utils import ( - SerializationType, +from model.driver.src.icon4py.model.driver.io_utils import ( read_decomp_info, read_icon_grid, ) @@ -57,10 +56,9 @@ ), ) def test_decomposition_info_masked( - dim, owned, total, caplog, download_data # noqa F811 + dim, owned, total, caplog, download_data, decomposition_info # noqa F811 ): my_rank = props.rank - decomposition_info = read_decomp_info(data_path, props, SerializationType.SB) all_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.ALL) my_total = total[my_rank] my_owned = owned[my_rank] @@ -88,10 +86,9 @@ def test_decomposition_info_masked( ), ) def test_decomposition_info_local_index( - dim, owned, total, caplog, download_data # noqa F811 + dim, owned, total, caplog, download_data, decomposition_info # noqa F811 ): my_rank = props.rank - decomposition_info = read_decomp_info(data_path, props, SerializationType.SB) all_indices = decomposition_info.local_index(dim, DecompositionInfo.EntryType.ALL) my_total = total[my_rank] my_owned = owned[my_rank] @@ -149,12 +146,8 @@ def test_domain_descriptor_id_are_globally_unique(num): props.comm_size not in (1, 2, 4), reason="input files only available for 1 or 2 nodes", ) -def test_decomposition_info_matches_gridsize(caplog, download_data): # noqa F811 - decomposition_info = read_decomp_info( - data_path, - props, - SerializationType.SB, - ) +def test_decomposition_info_matches_gridsize(caplog, download_data, decomposition_info): # noqa F811 + icon_grid = read_icon_grid(data_path, props.rank) assert ( decomposition_info.global_index( diff --git a/model/driver/.bumpversion.cfg b/model/driver/.bumpversion.cfg new file mode 100644 index 000000000..ae24af96d --- /dev/null +++ b/model/driver/.bumpversion.cfg @@ -0,0 +1,10 @@ +[bumpversion] +current_version = 0.0.6 +parse = (?P\d+)\.(?P\d+)(\.(?P\d+))? +serialize = + {major}.{minor}.{patch} + +[bumpversion:file:src/icon4py/model/driver/__init__.py] +parse = \"(?P\d+)\.(?P\d+)(\.(?P\d+))?\" +serialize = + {major}.{minor}.{patch} diff --git a/atm_dyn_iconam/src/icon4py/driver/README.md b/model/driver/README.md similarity index 100% rename from atm_dyn_iconam/src/icon4py/driver/README.md rename to model/driver/README.md diff --git a/model/driver/pyproject.toml b/model/driver/pyproject.toml new file mode 100644 index 000000000..e9b7f3e98 --- /dev/null +++ b/model/driver/pyproject.toml @@ -0,0 +1,121 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel>=0.40.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "icon4py-driver" +description = "ICON model driver." +readme = "README.md" +requires-python = ">=3.10" +license = {file = "LICENSE"} +authors = [ + {email = "gridtools@cscs.ch"}, + {name = "ETH Zurich"} +] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)", + "Operating System :: POSIX", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: Implementation :: CPython", + "Topic :: Scientific/Engineering :: Atmospheric Science", + "Topic :: Scientific/Engineering :: Mathematics", + "Topic :: Scientific/Engineering :: Physics" +] +dependencies = [ + "gt4py>=1.0.1", + "icon4py-common>=0.0.5", + "icon4py-atmosphere-dycore>=0.0.5", + "icon4py-atmosphere-diffusion>=0.0.5", +] +dynamic = ['version'] + +[project.urls] +repository = "https://github.com/C2SM/icon4py" + +[tool.black] +exclude = ''' +/( + \.git + | \.hg + | \.mypy_cache + | \.tox + | \.venv + | _build + | buck-out + | build + | dist +)/ +''' +include = '\.pyi?$' +line-length = 100 +target-version = ['py310'] + +[tool.coverage] + +[tool.coverage.html] +directory = 'tests/_reports/coverage_html' + +[tool.coverage.paths] +source = ['src/icon4py/model/'] + +[tool.coverage.report] +exclude_lines = [ + 'raise AssertionError', # Don't complain if tests don't hit defensive assertion code + 'raise NotImplementedError', # Don't complain if tests don't hit defensive assertion code + 'if 0:', # Don't complain if non-runnable code isn't run + 'if __name__ == .__main__.:' # Don't complain if non-runnable code isn't run +] +ignore_errors = true + +[tool.coverage.run] +parallel = true +branch = true +source_pkgs = ['driver'] + +[tool.isort] +lexicographical = true +line_length = 100 # It should be the same as in `tool.black.line-length` above +lines_after_imports = 2 +profile = 'black' +sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'FIRSTPARTY', 'TESTS', 'LOCALFOLDER'] +skip_gitignore = true +skip_glob = ['*.venv/**', '_local/**'] +known_first_party = ['icon4py.model'] +known_third_party = ['gt4py'] +multi_line_output = 3 +use_parentheses = true +include_trailing_comma = true +force_grid_wrap = 0 + +[tool.mypy] +install_types = true +non_interactive = true +exclude = [ + '^tests/*.py', +] +disallow_incomplete_defs = true +disallow_untyped_defs = true +ignore_missing_imports = false +implicit_reexport = true +warn_unused_configs = true +warn_unused_ignores = true +warn_redundant_casts = true +show_column_numbers = true +show_error_codes = true + +[tool.pytest] + +[tool.pytest.ini_options] +testpaths = 'tests' +markers = 'datatest: this test uses binary data' + +[tool.setuptools.dynamic] +version = {attr = 'icon4py.model.driver.__init__.__version__'} + +[tool.setuptools.package-data] +'icon4py.model.driver' = ['py.typed'] diff --git a/model/driver/requirements-dev.txt b/model/driver/requirements-dev.txt new file mode 100644 index 000000000..a42befbba --- /dev/null +++ b/model/driver/requirements-dev.txt @@ -0,0 +1,4 @@ +-r ../../../base-requirements-dev.txt +-e ../common +-e ../atmosphere +-e . diff --git a/model/driver/requirements.txt b/model/driver/requirements.txt new file mode 100644 index 000000000..b92b8313c --- /dev/null +++ b/model/driver/requirements.txt @@ -0,0 +1,5 @@ +-r ../../../base-requirements.txt +../common +../atmosphere/dycore +../atmosphere/diffusion +. diff --git a/atm_dyn_iconam/src/icon4py/driver/__init__.py b/model/driver/src/icon4py/model/driver/__init__.py similarity index 55% rename from atm_dyn_iconam/src/icon4py/driver/__init__.py rename to model/driver/src/icon4py/model/driver/__init__.py index 15dfdb009..1dcd287bc 100644 --- a/atm_dyn_iconam/src/icon4py/driver/__init__.py +++ b/model/driver/src/icon4py/model/driver/__init__.py @@ -10,3 +10,22 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +from typing import Final + +from packaging import version as pkg_version +__all__ = [ + "__author__", + "__copyright__", + "__license__", + "__version__", + "__version_info__", +] + + +__author__: Final = "ETH Zurich and individual contributors" +__copyright__: Final = "Copyright (c) 2014-2022 ETH Zurich" +__license__: Final = "GPL-3.0-or-later" + + +__version__: Final = "0.0.6" +__version_info__: Final = pkg_version.parse(__version__) diff --git a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py b/model/driver/src/icon4py/model/driver/dycore_driver.py similarity index 93% rename from atm_dyn_iconam/src/icon4py/driver/dycore_driver.py rename to model/driver/src/icon4py/model/driver/dycore_driver.py index e52802094..fc6220ddd 100644 --- a/atm_dyn_iconam/src/icon4py/driver/dycore_driver.py +++ b/model/driver/src/icon4py/model/driver/dycore_driver.py @@ -21,23 +21,23 @@ from gt4py.next import program, Field from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn -from icon4py.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim from icon4py.model.common.decomposition.decomposed import create_exchange from icon4py.model.common.decomposition.parallel_setup import ( ProcessProperties, get_processor_properties, ) -from model.atmosphere.diffusion.src.icon4py.model.atmosphere.diffusion import Diffusion, DiffusionParams -from model.atmosphere.diffusion.src.icon4py.model.atmosphere.diffusion import ( +from icon4py.model.common.test_utils import serialbox_utils as sb +from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.model.atmosphere.diffusion.diffusion_states import ( DiffusionDiagnosticState, PrognosticState, ) -from model.atmosphere.diffusion.src.icon4py.model.atmosphere.diffusion import _identity_c_k, _identity_e_k -from icon4py.driver.icon_configuration import IconRunConfig, read_config -from icon4py.driver.io_utils import ( +from icon4py.model.atmosphere.diffusion.diffusion_utils import _identity_c_k, _identity_e_k +from icon4py.model.driver.icon_configuration import IconRunConfig, read_config +from icon4py.model.driver.io_utils import ( SIMULATION_START_DATE, configure_logging, - import_testutils, read_decomp_info, read_geometry_fields, read_icon_grid, @@ -45,11 +45,6 @@ read_static_fields, ) - -helpers = import_testutils() -from helpers import serialbox_utils as sb_utils # noqa - - log = logging.getLogger(__name__) @@ -84,7 +79,7 @@ def _copy_diagnostic_and_prognostics( class DummyAtmoNonHydro: - def __init__(self, data_provider: sb_utils.IconSerialDataProvider): + def __init__(self, data_provider: sb.IconSerialDataProvider): self.config = None self.data_provider = data_provider self.simulation_date = datetime.fromisoformat(SIMULATION_START_DATE) diff --git a/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py b/model/driver/src/icon4py/model/driver/icon_configuration.py similarity index 96% rename from atm_dyn_iconam/src/icon4py/driver/icon_configuration.py rename to model/driver/src/icon4py/model/driver/icon_configuration.py index 63a42a191..5d67a0781 100644 --- a/atm_dyn_iconam/src/icon4py/driver/icon_configuration.py +++ b/model/driver/src/icon4py/model/driver/icon_configuration.py @@ -14,8 +14,7 @@ from dataclasses import dataclass from typing import Optional -from model.atmosphere.diffusion.src.icon4py.model.atmosphere.diffusion import DiffusionConfig - +from icon4py.model.atmosphere.diffusion.diffusion import DiffusionConfig n_substeps_reduced = 2 diff --git a/atm_dyn_iconam/src/icon4py/driver/io_utils.py b/model/driver/src/icon4py/model/driver/io_utils.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/driver/io_utils.py rename to model/driver/src/icon4py/model/driver/io_utils.py index 3e5c49494..c0d7ce181 100644 --- a/atm_dyn_iconam/src/icon4py/driver/io_utils.py +++ b/model/driver/src/icon4py/model/driver/io_utils.py @@ -10,50 +10,31 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -import importlib + import logging -import sys + from datetime import datetime from enum import Enum from pathlib import Path from icon4py.model.common.decomposition.decomposed import DecompositionInfo from icon4py.model.common.decomposition.parallel_setup import ParallelLogger, ProcessProperties -from model.atmosphere.diffusion.src.icon4py.model.atmosphere.diffusion import ( - DiffusionDiagnosticState, +from icon4py.model.atmosphere.diffusion.diffusion_states import (DiffusionDiagnosticState, DiffusionInterpolationState, DiffusionMetricState, - PrognosticState, -) + PrognosticState) + from icon4py.model.common.grid.horizontal import CellParams, EdgeParams from icon4py.model.common.grid.icon_grid import IconGrid from icon4py.model.common.grid.vertical import VerticalModelParams - +from icon4py.model.common.test_utils import serialbox_utils as sb SB_ONLY_MSG = "Only ser_type='sb' is implemented so far." SIMULATION_START_DATE = "2021-06-20T12:00:10.000" log = logging.getLogger(__name__) -# TODO(Magdalena): for preliminary version of the driver we need serialbox data which is in -# testutils, since that is no proper package we need to import it by hand here. -# remove once there is the testutils package again. -def import_testutils(): - testutils = ( - Path(__file__).parent.__str__() + "/../../../tests/test_utils/__init__.py" - ) - spec = importlib.util.spec_from_file_location("helpers", testutils) - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - return module - - -helpers = import_testutils() - -from helpers import serialbox_utils as sb # noqa F401 - class SerializationType(str, Enum): SB = "serialbox" diff --git a/model/driver/tests/conftest.py b/model/driver/tests/conftest.py new file mode 100644 index 000000000..79b96da84 --- /dev/null +++ b/model/driver/tests/conftest.py @@ -0,0 +1 @@ +from icon4py.model.common.test_utils.fixtures import setup_icon_data,datapath diff --git a/model/atmosphere/dycore/tests/test_dycore_driver.py b/model/driver/tests/test_dycore_driver.py similarity index 87% rename from model/atmosphere/dycore/tests/test_dycore_driver.py rename to model/driver/tests/test_dycore_driver.py index 0d112c78e..4ca4f649d 100644 --- a/model/atmosphere/dycore/tests/test_dycore_driver.py +++ b/model/driver/tests/test_dycore_driver.py @@ -13,11 +13,11 @@ import numpy as np -from icon4py.common.dimension import CellDim, EdgeDim, KDim -from icon4py.driver.dycore_driver import _copy_diagnostic_and_prognostics +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.driver.dycore_driver import _copy_diagnostic_and_prognostics -from .test_utils.helpers import random_field, zero_field -from .test_utils.simple_mesh import SimpleMesh +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def test_copy_diagnostic_and_prognostics(): diff --git a/model/atmosphere/dycore/tests/test_io_utils.py b/model/driver/tests/test_io_utils.py similarity index 96% rename from model/atmosphere/dycore/tests/test_io_utils.py rename to model/driver/tests/test_io_utils.py index fda2ab68e..f687609b4 100644 --- a/model/atmosphere/dycore/tests/test_io_utils.py +++ b/model/driver/tests/test_io_utils.py @@ -14,14 +14,9 @@ import pytest -from icon4py.driver.io_utils import ( - SerializationType, - read_geometry_fields, - read_icon_grid, - read_static_fields, -) from icon4py.model.common.grid.horizontal import CellParams, EdgeParams from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.driver.io_utils import read_geometry_fields, read_static_fields, read_icon_grid, SerializationType @pytest.mark.datatest diff --git a/model/requirements-dev.txt b/model/requirements-dev.txt index 1b7e16d5a..206e27f11 100644 --- a/model/requirements-dev.txt +++ b/model/requirements-dev.txt @@ -1,3 +1,5 @@ -r ../base-requirements-dev.txt -e ./atmosphere/dycore +-e ./atmosphere/diffusion -e ./common +-e ./driver diff --git a/model/requirements.txt b/model/requirements.txt index 2045fe905..519a20834 100644 --- a/model/requirements.txt +++ b/model/requirements.txt @@ -1,3 +1,5 @@ -r ../base-requirements.txt ./atmosphere/dycore +./atmosphere/diffusion ./common +./driver diff --git a/requirements-dev.txt b/requirements-dev.txt index ff9021871..256d7e850 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -4,8 +4,8 @@ -e ./model/atmosphere/dycore -e ./model/common -e ./model/atmosphere/diffusion +-e ./model/driver # icon4pytools -e ./tools -# TODO (magdalena) remove after merge --e ./py2f +-e ./py2f #TODO (magdalena) move to tools and remove diff --git a/requirements.txt b/requirements.txt index dfcf2e291..87afe68b2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,6 +4,7 @@ ./model/atmosphere/dycore ./model/atmosphere/diffusion ./model/common -./py2f +./model/driver +./py2f #TODO (magdalena) move to tools and remove # icon4pytools ./tools From 468fc2197d6a66cccafc2b42b386ff192957de53 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 22 Aug 2023 17:09:06 +0200 Subject: [PATCH 232/263] move field_managemange to common make mpi tests in common run --- ...sion.py => test_parallel_diffusion.py.txt} | 0 .../diffusion/tests/test_diffusion_utils.py | 3 +- model/common/pyproject.toml | 1 + .../field_management/interpolation_fields.py | 0 .../common/test_utils/parallel_fixtures.py | 81 ++++++++++++++++ .../common/test_utils/serialbox_utils.py | 3 +- model/common/tests/conftest.py | 7 +- model/common/tests/mpi_tests/common.py | 46 --------- .../common/tests/mpi_tests/test_decomposed.py | 95 ++++++++++--------- .../tests/test_interpolation_fields.py | 2 +- 10 files changed, 140 insertions(+), 98 deletions(-) rename model/atmosphere/diffusion/tests/mpi_tests/{test_parallel_diffusion.py => test_parallel_diffusion.py.txt} (100%) rename {atm_dyn_iconam/src/icon4py => model/common/src/icon4py/model/common}/field_management/interpolation_fields.py (100%) create mode 100644 model/common/src/icon4py/model/common/test_utils/parallel_fixtures.py delete mode 100644 model/common/tests/mpi_tests/common.py rename model/{atmosphere/dycore => common}/tests/test_interpolation_fields.py (95%) diff --git a/model/atmosphere/diffusion/tests/mpi_tests/test_parallel_diffusion.py b/model/atmosphere/diffusion/tests/mpi_tests/test_parallel_diffusion.py.txt similarity index 100% rename from model/atmosphere/diffusion/tests/mpi_tests/test_parallel_diffusion.py rename to model/atmosphere/diffusion/tests/mpi_tests/test_parallel_diffusion.py.txt diff --git a/model/atmosphere/diffusion/tests/test_diffusion_utils.py b/model/atmosphere/diffusion/tests/test_diffusion_utils.py index 9d36cfddb..14753b417 100644 --- a/model/atmosphere/diffusion/tests/test_diffusion_utils.py +++ b/model/atmosphere/diffusion/tests/test_diffusion_utils.py @@ -150,8 +150,9 @@ def test_set_zero_vertex_k(): @pytest.mark.datatest +@pytest.mark.parametrize("linit", [True]) def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( - diffusion_savepoint_init, r04b09_diffusion_config, icon_grid + diffusion_savepoint_init, r04b09_diffusion_config, icon_grid, linit ): savepoint = diffusion_savepoint_init config = r04b09_diffusion_config diff --git a/model/common/pyproject.toml b/model/common/pyproject.toml index 8fae7b1c1..83af37333 100644 --- a/model/common/pyproject.toml +++ b/model/common/pyproject.toml @@ -109,6 +109,7 @@ show_error_codes = true [tool.pytest.ini_options] testpaths = 'tests' +markers = 'datatest: this test uses binary data' [tool.setuptools.dynamic] version = {attr = 'icon4py.model.common.__init__.__version__'} diff --git a/atm_dyn_iconam/src/icon4py/field_management/interpolation_fields.py b/model/common/src/icon4py/model/common/field_management/interpolation_fields.py similarity index 100% rename from atm_dyn_iconam/src/icon4py/field_management/interpolation_fields.py rename to model/common/src/icon4py/model/common/field_management/interpolation_fields.py diff --git a/model/common/src/icon4py/model/common/test_utils/parallel_fixtures.py b/model/common/src/icon4py/model/common/test_utils/parallel_fixtures.py new file mode 100644 index 000000000..ba9cc09d8 --- /dev/null +++ b/model/common/src/icon4py/model/common/test_utils/parallel_fixtures.py @@ -0,0 +1,81 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + + +import pytest + +from icon4py.model.common.decomposition.parallel_setup import get_processor_properties, \ + ProcessProperties +from icon4py.model.common.test_utils.data_handling import download_and_extract +from icon4py.model.common.test_utils.fixtures import data_path, data_uris +from icon4py.model.common.test_utils.serialbox_utils import IconSerialDataProvider + + + + +@pytest.fixture(params=[False], scope="session") +def processor_props(request): + with_mpi = request.param + return get_processor_properties(with_mpi=with_mpi) + +@pytest.fixture(scope="session") +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def ranked_data_path(processor_props): + return data_path.absolute().joinpath(f"mpitask{processor_props.comm_size}") + +@pytest.fixture(scope="session") +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def download_data(request, processor_props, ranked_data_path): + """ + Get the binary ICON data from a remote server. + + Session scoped fixture which is a prerequisite of all the other fixtures in this file. + """ + try: + uri = data_uris[processor_props.comm_size] + + data_file = data_path.joinpath( + f"mch_ch_r04b09_dsl_mpitask{processor_props.comm_size}.tar.gz" + ).name + if processor_props.rank == 0: + download_and_extract(uri, ranked_data_path, data_file) + if processor_props.comm: + processor_props.comm.barrier() + except KeyError: + raise AssertionError( + f"no data for communicator of size {processor_props.comm_size} exists, use 1, 2 or 4" + ) + +@pytest.fixture +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def get_decomposition_info(processor_props, ranked_data_path): + local_path = ranked_data_path.joinpath("mch_ch_r04b09_dsl/ser_data") + sp = IconSerialDataProvider( + "icon_pydycore", str(local_path.absolute()), True, processor_props.rank + ) + return sp.from_savepoint_grid().construct_decomposition_info() + + +@pytest.fixture +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def get_icon_grid(processor_props, ranked_data_path): + local_path = ranked_data_path.joinpath("mch_ch_r04b09_dsl/ser_data") + return IconSerialDataProvider( + "icon_pydycore", str(local_path.absolute()), False, mpi_rank=processor_props.rank + ).from_savepoint_grid().construct_icon_grid() + + +def check_comm_size(props:ProcessProperties, sizes=[1,2,4]): + if props.comm_size not in sizes: + pytest.xfail(f"wrong comm size: {props.comm_size}: test only works for sizes: {sizes}") + 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 e642168dc..4a17c7542 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 @@ -18,6 +18,7 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import np_as_located_field +from icon4py.model.common import dimension from icon4py.model.common.dimension import ( C2E2CDim, C2E2CODim, @@ -208,7 +209,7 @@ def refin_ctrl(self, dim: Dimension): def num(self, dim: Dimension): return self.sizes[dim] - def _read_field_for_dim(self, field_name, read_func, dim): + def _read_field_for_dim(self, field_name, read_func, dim:Dimension): match (dim): case dimension.CellDim: return read_func(f"c_{field_name}") diff --git a/model/common/tests/conftest.py b/model/common/tests/conftest.py index 111377863..3bd7685c9 100644 --- a/model/common/tests/conftest.py +++ b/model/common/tests/conftest.py @@ -10,9 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -import importlib -import sys -from pathlib import Path + from icon4py.model.common.test_utils.fixtures import ( # noqa F401 damping_height, @@ -23,4 +21,7 @@ icon_grid, r04b09_dsl_gridfile, setup_icon_data, + + ) +from icon4py.model.common.test_utils.parallel_fixtures import (processor_props, ranked_data_path, get_decomposition_info, get_icon_grid, download_data) # noqa F401 diff --git a/model/common/tests/mpi_tests/common.py b/model/common/tests/mpi_tests/common.py deleted file mode 100644 index 9e9f8ccbf..000000000 --- a/model/common/tests/mpi_tests/common.py +++ /dev/null @@ -1,46 +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 pytest - -from icon4py.model.common.test_utils.data_handling import download_and_extract -from icon4py.model.common.test_utils.fixtures import data_path, data_uris -from icon4py.model.common.decomposition.parallel_setup import get_processor_properties - - -props = get_processor_properties(with_mpi=True) -path = data_path.joinpath(f"mpitask{props.comm_size}") -data_path = path.joinpath("mch_ch_r04b09_dsl/ser_data") - - -@pytest.fixture(scope="session") -def download_data(): - """ - Get the binary ICON data from a remote server. - - Session scoped fixture which is a prerequisite of all the other fixtures in this file. - """ - try: - uri = data_uris[props.comm_size] - - data_file = data_path.joinpath( - f"mch_ch_r04b09_dsl_mpitask{props.comm_size}.tar.gz" - ).name - if props.rank == 0: - download_and_extract(uri, path, data_file) - props.comm.barrier() - except KeyError: - raise AssertionError( - f"no data for communicator of size {props.comm_size} exists, use 1, 2 or 4" - ) diff --git a/model/common/tests/mpi_tests/test_decomposed.py b/model/common/tests/mpi_tests/test_decomposed.py index 5ab11d4b8..1649f5a08 100644 --- a/model/common/tests/mpi_tests/test_decomposed.py +++ b/model/common/tests/mpi_tests/test_decomposed.py @@ -14,11 +14,7 @@ import numpy as np import pytest -from icon4py.model.common.test_utils.parallel_fixtures import ( - data_path, - download_data, - props, -) + from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.model.common.decomposition.decomposed import ( DecompositionInfo, @@ -27,12 +23,7 @@ SingleNode, create_exchange, ) -from icon4py.model.common.decomposition.parallel_setup import ProcessProperties -from model.driver.src.icon4py.model.driver.io_utils import ( - read_decomp_info, - read_icon_grid, -) - +from icon4py.model.common.test_utils.parallel_fixtures import check_comm_size """ running tests with mpi: @@ -43,10 +34,12 @@ """ +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def test_props(processor_props): + assert processor_props.comm - -@pytest.mark.mpi -@pytest.mark.skipif(props.comm_size != 2, reason="runs on 2 nodes only") +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize("processor_props", [True], indirect=True) @pytest.mark.parametrize( ("dim, owned, total"), ( @@ -56,9 +49,11 @@ ), ) def test_decomposition_info_masked( - dim, owned, total, caplog, download_data, decomposition_info # noqa F811 + dim, owned, total, caplog, download_data, get_decomposition_info, processor_props, # noqa F811 ): - my_rank = props.rank + check_comm_size(processor_props, sizes=[2]) + my_rank = processor_props.rank + decomposition_info = get_decomposition_info all_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.ALL) my_total = total[my_rank] my_owned = owned[my_rank] @@ -76,7 +71,18 @@ def test_decomposition_info_masked( _assert_index_partitioning(all_indices, halo_indices, owned_indices) -@pytest.mark.skipif(props.comm_size != 2, reason="runs on 2 nodes only") +#@pytest.mark.skipif(props.comm_size != 2, reason="runs on 2 nodes only") +def _assert_index_partitioning(all_indices, halo_indices, owned_indices): + owned_list = owned_indices.tolist() + halos_list = halo_indices.tolist() + all_list = all_indices.tolist() + assert set(owned_list) & set(halos_list) == set() + assert set(owned_list) & set(all_list) == set(owned_list) + assert set(halos_list) & set(all_list) == set(halos_list) + assert set(halos_list) | set(owned_list) == set(all_list) + + +@pytest.mark.parametrize("processor_props", [True], indirect=True) @pytest.mark.parametrize( ("dim, owned, total"), ( @@ -85,10 +91,13 @@ def test_decomposition_info_masked( (VertexDim, (5373, 5290), (5455, 5456)), ), ) +@pytest.mark.mpi(min_size=2) def test_decomposition_info_local_index( - dim, owned, total, caplog, download_data, decomposition_info # noqa F811 + dim, owned, total, caplog, download_data, get_decomposition_info, processor_props # noqa F811 ): - my_rank = props.rank + check_comm_size(processor_props, sizes=[2]) + my_rank = processor_props.rank + decomposition_info = get_decomposition_info all_indices = decomposition_info.local_index(dim, DecompositionInfo.EntryType.ALL) my_total = total[my_rank] my_owned = owned[my_rank] @@ -109,20 +118,11 @@ def test_decomposition_info_local_index( _assert_index_partitioning(all_indices, halo_indices, owned_indices) -def _assert_index_partitioning(all_indices, halo_indices, owned_indices): - owned_list = owned_indices.tolist() - halos_list = halo_indices.tolist() - all_list = all_indices.tolist() - assert set(owned_list) & set(halos_list) == set() - assert set(owned_list) & set(all_list) == set(owned_list) - assert set(halos_list) & set(all_list) == set(halos_list) - assert set(halos_list) | set(owned_list) == set(all_list) - - @pytest.mark.mpi +@pytest.mark.parametrize("processor_props", [True], indirect=True) @pytest.mark.parametrize("num", [1, 2, 3]) -def test_domain_descriptor_id_are_globally_unique(num): - +def test_domain_descriptor_id_are_globally_unique(num, processor_props): + props = processor_props size = props.comm_size id_gen = DomainDescriptorIdGenerator(parallel_props=props) id1 = id_gen() @@ -142,13 +142,16 @@ def test_domain_descriptor_id_are_globally_unique(num): @pytest.mark.mpi -@pytest.mark.skipif( - props.comm_size not in (1, 2, 4), - reason="input files only available for 1 or 2 nodes", -) -def test_decomposition_info_matches_gridsize(caplog, download_data, decomposition_info): # noqa F811 - - icon_grid = read_icon_grid(data_path, props.rank) +#@pytest.mark.skipif( +# props.comm_size not in (1, 2, 4), +# reason="input files only available for 1 or 2 4nodes", +#) +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def test_decomposition_info_matches_gridsize(caplog, download_data, get_decomposition_info, get_icon_grid, processor_props, + ): # noqa F811 + decomposition_info = get_decomposition_info + check_comm_size(processor_props) + icon_grid = get_icon_grid assert ( decomposition_info.global_index( dim=CellDim, entry_type=DecompositionInfo.EntryType.ALL @@ -170,18 +173,18 @@ def test_decomposition_info_matches_gridsize(caplog, download_data, decompositio @pytest.mark.mpi -def test_create_multi_node_runtime_with_mpi(download_data): # noqa F811 - decomp_info = read_decomp_info(data_path, props) - exchange = create_exchange(props, decomp_info) +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def test_create_multi_node_runtime_with_mpi(download_data, get_decomposition_info, processor_props): # noqa F811 + props = processor_props + exchange = create_exchange(props, get_decomposition_info) if props.comm_size > 1: assert isinstance(exchange, GHexMultiNode) else: assert isinstance(exchange, SingleNode) - -def test_create_single_node_runtime_without_mpi(): - props = ProcessProperties.from_single_node() - decomp_info = read_decomp_info(data_path, props) - exchange = create_exchange(props, decomp_info) +@pytest.mark.parametrize("processor_props", [False], indirect=True) +def test_create_single_node_runtime_without_mpi(processor_props, get_decomposition_info, ranked_data_path): + props = processor_props + exchange = create_exchange(props, get_decomposition_info) assert isinstance(exchange, SingleNode) diff --git a/model/atmosphere/dycore/tests/test_interpolation_fields.py b/model/common/tests/test_interpolation_fields.py similarity index 95% rename from model/atmosphere/dycore/tests/test_interpolation_fields.py rename to model/common/tests/test_interpolation_fields.py index 37c0b3a31..b2c2bfeff 100644 --- a/model/atmosphere/dycore/tests/test_interpolation_fields.py +++ b/model/common/tests/test_interpolation_fields.py @@ -27,7 +27,7 @@ import pytest from icon4py.model.common.dimension import EdgeDim -from icon4py.field_management.interpolation_fields import compute_c_lin_e +from icon4py.model.common.field_management.interpolation_fields import compute_c_lin_e from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex From de5c6efc2291bde7d97acc4087d9f6be5f007a5f Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 23 Aug 2023 17:49:25 +0200 Subject: [PATCH 233/263] consolidate fixtures make mpi tests run --- model/atmosphere/diffusion/tests/conftest.py | 3 + .../diffusion/tests/mpi_tests/__init__.py | 0 .../mpi_tests/test_parallel_diffusion.py | 111 +++++++++++++++ .../mpi_tests/test_parallel_diffusion.py.txt | 127 ------------------ .../model/common/test_utils/fixtures.py | 33 +++-- ...rallel_fixtures.py => parallel_helpers.py} | 29 +--- model/common/tests/conftest.py | 6 +- .../common/tests/mpi_tests/test_decomposed.py | 25 ++-- model/driver/tests/conftest.py | 2 +- 9 files changed, 151 insertions(+), 185 deletions(-) create mode 100644 model/atmosphere/diffusion/tests/mpi_tests/__init__.py create mode 100644 model/atmosphere/diffusion/tests/mpi_tests/test_parallel_diffusion.py delete mode 100644 model/atmosphere/diffusion/tests/mpi_tests/test_parallel_diffusion.py.txt rename model/common/src/icon4py/model/common/test_utils/{parallel_fixtures.py => parallel_helpers.py} (56%) diff --git a/model/atmosphere/diffusion/tests/conftest.py b/model/atmosphere/diffusion/tests/conftest.py index 44f095ce8..5f2472e8b 100644 --- a/model/atmosphere/diffusion/tests/conftest.py +++ b/model/atmosphere/diffusion/tests/conftest.py @@ -9,7 +9,10 @@ datapath, grid_savepoint, icon_grid, + decomposition_info, setup_icon_data, + ranked_data_path, + processor_props, linit, ndyn_substeps, step_date_exit, diff --git a/model/atmosphere/diffusion/tests/mpi_tests/__init__.py b/model/atmosphere/diffusion/tests/mpi_tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/model/atmosphere/diffusion/tests/mpi_tests/test_parallel_diffusion.py b/model/atmosphere/diffusion/tests/mpi_tests/test_parallel_diffusion.py new file mode 100644 index 000000000..b47f02a8f --- /dev/null +++ b/model/atmosphere/diffusion/tests/mpi_tests/test_parallel_diffusion.py @@ -0,0 +1,111 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + + +import pytest + +from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.test_utils.parallel_helpers import check_comm_size +from ..test_diffusion import _verify_diffusion_fields +from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim +from icon4py.model.common.decomposition.decomposed import DecompositionInfo, create_exchange +from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams + + + +@pytest.mark.mpi +@pytest.mark.parametrize("ndyn_substeps", [2]) +@pytest.mark.parametrize("linit", [True, False]) +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def test_parallel_diffusion( + r04b09_diffusion_config, + step_date_init, + linit, + ndyn_substeps, + processor_props, + decomposition_info, + icon_grid, + diffusion_savepoint_init, + diffusion_savepoint_exit, + grid_savepoint, + metrics_savepoint, + interpolation_savepoint, + damping_height + +): + check_comm_size(processor_props) + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: inializing diffusion for experiment 'mch_ch_r04_b09_dsl" + ) + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: decomposition info : klevels = {decomposition_info.klevels}, " + f"local cells = {decomposition_info.global_index(CellDim, DecompositionInfo.EntryType.ALL).shape} " + f"local edges = {decomposition_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape} " + f"local vertices = {decomposition_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape}" + ) + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: GHEX context setup: from {processor_props.comm_name} with {processor_props.comm_size} nodes" + ) + + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: using local grid with {icon_grid.num_cells()} Cells, {icon_grid.num_edges()} Edges, {icon_grid.num_vertices()} Vertices" + ) + metric_state= metrics_savepoint.construct_metric_state_for_diffusion() + cell_geometry = grid_savepoint.construct_cell_geometry() + edge_geometry = grid_savepoint.construct_edge_geometry() + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() + + diffusion_params = DiffusionParams(r04b09_diffusion_config) + dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: setup: using {processor_props.comm_name} with {processor_props.comm_size} nodes" + ) + exchange = create_exchange(processor_props, decomposition_info) + + diffusion = Diffusion(exchange) + + diffusion.init( + grid=icon_grid, + config=r04b09_diffusion_config, + params=diffusion_params, + vertical_params=VerticalModelParams(grid_savepoint.vct_a(), damping_height), + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry, + ) + print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ") + diagnostic_state = diffusion_savepoint_init.construct_diagnostics_for_diffusion() + prognostic_state = diffusion_savepoint_init.construct_prognostics() + if linit: + diffusion.initial_run( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + ) + else: + diffusion.run( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + ) + print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ") + + _verify_diffusion_fields( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + diffusion_savepoint=diffusion_savepoint_exit, + ) + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: running diffusion step - using {processor_props.comm_name} with {processor_props.comm_size} nodes - DONE" + ) diff --git a/model/atmosphere/diffusion/tests/mpi_tests/test_parallel_diffusion.py.txt b/model/atmosphere/diffusion/tests/mpi_tests/test_parallel_diffusion.py.txt deleted file mode 100644 index 6a85e23b8..000000000 --- a/model/atmosphere/diffusion/tests/mpi_tests/test_parallel_diffusion.py.txt +++ /dev/null @@ -1,127 +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 pytest - -from icon4py.model.common.decomposition.parallel_setup import get_processor_properties -from icon4py.model.common.test_utils.parallel_fixtures import ( # noqa F401 - data_path, - download_data, - props, -) -from icon4py.model.atmosphere.diffusion.tests.test_diffusion import _verify_diffusion_fields -from icon4py.model.common.test_utils.serialbox_utils import IconSerialDataProvider -from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim -from icon4py.model.common.decomposition.decomposed import DecompositionInfo, create_exchange -from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams -from model.driver.src.icon4py.model.driver.io_utils import ( - read_decomp_info, - read_geometry_fields, - read_icon_grid, - read_static_fields, -) - -props = get_processor_properties(with_mpi=True) -@pytest.mark.mpi -@pytest.mark.parametrize("ndyn_substeps", [2]) -@pytest.mark.parametrize("linit", [True, False]) -@pytest.mark.skipif( - props.comm_size not in (1, 2, 4), - reason="input files only available for 1 or 2 nodes", -) -def test_parallel_diffusion( - r04b09_diffusion_config, - step_date_init, - linit, - ndyn_substeps, - download_data, # noqa: F811 -): - - print( - f"rank={props.rank}/{props.comm_size}: inializing diffusion for experiment 'mch_ch_r04_b09_dsl" - ) - decomp_info = read_decomp_info( - data_path, - props, - ) - print( - f"rank={props.rank}/{props.comm_size}: decomposition info : klevels = {decomp_info.klevels}, " - f"local cells = {decomp_info.global_index(CellDim, DecompositionInfo.EntryType.ALL).shape} " - f"local edges = {decomp_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape} " - f"local vertices = {decomp_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape}" - ) - print( - f"rank={props.rank}/{props.comm_size}: GHEX context setup: from {props.comm_name} with {props.comm_size} nodes" - ) - - icon_grid = read_icon_grid(data_path, rank=props.rank) - print( - f"rank={props.rank}/{props.comm_size}: using local grid with {icon_grid.num_cells()} Cells, {icon_grid.num_edges()} Edges, {icon_grid.num_vertices()} Vertices" - ) - diffusion_params = DiffusionParams(r04b09_diffusion_config) - - diffusion_initial_data = IconSerialDataProvider( - "icon_pydycore", str(data_path), True, mpi_rank=props.rank - ).from_savepoint_diffusion_init(linit=linit, date=step_date_init) - (edge_geometry, cell_geometry, vertical_geometry) = read_geometry_fields( - data_path, rank=props.rank - ) - (metric_state, interpolation_state) = read_static_fields(data_path, rank=props.rank) - - dtime = diffusion_initial_data.get_metadata("dtime").get("dtime") - print( - f"rank={props.rank}/{props.comm_size}: setup: using {props.comm_name} with {props.comm_size} nodes" - ) - exchange = create_exchange(props, decomp_info) - - diffusion = Diffusion(exchange) - - diffusion.init( - grid=icon_grid, - config=r04b09_diffusion_config, - params=diffusion_params, - vertical_params=vertical_geometry, - metric_state=metric_state, - interpolation_state=interpolation_state, - edge_params=edge_geometry, - cell_params=cell_geometry, - ) - print(f"rank={props.rank}/{props.comm_size}: diffusion initialized ") - diagnostic_state = diffusion_initial_data.construct_diagnostics_for_diffusion() - prognostic_state = diffusion_initial_data.construct_prognostics() - if linit: - diffusion.initial_run( - diagnostic_state=diagnostic_state, - prognostic_state=prognostic_state, - dtime=dtime, - ) - else: - diffusion.run( - diagnostic_state=diagnostic_state, - prognostic_state=prognostic_state, - dtime=dtime, - ) - print(f"rank={props.rank}/{props.comm_size}: diffusion run ") - - diffusion_savepoint_exit = IconSerialDataProvider( - "icon_pydycore", str(data_path), True, mpi_rank=props.rank - ).from_savepoint_diffusion_exit(linit=linit, date=step_date_init) - _verify_diffusion_fields( - diagnostic_state=diagnostic_state, - prognostic_state=prognostic_state, - diffusion_savepoint=diffusion_savepoint_exit, - ) - print( - f"rank={props.rank}/{props.comm_size}: running diffusion step - using {props.comm_name} with {props.comm_size} nodes - DONE" - ) diff --git a/model/common/src/icon4py/model/common/test_utils/fixtures.py b/model/common/src/icon4py/model/common/test_utils/fixtures.py index dc3b54d13..5cf874438 100644 --- a/model/common/src/icon4py/model/common/test_utils/fixtures.py +++ b/model/common/src/icon4py/model/common/test_utils/fixtures.py @@ -19,8 +19,7 @@ from .data_handling import download_and_extract from .serialbox_utils import IconSerialDataProvider from .simple_mesh import SimpleMesh - - +from ..decomposition.parallel_setup import get_processor_properties test_utils = Path(__file__).parent model = test_utils.parent.parent @@ -52,10 +51,19 @@ ).name -@pytest.fixture -def datapath(setup_icon_data): - local_path = "mpitask1/mch_ch_r04b09_dsl/ser_data" - return data_path.joinpath(local_path) +@pytest.fixture(params=[False], scope="session") +def processor_props(request): + with_mpi = request.param + return get_processor_properties(with_mpi=with_mpi) + + +@pytest.fixture(scope="session") +def ranked_data_path(processor_props): + return data_path.absolute().joinpath(f"mpitask{processor_props.comm_size}") + +@pytest.fixture(scope="session") +def datapath(setup_icon_data, ranked_data_path): + return ranked_data_path.joinpath("mch_ch_r04b09_dsl/ser_data") @pytest.fixture(scope="session") @@ -67,11 +75,10 @@ def setup_icon_data(): """ download_and_extract(data_uris[1], data_path, data_file) - -@pytest.fixture -def data_provider(setup_icon_data, datapath) -> IconSerialDataProvider: - path = str(datapath) - return IconSerialDataProvider("icon_pydycore", path, True) +# TODO (magdalena) dependency on setup_icon_data or download_data +@pytest.fixture(scope="session") +def data_provider(setup_icon_data, datapath, processor_props) -> IconSerialDataProvider: + return IconSerialDataProvider(fname_prefix="icon_pydycore", path=str(datapath), mpi_rank=processor_props.rank, do_print=True) @pytest.fixture @@ -89,6 +96,10 @@ def icon_grid(grid_savepoint): return grid_savepoint.construct_icon_grid() +@pytest.fixture +def decomposition_info(data_provider): + return data_provider.from_savepoint_grid().construct_decomposition_info() + @pytest.fixture def damping_height(): return 12500 diff --git a/model/common/src/icon4py/model/common/test_utils/parallel_fixtures.py b/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py similarity index 56% rename from model/common/src/icon4py/model/common/test_utils/parallel_fixtures.py rename to model/common/src/icon4py/model/common/test_utils/parallel_helpers.py index ba9cc09d8..4a7ebda58 100644 --- a/model/common/src/icon4py/model/common/test_utils/parallel_fixtures.py +++ b/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py @@ -14,25 +14,13 @@ import pytest -from icon4py.model.common.decomposition.parallel_setup import get_processor_properties, \ - ProcessProperties +from icon4py.model.common.decomposition.parallel_setup import ProcessProperties from icon4py.model.common.test_utils.data_handling import download_and_extract from icon4py.model.common.test_utils.fixtures import data_path, data_uris -from icon4py.model.common.test_utils.serialbox_utils import IconSerialDataProvider -@pytest.fixture(params=[False], scope="session") -def processor_props(request): - with_mpi = request.param - return get_processor_properties(with_mpi=with_mpi) - -@pytest.fixture(scope="session") -@pytest.mark.parametrize("processor_props", [True], indirect=True) -def ranked_data_path(processor_props): - return data_path.absolute().joinpath(f"mpitask{processor_props.comm_size}") - @pytest.fixture(scope="session") @pytest.mark.parametrize("processor_props", [True], indirect=True) def download_data(request, processor_props, ranked_data_path): @@ -56,23 +44,8 @@ def download_data(request, processor_props, ranked_data_path): f"no data for communicator of size {processor_props.comm_size} exists, use 1, 2 or 4" ) -@pytest.fixture -@pytest.mark.parametrize("processor_props", [True], indirect=True) -def get_decomposition_info(processor_props, ranked_data_path): - local_path = ranked_data_path.joinpath("mch_ch_r04b09_dsl/ser_data") - sp = IconSerialDataProvider( - "icon_pydycore", str(local_path.absolute()), True, processor_props.rank - ) - return sp.from_savepoint_grid().construct_decomposition_info() -@pytest.fixture -@pytest.mark.parametrize("processor_props", [True], indirect=True) -def get_icon_grid(processor_props, ranked_data_path): - local_path = ranked_data_path.joinpath("mch_ch_r04b09_dsl/ser_data") - return IconSerialDataProvider( - "icon_pydycore", str(local_path.absolute()), False, mpi_rank=processor_props.rank - ).from_savepoint_grid().construct_icon_grid() def check_comm_size(props:ProcessProperties, sizes=[1,2,4]): diff --git a/model/common/tests/conftest.py b/model/common/tests/conftest.py index 3bd7685c9..501d5e997 100644 --- a/model/common/tests/conftest.py +++ b/model/common/tests/conftest.py @@ -21,7 +21,9 @@ icon_grid, r04b09_dsl_gridfile, setup_icon_data, - + processor_props, +interpolation_savepoint, +ranked_data_path,decomposition_info ) -from icon4py.model.common.test_utils.parallel_fixtures import (processor_props, ranked_data_path, get_decomposition_info, get_icon_grid, download_data) # noqa F401 +from icon4py.model.common.test_utils.parallel_helpers import (download_data) # noqa F401 diff --git a/model/common/tests/mpi_tests/test_decomposed.py b/model/common/tests/mpi_tests/test_decomposed.py index 1649f5a08..80c66bff5 100644 --- a/model/common/tests/mpi_tests/test_decomposed.py +++ b/model/common/tests/mpi_tests/test_decomposed.py @@ -23,7 +23,7 @@ SingleNode, create_exchange, ) -from icon4py.model.common.test_utils.parallel_fixtures import check_comm_size +from icon4py.model.common.test_utils.parallel_helpers import check_comm_size """ running tests with mpi: @@ -49,11 +49,10 @@ def test_props(processor_props): ), ) def test_decomposition_info_masked( - dim, owned, total, caplog, download_data, get_decomposition_info, processor_props, # noqa F811 + dim, owned, total, caplog, download_data, decomposition_info, processor_props, # noqa F811 ): check_comm_size(processor_props, sizes=[2]) my_rank = processor_props.rank - decomposition_info = get_decomposition_info all_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.ALL) my_total = total[my_rank] my_owned = owned[my_rank] @@ -93,11 +92,10 @@ def _assert_index_partitioning(all_indices, halo_indices, owned_indices): ) @pytest.mark.mpi(min_size=2) def test_decomposition_info_local_index( - dim, owned, total, caplog, download_data, get_decomposition_info, processor_props # noqa F811 + dim, owned, total, caplog, download_data, decomposition_info, processor_props # noqa F811 ): check_comm_size(processor_props, sizes=[2]) my_rank = processor_props.rank - decomposition_info = get_decomposition_info all_indices = decomposition_info.local_index(dim, DecompositionInfo.EntryType.ALL) my_total = total[my_rank] my_owned = owned[my_rank] @@ -142,16 +140,11 @@ def test_domain_descriptor_id_are_globally_unique(num, processor_props): @pytest.mark.mpi -#@pytest.mark.skipif( -# props.comm_size not in (1, 2, 4), -# reason="input files only available for 1 or 2 4nodes", -#) @pytest.mark.parametrize("processor_props", [True], indirect=True) -def test_decomposition_info_matches_gridsize(caplog, download_data, get_decomposition_info, get_icon_grid, processor_props, +def test_decomposition_info_matches_gridsize(caplog, download_data, decomposition_info, icon_grid, processor_props, ): # noqa F811 - decomposition_info = get_decomposition_info + check_comm_size(processor_props) - icon_grid = get_icon_grid assert ( decomposition_info.global_index( dim=CellDim, entry_type=DecompositionInfo.EntryType.ALL @@ -174,17 +167,17 @@ def test_decomposition_info_matches_gridsize(caplog, download_data, get_decompos @pytest.mark.mpi @pytest.mark.parametrize("processor_props", [True], indirect=True) -def test_create_multi_node_runtime_with_mpi(download_data, get_decomposition_info, processor_props): # noqa F811 +def test_create_multi_node_runtime_with_mpi(download_data, decomposition_info, processor_props): # noqa F811 props = processor_props - exchange = create_exchange(props, get_decomposition_info) + exchange = create_exchange(props, decomposition_info) if props.comm_size > 1: assert isinstance(exchange, GHexMultiNode) else: assert isinstance(exchange, SingleNode) @pytest.mark.parametrize("processor_props", [False], indirect=True) -def test_create_single_node_runtime_without_mpi(processor_props, get_decomposition_info, ranked_data_path): +def test_create_single_node_runtime_without_mpi(processor_props, decomposition_info): props = processor_props - exchange = create_exchange(props, get_decomposition_info) + exchange = create_exchange(props, decomposition_info) assert isinstance(exchange, SingleNode) diff --git a/model/driver/tests/conftest.py b/model/driver/tests/conftest.py index 79b96da84..871d94970 100644 --- a/model/driver/tests/conftest.py +++ b/model/driver/tests/conftest.py @@ -1 +1 @@ -from icon4py.model.common.test_utils.fixtures import setup_icon_data,datapath +from icon4py.model.common.test_utils.fixtures import setup_icon_data,datapath, ranked_data_path,processor_props From 5fc58b69d45be7f97e436708731a312328ecd88c Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 23 Aug 2023 19:22:16 +0200 Subject: [PATCH 234/263] move grid related fixtures to common/conftest.py fix wrong url for grid files delete duplicate download fixture --- model/atmosphere/diffusion/tests/conftest.py | 2 +- model/atmosphere/dycore/tests/conftest.py | 7 +-- .../model/common/test_utils/fixtures.py | 63 +++++++------------ .../common/test_utils/parallel_helpers.py | 32 +--------- model/common/tests/conftest.py | 45 ++++++++++--- .../common/tests/mpi_tests/test_decomposed.py | 10 +-- model/driver/tests/conftest.py | 2 +- 7 files changed, 71 insertions(+), 90 deletions(-) diff --git a/model/atmosphere/diffusion/tests/conftest.py b/model/atmosphere/diffusion/tests/conftest.py index 5f2472e8b..ba01ee218 100644 --- a/model/atmosphere/diffusion/tests/conftest.py +++ b/model/atmosphere/diffusion/tests/conftest.py @@ -10,7 +10,7 @@ grid_savepoint, icon_grid, decomposition_info, - setup_icon_data, + download_ser_data, ranked_data_path, processor_props, linit, diff --git a/model/atmosphere/dycore/tests/conftest.py b/model/atmosphere/dycore/tests/conftest.py index 5133429bd..64e66161f 100644 --- a/model/atmosphere/dycore/tests/conftest.py +++ b/model/atmosphere/dycore/tests/conftest.py @@ -10,19 +10,16 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -import pytest from icon4py.model.common.test_utils.fixtures import ( # noqa F401 damping_height, data_provider, datapath, - get_grid_files, grid_savepoint, icon_grid, - r04b09_dsl_gridfile, - setup_icon_data, - linit,step_date_exit, step_date_init,mesh, backend + download_ser_data, + linit, step_date_exit, step_date_init, mesh, backend ) diff --git a/model/common/src/icon4py/model/common/test_utils/fixtures.py b/model/common/src/icon4py/model/common/test_utils/fixtures.py index 5cf874438..4bae1fdd2 100644 --- a/model/common/src/icon4py/model/common/test_utils/fixtures.py +++ b/model/common/src/icon4py/model/common/test_utils/fixtures.py @@ -31,24 +31,9 @@ 2: "https://polybox.ethz.ch/index.php/s/NUQjmJcMEoQxFiK/download", 4: "https://polybox.ethz.ch/index.php/s/QC7xt7xLT5xeVN5/download", } -mch_ch_r04b09_dsl_grid_uri = ( - "https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p/download" -) -r02b04_global_grid_uri = "https://polybox.ethz.ch/index.php/s/0EM8O8U53GKGsst/download" - -data_path = base_path.joinpath("ser_icondata") -data_file = data_path.joinpath("mch_ch_r04b09_dsl.tar.gz").name +ser_data_basepath = base_path.joinpath("ser_icondata") -grids_path = base_path.joinpath("grids") -r04b09_dsl_grid_path = grids_path.joinpath("mch_ch_r04b09_dsl") -r04b09_dsl_data_file = r04b09_dsl_grid_path.joinpath( - "mch_ch_r04b09_dsl_grids_v1.tar.gz" -).name -r02b04_global_grid_path = grids_path.joinpath("r02b04_global") -r02b04_global_data_file = r02b04_global_grid_path.joinpath( - "icon_grid_0013_R02B04_G.tar.gz" -).name @pytest.fixture(params=[False], scope="session") @@ -59,25 +44,40 @@ def processor_props(request): @pytest.fixture(scope="session") def ranked_data_path(processor_props): - return data_path.absolute().joinpath(f"mpitask{processor_props.comm_size}") + return ser_data_basepath.absolute().joinpath(f"mpitask{processor_props.comm_size}") @pytest.fixture(scope="session") -def datapath(setup_icon_data, ranked_data_path): +def datapath(ranked_data_path): return ranked_data_path.joinpath("mch_ch_r04b09_dsl/ser_data") @pytest.fixture(scope="session") -def setup_icon_data(): +def download_ser_data(request, processor_props, ranked_data_path): """ - Get the binary ICON data for single node run from a remote server. + Get the binary ICON data from a remote server. Session scoped fixture which is a prerequisite of all the other fixtures in this file. """ - download_and_extract(data_uris[1], data_path, data_file) + try: + uri = data_uris[processor_props.comm_size] + + data_file = ranked_data_path.joinpath( + f"mch_ch_r04b09_dsl_mpitask{processor_props.comm_size}.tar.gz" + ).name + if processor_props.rank == 0: + download_and_extract(uri, ranked_data_path, data_file) + if processor_props.comm: + processor_props.comm.barrier() + except KeyError: + raise AssertionError( + f"no data for communicator of size {processor_props.comm_size} exists, use 1, 2 or 4" + ) + + + -# TODO (magdalena) dependency on setup_icon_data or download_data @pytest.fixture(scope="session") -def data_provider(setup_icon_data, datapath, processor_props) -> IconSerialDataProvider: +def data_provider(download_ser_data, datapath, processor_props) -> IconSerialDataProvider: return IconSerialDataProvider(fname_prefix="icon_pydycore", path=str(datapath), mpi_rank=processor_props.rank, do_print=True) @@ -147,20 +147,6 @@ def step_date_exit(): return "2021-06-20T12:00:10.000" -@pytest.fixture(scope="session") -def get_grid_files(): - """ - Get the grid files used for testing. - - Session scoped fixture which is a prerequisite of all the other fixtures in this file. - """ - download_and_extract( - mch_ch_r04b09_dsl_grid_uri, r04b09_dsl_grid_path, r04b09_dsl_data_file - ) - download_and_extract( - r02b04_global_grid_uri, r02b04_global_grid_path, r02b04_global_data_file - ) - @pytest.fixture def interpolation_savepoint(data_provider): # noqa F811 @@ -175,9 +161,6 @@ def metrics_savepoint(data_provider): # noqa F811 -@pytest.fixture() -def r04b09_dsl_gridfile(get_grid_files): - return r04b09_dsl_grid_path.joinpath("grid.nc") BACKENDS = {"embedded": executor} MESHES = {"simple_mesh": SimpleMesh()} diff --git a/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py b/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py index 4a7ebda58..b8b1ed965 100644 --- a/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py @@ -15,40 +15,10 @@ import pytest from icon4py.model.common.decomposition.parallel_setup import ProcessProperties -from icon4py.model.common.test_utils.data_handling import download_and_extract -from icon4py.model.common.test_utils.fixtures import data_path, data_uris - - - - -@pytest.fixture(scope="session") -@pytest.mark.parametrize("processor_props", [True], indirect=True) -def download_data(request, processor_props, ranked_data_path): - """ - Get the binary ICON data from a remote server. - - Session scoped fixture which is a prerequisite of all the other fixtures in this file. - """ - try: - uri = data_uris[processor_props.comm_size] - - data_file = data_path.joinpath( - f"mch_ch_r04b09_dsl_mpitask{processor_props.comm_size}.tar.gz" - ).name - if processor_props.rank == 0: - download_and_extract(uri, ranked_data_path, data_file) - if processor_props.comm: - processor_props.comm.barrier() - except KeyError: - raise AssertionError( - f"no data for communicator of size {processor_props.comm_size} exists, use 1, 2 or 4" - ) - - def check_comm_size(props:ProcessProperties, sizes=[1,2,4]): if props.comm_size not in sizes: - pytest.xfail(f"wrong comm size: {props.comm_size}: test only works for sizes: {sizes}") + pytest.xfail(f"wrong comm size: {props.comm_size}: test only works for comm-sizes: {sizes}") diff --git a/model/common/tests/conftest.py b/model/common/tests/conftest.py index 501d5e997..764dfd166 100644 --- a/model/common/tests/conftest.py +++ b/model/common/tests/conftest.py @@ -10,20 +10,51 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import pytest - +from icon4py.model.common.test_utils.data_handling import download_and_extract from icon4py.model.common.test_utils.fixtures import ( # noqa F401 damping_height, data_provider, datapath, - get_grid_files, grid_savepoint, icon_grid, - r04b09_dsl_gridfile, - setup_icon_data, + download_ser_data, processor_props, -interpolation_savepoint, -ranked_data_path,decomposition_info + interpolation_savepoint, + ranked_data_path, decomposition_info, base_path + +) +grids_path = base_path.joinpath("grids") +r04b09_dsl_grid_path = grids_path.joinpath("mch_ch_r04b09_dsl") +r04b09_dsl_data_file = r04b09_dsl_grid_path.joinpath( + "mch_ch_r04b09_dsl_grids_v1.tar.gz" +).name +r02b04_global_grid_path = grids_path.joinpath("r02b04_global") +r02b04_global_data_file = r02b04_global_grid_path.joinpath( + "icon_grid_0013_R02B04_G.tar.gz" +).name + +mch_ch_r04b09_dsl_grid_uri = ( + "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download" ) -from icon4py.model.common.test_utils.parallel_helpers import (download_data) # noqa F401 +r02b04_global_grid_uri = "https://polybox.ethz.ch/index.php/s/0EM8O8U53GKGsst/download" + +@pytest.fixture() +def r04b09_dsl_gridfile(get_grid_files): + return r04b09_dsl_grid_path.joinpath("grid.nc") + +@pytest.fixture(scope="session") +def get_grid_files(): + """ + Get the grid files used for testing. + + Session scoped fixture which is a prerequisite of all the other fixtures in this file. + """ + download_and_extract( + mch_ch_r04b09_dsl_grid_uri, r04b09_dsl_grid_path, r04b09_dsl_data_file + ) + download_and_extract( + r02b04_global_grid_uri, r02b04_global_grid_path, r02b04_global_data_file + ) diff --git a/model/common/tests/mpi_tests/test_decomposed.py b/model/common/tests/mpi_tests/test_decomposed.py index 80c66bff5..9c2dac9ba 100644 --- a/model/common/tests/mpi_tests/test_decomposed.py +++ b/model/common/tests/mpi_tests/test_decomposed.py @@ -49,7 +49,7 @@ def test_props(processor_props): ), ) def test_decomposition_info_masked( - dim, owned, total, caplog, download_data, decomposition_info, processor_props, # noqa F811 + dim, owned, total, caplog, download_ser_data, decomposition_info, processor_props, # noqa F811 ): check_comm_size(processor_props, sizes=[2]) my_rank = processor_props.rank @@ -92,7 +92,7 @@ def _assert_index_partitioning(all_indices, halo_indices, owned_indices): ) @pytest.mark.mpi(min_size=2) def test_decomposition_info_local_index( - dim, owned, total, caplog, download_data, decomposition_info, processor_props # noqa F811 + dim, owned, total, caplog, download_ser_data, decomposition_info, processor_props # noqa F811 ): check_comm_size(processor_props, sizes=[2]) my_rank = processor_props.rank @@ -141,8 +141,8 @@ def test_domain_descriptor_id_are_globally_unique(num, processor_props): @pytest.mark.mpi @pytest.mark.parametrize("processor_props", [True], indirect=True) -def test_decomposition_info_matches_gridsize(caplog, download_data, decomposition_info, icon_grid, processor_props, - ): # noqa F811 +def test_decomposition_info_matches_gridsize(caplog, download_ser_data, decomposition_info, icon_grid, processor_props, + ): # noqa F811 check_comm_size(processor_props) assert ( @@ -167,7 +167,7 @@ def test_decomposition_info_matches_gridsize(caplog, download_data, decompositio @pytest.mark.mpi @pytest.mark.parametrize("processor_props", [True], indirect=True) -def test_create_multi_node_runtime_with_mpi(download_data, decomposition_info, processor_props): # noqa F811 +def test_create_multi_node_runtime_with_mpi(download_ser_data, decomposition_info, processor_props): # noqa F811 props = processor_props exchange = create_exchange(props, decomposition_info) if props.comm_size > 1: diff --git a/model/driver/tests/conftest.py b/model/driver/tests/conftest.py index 871d94970..5bb19df73 100644 --- a/model/driver/tests/conftest.py +++ b/model/driver/tests/conftest.py @@ -1 +1 @@ -from icon4py.model.common.test_utils.fixtures import setup_icon_data,datapath, ranked_data_path,processor_props +from icon4py.model.common.test_utils.fixtures import datapath, download_ser_data,ranked_data_path,processor_props From 5b9211b455a06ffe252004c98b2997c946fd7e2a Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 24 Aug 2023 10:34:03 +0200 Subject: [PATCH 235/263] rename diffusion/tests to diffusion/diffusion_tests --- model/atmosphere/diffusion/{tests => diffusion_tests}/__init__.py | 0 model/atmosphere/diffusion/{tests => diffusion_tests}/conftest.py | 0 .../diffusion/{tests => diffusion_tests}/mpi_tests/__init__.py | 0 .../mpi_tests/test_parallel_diffusion.py | 0 .../test_apply_nabla2_and_nabla4_global_to_vn.py | 0 .../test_apply_nabla2_and_nabla4_to_vn.py | 0 .../test_apply_nabla2_to_vn_in_lateral_boundary.py | 0 .../{tests => diffusion_tests}/test_apply_nabla2_to_w.py | 0 .../test_apply_nabla2_to_w_in_upper_damping_layer.py | 0 .../test_calculate_diagnostics_for_turbulence.py | 0 .../test_calculate_horizontal_gradients_for_turbulence.py | 0 .../test_calculate_nabla2_and_smag_coefficients_for_vn.py | 0 .../{tests => diffusion_tests}/test_calculate_nabla2_for_w.py | 0 .../{tests => diffusion_tests}/test_calculate_nabla2_for_z.py | 0 .../{tests => diffusion_tests}/test_calculate_nabla2_of_theta.py | 0 .../diffusion/{tests => diffusion_tests}/test_calculate_nabla4.py | 0 .../diffusion/{tests => diffusion_tests}/test_diffusion.py | 0 .../diffusion/{tests => diffusion_tests}/test_diffusion_states.py | 0 .../diffusion/{tests => diffusion_tests}/test_diffusion_utils.py | 0 ...est_enhance_diffusion_coefficient_for_grid_point_cold_pools.py | 0 .../test_temporary_field_for_grid_point_cold_pools_enhancement.py | 0 .../test_temporary_fields_for_turbulence_diagnostics.py | 0 ...truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py | 0 .../{tests => diffusion_tests}/test_update_theta_and_exner.py | 0 24 files changed, 0 insertions(+), 0 deletions(-) rename model/atmosphere/diffusion/{tests => diffusion_tests}/__init__.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/conftest.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/mpi_tests/__init__.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/mpi_tests/test_parallel_diffusion.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_apply_nabla2_and_nabla4_global_to_vn.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_apply_nabla2_and_nabla4_to_vn.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_apply_nabla2_to_vn_in_lateral_boundary.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_apply_nabla2_to_w.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_apply_nabla2_to_w_in_upper_damping_layer.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_calculate_diagnostics_for_turbulence.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_calculate_horizontal_gradients_for_turbulence.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_calculate_nabla2_and_smag_coefficients_for_vn.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_calculate_nabla2_for_w.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_calculate_nabla2_for_z.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_calculate_nabla2_of_theta.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_calculate_nabla4.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_diffusion.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_diffusion_states.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_diffusion_utils.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_temporary_field_for_grid_point_cold_pools_enhancement.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_temporary_fields_for_turbulence_diagnostics.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py (100%) rename model/atmosphere/diffusion/{tests => diffusion_tests}/test_update_theta_and_exner.py (100%) diff --git a/model/atmosphere/diffusion/tests/__init__.py b/model/atmosphere/diffusion/diffusion_tests/__init__.py similarity index 100% rename from model/atmosphere/diffusion/tests/__init__.py rename to model/atmosphere/diffusion/diffusion_tests/__init__.py diff --git a/model/atmosphere/diffusion/tests/conftest.py b/model/atmosphere/diffusion/diffusion_tests/conftest.py similarity index 100% rename from model/atmosphere/diffusion/tests/conftest.py rename to model/atmosphere/diffusion/diffusion_tests/conftest.py diff --git a/model/atmosphere/diffusion/tests/mpi_tests/__init__.py b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/__init__.py similarity index 100% rename from model/atmosphere/diffusion/tests/mpi_tests/__init__.py rename to model/atmosphere/diffusion/diffusion_tests/mpi_tests/__init__.py diff --git a/model/atmosphere/diffusion/tests/mpi_tests/test_parallel_diffusion.py b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py similarity index 100% rename from model/atmosphere/diffusion/tests/mpi_tests/test_parallel_diffusion.py rename to model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py diff --git a/model/atmosphere/diffusion/tests/test_apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_global_to_vn.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_apply_nabla2_and_nabla4_global_to_vn.py rename to model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_global_to_vn.py diff --git a/model/atmosphere/diffusion/tests/test_apply_nabla2_and_nabla4_to_vn.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_to_vn.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_apply_nabla2_and_nabla4_to_vn.py rename to model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_to_vn.py diff --git a/model/atmosphere/diffusion/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_apply_nabla2_to_vn_in_lateral_boundary.py rename to model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py diff --git a/model/atmosphere/diffusion/tests/test_apply_nabla2_to_w.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_apply_nabla2_to_w.py rename to model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w.py diff --git a/model/atmosphere/diffusion/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_apply_nabla2_to_w_in_upper_damping_layer.py rename to model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py diff --git a/model/atmosphere/diffusion/tests/test_calculate_diagnostics_for_turbulence.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_diagnostics_for_turbulence.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_calculate_diagnostics_for_turbulence.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_diagnostics_for_turbulence.py diff --git a/model/atmosphere/diffusion/tests/test_calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_horizontal_gradients_for_turbulence.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_calculate_horizontal_gradients_for_turbulence.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_horizontal_gradients_for_turbulence.py diff --git a/model/atmosphere/diffusion/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py diff --git a/model/atmosphere/diffusion/tests/test_calculate_nabla2_for_w.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_w.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_calculate_nabla2_for_w.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_w.py diff --git a/model/atmosphere/diffusion/tests/test_calculate_nabla2_for_z.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_z.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_calculate_nabla2_for_z.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_z.py diff --git a/model/atmosphere/diffusion/tests/test_calculate_nabla2_of_theta.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_of_theta.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_calculate_nabla2_of_theta.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_of_theta.py diff --git a/model/atmosphere/diffusion/tests/test_calculate_nabla4.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla4.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_calculate_nabla4.py rename to model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla4.py diff --git a/model/atmosphere/diffusion/tests/test_diffusion.py b/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_diffusion.py rename to model/atmosphere/diffusion/diffusion_tests/test_diffusion.py diff --git a/model/atmosphere/diffusion/tests/test_diffusion_states.py b/model/atmosphere/diffusion/diffusion_tests/test_diffusion_states.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_diffusion_states.py rename to model/atmosphere/diffusion/diffusion_tests/test_diffusion_states.py diff --git a/model/atmosphere/diffusion/tests/test_diffusion_utils.py b/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_diffusion_utils.py rename to model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py diff --git a/model/atmosphere/diffusion/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/diffusion_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py rename to model/atmosphere/diffusion/diffusion_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py diff --git a/model/atmosphere/diffusion/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py b/model/atmosphere/diffusion/diffusion_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py rename to model/atmosphere/diffusion/diffusion_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py diff --git a/model/atmosphere/diffusion/tests/test_temporary_fields_for_turbulence_diagnostics.py b/model/atmosphere/diffusion/diffusion_tests/test_temporary_fields_for_turbulence_diagnostics.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_temporary_fields_for_turbulence_diagnostics.py rename to model/atmosphere/diffusion/diffusion_tests/test_temporary_fields_for_turbulence_diagnostics.py diff --git a/model/atmosphere/diffusion/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/diffusion_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py rename to model/atmosphere/diffusion/diffusion_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py diff --git a/model/atmosphere/diffusion/tests/test_update_theta_and_exner.py b/model/atmosphere/diffusion/diffusion_tests/test_update_theta_and_exner.py similarity index 100% rename from model/atmosphere/diffusion/tests/test_update_theta_and_exner.py rename to model/atmosphere/diffusion/diffusion_tests/test_update_theta_and_exner.py From 6ed055a72f685fa643894d34a916360204bc2ce0 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 24 Aug 2023 10:34:48 +0200 Subject: [PATCH 236/263] pre-commit config --- model/.pre-commit-config.yaml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/model/.pre-commit-config.yaml b/model/.pre-commit-config.yaml index 04f4ab26a..fe7f761e1 100644 --- a/model/.pre-commit-config.yaml +++ b/model/.pre-commit-config.yaml @@ -5,7 +5,7 @@ default_language_version: # Remove frozen version once we migrated away from tsa node: 17.9.1 minimum_pre_commit_version: 2.20.0 -exclude: "tools/.*|model/common/.*" +exclude: "tools/.*" repos: - repo: meta @@ -66,7 +66,6 @@ repos: name: format INI config files args: [--autofix] # - id: pretty-format-toml - # args: [--autofix] - repo: https://github.com/pre-commit/mirrors-prettier rev: v2.2.1 @@ -120,7 +119,9 @@ repos: [ --config=model/.flake8, model/atmosphere/dycore/src/icon4py/, + model/atmosphere/diffusion/src/icon4py/, model/common/src/icon4py/, + model/driver/src/icon4py ] - repo: local From c92ff0c635a9658e6b69b10ce338f480bdcb582c Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 24 Aug 2023 11:11:46 +0200 Subject: [PATCH 237/263] pre-commit on model --- model/.pre-commit-config.yaml | 5 +- .../diffusion/diffusion_tests/__init__.py | 12 +++++ .../diffusion/diffusion_tests/conftest.py | 38 ++++++++++---- .../diffusion_tests/mpi_tests/__init__.py | 12 +++++ .../mpi_tests/test_parallel_diffusion.py | 27 ++++++---- ...st_apply_nabla2_and_nabla4_global_to_vn.py | 8 ++- .../test_apply_nabla2_and_nabla4_to_vn.py | 6 +-- .../diffusion_tests/test_apply_nabla2_to_w.py | 8 ++- ...st_calculate_diagnostics_for_turbulence.py | 8 ++- .../test_calculate_nabla2_of_theta.py | 4 +- .../diffusion_tests/test_calculate_nabla4.py | 4 +- .../diffusion_tests/test_diffusion.py | 11 ++-- .../diffusion_tests/test_diffusion_utils.py | 6 +-- ...orary_fields_for_turbulence_diagnostics.py | 5 +- ...fusion_nabla_of_theta_over_steep_points.py | 7 ++- .../model/atmosphere/diffusion/__init__.py | 3 +- .../model/atmosphere/diffusion/diffusion.py | 52 ++++++++++++------- .../atmosphere/diffusion/diffusion_utils.py | 2 - .../atmosphere/diffusion/stencils/__init__.py | 12 +++++ .../stencils/apply_diffusion_to_vn.py | 8 ++- ...ute_horizontal_gradients_for_turbulance.py | 12 +++-- .../apply_nabla2_and_nabla4_global_to_vn.py | 4 +- .../apply_nabla2_to_vn_in_lateral_boundary.py | 4 +- .../calculate_diagnostics_for_turbulence.py | 4 +- ..._coefficients_for_grid_point_cold_pools.py | 4 +- ...n_coefficient_for_grid_point_cold_pools.py | 4 +- ...orary_fields_for_turbulence_diagnostics.py | 10 +++- ...fusion_nabla_of_theta_over_steep_points.py | 15 ++++-- .../model/atmosphere/dycore/__init__.py | 1 - .../mo_advection_traj_btraj_compute_o1_dsl.py | 6 ++- ...lation_scalar_cells2verts_scalar_ri_dsl.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_02.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_10.py | 4 +- ...nonhydro_stencil_16_fused_btraj_traj_o1.py | 6 ++- .../dycore/mo_solve_nonhydro_stencil_17.py | 5 +- .../dycore/mo_solve_nonhydro_stencil_18.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_20.py | 10 +++- .../dycore/mo_solve_nonhydro_stencil_22.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_24.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_39.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_40.py | 12 +++-- .../dycore/mo_solve_nonhydro_stencil_42.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_55.py | 13 ++++- .../dycore/mo_solve_nonhydro_stencil_56_63.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_58.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_60.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_65.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_67.py | 4 +- .../dycore/mo_solve_nonhydro_stencil_68.py | 5 +- .../mo_velocity_advection_stencil_04.py | 4 +- .../mo_velocity_advection_stencil_07.py | 4 +- .../mo_velocity_advection_stencil_10.py | 4 +- .../mo_velocity_advection_stencil_14.py | 8 ++- .../mo_velocity_advection_stencil_16.py | 8 ++- .../mo_velocity_advection_stencil_18.py | 3 +- .../mo_velocity_advection_stencil_19.py | 5 +- .../mo_velocity_advection_stencil_20.py | 4 +- model/atmosphere/dycore/tests/conftest.py | 14 +++-- ..._mo_advection_traj_btraj_compute_o1_dsl.py | 18 ++++--- .../test_mo_solve_nonhydro_stencil_02.py | 4 +- .../test_mo_solve_nonhydro_stencil_04.py | 9 ++-- .../test_mo_solve_nonhydro_stencil_06.py | 4 +- .../test_mo_solve_nonhydro_stencil_15.py | 4 +- ...nonhydro_stencil_16_fused_btraj_traj_o1.py | 23 +++++--- .../test_mo_solve_nonhydro_stencil_19.py | 5 +- .../test_mo_solve_nonhydro_stencil_20.py | 8 ++- .../test_mo_solve_nonhydro_stencil_21.py | 4 +- .../test_mo_solve_nonhydro_stencil_25.py | 4 +- .../test_mo_solve_nonhydro_stencil_26.py | 4 +- .../test_mo_solve_nonhydro_stencil_28.py | 4 +- .../test_mo_solve_nonhydro_stencil_29.py | 4 +- .../test_mo_solve_nonhydro_stencil_42.py | 4 +- .../test_mo_solve_nonhydro_stencil_44.py | 4 +- .../test_mo_solve_nonhydro_stencil_47.py | 4 +- .../test_mo_solve_nonhydro_stencil_51.py | 4 +- .../test_mo_solve_nonhydro_stencil_52.py | 8 ++- .../test_mo_solve_nonhydro_stencil_54.py | 4 +- .../test_mo_solve_nonhydro_stencil_55.py | 5 +- .../test_mo_solve_nonhydro_stencil_58.py | 4 +- .../test_mo_solve_nonhydro_stencil_60.py | 4 +- .../test_mo_solve_nonhydro_stencil_62.py | 4 +- .../test_mo_solve_nonhydro_stencil_65.py | 4 +- .../test_mo_solve_nonhydro_stencil_68.py | 5 +- .../test_mo_velocity_advection_stencil_02.py | 8 ++- .../test_mo_velocity_advection_stencil_03.py | 8 ++- .../test_mo_velocity_advection_stencil_06.py | 8 ++- .../test_mo_velocity_advection_stencil_08.py | 4 +- .../test_mo_velocity_advection_stencil_09.py | 4 +- .../test_mo_velocity_advection_stencil_10.py | 4 +- .../test_mo_velocity_advection_stencil_13.py | 4 +- .../test_mo_velocity_advection_stencil_19.py | 5 +- .../model/common/decomposition/__init__.py | 12 +++++ .../model/common/decomposition/decomposed.py | 7 ++- .../common/decomposition/parallel_setup.py | 8 +-- .../icon4py/model/common/grid/grid_manager.py | 6 ++- .../icon4py/model/common/grid/icon_grid.py | 1 - .../model/common/test_utils/fixtures.py | 29 ++++++----- .../model/common/test_utils/helpers.py | 12 +++-- .../common/test_utils/parallel_helpers.py | 8 +-- .../common/test_utils/serialbox_utils.py | 29 +++++++---- .../common/src/icon4py/model/common/utils.py | 15 ++++++ model/common/tests/conftest.py | 13 +++-- model/common/tests/mpi_tests/README.md | 5 +- .../common/tests/mpi_tests/test_decomposed.py | 39 +++++++++++--- model/common/tests/test_grid_manager.py | 2 +- .../common/tests/test_interpolation_fields.py | 4 +- .../src/icon4py/model/driver/__init__.py | 2 + .../src/icon4py/model/driver/dycore_driver.py | 23 +++++--- .../model/driver/icon_configuration.py | 1 + .../src/icon4py/model/driver/io_utils.py | 18 ++++--- model/driver/tests/conftest.py | 20 ++++++- model/driver/tests/test_dycore_driver.py | 3 +- model/driver/tests/test_io_utils.py | 7 ++- py2f/src/icon4py/py2f/parsing.py | 1 - py2f/src/icon4py/py2f/py2fgen.py | 1 - .../py2f/wrappers/diffusion_wrapper.py | 15 +++--- py2f/tests/test_code_generation.py | 1 - 117 files changed, 669 insertions(+), 266 deletions(-) diff --git a/model/.pre-commit-config.yaml b/model/.pre-commit-config.yaml index fe7f761e1..1b7302dae 100644 --- a/model/.pre-commit-config.yaml +++ b/model/.pre-commit-config.yaml @@ -76,9 +76,10 @@ repos: types_or: [json, yaml, markdown] - repo: https://github.com/asottile/yesqa - rev: v1.3.0 + rev: v1.5.0 hooks: - id: yesqa + exclude: model/common/decomposition/parallel_setup.py - repo: https://github.com/psf/black rev: 22.3.0 @@ -121,7 +122,7 @@ repos: model/atmosphere/dycore/src/icon4py/, model/atmosphere/diffusion/src/icon4py/, model/common/src/icon4py/, - model/driver/src/icon4py + model/driver/src/icon4py, ] - repo: local diff --git a/model/atmosphere/diffusion/diffusion_tests/__init__.py b/model/atmosphere/diffusion/diffusion_tests/__init__.py index e69de29bb..15dfdb009 100644 --- a/model/atmosphere/diffusion/diffusion_tests/__init__.py +++ b/model/atmosphere/diffusion/diffusion_tests/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/model/atmosphere/diffusion/diffusion_tests/conftest.py b/model/atmosphere/diffusion/diffusion_tests/conftest.py index ba01ee218..2ec7b0911 100644 --- a/model/atmosphere/diffusion/diffusion_tests/conftest.py +++ b/model/atmosphere/diffusion/diffusion_tests/conftest.py @@ -1,27 +1,45 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + import pytest -from icon4py.model.atmosphere.diffusion.diffusion import DiffusionConfig, DiffusionType +from icon4py.model.atmosphere.diffusion.diffusion import ( + DiffusionConfig, + DiffusionType, +) from icon4py.model.common.test_utils.fixtures import ( # noqa F401 - interpolation_savepoint, - metrics_savepoint, + backend, damping_height, data_provider, datapath, - grid_savepoint, - icon_grid, decomposition_info, download_ser_data, - ranked_data_path, - processor_props, + grid_savepoint, + icon_grid, + interpolation_savepoint, linit, + mesh, + metrics_savepoint, ndyn_substeps, + processor_props, + ranked_data_path, step_date_exit, step_date_init, - mesh, - backend ) + + @pytest.fixture -def r04b09_diffusion_config(ndyn_substeps) -> DiffusionConfig: +def r04b09_diffusion_config(ndyn_substeps) -> DiffusionConfig: # noqa F811 """ Create DiffusionConfig matching MCH_CH_r04b09_dsl. diff --git a/model/atmosphere/diffusion/diffusion_tests/mpi_tests/__init__.py b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/__init__.py index e69de29bb..15dfdb009 100644 --- a/model/atmosphere/diffusion/diffusion_tests/mpi_tests/__init__.py +++ b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py index b47f02a8f..16edaf64c 100644 --- a/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py +++ b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py @@ -14,13 +14,19 @@ import pytest +from icon4py.model.atmosphere.diffusion.diffusion import ( + Diffusion, + DiffusionParams, +) +from icon4py.model.common.decomposition.decomposed import ( + DecompositionInfo, + create_exchange, +) +from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.model.common.grid.vertical import VerticalModelParams from icon4py.model.common.test_utils.parallel_helpers import check_comm_size -from ..test_diffusion import _verify_diffusion_fields -from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim -from icon4py.model.common.decomposition.decomposed import DecompositionInfo, create_exchange -from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams +from ..test_diffusion import _verify_diffusion_fields @pytest.mark.mpi @@ -40,8 +46,7 @@ def test_parallel_diffusion( grid_savepoint, metrics_savepoint, interpolation_savepoint, - damping_height - + damping_height, ): check_comm_size(processor_props) print( @@ -60,10 +65,12 @@ def test_parallel_diffusion( print( f"rank={processor_props.rank}/{processor_props.comm_size}: using local grid with {icon_grid.num_cells()} Cells, {icon_grid.num_edges()} Edges, {icon_grid.num_vertices()} Vertices" ) - metric_state= metrics_savepoint.construct_metric_state_for_diffusion() + metric_state = metrics_savepoint.construct_metric_state_for_diffusion() cell_geometry = grid_savepoint.construct_cell_geometry() edge_geometry = grid_savepoint.construct_edge_geometry() - interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() + interpolation_state = ( + interpolation_savepoint.construct_interpolation_state_for_diffusion() + ) diffusion_params = DiffusionParams(r04b09_diffusion_config) dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") @@ -84,7 +91,9 @@ def test_parallel_diffusion( edge_params=edge_geometry, cell_params=cell_geometry, ) - print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ") + print( + f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized " + ) diagnostic_state = diffusion_savepoint_init.construct_diagnostics_for_diffusion() prognostic_state = diffusion_savepoint_init.construct_prognostics() if linit: diff --git a/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_global_to_vn.py index 02e618fcd..1930df665 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_global_to_vn.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_global_to_vn.py @@ -44,8 +44,12 @@ def input_data(self, mesh): ) @staticmethod - def reference(mesh, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn): + def reference( + mesh, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn + ): area_edge = np.expand_dims(area_edge, axis=-1) diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0) - vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge) + vn = vn + area_edge * ( + kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge + ) return dict(vn=vn) diff --git a/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_to_vn.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_to_vn.py index a73086bf0..092a9d945 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_to_vn.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_to_vn.py @@ -14,12 +14,12 @@ import numpy as np import pytest -from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_to_vn import ( - apply_nabla2_and_nabla4_to_vn, -) from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_global_to_vn import ( apply_nabla2_and_nabla4_global_to_vn, ) +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_to_vn import ( + apply_nabla2_and_nabla4_to_vn, +) from icon4py.model.common.dimension import EdgeDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field diff --git a/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w.py index 147fd5c59..8cc82ea2c 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w.py @@ -15,7 +15,9 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w import apply_nabla2_to_w +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w import ( + apply_nabla2_to_w, +) from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field @@ -36,7 +38,9 @@ def reference( ) -> np.array: geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) area = np.expand_dims(area, axis=-1) - w = w - diff_multfac_w * area * area * np.sum(z_nabla2_c[mesh.c2e2cO] * geofac_n2s, axis=1) + w = w - diff_multfac_w * area * area * np.sum( + z_nabla2_c[mesh.c2e2cO] * geofac_n2s, axis=1 + ) return dict(w=w) @pytest.fixture diff --git a/model/atmosphere/diffusion/diffusion_tests/test_calculate_diagnostics_for_turbulence.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_diagnostics_for_turbulence.py index 5e4673053..f7e7831a2 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_calculate_diagnostics_for_turbulence.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_diagnostics_for_turbulence.py @@ -36,7 +36,9 @@ def reference( kc_offset_1 = np.roll(kh_c, shift=1, axis=1) div_offset_1 = np.roll(div, shift=1, axis=1) div_ic[:, 1:] = (wgtfac_c * div + (1.0 - wgtfac_c) * div_offset_1)[:, 1:] - hdef_ic[:, 1:] = ((wgtfac_c * kh_c + (1.0 - wgtfac_c) * kc_offset_1) ** 2)[:, 1:] + hdef_ic[:, 1:] = ((wgtfac_c * kh_c + (1.0 - wgtfac_c) * kc_offset_1) ** 2)[ + :, 1: + ] return dict(div_ic=div_ic, hdef_ic=hdef_ic) @pytest.fixture @@ -46,4 +48,6 @@ def input_data(self, mesh): kh_c = random_field(mesh, CellDim, KDim) div_ic = zero_field(mesh, CellDim, KDim) hdef_ic = zero_field(mesh, CellDim, KDim) - return dict(wgtfac_c=wgtfac_c, div=div, kh_c=kh_c, div_ic=div_ic, hdef_ic=hdef_ic) + return dict( + wgtfac_c=wgtfac_c, div=div, kh_c=kh_c, div_ic=div_ic, hdef_ic=hdef_ic + ) diff --git a/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_of_theta.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_of_theta.py index 79cbb2ba5..0bc301f40 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_of_theta.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_of_theta.py @@ -43,7 +43,9 @@ def test_calculate_nabla2_of_theta(): out = zero_field(mesh, CellDim, KDim) - ref = calculate_nabla2_of_theta_numpy(mesh.c2e, np.asarray(z_nabla2_e), np.asarray(geofac_div)) + ref = calculate_nabla2_of_theta_numpy( + mesh.c2e, np.asarray(z_nabla2_e), np.asarray(geofac_div) + ) calculate_nabla2_of_theta( z_nabla2_e, geofac_div_new, diff --git a/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla4.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla4.py index 5bc20fe4a..a29105933 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla4.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla4.py @@ -14,7 +14,9 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import calculate_nabla4 +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import ( + calculate_nabla4, +) from icon4py.model.common.dimension import ( E2C2VDim, ECVDim, diff --git a/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py b/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py index f9a8a3e96..2a57d665f 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py @@ -14,11 +14,10 @@ import numpy as np import pytest -from icon4py.model.common.test_utils.serialbox_utils import ( - IconDiffusionExitSavepoint, - IconDiffusionInitSavepoint, +from icon4py.model.atmosphere.diffusion.diffusion import ( + Diffusion, + DiffusionParams, ) -from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.model.atmosphere.diffusion.diffusion_states import ( DiffusionDiagnosticState, PrognosticState, @@ -26,6 +25,10 @@ from icon4py.model.atmosphere.diffusion.diffusion_utils import scale_k from icon4py.model.common.grid.horizontal import CellParams, EdgeParams from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.test_utils.serialbox_utils import ( + IconDiffusionExitSavepoint, + IconDiffusionInitSavepoint, +) from .test_diffusion_utils import ( diff_multfac_vn_numpy, diff --git a/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py b/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py index 14753b417..733243bb7 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py @@ -14,9 +14,6 @@ import numpy as np import pytest -from icon4py.model.common.test_utils.helpers import random_field, zero_field -from icon4py.model.common.test_utils.simple_mesh import SimpleMesh -from icon4py.model.common.dimension import KDim, VertexDim from icon4py.model.atmosphere.diffusion.diffusion import DiffusionParams from icon4py.model.atmosphere.diffusion.diffusion_utils import ( _en_smag_fac_for_zero_nshift, @@ -26,6 +23,9 @@ set_zero_v_k, setup_fields_for_initial_step, ) +from icon4py.model.common.dimension import KDim, VertexDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh def diff_multfac_vn_numpy(shape, k4, substeps): diff --git a/model/atmosphere/diffusion/diffusion_tests/test_temporary_fields_for_turbulence_diagnostics.py b/model/atmosphere/diffusion/diffusion_tests/test_temporary_fields_for_turbulence_diagnostics.py index 521214d70..6b044d2fb 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_temporary_fields_for_turbulence_diagnostics.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_temporary_fields_for_turbulence_diagnostics.py @@ -17,13 +17,12 @@ from icon4py.model.atmosphere.diffusion.stencils.temporary_fields_for_turbulence_diagnostics import ( temporary_fields_for_turbulence_diagnostics, ) -from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim, CEDim +from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim from icon4py.model.common.test_utils.helpers import ( StencilTest, + as_1D_sparse_field, random_field, zero_field, - as_1D_sparse_field - ) diff --git a/model/atmosphere/diffusion/diffusion_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/diffusion_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index 3688d917e..b0cf9d517 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -56,12 +56,15 @@ def mo_nh_diffusion_stencil_15_numpy( ] sum_over = np.sum( - geofac_n2s_nbh * (vcoef * theta_v_at_zd_vertidx + (1.0 - vcoef) * theta_v_at_zd_vertidx_p1), + geofac_n2s_nbh + * (vcoef * theta_v_at_zd_vertidx + (1.0 - vcoef) * theta_v_at_zd_vertidx_p1), axis=1, ) geofac_n2s_c = np.expand_dims(geofac_n2s_c, axis=1) # add KDim - return np.where(mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum_over), z_temp) + return np.where( + mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum_over), z_temp + ) def test_mo_nh_diffusion_stencil_15(): diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/__init__.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/__init__.py index 5c4011cc5..49c96a94c 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/__init__.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/__init__.py @@ -14,6 +14,8 @@ from typing import Final from packaging import version as pkg_version + + __all__ = [ "__author__", "__copyright__", @@ -30,4 +32,3 @@ __version__: Final = "0.0.6" __version_info__: Final = pkg_version.parse(__version__) - diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py index 95ba73c96..4a8d8a110 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -29,7 +29,23 @@ run_gtfn_imperative, ) -from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_vn import apply_diffusion_to_vn +from icon4py.model.atmosphere.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, + PrognosticState, +) +from icon4py.model.atmosphere.diffusion.diffusion_utils import ( + copy_field, + init_diffusion_local_fields_for_regular_timestep, + init_nabla2_factor_in_upper_damping_zone, + scale_k, + setup_fields_for_initial_step, + zero_field, +) +from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_vn import ( + apply_diffusion_to_vn, +) from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance, ) @@ -45,36 +61,32 @@ from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_theta import ( calculate_nabla2_for_theta, ) +from icon4py.model.atmosphere.diffusion.stencils.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( + truly_horizontal_diffusion_nabla_of_theta_over_steep_points, +) +from icon4py.model.atmosphere.diffusion.stencils.update_theta_and_exner import ( + update_theta_and_exner, +) + # TODO (magdalena) should go to common.math from icon4py.model.atmosphere.dycore.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) -from icon4py.model.atmosphere.diffusion.stencils.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( - truly_horizontal_diffusion_nabla_of_theta_over_steep_points, -) -from icon4py.model.atmosphere.diffusion.stencils.update_theta_and_exner import update_theta_and_exner from icon4py.model.common.constants import ( CPD, DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO, GAS_CONSTANT_DRY_AIR, ) -from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim -from icon4py.model.common.decomposition.decomposed import ExchangeRuntime, SingleNode -from icon4py.model.atmosphere.diffusion.diffusion_states import ( - DiffusionDiagnosticState, - DiffusionInterpolationState, - DiffusionMetricState, - PrognosticState, +from icon4py.model.common.decomposition.decomposed import ( + ExchangeRuntime, + SingleNode, ) -from icon4py.model.atmosphere.diffusion.diffusion_utils import ( - copy_field, - init_diffusion_local_fields_for_regular_timestep, - init_nabla2_factor_in_upper_damping_zone, - scale_k, - setup_fields_for_initial_step, - zero_field, +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim +from icon4py.model.common.grid.horizontal import ( + CellParams, + EdgeParams, + HorizontalMarkerIndex, ) -from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalMarkerIndex from icon4py.model.common.grid.icon_grid import IconGrid from icon4py.model.common.grid.vertical import VerticalModelParams diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py index bec8a87ae..fc17ed567 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py @@ -242,5 +242,3 @@ def init_nabla2_factor_in_upper_damping_zone( ** 4 ) return np_as_located_field(KDim)(buffer) - - diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/__init__.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/__init__.py index e69de29bb..15dfdb009 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/__init__.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py index 88b74a92c..5fa0d5e50 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py @@ -23,7 +23,9 @@ from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_vn_in_lateral_boundary import ( _apply_nabla2_to_vn_in_lateral_boundary, ) -from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import _calculate_nabla4 +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import ( + _calculate_nabla4, +) from icon4py.model.common.dimension import ECVDim, EdgeDim, KDim, VertexDim @@ -71,7 +73,9 @@ def _apply_diffusion_to_vn( vn, nudgezone_diff, ), - _apply_nabla2_to_vn_in_lateral_boundary(z_nabla2_e, area_edge, vn, fac_bdydiff_v), + _apply_nabla2_to_vn_in_lateral_boundary( + z_nabla2_e, area_edge, vn, fac_bdydiff_v + ), ) if limited_area else where( diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py index dd4f0a264..44f4ab69f 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py @@ -14,7 +14,9 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where -from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w import _apply_nabla2_to_w +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w import ( + _apply_nabla2_to_w, +) from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w_in_upper_damping_layer import ( _apply_nabla2_to_w_in_upper_damping_layer, ) @@ -52,7 +54,9 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( dwdx, dwdy = where( int32(0) < vert_idx, - _calculate_horizontal_gradients_for_turbulence(w_old, geofac_grg_x, geofac_grg_y), + _calculate_horizontal_gradients_for_turbulence( + w_old, geofac_grg_x, geofac_grg_y + ), (dwdx, dwdy), ) @@ -69,7 +73,9 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( & (vert_idx < nrdmax) & (interior_idx <= horz_idx) & (horz_idx < halo_idx), - _apply_nabla2_to_w_in_upper_damping_layer(w, diff_multfac_n2w, area, z_nabla2_c), + _apply_nabla2_to_w_in_upper_damping_layer( + w, diff_multfac_n2w, area, z_nabla2_c + ), w, ) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py index 313e9ac50..aa76bed1b 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py @@ -27,7 +27,9 @@ def _apply_nabla2_and_nabla4_global_to_vn( diff_multfac_vn: Field[[KDim], float], vn: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge) + vn = vn + area_edge * ( + kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge + ) return vn diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py index c2f107489..1710922bd 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py @@ -35,4 +35,6 @@ def apply_nabla2_to_vn_in_lateral_boundary( vn: Field[[EdgeDim, KDim], float], fac_bdydiff_v: float, ): - _apply_nabla2_to_vn_in_lateral_boundary(z_nabla2_e, area_edge, vn, fac_bdydiff_v, out=vn) + _apply_nabla2_to_vn_in_lateral_boundary( + z_nabla2_e, area_edge, vn, fac_bdydiff_v, out=vn + ) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py index 1c686fe58..60c3126a2 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py @@ -38,4 +38,6 @@ def calculate_diagnostics_for_turbulence( div_ic: Field[[CellDim, KDim], float], hdef_ic: Field[[CellDim, KDim], float], ): - _calculate_diagnostics_for_turbulence(div, kh_c, wgtfac_c, out=(div_ic[:, 1:], hdef_ic[:, 1:])) + _calculate_diagnostics_for_turbulence( + div, kh_c, wgtfac_c, out=(div_ic[:, 1:], hdef_ic[:, 1:]) + ) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py index 6a997cd21..cc7eb2642 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py @@ -33,7 +33,9 @@ def _calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( enh_diffu_3d = _temporary_field_for_grid_point_cold_pools_enhancement( theta_v, theta_ref_mc, thresh_tdiff ) - kh_smag_e = _enhance_diffusion_coefficient_for_grid_point_cold_pools(kh_smag_e, enh_diffu_3d) + kh_smag_e = _enhance_diffusion_coefficient_for_grid_point_cold_pools( + kh_smag_e, enh_diffu_3d + ) return kh_smag_e diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py index 37430b544..ca00e06c5 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py @@ -32,4 +32,6 @@ def enhance_diffusion_coefficient_for_grid_point_cold_pools( kh_smag_e: Field[[EdgeDim, KDim], float], enh_diffu_3d: Field[[CellDim, KDim], float], ): - _enhance_diffusion_coefficient_for_grid_point_cold_pools(kh_smag_e, enh_diffu_3d, out=kh_smag_e) + _enhance_diffusion_coefficient_for_grid_point_cold_pools( + kh_smag_e, enh_diffu_3d, out=kh_smag_e + ) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py index 049c1202a..e1e38351a 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py @@ -15,7 +15,15 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.model.common.dimension import C2CE, CEDim, C2E, C2EDim, CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import ( + C2CE, + C2E, + C2EDim, + CEDim, + CellDim, + EdgeDim, + KDim, +) @field_operator diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index b4564f007..0984444ec 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -46,11 +46,20 @@ def _truly_horizontal_diffusion_nabla_of_theta_over_steep_points( sum_over_neighbors = ( geofac_n2s_nbh(C2CEC[0]) - * (vcoef(C2CEC[0]) * theta_v_0(C2E2C[0]) + (1.0 - vcoef(C2CEC[0])) * theta_v_0_m1(C2E2C[0])) + * ( + vcoef(C2CEC[0]) * theta_v_0(C2E2C[0]) + + (1.0 - vcoef(C2CEC[0])) * theta_v_0_m1(C2E2C[0]) + ) + geofac_n2s_nbh(C2CEC[1]) - * (vcoef(C2CEC[1]) * theta_v_1(C2E2C[1]) + (1.0 - vcoef(C2CEC[1])) * theta_v_1_m1(C2E2C[1])) + * ( + vcoef(C2CEC[1]) * theta_v_1(C2E2C[1]) + + (1.0 - vcoef(C2CEC[1])) * theta_v_1_m1(C2E2C[1]) + ) + geofac_n2s_nbh(C2CEC[2]) - * (vcoef(C2CEC[2]) * theta_v_2(C2E2C[2]) + (1.0 - vcoef(C2CEC[2])) * theta_v_2_m1(C2E2C[2])) + * ( + vcoef(C2CEC[2]) * theta_v_2(C2E2C[2]) + + (1.0 - vcoef(C2CEC[2])) * theta_v_2_m1(C2E2C[2]) + ) ) z_temp = where( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/__init__.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/__init__.py index dab708955..6cfd5ac92 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/__init__.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/__init__.py @@ -14,7 +14,6 @@ from packaging import version as pkg_version - __all__ = [ "__author__", "__copyright__", diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py index cd8a2e9bf..06a94e5ec 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py @@ -43,11 +43,13 @@ def _mo_advection_traj_btraj_compute_o1_dsl( p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) z_ntdistv_bary_1 = -( - p_vn * p_dthalf + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) + p_vn * p_dthalf + + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) + p_vt * p_dthalf + + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) ) p_distv_bary_1 = where( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py index 295be5cab..06a4125d8 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py @@ -33,4 +33,6 @@ def mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl( c_intp: Field[[VertexDim, V2CDim], float], p_vert_out: Field[[VertexDim, KDim], float], ): - _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl(p_cell_in, c_intp, out=p_vert_out) + _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl( + p_cell_in, c_intp, out=p_vert_out + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py index ca293f5a9..509a9e615 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py @@ -24,7 +24,9 @@ def _mo_solve_nonhydro_stencil_02( exner_ref_mc: Field[[CellDim, KDim], float], exner_pr: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - z_exner_ex_pr = (1.0 + exner_exfac) * (exner - exner_ref_mc) - exner_exfac * exner_pr + z_exner_ex_pr = (1.0 + exner_exfac) * ( + exner - exner_ref_mc + ) - exner_exfac * exner_pr exner_pr = exner - exner_ref_mc return z_exner_ex_pr, exner_pr diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py index b70912057..bfb6b25f6 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py @@ -43,7 +43,9 @@ def _mo_solve_nonhydro_stencil_10( ]: z_w_backtraj = -(w - w_concorr_c) * dtime * 0.5 / ddqz_z_half z_rho_tavg_m1 = wgt_nnow_rth * rho_now(Koff[-1]) + wgt_nnew_rth * rho_var(Koff[-1]) - z_theta_tavg_m1 = wgt_nnow_rth * theta_now(Koff[-1]) + wgt_nnew_rth * theta_var(Koff[-1]) + z_theta_tavg_m1 = wgt_nnow_rth * theta_now(Koff[-1]) + wgt_nnew_rth * theta_var( + Koff[-1] + ) z_rho_tavg = wgt_nnow_rth * rho_now + wgt_nnew_rth * rho_var z_theta_tavg = wgt_nnow_rth * theta_now + wgt_nnew_rth * theta_var rho_ic = ( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py index eb9ac90b6..a2205e3cb 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py @@ -40,11 +40,13 @@ def _compute_btraj( lvn_pos = where(p_vn > 0.0, True, False) z_ntdistv_bary_1 = -( - p_vn * p_dthalf + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) + p_vn * p_dthalf + + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) + p_vt * p_dthalf + + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) ) p_distv_bary_1 = where( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py index 4c628b66c..f7841c6d0 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py @@ -28,7 +28,10 @@ def _mo_solve_nonhydro_stencil_17( ) -> Field[[EdgeDim, KDim], float]: scalfac_dd3d = broadcast(scalfac_dd3d, (EdgeDim, KDim)) z_graddiv_vn = z_graddiv_vn + ( - hmask_dd3d * scalfac_dd3d * inv_dual_edge_length * (z_dwdz_dd(E2C[1]) - z_dwdz_dd(E2C[0])) + hmask_dd3d + * scalfac_dd3d + * inv_dual_edge_length + * (z_dwdz_dd(E2C[1]) - z_dwdz_dd(E2C[0])) ) return z_graddiv_vn diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py index 5046f22bf..9d99e03ba 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py @@ -23,7 +23,9 @@ def _mo_solve_nonhydro_stencil_18( inv_dual_edge_length: Field[[EdgeDim], float], z_exner_ex_pr: Field[[CellDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - z_gradh_exner = inv_dual_edge_length * (z_exner_ex_pr(E2C[1]) - z_exner_ex_pr(E2C[0])) + z_gradh_exner = inv_dual_edge_length * ( + z_exner_ex_pr(E2C[1]) - z_exner_ex_pr(E2C[0]) + ) return z_gradh_exner diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py index 0477ff1fc..453dfec9c 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py @@ -49,12 +49,18 @@ def _mo_solve_nonhydro_stencil_20( ( z_exner_ex_pr_1(E2C[1]) + zdiff_gradp(E2EC[1]) - * (z_dexner_dz_c1_1(E2C[1]) + zdiff_gradp(E2EC[1]) * z_dexner_dz_c2_1(E2C[1])) + * ( + z_dexner_dz_c1_1(E2C[1]) + + zdiff_gradp(E2EC[1]) * z_dexner_dz_c2_1(E2C[1]) + ) ) - ( z_exner_ex_pr_0(E2C[0]) + zdiff_gradp(E2EC[0]) - * (z_dexner_dz_c1_0(E2C[0]) + zdiff_gradp(E2EC[0]) * z_dexner_dz_c2_0(E2C[0])) + * ( + z_dexner_dz_c1_0(E2C[0]) + + zdiff_gradp(E2EC[0]) * z_dexner_dz_c2_0(E2C[0]) + ) ) ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py index 84fb050f6..f1919b193 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py @@ -25,7 +25,9 @@ def _mo_solve_nonhydro_stencil_22( z_hydro_corr: Field[[EdgeDim], float], z_gradh_exner: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - z_gradh_exner = where(ipeidx_dsl, z_gradh_exner + z_hydro_corr * pg_exdist, z_gradh_exner) + z_gradh_exner = where( + ipeidx_dsl, z_gradh_exner + z_hydro_corr * pg_exdist, z_gradh_exner + ) return z_gradh_exner diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py index a23544d44..66f3c2234 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py @@ -28,7 +28,9 @@ def _mo_solve_nonhydro_stencil_24( dtime: float, cpd: float, ) -> Field[[EdgeDim, KDim], float]: - vn_nnew = vn_nnow + dtime * (ddt_vn_adv_ntl1 + ddt_vn_phy - cpd * z_theta_v_e * z_gradh_exner) + vn_nnew = vn_nnow + dtime * ( + ddt_vn_adv_ntl1 + ddt_vn_phy - cpd * z_theta_v_e * z_gradh_exner + ) return vn_nnew diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py index 84abeee0e..de3b9253b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py @@ -32,7 +32,9 @@ def _mo_solve_nonhydro_stencil_39( wgtfac_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: z_w_concorr_me_offset_1 = z_w_concorr_me(Koff[-1]) - z_w_concorr_mc_m1 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim) + z_w_concorr_mc_m1 = neighbor_sum( + e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim + ) z_w_concorr_mc_m0 = neighbor_sum(e_bln_c_s * z_w_concorr_me(C2E), axis=C2EDim) w_concorr_c = wgtfac_c * z_w_concorr_mc_m0 + (1.0 - wgtfac_c) * z_w_concorr_mc_m1 return w_concorr_c diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py index 6c4016960..a0adf8bb9 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py @@ -35,9 +35,15 @@ def _mo_solve_nonhydro_stencil_40( z_w_concorr_me_offset_2 = z_w_concorr_me(Koff[-2]) z_w_concorr_me_offset_3 = z_w_concorr_me(Koff[-3]) - z_w_concorr_mc_m1 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim) - z_w_concorr_mc_m2 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_2(C2E), axis=C2EDim) - z_w_concorr_mc_m3 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_3(C2E), axis=C2EDim) + z_w_concorr_mc_m1 = neighbor_sum( + e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim + ) + z_w_concorr_mc_m2 = neighbor_sum( + e_bln_c_s * z_w_concorr_me_offset_2(C2E), axis=C2EDim + ) + z_w_concorr_mc_m3 = neighbor_sum( + e_bln_c_s * z_w_concorr_me_offset_3(C2E), axis=C2EDim + ) return ( wgtfacq_c(Koff[-1]) * z_w_concorr_mc_m1 diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py index 5b6b0d51e..8606aa719 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py @@ -33,7 +33,9 @@ def _mo_solve_nonhydro_stencil_42( cpd: float, ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: z_w_expl = w_nnow + dtime * ( - wgt_nnow_vel * ddt_w_adv_ntl1 + wgt_nnew_vel * ddt_w_adv_ntl2 - cpd * z_th_ddz_exner_c + wgt_nnow_vel * ddt_w_adv_ntl1 + + wgt_nnew_vel * ddt_w_adv_ntl2 + - cpd * z_th_ddz_exner_c ) z_contr_w_fl_l = rho_ic * (-w_concorr_c + vwind_expl_wgt * w_nnow) return z_w_expl, z_contr_w_fl_l diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py index a354730e7..d55620484 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py @@ -42,8 +42,17 @@ def _mo_solve_nonhydro_stencil_55( rho_new = z_rho_expl - vwind_impl_wgt * dtime * inv_ddqz_z_full * ( rho_ic * w - rho_ic(Koff[1]) * w(Koff[1]) ) - exner_new = z_exner_expl + exner_ref_mc - z_beta * (z_alpha * w - z_alpha(Koff[1]) * w(Koff[1])) - theta_v_new = rho_now * theta_v_now * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) / rho_new + exner_new = ( + z_exner_expl + + exner_ref_mc + - z_beta * (z_alpha * w - z_alpha(Koff[1]) * w(Koff[1])) + ) + theta_v_new = ( + rho_now + * theta_v_now + * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) + / rho_new + ) return rho_new, exner_new, theta_v_new diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py index e5e70e7ba..b1f718731 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py @@ -24,7 +24,9 @@ def _mo_solve_nonhydro_stencil_56_63( w: Field[[CellDim, KDim], float], w_concorr_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - z_dwdz_dd = inv_ddqz_z_full * ((w - w(Koff[1])) - (w_concorr_c - w_concorr_c(Koff[1]))) + z_dwdz_dd = inv_ddqz_z_full * ( + (w - w(Koff[1])) - (w_concorr_c - w_concorr_c(Koff[1])) + ) return z_dwdz_dd diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py index 17879f2db..8b5347ada 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py @@ -27,7 +27,9 @@ def _mo_solve_nonhydro_stencil_58( mass_flx_ic: Field[[CellDim, KDim], float], r_nsubsteps: float, ) -> Field[[CellDim, KDim], float]: - mass_flx_ic = mass_flx_ic + (r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w)) + mass_flx_ic = mass_flx_ic + ( + r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w) + ) return mass_flx_ic diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py index 995322d72..d0435fdde 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py @@ -26,7 +26,9 @@ def _mo_solve_nonhydro_stencil_60( ndyn_substeps_var: float, dtime: float, ) -> Field[[CellDim, KDim], float]: - exner_dyn_incr = exner - (exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy) + exner_dyn_incr = exner - ( + exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy + ) return exner_dyn_incr diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py index eea826bad..3ff46a988 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py @@ -30,7 +30,9 @@ def _mo_solve_nonhydro_stencil_65( r_nsubsteps: float, ) -> Field[[CellDim, KDim], float]: mass_flx_ic = mass_flx_ic + ( - r_nsubsteps * rho_ic * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) + r_nsubsteps + * rho_ic + * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) ) return mass_flx_ic diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py index 6f46f314d..0c949bd69 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py @@ -39,4 +39,6 @@ def mo_solve_nonhydro_stencil_67( rd_o_cvd: float, rd_o_p0ref: float, ): - _mo_solve_nonhydro_stencil_67(rho, theta_v, exner, rd_o_cvd, rd_o_p0ref, out=(theta_v, exner)) + _mo_solve_nonhydro_stencil_67( + rho, theta_v, exner, rd_o_cvd, rd_o_p0ref, out=(theta_v, exner) + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py index 15230b8af..8725c196c 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py @@ -31,7 +31,10 @@ def _mo_solve_nonhydro_stencil_68( ) -> Field[[CellDim, KDim], float]: theta_v_new = where( mask_prog_halo_c, - rho_now * theta_v_now * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) / rho_new, + rho_now + * theta_v_now + * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) + / rho_new, theta_v_new, ) return theta_v_new diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py index c7c051da0..72ad7e547 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py @@ -37,4 +37,6 @@ def mo_velocity_advection_stencil_04( vt: Field[[EdgeDim, KDim], float], z_w_concorr_me: Field[[EdgeDim, KDim], float], ): - _mo_velocity_advection_stencil_04(vn, ddxn_z_full, ddxt_z_full, vt, out=z_w_concorr_me) + _mo_velocity_advection_stencil_04( + vn, ddxn_z_full, ddxt_z_full, vt, out=z_w_concorr_me + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py index a08de19b3..4b965b6e8 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py @@ -37,7 +37,9 @@ def _mo_velocity_advection_stencil_07( ) -> Field[[EdgeDim, KDim], float]: return vn_ie * inv_dual_edge_length * ( w(E2C[0]) - w(E2C[1]) - ) + z_vt_ie * inv_primal_edge_length * tangent_orientation * (z_w_v(E2V[0]) - z_w_v(E2V[1])) + ) + z_vt_ie * inv_primal_edge_length * tangent_orientation * ( + z_w_v(E2V[0]) - z_w_v(E2V[1]) + ) @program(grid_type=GridType.UNSTRUCTURED) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py index 4dd324603..4cdbc7d73 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py @@ -23,7 +23,9 @@ def _mo_velocity_advection_stencil_10( z_w_concorr_mc: Field[[CellDim, KDim], float], wgtfac_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc(Koff[-1]) + w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc( + Koff[-1] + ) return w_concorr_c diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py index 3bce7d7ef..f413dceae 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py @@ -43,9 +43,13 @@ def _mo_velocity_advection_stencil_14( vcfl = where(cfl_clipping, z_w_con_c * dtime / ddqz_z_half, 0.0) - z_w_con_c = where((cfl_clipping) & (vcfl < -0.85), -0.85 * ddqz_z_half / dtime, z_w_con_c) + z_w_con_c = where( + (cfl_clipping) & (vcfl < -0.85), -0.85 * ddqz_z_half / dtime, z_w_con_c + ) - z_w_con_c = where((cfl_clipping) & (vcfl > 0.85), 0.85 * ddqz_z_half / dtime, z_w_con_c) + z_w_con_c = where( + (cfl_clipping) & (vcfl > 0.85), 0.85 * ddqz_z_half / dtime, z_w_con_c + ) return cfl_clipping, pre_levelmask, vcfl, z_w_con_c diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py index a1823ac6b..38b60dd96 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py @@ -26,7 +26,9 @@ def _mo_velocity_advection_stencil_16( coeff2_dwdz: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: ddt_w_adv = -z_w_con_c * ( - w(Koff[-1]) * coeff1_dwdz - w(Koff[1]) * coeff2_dwdz + w * (coeff2_dwdz - coeff1_dwdz) + w(Koff[-1]) * coeff1_dwdz + - w(Koff[1]) * coeff2_dwdz + + w * (coeff2_dwdz - coeff1_dwdz) ) return ddt_w_adv @@ -39,4 +41,6 @@ def mo_velocity_advection_stencil_16( coeff2_dwdz: Field[[CellDim, KDim], float], ddt_w_adv: Field[[CellDim, KDim], float], ): - _mo_velocity_advection_stencil_16(z_w_con_c, w, coeff1_dwdz, coeff2_dwdz, out=ddt_w_adv[:, 1:]) + _mo_velocity_advection_stencil_16( + z_w_con_c, w, coeff1_dwdz, coeff2_dwdz, out=ddt_w_adv[:, 1:] + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_18.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_18.py index c4e8d32d2..9499469aa 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_18.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_18.py @@ -45,7 +45,8 @@ def _mo_velocity_advection_stencil_18( ddt_w_adv = where( levelmask & cfl_clipping & owner_mask, - ddt_w_adv + difcoef * area * neighbor_sum(w(C2E2CO) * geofac_n2s, axis=C2E2CODim), + ddt_w_adv + + difcoef * area * neighbor_sum(w(C2E2CO) * geofac_n2s, axis=C2E2CODim), ddt_w_adv, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py index 621e6befe..6476e2a57 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py @@ -45,7 +45,10 @@ def _mo_velocity_advection_stencil_19( ): ddt_vn_adv = -( (coeff_gradekin(E2EC[0]) - coeff_gradekin(E2EC[1])) * z_kin_hor_e - + (-coeff_gradekin(E2EC[0]) * z_ekinh(E2C[0]) + coeff_gradekin(E2EC[1]) * z_ekinh(E2C[1])) + + ( + -coeff_gradekin(E2EC[0]) * z_ekinh(E2C[0]) + + coeff_gradekin(E2EC[1]) * z_ekinh(E2C[1]) + ) + vt * (f_e + 0.5 * neighbor_sum(zeta(E2V), axis=E2VDim)) + neighbor_sum(z_w_con_c_full(E2C) * c_lin_e, axis=E2CDim) * (vn_ie - vn_ie(Koff[1])) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_20.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_20.py index 173bd2e9b..abcc04c01 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_20.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_20.py @@ -76,7 +76,9 @@ def _mo_velocity_advection_stencil_20( (levelmask | levelmask(Koff[1])) & (abs(w_con_e) > cfl_w_limit * ddqz_z_full_e), ddt_vn_adv + difcoef * area_edge * neighbor_sum(geofac_grdiv * vn(E2C2EO), axis=E2C2EODim) - + tangent_orientation * inv_primal_edge_length * neighbor_sum(zeta(E2V), axis=E2VDim), + + tangent_orientation + * inv_primal_edge_length + * neighbor_sum(zeta(E2V), axis=E2VDim), ddt_vn_adv, ) return ddt_vn_adv diff --git a/model/atmosphere/dycore/tests/conftest.py b/model/atmosphere/dycore/tests/conftest.py index 64e66161f..8e6147e6d 100644 --- a/model/atmosphere/dycore/tests/conftest.py +++ b/model/atmosphere/dycore/tests/conftest.py @@ -13,17 +13,15 @@ from icon4py.model.common.test_utils.fixtures import ( # noqa F401 + backend, damping_height, data_provider, datapath, + download_ser_data, grid_savepoint, icon_grid, - download_ser_data, - linit, step_date_exit, step_date_init, mesh, backend + linit, + mesh, + step_date_exit, + step_date_init, ) - - - - - - diff --git a/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py index a38c366a5..dff37381e 100644 --- a/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py @@ -54,22 +54,28 @@ def mo_advection_traj_btraj_compute_o1_dsl_numpy( p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) z_ntdistv_bary_1 = -( - p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) + p_vn * p_dthalf + + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) + p_vt * p_dthalf + + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) ) p_distv_bary_1 = np.where( lvn_pos, - z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], - z_ntdistv_bary_1 * primal_normal_cell_1[:, 1] + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], + z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] + + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], + z_ntdistv_bary_1 * primal_normal_cell_1[:, 1] + + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], ) p_distv_bary_2 = np.where( lvn_pos, - z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], - z_ntdistv_bary_1 * primal_normal_cell_2[:, 1] + z_ntdistv_bary_2 * dual_normal_cell_2[:, 1], + z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] + + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], + z_ntdistv_bary_1 * primal_normal_cell_2[:, 1] + + z_ntdistv_bary_2 * dual_normal_cell_2[:, 1], ) return p_cell_idx, p_cell_blk, p_distv_bary_1, p_distv_bary_2 diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py index 4238b31fd..0d5c38403 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py @@ -38,7 +38,9 @@ def reference( exner_exfac: np.array, **kwargs, ) -> dict: - z_exner_ex_pr = (1 + exner_exfac) * (exner - exner_ref_mc) - exner_exfac * exner_pr + z_exner_ex_pr = (1 + exner_exfac) * ( + exner - exner_ref_mc + ) - exner_exfac * exner_pr exner_pr = exner - exner_ref_mc return dict(z_exner_ex_pr=z_exner_ex_pr, exner_pr=exner_pr) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py index 29cb23efe..04df89bb9 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py @@ -37,9 +37,12 @@ def reference( z_exner_ic: np.array, ) -> np.array: z_exner_ic[:, 3:] = ( - np.roll(wgtfacq_c, shift=1, axis=1) * np.roll(z_exner_ex_pr, shift=1, axis=1) - + np.roll(wgtfacq_c, shift=2, axis=1) * np.roll(z_exner_ex_pr, shift=2, axis=1) - + np.roll(wgtfacq_c, shift=3, axis=1) * np.roll(z_exner_ex_pr, shift=3, axis=1) + np.roll(wgtfacq_c, shift=1, axis=1) + * np.roll(z_exner_ex_pr, shift=1, axis=1) + + np.roll(wgtfacq_c, shift=2, axis=1) + * np.roll(z_exner_ex_pr, shift=2, axis=1) + + np.roll(wgtfacq_c, shift=3, axis=1) + * np.roll(z_exner_ex_pr, shift=3, axis=1) )[:, 3:] return {"z_exner_ic": z_exner_ic} diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py index 9a86409fa..2ba664c62 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py @@ -30,7 +30,9 @@ class TestMoSolveNonhydroStencil06(StencilTest): OUTPUTS = ("z_dexner_dz_c_1",) @staticmethod - def reference(mesh, z_exner_ic: np.array, inv_ddqz_z_full: np.array, **kwargs) -> np.array: + def reference( + mesh, z_exner_ic: np.array, inv_ddqz_z_full: np.array, **kwargs + ) -> np.array: z_dexner_dz_c_1 = (z_exner_ic[:, :-1] - z_exner_ic[:, 1:]) * inv_ddqz_z_full return dict(z_dexner_dz_c_1=z_dexner_dz_c_1) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py index bb92f1fbe..40f10d203 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py @@ -26,7 +26,9 @@ class TestMoSolveNonhydroStencil15(StencilTest): OUTPUTS = ("z_rho_e", "z_theta_v_e") @staticmethod - def reference(mesh, z_rho_e: np.array, z_theta_v_e: np.array, **kwargs) -> tuple[np.array]: + def reference( + mesh, z_rho_e: np.array, z_theta_v_e: np.array, **kwargs + ) -> tuple[np.array]: z_rho_e = np.zeros_like(z_rho_e) z_theta_v_e = np.zeros_like(z_theta_v_e) return dict(z_rho_e=z_rho_e, z_theta_v_e=z_theta_v_e) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py index 8282ea91e..4e5e50ed2 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py @@ -45,22 +45,28 @@ def compute_btraj_numpy( dual_normal_cell_2 = np.expand_dims(dual_normal_cell_2, axis=-1) z_ntdistv_bary_1 = -( - p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) + p_vn * p_dthalf + + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) + p_vt * p_dthalf + + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) ) p_distv_bary_1 = np.where( lvn_pos, - z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], - z_ntdistv_bary_1 * primal_normal_cell_1[:, 1] + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], + z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] + + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], + z_ntdistv_bary_1 * primal_normal_cell_1[:, 1] + + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], ) p_distv_bary_2 = np.where( lvn_pos, - z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], - z_ntdistv_bary_1 * primal_normal_cell_2[:, 1] + z_ntdistv_bary_2 * dual_normal_cell_2[:, 1], + z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] + + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], + z_ntdistv_bary_1 * primal_normal_cell_2[:, 1] + + z_ntdistv_bary_2 * dual_normal_cell_2[:, 1], ) return p_distv_bary_1, p_distv_bary_2 @@ -194,7 +200,10 @@ def test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1(): z_rho_e = random_field(mesh, EdgeDim, KDim) z_theta_v_e = random_field(mesh, EdgeDim, KDim) - (z_rho_e_ref, z_theta_v_e_ref,) = mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1_numpy( + ( + z_rho_e_ref, + z_theta_v_e_ref, + ) = mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1_numpy( mesh.e2c, np.asarray(p_vn), np.asarray(p_vt), diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py index 5a782f7c5..56d05c7b9 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py @@ -41,8 +41,9 @@ def reference( z_exner_ex_pr_e2c = z_exner_ex_pr[mesh.e2c] z_exner_ex_weighted = z_exner_ex_pr_e2c[:, 1] - z_exner_ex_pr_e2c[:, 0] - z_gradh_exner = inv_dual_edge_length * z_exner_ex_weighted - ddxn_z_full * np.sum( - c_lin_e * z_dexner_dz_c_1[mesh.e2c], axis=1 + z_gradh_exner = ( + inv_dual_edge_length * z_exner_ex_weighted + - ddxn_z_full * np.sum(c_lin_e * z_dexner_dz_c_1[mesh.e2c], axis=1) ) return dict(z_gradh_exner=z_gradh_exner) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py index 349355cd7..8e5bc613b 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py @@ -51,8 +51,12 @@ def _apply_index_field(shape, to_index, neighbor_table, offset_field): inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, -1) z_exner_ex_pr_at_kidx = _apply_index_field(full_shape, z_exner_ex_pr, e2c, ikoffset) - z_dexner_dz_c_1_at_kidx = _apply_index_field(full_shape, z_dexner_dz_c_1, e2c, ikoffset) - z_dexner_dz_c_2_at_kidx = _apply_index_field(full_shape, z_dexner_dz_c_2, e2c, ikoffset) + z_dexner_dz_c_1_at_kidx = _apply_index_field( + full_shape, z_dexner_dz_c_1, e2c, ikoffset + ) + z_dexner_dz_c_2_at_kidx = _apply_index_field( + full_shape, z_dexner_dz_c_2, e2c, ikoffset + ) def at_neighbor(i): return z_exner_ex_pr_at_kidx[:, i, :] + zdiff_gradp[:, i, :] * ( diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py index e0f31fb5e..4423ac168 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py @@ -61,7 +61,9 @@ def _apply_index_field(shape, to_index, neighbor_table, offset_field): full_shape, theta_v_ic, e2c, ikoffset ) - inv_ddqz_z_full_at_kidx, _ = _apply_index_field(full_shape, inv_ddqz_z_full, e2c, ikoffset) + inv_ddqz_z_full_at_kidx, _ = _apply_index_field( + full_shape, inv_ddqz_z_full, e2c, ikoffset + ) z_theta1 = ( theta_v_at_kidx[:, 0, :] diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py index 18623f00d..cac9fbecc 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py @@ -30,7 +30,9 @@ class TestMoSolveNonhydroStencil25(StencilTest): OUTPUTS = ("z_graddiv2_vn",) @staticmethod - def reference(mesh, geofac_grdiv: np.array, z_graddiv_vn: np.array, **kwargs) -> np.array: + def reference( + mesh, geofac_grdiv: np.array, z_graddiv_vn: np.array, **kwargs + ) -> np.array: geofac_grdiv = np.expand_dims(geofac_grdiv, axis=-1) z_graddiv2_vn = np.sum(z_graddiv_vn[mesh.e2c2eO] * geofac_grdiv, axis=1) return dict(z_graddiv2_vn=z_graddiv2_vn) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py index ec2389a8d..dded703d1 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py @@ -26,7 +26,9 @@ class TestMoSolveNonhydroStencil26(StencilTest): OUTPUTS = ("vn",) @staticmethod - def reference(mesh, z_graddiv_vn: np.array, vn: np.array, scal_divdamp_o2, **kwargs) -> dict: + def reference( + mesh, z_graddiv_vn: np.array, vn: np.array, scal_divdamp_o2, **kwargs + ) -> dict: vn = vn + (scal_divdamp_o2 * z_graddiv_vn) return dict(vn=vn) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py index 17234d448..2eeba98f4 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py @@ -26,7 +26,9 @@ class TestMoSolveNonhydroStencil28(StencilTest): OUTPUTS = ("vn",) @staticmethod - def reference(mesh, vn_incr: np.array, vn: np.array, iau_wgt_dyn, **kwargs) -> np.array: + def reference( + mesh, vn_incr: np.array, vn: np.array, iau_wgt_dyn, **kwargs + ) -> np.array: vn = vn + (iau_wgt_dyn * vn_incr) return dict(vn=vn) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py index da69d23da..deb5bf35c 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py @@ -30,7 +30,9 @@ class TestMoSolveNonhydroStencil29(StencilTest): OUTPUTS = ("vn_new",) @staticmethod - def reference(mesh, grf_tend_vn: np.array, vn_now: np.array, dtime, **kwargs) -> dict: + def reference( + mesh, grf_tend_vn: np.array, vn_now: np.array, dtime, **kwargs + ) -> dict: vn_new = vn_now + dtime * grf_tend_vn return dict(vn_new=vn_new) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py index 7e0ea0bdd..7af31248f 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py @@ -46,7 +46,9 @@ def reference( **kwargs, ) -> tuple[np.array]: z_w_expl = w_nnow + dtime * ( - wgt_nnow_vel * ddt_w_adv_ntl1 + wgt_nnew_vel * ddt_w_adv_ntl2 - cpd * z_th_ddz_exner_c + wgt_nnow_vel * ddt_w_adv_ntl1 + + wgt_nnew_vel * ddt_w_adv_ntl2 + - cpd * z_th_ddz_exner_c ) vwind_expl_wgt = np.expand_dims(vwind_expl_wgt, axis=-1) z_contr_w_fl_l = rho_ic * (-w_concorr_c + vwind_expl_wgt * w_nnow) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py index 2660d320e..6624930d7 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py @@ -44,7 +44,9 @@ def reference( cvd, **kwargs, ) -> dict: - z_beta = dtime * rd * exner_nnow / (cvd * rho_nnow * theta_v_nnow) * inv_ddqz_z_full + z_beta = ( + dtime * rd * exner_nnow / (cvd * rho_nnow * theta_v_nnow) * inv_ddqz_z_full + ) vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) z_alpha = vwind_impl_wgt * theta_v_ic * rho_ic diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py index 60253bc78..3f477369d 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py @@ -30,7 +30,9 @@ class TestMoSolveNonhydroStencil47(StencilTest): OUTPUTS = ("w_nnew", "z_contr_w_fl_l") @staticmethod - def reference(mesh, w_concorr_c: np.array, z_contr_w_fl_l: np.array, **kwargs) -> dict: + def reference( + mesh, w_concorr_c: np.array, z_contr_w_fl_l: np.array, **kwargs + ) -> dict: w_nnew = w_concorr_c z_contr_w_fl_l = np.zeros_like(z_contr_w_fl_l) return dict(w_nnew=w_nnew, z_contr_w_fl_l=z_contr_w_fl_l) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py index f7e12c96f..d1cb07a66 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py @@ -57,7 +57,9 @@ def mo_solve_nonhydro_stencil_51_numpy( z_c = -z_gamma * z_beta * z_alpha_k_plus_1 z_b = 1.0 + z_gamma * z_alpha[:, :-1] * (z_beta_k_minus_1 + z_beta) z_q = mo_solve_nonhydro_stencil_51_z_q_numpy(z_c, z_b) - w_nnew = mo_solve_nonhydro_stencil_51_w_nnew_numpy(z_gamma, z_b, z_w_expl, z_exner_expl) + w_nnew = mo_solve_nonhydro_stencil_51_w_nnew_numpy( + z_gamma, z_b, z_w_expl, z_exner_expl + ) return z_q, w_nnew diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py index 6c0cc28e0..85b153189 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py @@ -49,11 +49,15 @@ def mo_solve_nonhydro_stencil_52_numpy( for k in range(1, k_size): z_a[:, k] = -z_gamma[:, k] * z_beta[:, k - 1] * z_alpha[:, k - 1] z_c[:, k] = -z_gamma[:, k] * z_beta[:, k] * z_alpha[:, k + 1] - z_b[:, k] = 1.0 + z_gamma[:, k] * z_alpha[:, k] * (z_beta[:, k - 1] + z_beta[:, k]) + z_b[:, k] = 1.0 + z_gamma[:, k] * z_alpha[:, k] * ( + z_beta[:, k - 1] + z_beta[:, k] + ) z_g[:, k] = 1.0 / (z_b[:, k] + z_a[:, k] * z_q[:, k - 1]) z_q[:, k] = -z_c[:, k] * z_g[:, k] - w[:, k] = z_w_expl[:, k] - z_gamma[:, k] * (z_exner_expl[:, k - 1] - z_exner_expl[:, k]) + w[:, k] = z_w_expl[:, k] - z_gamma[:, k] * ( + z_exner_expl[:, k - 1] - z_exner_expl[:, k] + ) w[:, k] = (w[:, k] - z_a[:, k] * w[:, k - 1]) * z_g[:, k] return z_q, w diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py index eeb17c2cc..bc2e501a2 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py @@ -26,7 +26,9 @@ class TestMoSolveNonhydroStencil54(StencilTest): OUTPUTS = ("w",) @staticmethod - def reference(mesh, z_raylfac: np.array, w_1: np.array, w: np.array, **kwargs) -> np.array: + def reference( + mesh, z_raylfac: np.array, w_1: np.array, w: np.array, **kwargs + ) -> np.array: z_raylfac = np.expand_dims(z_raylfac, axis=0) w_1 = np.expand_dims(w_1, axis=-1) w = z_raylfac * w + (1.0 - z_raylfac) * w_1 diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py index 1d9e12649..cd6a2683c 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py @@ -62,7 +62,10 @@ def reference( - z_beta * (z_alpha[:, :-1] * w_offset_0 - z_alpha_offset_1 * w_offset_1) ) theta_v_new = ( - rho_now * theta_v_now * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) / rho_new + rho_now + * theta_v_now + * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) + / rho_new ) return dict(rho_new=rho_new, exner_new=exner_new, theta_v_new=theta_v_new) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py index dff001052..93e8c6a95 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py @@ -37,7 +37,9 @@ def reference( **kwargs, ) -> dict: vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) - mass_flx_ic = mass_flx_ic + (r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w)) + mass_flx_ic = mass_flx_ic + ( + r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w) + ) return dict(mass_flx_ic=mass_flx_ic) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py index 634446cb3..407dc06e9 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py @@ -35,7 +35,9 @@ def reference( dtime: float, **kwargs, ) -> np.array: - exner_dyn_incr = exner - (exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy) + exner_dyn_incr = exner - ( + exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy + ) return dict(exner_dyn_incr=exner_dyn_incr) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py index 41ef7c943..8ff168cd3 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py @@ -30,7 +30,9 @@ class TestMoSolveNonhydroStencil62(StencilTest): OUTPUTS = ("w_new",) @staticmethod - def reference(mesh, w_now: np.array, grf_tend_w: np.array, dtime: float, **kwargs) -> np.array: + def reference( + mesh, w_now: np.array, grf_tend_w: np.array, dtime: float, **kwargs + ) -> np.array: w_new = w_now + dtime * grf_tend_w return dict(w_new=w_new) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py index b5b5b9c0c..77e46f592 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py @@ -41,7 +41,9 @@ def reference( vwind_expl_wgt = np.expand_dims(vwind_expl_wgt, axis=-1) vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) mass_flx_ic = mass_flx_ic + ( - r_nsubsteps * rho_ic * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) + r_nsubsteps + * rho_ic + * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) ) return dict(mass_flx_ic=mass_flx_ic) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py index dd54adca2..297ec945d 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py @@ -46,7 +46,10 @@ def reference( theta_v_new = np.where( mask_prog_halo_c, - rho_now * theta_v_now * ((exner_new / exner_now - 1) * cvd_o_rd + 1.0) / rho_new, + rho_now + * theta_v_now + * ((exner_new / exner_now - 1) * cvd_o_rd + 1.0) + / rho_new, theta_v_new, ) return dict(theta_v_new=theta_v_new) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py index 9bdc11bf2..aa64258d7 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py @@ -21,13 +21,17 @@ from icon4py.model.common.test_utils.simple_mesh import SimpleMesh -def mo_velocity_advection_stencil_02_vn_ie_numpy(wgtfac_e: np.array, vn: np.array) -> np.array: +def mo_velocity_advection_stencil_02_vn_ie_numpy( + wgtfac_e: np.array, vn: np.array +) -> np.array: vn_ie_k_minus_1 = np.roll(vn, shift=1, axis=1) vn_ie = wgtfac_e * vn + (1.0 - wgtfac_e) * vn_ie_k_minus_1 return vn_ie -def mo_velocity_advection_stencil_02_z_kin_hor_e_numpy(vn: np.array, vt: np.array) -> np.array: +def mo_velocity_advection_stencil_02_z_kin_hor_e_numpy( + vn: np.array, vt: np.array +) -> np.array: z_kin_hor_e = 0.5 * (vn * vn + vt * vt) return z_kin_hor_e diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_03.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_03.py index f32c87b1b..d7c3f4189 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_03.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_03.py @@ -21,7 +21,9 @@ from icon4py.model.common.test_utils.simple_mesh import SimpleMesh -def mo_velocity_advection_stencil_03_numpy(wgtfac_e: np.array, vt: np.array) -> np.array: +def mo_velocity_advection_stencil_03_numpy( + wgtfac_e: np.array, vt: np.array +) -> np.array: vt_k_minus_1 = np.roll(vt, shift=1, axis=1) z_vt_ie = wgtfac_e * vt + (1.0 - wgtfac_e) * vt_k_minus_1 @@ -36,7 +38,9 @@ def test_mo_velocity_advection_stencil_03(): z_vt_ie = zero_field(mesh, EdgeDim, KDim) - z_vt_ie_ref = mo_velocity_advection_stencil_03_numpy(np.asarray(wgtfac_e), np.asarray(vt)) + z_vt_ie_ref = mo_velocity_advection_stencil_03_numpy( + np.asarray(wgtfac_e), np.asarray(vt) + ) mo_velocity_advection_stencil_03( wgtfac_e, vt, diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_06.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_06.py index 1acd23aeb..f1eefbb0c 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_06.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_06.py @@ -21,7 +21,9 @@ from icon4py.model.common.test_utils.simple_mesh import SimpleMesh -def mo_velocity_advection_stencil_06_numpy(wgtfacq_e: np.array, vn: np.array) -> np.array: +def mo_velocity_advection_stencil_06_numpy( + wgtfacq_e: np.array, vn: np.array +) -> np.array: vn_k_minus_1 = np.roll(vn, shift=1, axis=1) vn_k_minus_2 = np.roll(vn, shift=2, axis=1) vn_k_minus_3 = np.roll(vn, shift=3, axis=1) @@ -45,7 +47,9 @@ def test_mo_velocity_advection_stencil_06(): vn_ie = zero_field(mesh, EdgeDim, KDim) - vn_ie_ref = mo_velocity_advection_stencil_06_numpy(np.asarray(wgtfacq_e), np.asarray(vn)) + vn_ie_ref = mo_velocity_advection_stencil_06_numpy( + np.asarray(wgtfacq_e), np.asarray(vn) + ) mo_velocity_advection_stencil_06( wgtfacq_e, vn, diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py index 192de3f3c..80b8a8b93 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py @@ -30,7 +30,9 @@ class TestMoVelocityAdvectionStencil08(StencilTest): OUTPUTS = ("z_ekinh",) @staticmethod - def reference(mesh, z_kin_hor_e: np.array, e_bln_c_s: np.array, **kwargs) -> np.array: + def reference( + mesh, z_kin_hor_e: np.array, e_bln_c_s: np.array, **kwargs + ) -> np.array: e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) z_ekinh = np.sum(z_kin_hor_e[mesh.c2e] * e_bln_c_s, axis=1) return dict(z_ekinh=z_ekinh) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py index 664b16856..e02b7c37e 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py @@ -30,7 +30,9 @@ class TestMoVelocityAdvectionStencil09(StencilTest): OUTPUTS = ("z_w_concorr_mc",) @staticmethod - def reference(mesh, z_w_concorr_me: np.array, e_bln_c_s: np.array, **kwargs) -> np.array: + def reference( + mesh, z_w_concorr_me: np.array, e_bln_c_s: np.array, **kwargs + ) -> np.array: e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) z_w_concorr_mc = np.sum(z_w_concorr_me[mesh.c2e] * e_bln_c_s, axis=1) return dict(z_w_concorr_mc=z_w_concorr_mc) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py index 33c9bfe84..19f5e0f32 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py @@ -25,7 +25,9 @@ def mo_velocity_advection_stencil_10_numpy( wgtfac_c: np.array, z_w_concorr_mc: np.array ) -> np.array: z_w_concorr_mc_k_minus_1 = np.roll(z_w_concorr_mc, shift=1, axis=1) - w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc_k_minus_1 + w_concorr_c = ( + wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc_k_minus_1 + ) return w_concorr_c diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py index eb7fb3fc4..336ce24d3 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py @@ -26,7 +26,9 @@ class TestMoVelocityAdvectionStencil13(StencilTest): OUTPUTS = ("z_w_con_c",) @staticmethod - def reference(mesh, w_concorr_c: np.array, z_w_con_c: np.array, **kwargs) -> np.array: + def reference( + mesh, w_concorr_c: np.array, z_w_con_c: np.array, **kwargs + ) -> np.array: z_w_con_c = z_w_con_c - w_concorr_c return dict(z_w_con_c=z_w_con_c) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py index eaa6532be..48d9cb6b3 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py @@ -66,7 +66,10 @@ def mo_velocity_advection_stencil_19_numpy( ddt_vn_adv = -( (coeff_gradekin[:, 0] - coeff_gradekin[:, 1]) * z_kin_hor_e - + (-coeff_gradekin[:, 0] * z_ekinh_e2c[:, 0] + coeff_gradekin[:, 1] * z_ekinh_e2c[:, 1]) + + ( + -coeff_gradekin[:, 0] * z_ekinh_e2c[:, 0] + + coeff_gradekin[:, 1] * z_ekinh_e2c[:, 1] + ) + vt * (f_e + 0.5 * np.sum(zeta[e2v], axis=1)) + np.sum(z_w_con_c_full[e2c] * c_lin_e, axis=1) * (vn_ie[:, :-1] - vn_ie[:, 1:]) diff --git a/model/common/src/icon4py/model/common/decomposition/__init__.py b/model/common/src/icon4py/model/common/decomposition/__init__.py index e69de29bb..15dfdb009 100644 --- a/model/common/src/icon4py/model/common/decomposition/__init__.py +++ b/model/common/src/icon4py/model/common/decomposition/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/model/common/src/icon4py/model/common/decomposition/decomposed.py b/model/common/src/icon4py/model/common/decomposition/decomposed.py index d19fdd915..d46e52a15 100644 --- a/model/common/src/icon4py/model/common/decomposition/decomposed.py +++ b/model/common/src/icon4py/model/common/decomposition/decomposed.py @@ -21,8 +21,13 @@ import numpy.ma as ma from gt4py.next import Dimension -from icon4py.model.common.dimension import CellDim, DimensionKind, EdgeDim, VertexDim from icon4py.model.common.decomposition.parallel_setup import ProcessProperties +from icon4py.model.common.dimension import ( + CellDim, + DimensionKind, + EdgeDim, + VertexDim, +) from icon4py.model.common.utils import builder diff --git a/model/common/src/icon4py/model/common/decomposition/parallel_setup.py b/model/common/src/icon4py/model/common/decomposition/parallel_setup.py index f848edd5d..877885002 100644 --- a/model/common/src/icon4py/model/common/decomposition/parallel_setup.py +++ b/model/common/src/icon4py/model/common/decomposition/parallel_setup.py @@ -85,12 +85,12 @@ def from_single_node(cls): class ParallelLogger(logging.Filter): - def __init__(self, processProperties: ProcessProperties = None): + def __init__(self, process_properties: ProcessProperties = None): super().__init__() self._rank_info = "" - if processProperties and processProperties.comm_size > 1: - self._rank_info = f"rank={processProperties.rank}/{processProperties.comm_size} [{processProperties.comm_name}] " + if process_properties and process_properties.comm_size > 1: + self._rank_info = f"rank={process_properties.rank}/{process_properties.comm_size} [{process_properties.comm_name}] " - def filter(self, record: logging.LogRecord) -> bool: + def filter(self, record: logging.LogRecord) -> bool: # noqa: A003 record.rank = self._rank_info return True diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py index 3983cb423..56bc1dcaf 100644 --- a/model/common/src/icon4py/model/common/grid/grid_manager.py +++ b/model/common/src/icon4py/model/common/grid/grid_manager.py @@ -36,7 +36,11 @@ VertexDim, ) from icon4py.model.common.grid.horizontal import HorizontalGridSize -from icon4py.model.common.grid.icon_grid import GridConfig, IconGrid, VerticalGridSize +from icon4py.model.common.grid.icon_grid import ( + GridConfig, + IconGrid, + VerticalGridSize, +) class GridFileName(str, Enum): diff --git a/model/common/src/icon4py/model/common/grid/icon_grid.py b/model/common/src/icon4py/model/common/grid/icon_grid.py index 4153486cb..3332ef405 100644 --- a/model/common/src/icon4py/model/common/grid/icon_grid.py +++ b/model/common/src/icon4py/model/common/grid/icon_grid.py @@ -59,7 +59,6 @@ def num_cells(self): return self.horizontal_config.num_cells - class IconGrid: def __init__(self): self.config: GridConfig = None diff --git a/model/common/src/icon4py/model/common/test_utils/fixtures.py b/model/common/src/icon4py/model/common/test_utils/fixtures.py index 4bae1fdd2..66921a308 100644 --- a/model/common/src/icon4py/model/common/test_utils/fixtures.py +++ b/model/common/src/icon4py/model/common/test_utils/fixtures.py @@ -16,10 +16,11 @@ import pytest from gt4py.next.program_processors.runners.roundtrip import executor +from ..decomposition.parallel_setup import get_processor_properties from .data_handling import download_and_extract from .serialbox_utils import IconSerialDataProvider from .simple_mesh import SimpleMesh -from ..decomposition.parallel_setup import get_processor_properties + test_utils = Path(__file__).parent model = test_utils.parent.parent @@ -35,7 +36,6 @@ ser_data_basepath = base_path.joinpath("ser_icondata") - @pytest.fixture(params=[False], scope="session") def processor_props(request): with_mpi = request.param @@ -44,7 +44,8 @@ def processor_props(request): @pytest.fixture(scope="session") def ranked_data_path(processor_props): - return ser_data_basepath.absolute().joinpath(f"mpitask{processor_props.comm_size}") + return ser_data_basepath.absolute().joinpath(f"mpitask{processor_props.comm_size}") + @pytest.fixture(scope="session") def datapath(ranked_data_path): @@ -74,11 +75,16 @@ def download_ser_data(request, processor_props, ranked_data_path): ) - - @pytest.fixture(scope="session") -def data_provider(download_ser_data, datapath, processor_props) -> IconSerialDataProvider: - return IconSerialDataProvider(fname_prefix="icon_pydycore", path=str(datapath), mpi_rank=processor_props.rank, do_print=True) +def data_provider( + download_ser_data, datapath, processor_props +) -> IconSerialDataProvider: + return IconSerialDataProvider( + fname_prefix="icon_pydycore", + path=str(datapath), + mpi_rank=processor_props.rank, + do_print=True, + ) @pytest.fixture @@ -100,12 +106,12 @@ def icon_grid(grid_savepoint): def decomposition_info(data_provider): return data_provider.from_savepoint_grid().construct_decomposition_info() + @pytest.fixture def damping_height(): return 12500 - @pytest.fixture def ndyn_substeps(): """ @@ -147,21 +153,18 @@ def step_date_exit(): return "2021-06-20T12:00:10.000" - @pytest.fixture -def interpolation_savepoint(data_provider): # noqa F811 +def interpolation_savepoint(data_provider): # F811 """Load data from ICON interplation state savepoint.""" return data_provider.from_interpolation_savepoint() @pytest.fixture -def metrics_savepoint(data_provider): # noqa F811 +def metrics_savepoint(data_provider): # F811 """Load data from ICON mestric state savepoint.""" return data_provider.from_metrics_savepoint() - - BACKENDS = {"embedded": executor} MESHES = {"simple_mesh": SimpleMesh()} diff --git a/model/common/src/icon4py/model/common/test_utils/helpers.py b/model/common/src/icon4py/model/common/test_utils/helpers.py index eee6ac3cf..9060bb89c 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -11,8 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from typing import ClassVar -from typing import Optional +from typing import ClassVar, Optional import numpy as np import numpy.typing as npt @@ -21,6 +20,7 @@ from gt4py.next.ffront.decorator import Program from gt4py.next.iterator import embedded as it_embedded + try: import pytest_benchmark except ModuleNotFoundError: @@ -66,7 +66,9 @@ def random_field( extend: Optional[dict[gt_common.Dimension, int]] = None, ) -> it_embedded.MutableLocatedField: return it_embedded.np_as_located_field(*dims)( - np.random.default_rng().uniform(low=low, high=high, size=_shape(mesh, *dims, extend=extend)) + np.random.default_rng().uniform( + low=low, high=high, size=_shape(mesh, *dims, extend=extend) + ) ) @@ -113,7 +115,9 @@ def flatten_first_two_dims( def _test_validation(self, mesh, backend, input_data): - reference_outputs = self.reference(mesh, **{k: np.array(v) for k, v in input_data.items()}) + reference_outputs = self.reference( + mesh, **{k: np.array(v) for k, v in input_data.items()} + ) self.PROGRAM.with_backend(backend)( **input_data, offset_provider=mesh.get_offset_provider(), diff --git a/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py b/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py index b8b1ed965..7fce8ba16 100644 --- a/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py @@ -17,8 +17,8 @@ from icon4py.model.common.decomposition.parallel_setup import ProcessProperties - -def check_comm_size(props:ProcessProperties, sizes=[1,2,4]): +def check_comm_size(props: ProcessProperties, sizes=(1, 2, 4)): if props.comm_size not in sizes: - pytest.xfail(f"wrong comm size: {props.comm_size}: test only works for comm-sizes: {sizes}") - + pytest.xfail( + f"wrong comm size: {props.comm_size}: test only works for comm-sizes: {sizes}" + ) 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 4a17c7542..b4e50b4f6 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 @@ -18,7 +18,15 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import np_as_located_field +from icon4py.model.atmosphere.diffusion.diffusion import VectorTuple +from icon4py.model.atmosphere.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, + PrognosticState, +) from icon4py.model.common import dimension +from icon4py.model.common.decomposition.decomposed import DecompositionInfo from icon4py.model.common.dimension import ( C2E2CDim, C2E2CODim, @@ -35,17 +43,16 @@ V2EDim, VertexDim, ) -from icon4py.model.common.decomposition.decomposed import DecompositionInfo -from icon4py.model.atmosphere.diffusion.diffusion import VectorTuple -from icon4py.model.atmosphere.diffusion.diffusion_states import ( - DiffusionDiagnosticState, - DiffusionInterpolationState, - DiffusionMetricState, - PrognosticState, +from icon4py.model.common.grid.horizontal import ( + CellParams, + EdgeParams, + HorizontalGridSize, +) +from icon4py.model.common.grid.icon_grid import ( + GridConfig, + IconGrid, + VerticalGridSize, ) -from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalGridSize -from icon4py.model.common.grid.icon_grid import GridConfig, IconGrid, VerticalGridSize - from icon4py.model.common.test_utils.helpers import as_1D_sparse_field @@ -209,7 +216,7 @@ def refin_ctrl(self, dim: Dimension): def num(self, dim: Dimension): return self.sizes[dim] - def _read_field_for_dim(self, field_name, read_func, dim:Dimension): + def _read_field_for_dim(self, field_name, read_func, dim: Dimension): match (dim): case dimension.CellDim: return read_func(f"c_{field_name}") diff --git a/model/common/src/icon4py/model/common/utils.py b/model/common/src/icon4py/model/common/utils.py index 6b24ee54e..403bf7a68 100644 --- a/model/common/src/icon4py/model/common/utils.py +++ b/model/common/src/icon4py/model/common/utils.py @@ -1,5 +1,20 @@ +# 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 + + def builder(func): """Use as decorator on builder functions.""" + def wrapper(self, *args, **kwargs): func(self, *args, **kwargs) return self diff --git a/model/common/tests/conftest.py b/model/common/tests/conftest.py index 764dfd166..684c0c381 100644 --- a/model/common/tests/conftest.py +++ b/model/common/tests/conftest.py @@ -14,17 +14,20 @@ from icon4py.model.common.test_utils.data_handling import download_and_extract from icon4py.model.common.test_utils.fixtures import ( # noqa F401 + base_path, damping_height, data_provider, datapath, + decomposition_info, + download_ser_data, grid_savepoint, icon_grid, - download_ser_data, - processor_props, interpolation_savepoint, - ranked_data_path, decomposition_info, base_path - + processor_props, + ranked_data_path, ) + + grids_path = base_path.joinpath("grids") r04b09_dsl_grid_path = grids_path.joinpath("mch_ch_r04b09_dsl") r04b09_dsl_data_file = r04b09_dsl_grid_path.joinpath( @@ -41,10 +44,12 @@ ) r02b04_global_grid_uri = "https://polybox.ethz.ch/index.php/s/0EM8O8U53GKGsst/download" + @pytest.fixture() def r04b09_dsl_gridfile(get_grid_files): return r04b09_dsl_grid_path.joinpath("grid.nc") + @pytest.fixture(scope="session") def get_grid_files(): """ diff --git a/model/common/tests/mpi_tests/README.md b/model/common/tests/mpi_tests/README.md index 0020ffbdb..a3e870bb3 100644 --- a/model/common/tests/mpi_tests/README.md +++ b/model/common/tests/mpi_tests/README.md @@ -33,7 +33,6 @@ Run the tests with mpirun -np 2 pytest -v -s --with-mpi tests/mpi_tests/ ``` -Test using serialized test data can be run with either 1, 2 or 4 nodes. +Test using serialized test data can be run with either 1, 2 or 4 nodes. -When you run the tests for the first time -the serialized test data will be downloaded from an online storage. That make some time... +When you run the tests for the first time the serialized test data will be downloaded from an online storage. That make some time... diff --git a/model/common/tests/mpi_tests/test_decomposed.py b/model/common/tests/mpi_tests/test_decomposed.py index 9c2dac9ba..27c5be973 100644 --- a/model/common/tests/mpi_tests/test_decomposed.py +++ b/model/common/tests/mpi_tests/test_decomposed.py @@ -14,8 +14,6 @@ import numpy as np import pytest - -from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.model.common.decomposition.decomposed import ( DecompositionInfo, DomainDescriptorIdGenerator, @@ -23,8 +21,10 @@ SingleNode, create_exchange, ) +from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.model.common.test_utils.parallel_helpers import check_comm_size + """ running tests with mpi: @@ -34,10 +34,13 @@ """ + + @pytest.mark.parametrize("processor_props", [True], indirect=True) def test_props(processor_props): assert processor_props.comm + @pytest.mark.mpi(min_size=2) @pytest.mark.parametrize("processor_props", [True], indirect=True) @pytest.mark.parametrize( @@ -49,7 +52,13 @@ def test_props(processor_props): ), ) def test_decomposition_info_masked( - dim, owned, total, caplog, download_ser_data, decomposition_info, processor_props, # noqa F811 + dim, + owned, + total, + caplog, + download_ser_data, + decomposition_info, + processor_props, # F811 ): check_comm_size(processor_props, sizes=[2]) my_rank = processor_props.rank @@ -70,7 +79,7 @@ def test_decomposition_info_masked( _assert_index_partitioning(all_indices, halo_indices, owned_indices) -#@pytest.mark.skipif(props.comm_size != 2, reason="runs on 2 nodes only") +# @pytest.mark.skipif(props.comm_size != 2, reason="runs on 2 nodes only") def _assert_index_partitioning(all_indices, halo_indices, owned_indices): owned_list = owned_indices.tolist() halos_list = halo_indices.tolist() @@ -92,7 +101,13 @@ def _assert_index_partitioning(all_indices, halo_indices, owned_indices): ) @pytest.mark.mpi(min_size=2) def test_decomposition_info_local_index( - dim, owned, total, caplog, download_ser_data, decomposition_info, processor_props # noqa F811 + dim, + owned, + total, + caplog, + download_ser_data, + decomposition_info, + processor_props, # F811 ): check_comm_size(processor_props, sizes=[2]) my_rank = processor_props.rank @@ -141,8 +156,13 @@ def test_domain_descriptor_id_are_globally_unique(num, processor_props): @pytest.mark.mpi @pytest.mark.parametrize("processor_props", [True], indirect=True) -def test_decomposition_info_matches_gridsize(caplog, download_ser_data, decomposition_info, icon_grid, processor_props, - ): # noqa F811 +def test_decomposition_info_matches_gridsize( + caplog, + download_ser_data, + decomposition_info, + icon_grid, + processor_props, +): # F811 check_comm_size(processor_props) assert ( @@ -167,7 +187,9 @@ def test_decomposition_info_matches_gridsize(caplog, download_ser_data, decompos @pytest.mark.mpi @pytest.mark.parametrize("processor_props", [True], indirect=True) -def test_create_multi_node_runtime_with_mpi(download_ser_data, decomposition_info, processor_props): # noqa F811 +def test_create_multi_node_runtime_with_mpi( + download_ser_data, decomposition_info, processor_props +): # F811 props = processor_props exchange = create_exchange(props, decomposition_info) if props.comm_size > 1: @@ -175,6 +197,7 @@ def test_create_multi_node_runtime_with_mpi(download_ser_data, decomposition_inf else: assert isinstance(exchange, SingleNode) + @pytest.mark.parametrize("processor_props", [False], indirect=True) def test_create_single_node_runtime_without_mpi(processor_props, decomposition_info): props = processor_props diff --git a/model/common/tests/test_grid_manager.py b/model/common/tests/test_grid_manager.py index 7bdf95d2d..a31f08469 100644 --- a/model/common/tests/test_grid_manager.py +++ b/model/common/tests/test_grid_manager.py @@ -14,7 +14,6 @@ from pathlib import Path from uuid import uuid4 -import conftest import numpy as np import pytest from netCDF4 import Dataset @@ -31,6 +30,7 @@ from icon4py.model.common.grid.icon_grid import VerticalGridSize from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + SIMPLE_MESH_NC = "./simple_mesh_grid.nc" diff --git a/model/common/tests/test_interpolation_fields.py b/model/common/tests/test_interpolation_fields.py index b2c2bfeff..7c76c446e 100644 --- a/model/common/tests/test_interpolation_fields.py +++ b/model/common/tests/test_interpolation_fields.py @@ -27,7 +27,9 @@ import pytest from icon4py.model.common.dimension import EdgeDim -from icon4py.model.common.field_management.interpolation_fields import compute_c_lin_e +from icon4py.model.common.field_management.interpolation_fields import ( + compute_c_lin_e, +) from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex diff --git a/model/driver/src/icon4py/model/driver/__init__.py b/model/driver/src/icon4py/model/driver/__init__.py index 1dcd287bc..dab708955 100644 --- a/model/driver/src/icon4py/model/driver/__init__.py +++ b/model/driver/src/icon4py/model/driver/__init__.py @@ -13,6 +13,8 @@ from typing import Final from packaging import version as pkg_version + + __all__ = [ "__author__", "__copyright__", diff --git a/model/driver/src/icon4py/model/driver/dycore_driver.py b/model/driver/src/icon4py/model/driver/dycore_driver.py index fc6220ddd..09ebc2f5f 100644 --- a/model/driver/src/icon4py/model/driver/dycore_driver.py +++ b/model/driver/src/icon4py/model/driver/dycore_driver.py @@ -18,22 +18,28 @@ import click import pytz from devtools import Timer -from gt4py.next import program, Field +from gt4py.next import Field, program from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn -from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.atmosphere.diffusion.diffusion import ( + Diffusion, + DiffusionParams, +) +from icon4py.model.atmosphere.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + PrognosticState, +) +from icon4py.model.atmosphere.diffusion.diffusion_utils import ( + _identity_c_k, + _identity_e_k, +) from icon4py.model.common.decomposition.decomposed import create_exchange from icon4py.model.common.decomposition.parallel_setup import ( ProcessProperties, get_processor_properties, ) +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim from icon4py.model.common.test_utils import serialbox_utils as sb -from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams -from icon4py.model.atmosphere.diffusion.diffusion_states import ( - DiffusionDiagnosticState, - PrognosticState, -) -from icon4py.model.atmosphere.diffusion.diffusion_utils import _identity_c_k, _identity_e_k from icon4py.model.driver.icon_configuration import IconRunConfig, read_config from icon4py.model.driver.io_utils import ( SIMULATION_START_DATE, @@ -45,6 +51,7 @@ read_static_fields, ) + log = logging.getLogger(__name__) diff --git a/model/driver/src/icon4py/model/driver/icon_configuration.py b/model/driver/src/icon4py/model/driver/icon_configuration.py index 5d67a0781..54d92833f 100644 --- a/model/driver/src/icon4py/model/driver/icon_configuration.py +++ b/model/driver/src/icon4py/model/driver/icon_configuration.py @@ -16,6 +16,7 @@ from icon4py.model.atmosphere.diffusion.diffusion import DiffusionConfig + n_substeps_reduced = 2 diff --git a/model/driver/src/icon4py/model/driver/io_utils.py b/model/driver/src/icon4py/model/driver/io_utils.py index c0d7ce181..bb69fef9f 100644 --- a/model/driver/src/icon4py/model/driver/io_utils.py +++ b/model/driver/src/icon4py/model/driver/io_utils.py @@ -12,30 +12,32 @@ # SPDX-License-Identifier: GPL-3.0-or-later import logging - from datetime import datetime from enum import Enum from pathlib import Path -from icon4py.model.common.decomposition.decomposed import DecompositionInfo -from icon4py.model.common.decomposition.parallel_setup import ParallelLogger, ProcessProperties -from icon4py.model.atmosphere.diffusion.diffusion_states import (DiffusionDiagnosticState, +from icon4py.model.atmosphere.diffusion.diffusion_states import ( + DiffusionDiagnosticState, DiffusionInterpolationState, DiffusionMetricState, - PrognosticState) - + PrognosticState, +) +from icon4py.model.common.decomposition.decomposed import DecompositionInfo +from icon4py.model.common.decomposition.parallel_setup import ( + ParallelLogger, + ProcessProperties, +) from icon4py.model.common.grid.horizontal import CellParams, EdgeParams from icon4py.model.common.grid.icon_grid import IconGrid from icon4py.model.common.grid.vertical import VerticalModelParams - from icon4py.model.common.test_utils import serialbox_utils as sb + SB_ONLY_MSG = "Only ser_type='sb' is implemented so far." SIMULATION_START_DATE = "2021-06-20T12:00:10.000" log = logging.getLogger(__name__) - class SerializationType(str, Enum): SB = "serialbox" NC = "netcdf" diff --git a/model/driver/tests/conftest.py b/model/driver/tests/conftest.py index 5bb19df73..777fbf547 100644 --- a/model/driver/tests/conftest.py +++ b/model/driver/tests/conftest.py @@ -1 +1,19 @@ -from icon4py.model.common.test_utils.fixtures import datapath, download_ser_data,ranked_data_path,processor_props +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from icon4py.model.common.test_utils.fixtures import ( # noqa F401 + datapath, + download_ser_data, + processor_props, + ranked_data_path, +) diff --git a/model/driver/tests/test_dycore_driver.py b/model/driver/tests/test_dycore_driver.py index 4ca4f649d..88c4fb95d 100644 --- a/model/driver/tests/test_dycore_driver.py +++ b/model/driver/tests/test_dycore_driver.py @@ -14,10 +14,9 @@ import numpy as np from icon4py.model.common.dimension import CellDim, EdgeDim, KDim -from icon4py.model.driver.dycore_driver import _copy_diagnostic_and_prognostics - from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh +from icon4py.model.driver.dycore_driver import _copy_diagnostic_and_prognostics def test_copy_diagnostic_and_prognostics(): diff --git a/model/driver/tests/test_io_utils.py b/model/driver/tests/test_io_utils.py index f687609b4..b3b7434e6 100644 --- a/model/driver/tests/test_io_utils.py +++ b/model/driver/tests/test_io_utils.py @@ -16,7 +16,12 @@ from icon4py.model.common.grid.horizontal import CellParams, EdgeParams from icon4py.model.common.grid.vertical import VerticalModelParams -from icon4py.model.driver.io_utils import read_geometry_fields, read_static_fields, read_icon_grid, SerializationType +from icon4py.model.driver.io_utils import ( + SerializationType, + read_geometry_fields, + read_icon_grid, + read_static_fields, +) @pytest.mark.datatest diff --git a/py2f/src/icon4py/py2f/parsing.py b/py2f/src/icon4py/py2f/parsing.py index 0bf398d74..875df134d 100644 --- a/py2f/src/icon4py/py2f/parsing.py +++ b/py2f/src/icon4py/py2f/parsing.py @@ -16,7 +16,6 @@ from gt4py.next.type_system.type_specifications import ScalarType from gt4py.next.type_system.type_translation import from_type_hint - from icon4py.py2f.cffi_utils import CffiMethod from icon4py.py2f.codegen import CffiPlugin, DimensionType, Func, FuncParameter diff --git a/py2f/src/icon4py/py2f/py2fgen.py b/py2f/src/icon4py/py2f/py2fgen.py index f638f6e28..0ad0ca780 100644 --- a/py2f/src/icon4py/py2f/py2fgen.py +++ b/py2f/src/icon4py/py2f/py2fgen.py @@ -14,7 +14,6 @@ import pathlib import click - from icon4py.f2py.cffi_utils import generate_and_compile_cffi_plugin from icon4py.f2py.codegen import ( generate_and_write_f90_interface, diff --git a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py index 54aef39ff..0da0d4b3a 100644 --- a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py +++ b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py @@ -13,7 +13,6 @@ # flake8: noqa from gt4py.next.common import Field - from icon4py.common.dimension import ( C2E2CDim, C2E2CODim, @@ -25,22 +24,20 @@ V2EDim, VertexDim, ) + +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams +from icon4py.model.common.grid.icon_grid import IconGrid +from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.py2f.cffi_utils import CffiMethod from model.atmosphere.diffusion.src.icon4py.model.atmosphere.diffusion import ( Diffusion, DiffusionConfig, - DiffusionParams, -) -from model.atmosphere.diffusion.src.icon4py.model.atmosphere.diffusion import ( DiffusionDiagnosticState, DiffusionInterpolationState, DiffusionMetricState, + DiffusionParams, PrognosticState, ) -from icon4py.model.common.grid.horizontal import CellParams, EdgeParams -from icon4py.model.common.grid.icon_grid import IconGrid -from icon4py.model.common.grid.vertical import VerticalModelParams -from icon4py.py2f.cffi_utils import CffiMethod - diffusion: Diffusion() diff --git a/py2f/tests/test_code_generation.py b/py2f/tests/test_code_generation.py index bfe72b953..0ae91102a 100644 --- a/py2f/tests/test_code_generation.py +++ b/py2f/tests/test_code_generation.py @@ -14,7 +14,6 @@ import pytest from gt4py.next.type_system.type_specifications import ScalarKind - from icon4py.py2f.codegen import ( CffiPlugin, CHeaderGenerator, From 5648609f488be786da3c6456e9db6fdd83f90b27 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 24 Aug 2023 12:16:27 +0200 Subject: [PATCH 238/263] tox.ini - exclude datatest marker update README.md: running of parallel tests clean up dependencies: - fix path in driver/requirements.txt - add dependencies in pyproject.toml - remove pypi dependencies from requirements-dev.txt --- README.md | 32 ++++++++++++++++++- base-requirements-dev.txt | 6 ++-- base-requirements.txt | 6 ++-- model/atmosphere/diffusion/pyproject.toml | 8 +++-- model/common/pyproject.toml | 3 ++ model/common/tests/mpi_tests/README.md | 38 ----------------------- model/common/tests/test_grid_manager.py | 31 +++++++++--------- model/driver/README.md | 32 ++++++++++--------- model/driver/pyproject.toml | 7 +++++ model/driver/requirements-dev.txt | 2 +- model/driver/requirements.txt | 2 +- model/tox.ini | 4 +-- tools/pyproject.toml | 6 ++-- tools/requirements-dev.txt | 1 + 14 files changed, 91 insertions(+), 87 deletions(-) delete mode 100644 model/common/tests/mpi_tests/README.md diff --git a/README.md b/README.md index 2b4738ba6..64aab8c2f 100644 --- a/README.md +++ b/README.md @@ -112,9 +112,39 @@ tox # Run test suite in a specific environment (use `tox -a` to see list of envs) tox -e py310 ``` - The default `tox` environment is configured to generate HTML test coverage reports in `_reports/coverage_html/`. +#### data dependent tests +Some test depend on serialized data generated by a full model run. Those test are marked with +`pytest.mark.datatest` and can be skipped by specifying +```bash +pytest -v -m "not datatest" +``` + + +#### testing parallel code +Test for parallel code is collected in a package's `tests/mpi_tests` folder. +All parallel tests are marked with `@pytest.mark.mpi` and are skipped if +the `--with-mpi` is not passed option is not passed to `pytest` +In order to run them, you need a MPI installation on your system: +On Debian-Linux do +```bash +sudo apt-get install libopenmpi-dev +``` +or +```bash +sudo apt-get install mpich +``` +on MacOs +```bash +brew install mpich +``` +Then you can run the tests with +```bash +mpirun -np 2 pytest -v -s --with-mpi path/to/test/folder/mpi_tests +``` + + ### Benchmarking We use [`pytest-benchmark`](https://pytest-benchmark.readthedocs.io/en/latest/) to benchmark the execution time of stencils in icon4py. To disable benchmarking during testing you can use `--benchmark-disable` when invoking `pytest`. diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index c38e53e47..021c88962 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -24,13 +24,11 @@ pytest-cache>=1.0 pytest-cov>=2.8 pytest-factoryboy>=2.0 pytest-xdist[psutil]>=2.2 +pytest-mpi>=0.6 setuptools>=40.8.0 wheel>=0.37.1 tox>=3.25 wget>=3.2 -netcdf4>=1.6.0 -pytz>=2023.2 -mpi4py<=3.1.4 -pytest-mpi>=0.6 + diff --git a/base-requirements.txt b/base-requirements.txt index d587738bd..7b7d94fac 100644 --- a/base-requirements.txt +++ b/base-requirements.txt @@ -1,7 +1,5 @@ # VCS gt4py @ git+https://github.com/GridTools/gt4py.git@main pyghex @ git+https://github.com/boeschf/GHEX.git@pipify#subdirectory=bindings/python -# PyPI -netcdf4>=1.6.0 -pytz>=2023.2 -mpi4py<=3.1.4 + + diff --git a/model/atmosphere/diffusion/pyproject.toml b/model/atmosphere/diffusion/pyproject.toml index 6d03d11e1..a797dfae8 100644 --- a/model/atmosphere/diffusion/pyproject.toml +++ b/model/atmosphere/diffusion/pyproject.toml @@ -28,7 +28,9 @@ classifiers = [ ] dependencies = [ "gt4py>=1.0.1", - "icon4py_common>=0.0.5", + "icon4py-common>=0.0.5", + "mpi4py<=3.1.4", + "pyghex>=0.3.0", ] dynamic = ['version'] @@ -56,7 +58,7 @@ target-version = ['py310'] [tool.coverage] [tool.coverage.html] -directory = 'tests/_reports/coverage_html' +directory = 'diffusion_tests/_reports/coverage_html' [tool.coverage.paths] source = ['src/icon4py/model/'] @@ -109,7 +111,7 @@ show_error_codes = true [tool.pytest] [tool.pytest.ini_options] -testpaths = 'tests' +testpaths = 'diffusion_tests' markers = 'datatest: this test uses binary data' [tool.setuptools.dynamic] diff --git a/model/common/pyproject.toml b/model/common/pyproject.toml index 83af37333..f5485f947 100644 --- a/model/common/pyproject.toml +++ b/model/common/pyproject.toml @@ -28,6 +28,9 @@ classifiers = [ ] dependencies = [ "gt4py>=1.0.1", + "mpi4py<=3.1.4", + "pyghex>=0.3.0", + "netcdf4>=1.6.0" ] dynamic = ['version'] diff --git a/model/common/tests/mpi_tests/README.md b/model/common/tests/mpi_tests/README.md deleted file mode 100644 index a3e870bb3..000000000 --- a/model/common/tests/mpi_tests/README.md +++ /dev/null @@ -1,38 +0,0 @@ -## Running parallel version of diffusion - -### Installation - -The parallelized code uses [GHEX](https://github.com/ghex-org/GHEX) with MPI for halo exchanges. The GHEX python bindings can be installed with pip from VCS. However it has some system dependencies - -1.You need a running MPI installation in the system. On linux (apt base system) do - -```bash -sudo apt-get install libopenmpi-dev -``` - -on MacOs - -```bash -brew install mpich -``` - -2. You need to have boost (headers) installed in the system -3. clone and install `icon4py` as described in the [Readme](../../../README.md) install python dependencies In an existing local `icon4py` clone rerun `pip` install: - -``` -pip install --src _external_src -r requirements-dev.txt -``` - -### Run parallel test - -The project includes the `pytest-mpi` utility which augments `pytest` with some MPI specific utilities. MPI dependent tests are marked with `@pytest.mark.mpi` and are skipped when not run under mpi (`--with-mpi`) All MPI related tests are in the folder `icon4py/atm_dyn_iconam/tests/mpi_tests` tests depending on MPI are marked with `@pytest.mark.mpi` are therefore skipped by any serial pytest run. - -Run the tests with - -```bash -mpirun -np 2 pytest -v -s --with-mpi tests/mpi_tests/ -``` - -Test using serialized test data can be run with either 1, 2 or 4 nodes. - -When you run the tests for the first time the serialized test data will be downloaded from an online storage. That make some time... diff --git a/model/common/tests/test_grid_manager.py b/model/common/tests/test_grid_manager.py index a31f08469..d87acbf5a 100644 --- a/model/common/tests/test_grid_manager.py +++ b/model/common/tests/test_grid_manager.py @@ -31,18 +31,14 @@ from icon4py.model.common.test_utils.simple_mesh import SimpleMesh -SIMPLE_MESH_NC = "./simple_mesh_grid.nc" +SIMPLE_MESH_NC = "simple_mesh_grid.nc" @pytest.fixture -def simple_mesh_path(simple_mesh_data): - return Path(SIMPLE_MESH_NC).absolute() - - -@pytest.fixture(scope="session") -def simple_mesh_data(): +def simple_mesh_gridfile(tmp_path): + path = tmp_path.joinpath(SIMPLE_MESH_NC).absolute() mesh = SimpleMesh() - dataset = Dataset(SIMPLE_MESH_NC, "w", format="NETCDF4") + dataset = Dataset(path, "w", format="NETCDF4") dataset.setncattr(GridFile.PropertyName.GRID_ID, str(uuid4())) dataset.createDimension(GridFile.DimensionName.VERTEX_NAME, size=mesh.n_vertices) @@ -194,6 +190,8 @@ def simple_mesh_data(): (GridFile.DimensionName.MAX_CHILD_DOMAINS, GridFile.DimensionName.VERTEX_GRF), ) dataset.close() + yield path + path.unlink() def _add_to_dataset( @@ -206,9 +204,8 @@ def _add_to_dataset( var[:] = np.transpose(data)[:] -def test_gridparser_dimension(simple_mesh_data): - - data = Dataset(SIMPLE_MESH_NC, "r") +def test_gridparser_dimension(simple_mesh_gridfile): + data = Dataset(simple_mesh_gridfile, "r") grid_parser = GridFile(data) mesh = SimpleMesh() assert grid_parser.dimension(GridFile.DimensionName.CELL_NAME) == mesh.n_cells @@ -232,9 +229,9 @@ def test_gridfile_vertex_cell_edge_dimensions(grid_savepoint, r04b09_dsl_gridfil ) == grid_savepoint.num(VertexDim) -def test_grid_parser_index_fields(simple_mesh_data, caplog): +def test_grid_parser_index_fields(simple_mesh_gridfile, caplog): caplog.set_level(logging.DEBUG) - data = Dataset(SIMPLE_MESH_NC, "r") + data = Dataset(simple_mesh_gridfile, "r") mesh = SimpleMesh() grid_parser = GridFile(data) @@ -406,20 +403,20 @@ def init_grid_manager(fname): @pytest.mark.parametrize("dim, size", [(CellDim, 18), (EdgeDim, 27), (VertexDim, 9)]) -def test_grid_manager_getsize(simple_mesh_data, simple_mesh_path, dim, size, caplog): +def test_grid_manager_getsize(simple_mesh_gridfile, dim, size, caplog): caplog.set_level(logging.DEBUG) gm = GridManager( - IndexTransformation(), simple_mesh_path, VerticalGridSize(num_lev=80) + IndexTransformation(), simple_mesh_gridfile, VerticalGridSize(num_lev=80) ) gm() assert size == gm.get_size(dim) -def test_grid_manager_diamond_offset(simple_mesh_path): +def test_grid_manager_diamond_offset(simple_mesh_gridfile): mesh = SimpleMesh() gm = GridManager( IndexTransformation(), - simple_mesh_path, + simple_mesh_gridfile, VerticalGridSize(num_lev=mesh.k_level), ) gm() diff --git a/model/driver/README.md b/model/driver/README.md index 069362ae3..ceeadcae1 100644 --- a/model/driver/README.md +++ b/model/driver/README.md @@ -1,35 +1,39 @@ -# dummy driver for python dycore code +# Dummy driver for Python ICON port -`dycore_driver.py` contains a simple python program to run the experimental ICON python port. The code mostly draws on serialized ICON data. It initializes the grid from serialized data from a `mch_ch_r04b09_dsl` run and configures a timeloop functionality based on that configuration. +`dycore_driver.py` contains a simple python program to run the experimental ICON python port. So far +code mostly draws on serialized ICON data until we increasingly can initialize and run the model by independently. -Currently, there is does no real timestepping instead it calls a dummy timestep that serves a batch of new serialized input fields from ICON. +It initializes the grid from serialized data from a `mch_ch_r04b09_dsl` run and configures a +timeloop functionality based on that configuration. + +Currently, there is does *no real timestepping*, instead it calls a dummy timestep that +serves a batch of new serialized input fields from ICON. The code is meant to be changed and enlarged as we port new parts of the model. -## usage +It runs single node or parallel versions. For parallel runs the domain has to be decomposed previousely +through a full ICON run that generates the necessary serialized data. Test data for runs with 1, 2, 4 nodes +are available. + +## Usage ```bash -cd atm_dyn_iconam/src/icon4py -python driver/dycore_driver.py ../../../testdata/ser_icondata/mpitask1/mch_ch_r04b09_dsl/ser_data --n_steps=2 --run_path=/home/magdalena/temp/icon +export ICON4PY_ROOT= +python driver/dycore_driver.py $ICON4PY_ROOT/testdata/ser_icondata/mpitask1/mch_ch_r04b09_dsl/ser_data --n_steps=2 --run_path=/home/magdalena/temp/icon ``` or if running in parallel ```bash -cd atm_dyn_iconam/src/icon4py -mpirun -np 2 python driver/dycore_driver.py ../../../testdata/ser_icondata/mpitask2/mch_ch_r04b09_dsl/ser_data --mpi=True --n_steps=2 --run_path=/home/magdalena/temp/icon +mpirun -np 2 python driver/dycore_driver.py $ICON4PY_ROOT/testdata/ser_icondata/mpitask2/mch_ch_r04b09_dsl/ser_data --mpi=True --n_steps=2 --run_path=/home/magdalena/temp/icon ``` #### remarks -- First (required) arg is the folder where the serialized input data is stored. The same input data s for the unit tests is used. It can be obtained from [polybox](https://polybox.ethz.ch/index.php/s/rzuvPf7p9sM801I/download) -- The driver heavily relies on serialized input data. The paths in the example are where the data is put when downloaded via the unit tests. +- First (required) arg is the folder where the serialized input data is stored. The same input data is for the unit tests is used. The path in the example is where the data is put when downloaded via the unit tests. - data for a serial (single node) run can be downloaded from https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p. - parallel runs are possible if corresponding data is provided, which is currently available for test with 2 or 4 MPI processes: - - - 2 processes: https://polybox.ethz.ch/index.php/s/NUQjmJcMEoQxFiK - - 4 processes: https://polybox.ethz.ch/index.php/s/QC7xt7xLT5xeVN5 - +check [fixtures.py](../common/src/icon4py/model/common/test_utils/fixtures.py) for download urls. - The serialized data used contains only 5 timesteps so `--n_steps > 2` will throw an exception. - The code logs to file and to console. Debug logging is only going to file. The log directory can be changed with the --run_path option. diff --git a/model/driver/pyproject.toml b/model/driver/pyproject.toml index e9b7f3e98..188d04d2d 100644 --- a/model/driver/pyproject.toml +++ b/model/driver/pyproject.toml @@ -28,12 +28,19 @@ classifiers = [ ] dependencies = [ "gt4py>=1.0.1", + "mpi4py<=3.1.4", + "pyghex>=0.3.0", + "pytz>=2023.2", "icon4py-common>=0.0.5", "icon4py-atmosphere-dycore>=0.0.5", "icon4py-atmosphere-diffusion>=0.0.5", ] dynamic = ['version'] +[project.scripts] +dycore_driver = "icon4py.model.driver.dycore_driver:main" + + [project.urls] repository = "https://github.com/C2SM/icon4py" diff --git a/model/driver/requirements-dev.txt b/model/driver/requirements-dev.txt index a42befbba..3f54c69af 100644 --- a/model/driver/requirements-dev.txt +++ b/model/driver/requirements-dev.txt @@ -1,4 +1,4 @@ --r ../../../base-requirements-dev.txt +-r ../../base-requirements-dev.txt -e ../common -e ../atmosphere -e . diff --git a/model/driver/requirements.txt b/model/driver/requirements.txt index b92b8313c..18a81dc7c 100644 --- a/model/driver/requirements.txt +++ b/model/driver/requirements.txt @@ -1,4 +1,4 @@ --r ../../../base-requirements.txt +-r ../../base-requirements.txt ../common ../atmosphere/dycore ../atmosphere/diffusion diff --git a/model/tox.ini b/model/tox.ini index 27fb18d61..6714d9033 100644 --- a/model/tox.ini +++ b/model/tox.ini @@ -14,8 +14,8 @@ passenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src common/src - pytest -v -s -n auto --cov --cov-append --benchmark-disable + -pytest -v -m "not datatest" -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src common/src driver/src + pytest -v -m "not datatest" -s -n auto --cov --cov-append --benchmark-disable commands_post = rm -rf tests/_reports/coverage_html -coverage html diff --git a/tools/pyproject.toml b/tools/pyproject.toml index d6978a938..6398383f3 100644 --- a/tools/pyproject.toml +++ b/tools/pyproject.toml @@ -27,11 +27,13 @@ classifiers = [ "Topic :: Scientific/Engineering :: Physics" ] dependencies = [ + "icon4py-atmosphere-diffusion", "icon4py-atmosphere-dycore", "gt4py>=1.0.1", - "icon4py_common", + "icon4py-common", "tabulate>=0.8.9", - "fprettify>=0.3.7" + "fprettify>=0.3.7", + "cffi>=1.5" ] dynamic = ['version'] diff --git a/tools/requirements-dev.txt b/tools/requirements-dev.txt index f3d56fb5c..d4652678b 100644 --- a/tools/requirements-dev.txt +++ b/tools/requirements-dev.txt @@ -1,4 +1,5 @@ -r ../base-requirements-dev.txt -e ../model/atmosphere/dycore +-e ../model/atmosphere/diffusion -e ../model/common -e . From b129b34805ddbd05b760eb84c7f2ee3fe98e6132 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 24 Aug 2023 15:06:58 +0200 Subject: [PATCH 239/263] move py2f to tools --- .../src/icon4py/model/driver/dycore_driver.py | 4 +- py2f/README.md | 20 -- py2f/requirements-dev.txt | 4 - py2f/requirements.txt | 4 - py2f/setup.cfg | 58 ---- py2f/src/icon4py/py2f/cffi_utils.py | 89 ------ .../py2f/wrappers/diffusion_wrapper.py | 158 ----------- requirements-dev.txt | 1 - requirements.txt | 1 - tools/README.md | 36 +++ tools/pyproject.toml | 5 +- .../src/icon4pytools}/py2f/__init__.py | 0 tools/src/icon4pytools/py2f/cffi_utils.py | 181 ++++++++++++ .../src/icon4pytools}/py2f/codegen.py | 1 + .../src/icon4pytools}/py2f/parsing.py | 19 +- .../src/icon4pytools}/py2f/py2fgen.py | 20 +- tools/src/icon4pytools/py2f/typing_utils.py | 27 ++ .../icon4pytools}/py2f/wrappers/__init__.py | 0 .../py2f/wrappers/diffusion_wrapper.py | 260 ++++++++++++++++++ tools/tests/py2f/test_cffi_utils.py | 84 ++++++ .../tests/py2f}/test_code_generation.py | 3 +- .../tests/py2f}/test_parsing_wrapper.py | 6 +- .../tests/py2f/test_py2fgen.py | 14 +- 23 files changed, 626 insertions(+), 369 deletions(-) delete mode 100644 py2f/README.md delete mode 100644 py2f/requirements-dev.txt delete mode 100644 py2f/requirements.txt delete mode 100644 py2f/setup.cfg delete mode 100644 py2f/src/icon4py/py2f/cffi_utils.py delete mode 100644 py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py rename {py2f/src/icon4py => tools/src/icon4pytools}/py2f/__init__.py (100%) create mode 100644 tools/src/icon4pytools/py2f/cffi_utils.py rename {py2f/src/icon4py => tools/src/icon4pytools}/py2f/codegen.py (99%) rename {py2f/src/icon4py => tools/src/icon4pytools}/py2f/parsing.py (67%) rename {py2f/src/icon4py => tools/src/icon4pytools}/py2f/py2fgen.py (74%) create mode 100644 tools/src/icon4pytools/py2f/typing_utils.py rename {py2f/src/icon4py => tools/src/icon4pytools}/py2f/wrappers/__init__.py (100%) create mode 100644 tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py create mode 100644 tools/tests/py2f/test_cffi_utils.py rename {py2f/tests => tools/tests/py2f}/test_code_generation.py (98%) rename {py2f/tests => tools/tests/py2f}/test_parsing_wrapper.py (86%) rename py2f/setup.py => tools/tests/py2f/test_py2fgen.py (60%) diff --git a/model/driver/src/icon4py/model/driver/dycore_driver.py b/model/driver/src/icon4py/model/driver/dycore_driver.py index 09ebc2f5f..f9f50b2ad 100644 --- a/model/driver/src/icon4py/model/driver/dycore_driver.py +++ b/model/driver/src/icon4py/model/driver/dycore_driver.py @@ -263,7 +263,7 @@ def initialize(n_time_steps, file_path: Path, props: ProcessProperties): "--n_steps", default=5, help="number of time steps to run, max 5 is supported" ) @click.option("--mpi", default=False, help="whether or not you are running with mpi") -def run(input_path, run_path, n_steps, mpi): +def main(input_path, run_path, n_steps, mpi): """ Run the driver. @@ -300,4 +300,4 @@ def run(input_path, run_path, n_steps, mpi): if __name__ == "__main__": - run() + main() diff --git a/py2f/README.md b/py2f/README.md deleted file mode 100644 index af6792103..000000000 --- a/py2f/README.md +++ /dev/null @@ -1,20 +0,0 @@ -# icon4py py2f90 - -## Description - -Python utilities for generating a C library and Fortran interface to call Python icon4py modules. The library [embeds python via CFFI ](https://cffi.readthedocs.io/en/latest/embedding.html) - -### py2f90gen - -Generates a C header file and a Fortran interface and compiles python functions into a C library embedding python. - -The functions need to be decorated with `CffiMethod.register` and have a signature with scalar arguments of gt4py fields: - -``` -@CffiMethod.register -def foo(i:int, param:float, field1: Field[[VertexDim, KDim], float], field2: Field[CellDim, KDim], float]) -``` - -## Installation instructions - -Check `README.md` file in the root of the repository. diff --git a/py2f/requirements-dev.txt b/py2f/requirements-dev.txt deleted file mode 100644 index 217c64d71..000000000 --- a/py2f/requirements-dev.txt +++ /dev/null @@ -1,4 +0,0 @@ --r ../base-requirements-dev.txt --e ../common --e ../atm_dyn_iconam --e . diff --git a/py2f/requirements.txt b/py2f/requirements.txt deleted file mode 100644 index b6eac2aa8..000000000 --- a/py2f/requirements.txt +++ /dev/null @@ -1,4 +0,0 @@ --r ../base-requirements.txt -../common -../atm_dyn_iconam -. diff --git a/py2f/setup.cfg b/py2f/setup.cfg deleted file mode 100644 index f7d5ac4cd..000000000 --- a/py2f/setup.cfg +++ /dev/null @@ -1,58 +0,0 @@ -# This file is mainly used to configure package creation with setuptools. -# Documentation: -# http://setuptools.readthedocs.io/en/latest/setuptools.html#configuring-setup-using-setup-cfg-files -# -[metadata] -name = icon4py_py2f -description = Icon inspired code in Python and GT4Py -long_description = file: README.md -long_description_content_type = text/markdown -url = https://github.com/C2SM/icon4py -author = ETH Zurich -author_email = gridtools@cscs.ch -license = gpl3 -license_files = LICENSE -platforms = Linux, Mac -classifiers = - Development Status :: 3 - Alpha - Intended Audience :: Science/Research - License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+) - Operating System :: POSIX - Programming Language :: Python - Programming Language :: Python :: 3 - Programming Language :: Python :: 3 :: Only - Programming Language :: Python :: 3.10 - Programming Language :: Python :: Implementation :: CPython - Topic :: Scientific/Engineering :: Atmospheric Science - Topic :: Scientific/Engineering :: Mathematics - Topic :: Scientific/Engineering :: Physics -project_urls = - Source Code = https://github.com/C2SM/icon4py - -[options] -packages = find_namespace: -install_requires = - icon4py-common - icon4py-atm_dyn_iconam - cffi>=1.5 - tabulate>=0.8.9 - fprettify>=0.3.7 -python_requires = >=3.10 -package_dir = - = src -zip_safe = False - -[options.package_data] -# References: -# https://setuptools.pypa.io/en/latest/userguide/datafiles.html -# https://github.com/abravalheri/experiment-setuptools-package-data -* = *.md, *.rst, *.toml, *.txt, py.typed - -[options.packages.find] -where = src -exclude = - tests - -[options.entry_points] -console_scripts = - py2fgen = icon4py.py2f.wrapper.py2fgen:main diff --git a/py2f/src/icon4py/py2f/cffi_utils.py b/py2f/src/icon4py/py2f/cffi_utils.py deleted file mode 100644 index e66a14253..000000000 --- a/py2f/src/icon4py/py2f/cffi_utils.py +++ /dev/null @@ -1,89 +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 functools - -import cffi - - -FFI_DEF_EXTERN_DECORATOR = "@ffi.def_extern()" - -CFFI_GEN_DECORATOR = "@CffiMethod.register" - - -class CffiMethod: - _registry = {} - - @classmethod - def register(cls, func): - try: - cls._registry[func.__module__].append(func.__name__) - except KeyError: - cls._registry[func.__module__] = [func.__name__] - - @functools.wraps(func) - def wrapper(*args, **kwargs): - return func(*args, **kwargs) - - return wrapper - - @classmethod - def get(cls, name: str): - return cls._registry[name] - - -def with_cffi_gen(func): - @functools.wraps(func) - def _cffi_gen(*args, **kwargs): - return func(*args, **kwargs) - - return _cffi_gen - - -def generate_and_compile_cffi_plugin( - plugin_name: str, c_header: str, module_name: str, build_path="." -): - """ - Create C shared library. - - Create a linkable C library and F90 interface for the functions in the python module - {module_name} that are decorated with '@CffiMethod.register'. - - Args: - plugin_name: name of the plugin, a linkable C library with the name - 'lib{plugin_name}.so' will be created in the {build_path} folder' - c_header: C type header signature for the python functions. - module_name: python module name that contains python functions corresponding - to the signature in the '{c_header}' string, these functions must be decorated - with @CffiMethod.register and the file must contain the import - build_path: *optional* path to build directory - - """ - python_src_file = f"{module_name.split('.')[-1]}.py" - c_header_file = plugin_name + ".h" - with open("/".join([build_path, c_header_file]), "w") as f: - f.write(c_header) - - builder = cffi.FFI() - - builder.embedding_api(c_header) - builder.set_source(plugin_name, f'#include "{c_header_file}"') - - with open(python_src_file) as f: - module = f.read() - - module = f"from {plugin_name} import ffi\n{module}".replace( - CFFI_GEN_DECORATOR, FFI_DEF_EXTERN_DECORATOR - ) - - builder.embedding_init_code(module) - builder.compile(tmpdir=build_path, target=f"lib{plugin_name}.*", verbose=True) diff --git a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py b/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py deleted file mode 100644 index 0da0d4b3a..000000000 --- a/py2f/src/icon4py/py2f/wrappers/diffusion_wrapper.py +++ /dev/null @@ -1,158 +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 - -# flake8: noqa -from gt4py.next.common import Field -from icon4py.common.dimension import ( - C2E2CDim, - C2E2CODim, - C2EDim, - CellDim, - ECVDim, - EdgeDim, - KDim, - V2EDim, - VertexDim, -) - -from icon4py.model.common.grid.horizontal import CellParams, EdgeParams -from icon4py.model.common.grid.icon_grid import IconGrid -from icon4py.model.common.grid.vertical import VerticalModelParams -from icon4py.py2f.cffi_utils import CffiMethod -from model.atmosphere.diffusion.src.icon4py.model.atmosphere.diffusion import ( - Diffusion, - DiffusionConfig, - DiffusionDiagnosticState, - DiffusionInterpolationState, - DiffusionMetricState, - DiffusionParams, - PrognosticState, -) - -diffusion: Diffusion() - - -@CffiMethod.register -def diffusion_init( - vct_a: Field[[KDim], float], - nrdmax: float, # in config - theta_ref_mc: Field[[CellDim, KDim], float], - wgtfac_c: Field[[CellDim, KDim], float], - mask_hdiff: Field[[CellDim, KDim], int], - zd_vertidx: Field[[CellDim, C2E2CDim, KDim], int], - zd_diffcoef: Field[[CellDim, KDim], float], - zd_intcoef: Field[[CellDim, C2E2CDim, KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - rbf_coeff_1: Field[[VertexDim, V2EDim], float], - rbf_coeff_2: Field[[VertexDim, V2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], - geofac_n2s: Field[[CellDim, C2E2CODim], float], - geofac_grg_x: Field[[CellDim, C2E2CODim], float], - geofac_grg_y: Field[[CellDim, C2E2CODim], float], - nudgecoeff_e: Field[[EdgeDim], float], - tangent_orientation: Field[[EdgeDim], float], - primal_edge_lengths: Field[[EdgeDim], float], - inverse_primal_edge_lengths: Field[[EdgeDim], float], - dual_edge_lengths: Field[[EdgeDim], float], - inverse_dual_edge_lengths: Field[[EdgeDim], float], - inverse_vertex_vertex_lengths: Field[[EdgeDim], float], - primal_normal_vert_1: Field[[ECVDim], float], - primal_normal_vert_2: Field[[ECVDim], float], - dual_normal_vert_1: Field[[ECVDim], float], - dual_normal_vert_2: Field[[ECVDim], float], - edge_areas: Field[[EdgeDim], float], - cell_areas: Field[[CellDim], float], -): - """ - Instantiate and Initialize the diffusion object. - - should only accept simple fields as arguments for compatibility with the standalone - Fortran ICON Diffusion component (aka Diffusion granule) - - """ - grid = IconGrid() # TODO(Magdalena) where to get this from - edge_params = EdgeParams( - tangent_orientation=tangent_orientation, - primal_edge_lengths=primal_edge_lengths, - inverse_primal_edge_lengths=inverse_primal_edge_lengths, - dual_normal_vert=(dual_normal_vert_1, dual_normal_vert_2), - dual_edge_lengths=dual_edge_lengths, - inverse_dual_edge_lengths=inverse_dual_edge_lengths, - inverse_vertex_vertex_lengths=inverse_vertex_vertex_lengths, - primal_normal_vert=(primal_normal_vert_1, primal_normal_vert_2), - edge_areas=edge_areas, - ) - cell_params = CellParams(cell_areas) - vertical_params = VerticalModelParams(vct_a=vct_a, rayleigh_damping_height=nrdmax) - config: DiffusionConfig = DiffusionConfig(grid, vertical_params) - derived_diffusion_params = DiffusionParams(config) - metric_state = DiffusionMetricState( - theta_ref_mc, wgtfac_c, mask_hdiff, zd_vertidx, zd_diffcoef, zd_intcoef - ) - interpolation_state = DiffusionInterpolationState( - e_bln_c_s, - rbf_coeff_1, - rbf_coeff_2, - geofac_div, - geofac_n2s, - geofac_grg_x, - geofac_grg_y, - nudgecoeff_e, - ) - - diffusion.init( - grid=grid, - config=config, - params=derived_diffusion_params, - vertical_params=vertical_params, - metric_state=metric_state, - interpolation_state=interpolation_state, - edge_params=edge_params, - cell_params=cell_params, - ) - - -@CffiMethod.register -def diffusion_run( - dtime: float, - linit: bool, - vn: Field[[EdgeDim, KDim], float], - w: Field[[CellDim, KDim], float], - theta_v: Field[[CellDim, KDim], float], - exner: Field[[CellDim, KDim], float], - div_ic: Field[[CellDim, KDim], float], - hdef_ic: Field[[CellDim, KDim], float], - dwdx: Field[[CellDim, KDim], float], - dwdy: Field[[CellDim, KDim], float], -): - diagnostic_state = DiffusionDiagnosticState(hdef_ic, div_ic, dwdx, dwdy) - prognostic_state = PrognosticState( - w=w, - vn=vn, - exner_pressure=exner, - theta_v=theta_v, - ) - if linit: - diffusion.initial_run( - diagnostic_state, - prognostic_state, - dtime, - ) - else: - diffusion.run(diagnostic_state, prognostic_state, dtime) - - -class DuplicateInitializationException(Exception): - """Raised if the component is already initilalized.""" - - pass diff --git a/requirements-dev.txt b/requirements-dev.txt index 256d7e850..ca807eac7 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -8,4 +8,3 @@ # icon4pytools -e ./tools --e ./py2f #TODO (magdalena) move to tools and remove diff --git a/requirements.txt b/requirements.txt index 87afe68b2..803c4f811 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,6 +5,5 @@ ./model/atmosphere/diffusion ./model/common ./model/driver -./py2f #TODO (magdalena) move to tools and remove # icon4pytools ./tools diff --git a/tools/README.md b/tools/README.md index fdb881588..41f040e0d 100644 --- a/tools/README.md +++ b/tools/README.md @@ -232,3 +232,39 @@ OUTPUT_FILEPATH A path to the output Fortran source file to be generated. ``` **Note:** The output of f2ser still has to be preprocessed using `pp_ser.py`, which then yields a compilable unit. The serialised files will have `f2ser` as their prefix in the default folder location of the experiment. + +## `py2f` + +Python utility for generating a C library and Fortran interface to call Python icon4py modules. The library [embeds python via CFFI ](https://cffi.readthedocs.io/en/latest/embedding.) + +This is **highly experimental** and has not been tested from within Fortran code! + +### usage + +`py2fgen`: Generates a C header file and a Fortran interface and compiles python functions into a C library embedding python. The functions need to be decorated with `CffiMethod.register` and have a signature with scalar arguments or `GT4Py` fields, for example: + +``` +@CffiMethod.register +def foo(i:int, param:float, field1: Field[[VertexDim, KDim], float], field2: Field[CellDim, KDim], float]) +``` + +see `src/icon4pytools/py2f/wrappers` for examples. + +```bash +py2fgen icon4pytools.py2f.wrappers.diffusion_wrapper py2f_build +``` + +where the first argument is the python module to parse and the second a build directory. The call above will generate the following in `py2f_build`: + +```bash + ls py2f_build/ +total 204K +drwxrwxr-x 2 magdalena magdalena 4.0K Aug 24 17:01 . +drwxrwxr-x 9 magdalena magdalena 4.0K Aug 24 17:01 .. +-rw-rw-r-- 1 magdalena magdalena 74K Aug 24 16:58 diffusion_wrapper.c +-rw-rw-r-- 1 magdalena magdalena 3.5K Aug 24 17:01 diffusion_wrapper.f90 +-rw-rw-r-- 1 magdalena magdalena 955 Aug 24 17:01 diffusion_wrapper.h +-rw-rw-r-- 1 magdalena magdalena 58K Aug 24 17:01 diffusion_wrapper.o +-rwxrwxr-x 1 magdalena magdalena 49K Aug 24 17:01 libdiffusion_wrapper.so + +``` diff --git a/tools/pyproject.toml b/tools/pyproject.toml index 6398383f3..030ff3df1 100644 --- a/tools/pyproject.toml +++ b/tools/pyproject.toml @@ -33,7 +33,7 @@ dependencies = [ "icon4py-common", "tabulate>=0.8.9", "fprettify>=0.3.7", - "cffi>=1.5" + "cffi>=1.5", ] dynamic = ['version'] @@ -41,6 +41,7 @@ dynamic = ['version'] icon_liskov = "icon4pytools.liskov.cli:main" icon4pygen = "icon4pytools.icon4pygen.cli:main" f2ser = "icon4pytools.f2ser.cli:main" +py2fgen = "icon4pytools.py2f.py2fgen:main" [project.urls] repository = "https://github.com/C2SM/icon4py" @@ -106,6 +107,7 @@ exclude = [ '^tests/f2ser/*.py', '^tests/icon4pygen/*.py', '^tests/liskov/*.py', + '^tests/py2f/*.py', ] disallow_incomplete_defs = true disallow_untyped_defs = true @@ -127,3 +129,4 @@ version = {attr = 'icon4pytools.__init__.__version__'} [tool.setuptools.package-data] 'icon4pytools' = ['py.typed'] + diff --git a/py2f/src/icon4py/py2f/__init__.py b/tools/src/icon4pytools/py2f/__init__.py similarity index 100% rename from py2f/src/icon4py/py2f/__init__.py rename to tools/src/icon4pytools/py2f/__init__.py diff --git a/tools/src/icon4pytools/py2f/cffi_utils.py b/tools/src/icon4pytools/py2f/cffi_utils.py new file mode 100644 index 000000000..bcceccec2 --- /dev/null +++ b/tools/src/icon4pytools/py2f/cffi_utils.py @@ -0,0 +1,181 @@ +# 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 functools +import inspect +from collections import OrderedDict +from importlib.resources import files +from types import MappingProxyType +from typing import Any + +import cffi +import numpy as np +from gt4py.next.common import Dimension, DimensionKind +from gt4py.next.iterator.embedded import np_as_located_field + +from icon4pytools.py2f.typing_utils import parse_annotation + + +FFI_DEF_EXTERN_DECORATOR = "@ffi.def_extern()" + +CFFI_GEN_DECORATOR = "@CffiMethod.register" + + +class CffiMethod: + _registry = {} + + @classmethod + def register(cls, func): + try: + cls._registry[func.__module__].append(func.__name__) + except KeyError: + cls._registry[func.__module__] = [func.__name__] + + @functools.wraps(func) + def wrapper(*args, **kwargs): + return func(*args, **kwargs) + + return wrapper + + @classmethod + def get(cls, name: str): + return cls._registry[name] + + +def generate_and_compile_cffi_plugin( + plugin_name: str, c_header: str, module_name: str, build_path="." +): + """ + Create C shared library. + + Create a linkable C library and F90 interface for the functions in the python module + {module_name} that are decorated with '@CffiMethod.register'. + + Args: + plugin_name: name of the plugin, a linkable C library with the name + 'lib{plugin_name}.so' will be created in the {build_path} folder' + c_header: C type header signature for the python functions. + module_name: python module name that contains python functions corresponding + to the signature in the '{c_header}' string, these functions must be decorated + with @CffiMethod.register and the file must contain the import + build_path: *optional* path to build directory + + """ + module_split = module_name.split(".") + python_src_file = "/".join(module_split[1:]) + ".py" + python_package = module_split[0] + + c_header_file = plugin_name + ".h" + with open("/".join([build_path, c_header_file]), "w") as f: + f.write(c_header) + + builder = cffi.FFI() + + builder.embedding_api(c_header) + builder.set_source(plugin_name, f'#include "{c_header_file}"') + + module = files(python_package).joinpath(python_src_file).read_text() + + module = f"from {plugin_name} import ffi\n{module}".replace( + CFFI_GEN_DECORATOR, FFI_DEF_EXTERN_DECORATOR + ) + + builder.embedding_init_code(module) + builder.compile(tmpdir=build_path, target=f"lib{plugin_name}.*", verbose=True) + + +class UnknownDimensionException(Exception): + """Raised if a Dimension is unknown to the interface generation.""" + + pass + + +def to_fields( + dim_sizes: dict[Dimension, int] +): + """ + Pack/Unpack Fortran 2d arrays to numpy arrays with using CFFI frombuffer. + #TODO (magdalena) handle in a better way? + Args: + dim_sizes: dictionary containing the sizes of the dimension. + + """ + ffi = cffi.FFI() + dim_sizes = dim_sizes + + def _dim_sizes(dims: list[Dimension]) -> tuple[int, int]: + """Extract the size of dimension from a dictionary.""" + v_size = None + h_size = None + for d in dims: + if d not in dim_sizes.keys(): + raise UnknownDimensionException( + f"size of dimension '{d}' not defined in '{dim_sizes.keys()}'" + ) + if d.kind == DimensionKind.VERTICAL or d.kind == DimensionKind.LOCAL: + v_size = dim_sizes[d] + elif d.kind == DimensionKind.HORIZONTAL: + h_size = dim_sizes[d] + return h_size, v_size + + def _unpack(ptr, size_h, size_v, dtype) -> np.ndarray: + """ + Unpack a 2d c/fortran field into a numpy array. + + :param dtype: expected type of the fields + :param ptr: c_pointer to the field + :param size_h: length of horizontal dimension + :param size_v: length of vertical dimension + :return: a numpy array with shape=(size_h, size_v) + and dtype = ctype of the pointer + """ + shape = (size_h, size_v) + length = np.prod(shape) + c_type = ffi.getctype(ffi.typeof(ptr).item) + # TODO (magdalena) fix dtype handling use SCALARTYPE? + mem_size = ffi.sizeof(c_type) + mem_size = np.dtype(c_type).itemsize + ar = np.frombuffer( + ffi.buffer(ptr, length * mem_size), + dtype=np.dtype(c_type), + count=-1, + offset=0, + ).reshape(shape) + return ar + + def _to_fields_decorator(func): + @functools.wraps(func) + def _wrapper( + *args, **kwargs + ): # these are the args of the decorated function ie to_fields(func(*args, **kwargs)) + signature = inspect.signature(func) + parameters: MappingProxyType[str, inspect.Parameter] = signature.parameters + arguments: OrderedDict[str, Any] = signature.bind(*args, **kwargs).arguments + f_args = [] + for name, argument in arguments.items(): + ar = _transform_arg(argument, name, parameters) + f_args.append(ar) + return func(*f_args, **kwargs) + + def _transform_arg(argument, name, parameters): + dims, dtype = parse_annotation(parameters[name].annotation) + if dims: + (size_h, size_v) = _dim_sizes(dims) + ar = _unpack(argument, size_h, size_v, dtype) + ar = np_as_located_field(*dims)(ar) + else: + ar = argument + return ar + + return _wrapper + + return _to_fields_decorator diff --git a/py2f/src/icon4py/py2f/codegen.py b/tools/src/icon4pytools/py2f/codegen.py similarity index 99% rename from py2f/src/icon4py/py2f/codegen.py rename to tools/src/icon4pytools/py2f/codegen.py index 6bb9a3a3c..502b20492 100644 --- a/py2f/src/icon4py/py2f/codegen.py +++ b/tools/src/icon4pytools/py2f/codegen.py @@ -17,6 +17,7 @@ from gt4py.eve.codegen import JinjaTemplate as as_jinja from gt4py.eve.codegen import TemplatedGenerator from gt4py.next.type_system.type_specifications import ScalarKind + from icon4pytools.icon4pygen.bindings.codegen.type_conversion import ( BUILTIN_TO_CPP_TYPE, BUILTIN_TO_ISO_C_TYPE, diff --git a/py2f/src/icon4py/py2f/parsing.py b/tools/src/icon4pytools/py2f/parsing.py similarity index 67% rename from py2f/src/icon4py/py2f/parsing.py rename to tools/src/icon4pytools/py2f/parsing.py index 875df134d..847285615 100644 --- a/py2f/src/icon4py/py2f/parsing.py +++ b/tools/src/icon4pytools/py2f/parsing.py @@ -14,10 +14,9 @@ import importlib from inspect import signature, unwrap -from gt4py.next.type_system.type_specifications import ScalarType -from gt4py.next.type_system.type_translation import from_type_hint -from icon4py.py2f.cffi_utils import CffiMethod -from icon4py.py2f.codegen import CffiPlugin, DimensionType, Func, FuncParameter +from icon4pytools.py2f.cffi_utils import CffiMethod +from icon4pytools.py2f.codegen import CffiPlugin, DimensionType, Func, FuncParameter +from icon4pytools.py2f.typing_utils import parse_annotation def parse_functions_from_module(module_name: str) -> CffiPlugin: @@ -38,11 +37,7 @@ def _parse_function(module, s): def _parse_params(params, s): - type_spec = from_type_hint(params[s].annotation) - if isinstance(type_spec, ScalarType): - dtype = type_spec.kind - dims = [] - else: - dtype = type_spec.dtype.kind - dims = [DimensionType(name=d.value, length=10) for d in type_spec.dims] - return FuncParameter(name=s, d_type=dtype, dimensions=dims) + annotation = params[s].annotation + dims, dtype = parse_annotation(annotation) + dim_types = [DimensionType(name=d.value, length=10) for d in dims] + return FuncParameter(name=s, d_type=dtype, dimensions=dim_types) diff --git a/py2f/src/icon4py/py2f/py2fgen.py b/tools/src/icon4pytools/py2f/py2fgen.py similarity index 74% rename from py2f/src/icon4py/py2f/py2fgen.py rename to tools/src/icon4pytools/py2f/py2fgen.py index 0ad0ca780..7bdb799d1 100644 --- a/py2f/src/icon4py/py2f/py2fgen.py +++ b/tools/src/icon4pytools/py2f/py2fgen.py @@ -14,16 +14,14 @@ import pathlib import click -from icon4py.f2py.cffi_utils import generate_and_compile_cffi_plugin -from icon4py.f2py.codegen import ( - generate_and_write_f90_interface, - generate_c_header, -) -from icon4py.f2py.parsing import parse_functions_from_module + +from icon4pytools.py2f.cffi_utils import generate_and_compile_cffi_plugin +from icon4pytools.py2f.codegen import generate_and_write_f90_interface, generate_c_header +from icon4pytools.py2f.parsing import parse_functions_from_module @click.command( - "py2f90gen", + "py2fgen", ) @click.argument("module", type=str) @click.argument( @@ -37,7 +35,7 @@ def main(module: str, build_path: pathlib.Path) -> None: Args: - module: name of the python module containing the methods to be embedded. Those - methods have to be decoratoed with CffiMethod.register + methods have to be decorated with CffiMethod.register - build_path: directory where the generated code and compiled libraries are to be found. """ @@ -45,10 +43,8 @@ def main(module: str, build_path: pathlib.Path) -> None: build_path.mkdir(exist_ok=True, parents=True) plugin = parse_functions_from_module(module_name) c_header = generate_c_header(plugin) - generate_and_compile_cffi_plugin( - plugin.name, c_header, module_name, str(build_path) - ) - generate_and_write_f90_interface(plugin) + generate_and_compile_cffi_plugin(plugin.name, c_header, module_name, str(build_path)) + generate_and_write_f90_interface(build_path, plugin) if __name__ == "__main__": diff --git a/tools/src/icon4pytools/py2f/typing_utils.py b/tools/src/icon4pytools/py2f/typing_utils.py new file mode 100644 index 000000000..72ba5fa28 --- /dev/null +++ b/tools/src/icon4pytools/py2f/typing_utils.py @@ -0,0 +1,27 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Dimension +from gt4py.next.type_system.type_specifications import ScalarType +from gt4py.next.type_system.type_translation import from_type_hint + + +def parse_annotation(annotation) -> tuple[list[Dimension], ScalarType]: + type_spec = from_type_hint(annotation) + if isinstance(type_spec, ScalarType): + dtype = type_spec.kind + dims = [] + else: + dtype = type_spec.dtype.kind + dims = type_spec.dims + return dims, dtype diff --git a/py2f/src/icon4py/py2f/wrappers/__init__.py b/tools/src/icon4pytools/py2f/wrappers/__init__.py similarity index 100% rename from py2f/src/icon4py/py2f/wrappers/__init__.py rename to tools/src/icon4pytools/py2f/wrappers/__init__.py diff --git a/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py new file mode 100644 index 000000000..adb517b3e --- /dev/null +++ b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py @@ -0,0 +1,260 @@ +# 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 +# flake8: noqa +import numpy as np +from gt4py.next.common import Field +from gt4py.next.ffront.fbuiltins import int32 +from icon4py.model.atmosphere.diffusion.diffusion import ( + Diffusion, + DiffusionConfig, + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, + DiffusionParams, + DiffusionType, + PrognosticState, +) +from icon4py.model.common.dimension import ( + C2E2CDim, + C2E2CODim, + C2EDim, + CECDim, + CellDim, + E2C2VDim, + E2CDim, + E2VDim, + ECVDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalGridSize +from icon4py.model.common.grid.icon_grid import GridConfig, IconGrid +from icon4py.model.common.grid.vertical import VerticalGridSize, VerticalModelParams + +from icon4pytools.py2f.cffi_utils import CffiMethod, to_fields +""" +Wrapper module for diffusion granule. + +Module contains a diffusion_init and diffusion_run function that follow the architecture of +Fortran granule interfaces: +- all arguments needed from external sources are passed. +- passing of scalar types or fields of simple types + + +""" +# TODO (magdalena) Revise interface architecture with Fortran granules: +# The module variable to match the Fortran interface: where only fields are passed. +# We should rather instantiate the object init and return it. +diffusion: Diffusion() + +nproma = 50000 +field_sizes = {EdgeDim: nproma, CellDim: nproma, VertexDim: nproma} + +@to_fields(dim_sizes=field_sizes) +@CffiMethod.register +def diffusion_init( + nproma: int, + nlev: int, + n_shift_total: int, + nrdmax: float, + n_dyn_substeps: int, + nudge_max_coeff: float, + denom_diffu_v: float, + hdiff_smag_fac: float, + hdiff_order: int, + hdiff_efdt_ratio: float, + itype_sher: int, + itype_vn_diffu: int, + itype_t_diffu: int, + lhdiff_rcf: bool, + lhdiff_w: bool, + lhdiff_temp: bool, + l_limited_area: bool, + l_zdiffu_t: bool, + lsmag_3d: bool, + vct_a: Field[[KDim], float], + e_bln_c_s: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + geofac_n2s: Field[[CellDim, C2E2CODim], float], + geofac_grg_x: Field[[CellDim, C2E2CODim], float], + geofac_grg_y: Field[[CellDim, C2E2CODim], float], + nudgecoeff_e: Field[[EdgeDim], float], + zd_intcoef: Field[ + [CellDim, C2E2CDim, KDim], float + ], # special DSL field: zd_intcoef_dsl in mo_vertical_grid.f90 + zd_diffcoef: Field[ + [CellDim, KDim], float + ], # special DSL field mask instead of list: zd_diffcoef_dsl in mo_vertical_grid.f90 + wgtfac_c: Field[[CellDim, KDim], float], + theta_ref_mc: Field[[CellDim, KDim], float], + edges_vertex_idx: Field[[EdgeDim, E2VDim], int32], + edges_cell_idx: Field[[EdgeDim, E2CDim], int32], + edges_tangent_orientation: Field[[EdgeDim], float], + edges_primal_normal_vert_1: Field[[ECVDim], float], # shallow derived type in Fortran + edges_primal_normal_vert_2: Field[[ECVDim], float], # shallow derived type in Fortran + edges_dual_normal_vert_1: Field[[ECVDim], float], # shallow derived type in Fortran + edges_dual_normal_vert_2: Field[[ECVDim], float], # shallow derived type in Fortran + edges_inv_vert_vert_lengths: Field[[EdgeDim], float], + edges_inv_primal_edge_length: Field[[EdgeDim], float], + edges_inv_dual_edge_length: Field[[EdgeDim], float], + edges_area_edge: Field[[EdgeDim], float], + cells_neighbor_idx: Field[[CellDim, C2E2CDim], int32], + cells_edge_idx: Field[[CellDim, C2EDim], int32], + cells_area: Field[[CellDim], float], + # dsl specific additional args + mask_hdiff: Field[[CellDim, KDim], bool], + zd_vertoffset: Field[[CECDim], int32], # vertical offsets used in DSL for absolute indices zd_vertidx in mo_vertical_grid.f90 + rbf_coeff_1: Field[[VertexDim, V2EDim], float], # -> used in rbf_vec_interpol_vertex + rbf_coeff_2: Field[[VertexDim, V2EDim], float], # -> used in rbf_vec_interpol_vertex + verts_edge_idx: Field[[VertexDim, V2EDim], int32], # -> mo_intp_rbf_rbf_vec_interpol_vertex +): + """ + Instantiate and Initialize the diffusion object. + + should only accept simple fields as arguments for compatibility with the standalone + Fortran ICON Diffusion component (aka Diffusion granule) + + """ + if diffusion.initialized: + raise DuplicateInitializationException("Diffusion has already been already initialized") + + horizontal_size = HorizontalGridSize(num_cells=nproma, num_vertices=nproma, num_edges=nproma) + vertical_size = VerticalGridSize(num_lev=nlev) + mesh_config = GridConfig( + horizontal_config=horizontal_size, + vertical_config=vertical_size, + limited_area=l_limited_area, + n_shift_total=n_shift_total, + ) + # we need the start, end indices in order for the grid to be functional those are not passed + # to init in the Fortran diffusion granule since they are hidden away in the get_indices_[c,e,v] + # diffusion_run does take p_patch in order to pass it on to other subroutines (interpolation, get_indices... + + c2e2c0 = np.column_stack( + ( + (np.asarray(cells_neighbor_idx)), + (np.asarray(range(np.asarray(cells_neighbor_idx).shape[0]))), + ) + ) + e2c2v = np.asarray(edges_vertex_idx) + e2v = e2c2v[:, 0:2] + connectivities = { + E2CDim: np.asarray(edges_cell_idx), + E2C2VDim: e2c2v, + E2VDim: e2v, + C2EDim: np.asarray(cells_edge_idx), + C2E2CDim: (np.asarray(cells_neighbor_idx)), + C2E2CODim: c2e2c0, + V2EDim: np.asarray(verts_edge_idx), + } + # TODO (Magdalena) we need start_index, end_index: those are not passed in the fortran granules, + # because they are used through get_indices only + grid = IconGrid().with_config(mesh_config).with_connectivities(connectivities) + edge_params = EdgeParams( + tangent_orientation=edges_tangent_orientation, + inverse_primal_edge_lengths=edges_inv_primal_edge_length, + dual_normal_vert_x=edges_dual_normal_vert_1, + dual_normal_vert_y=edges_dual_normal_vert_2, + inverse_dual_edge_lengths=edges_inv_dual_edge_length, + inverse_vertex_vertex_lengths=edges_inv_vert_vert_lengths, + primal_normal_vert_x=edges_primal_normal_vert_1, + primal_normal_vert_y=edges_primal_normal_vert_2, + edge_areas=edges_area_edge, + ) + cell_params = CellParams(cells_area) + config: DiffusionConfig = DiffusionConfig( + diffusion_type=DiffusionType(hdiff_order), + hdiff_w=lhdiff_w, + hdiff_temp=lhdiff_temp, + type_vn_diffu=itype_vn_diffu, + smag_3d=lsmag_3d, + type_t_diffu=itype_t_diffu, + hdiff_efdt_ratio=hdiff_efdt_ratio, + hdiff_w_efdt_ratio=hdiff_efdt_ratio, + smagorinski_scaling_factor=hdiff_smag_fac, + n_substeps=n_dyn_substeps, + zdiffu_t=l_zdiffu_t, + hdiff_rcf=lhdiff_rcf, + velocity_boundary_diffusion_denom=denom_diffu_v, + max_nudging_coeff=nudge_max_coeff, + ) + vertical_params = VerticalModelParams(vct_a=vct_a, rayleigh_damping_height=nrdmax) + + derived_diffusion_params = DiffusionParams(config) + metric_state = DiffusionMetricState( + theta_ref_mc=theta_ref_mc, + wgtfac_c=wgtfac_c, + mask_hdiff=mask_hdiff, + zd_vertoffset=zd_vertoffset, + zd_diffcoef=zd_diffcoef, + zd_intcoef=zd_intcoef, + ) + interpolation_state = DiffusionInterpolationState( + e_bln_c_s, + rbf_coeff_1, + rbf_coeff_2, + geofac_div, + geofac_n2s, + geofac_grg_x, + geofac_grg_y, + nudgecoeff_e, + ) + + diffusion.init( + grid=grid, + config=config, + params=derived_diffusion_params, + vertical_params=vertical_params, + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_params, + cell_params=cell_params, + ) + + +@CffiMethod.register +def diffusion_run( + dtime: float, + linit: bool, + vn: Field[[EdgeDim, KDim], float], + w: Field[[CellDim, KDim], float], + theta_v: Field[[CellDim, KDim], float], + exner: Field[[CellDim, KDim], float], + div_ic: Field[[CellDim, KDim], float], + hdef_ic: Field[[CellDim, KDim], float], + dwdx: Field[[CellDim, KDim], float], + dwdy: Field[[CellDim, KDim], float], +): + diagnostic_state = DiffusionDiagnosticState(hdef_ic, div_ic, dwdx, dwdy) + prognostic_state = PrognosticState( + w=w, + vn=vn, + exner_pressure=exner, + theta_v=theta_v, + ) + if linit: + diffusion.initial_run( + diagnostic_state, + prognostic_state, + dtime, + ) + else: + diffusion.run(diagnostic_state, prognostic_state, dtime) + + +class DuplicateInitializationException(Exception): + """Raised if the component is already initilalized.""" + + pass diff --git a/tools/tests/py2f/test_cffi_utils.py b/tools/tests/py2f/test_cffi_utils.py new file mode 100644 index 000000000..d00571682 --- /dev/null +++ b/tools/tests/py2f/test_cffi_utils.py @@ -0,0 +1,84 @@ +# 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 cffi +import numpy as np +import pytest +from gt4py.next.common import Field +from gt4py.next.ffront.fbuiltins import float32, float64, int32, int64 + +from icon4py.model.common.dimension import E2CDim, EdgeDim, KDim, VertexDim +from icon4pytools.py2f.cffi_utils import UnknownDimensionException, to_fields + + +n_vertices = 9 +n_edges = 27 +levels = 12 +e2c_sparse_size = 2 + + +def random(*sizes): + return np.random.default_rng().uniform(size=sizes) + + +@pytest.mark.parametrize("pointer_type", ["float*", "double*"]) +def test_unpack_from_buffer_to_field(pointer_type: str): + @to_fields(dim_sizes={VertexDim: n_vertices, KDim: levels}) + def identity( + f1: Field[[VertexDim, KDim], float], factor: float + ) -> tuple[float, Field[[VertexDim, KDim], float]]: + return factor, f1 + + ffi = cffi.FFI() + input_array = random(n_vertices, levels) + input_factor = 0.5 + res_factor, result_field = identity( + ffi.from_buffer(pointer_type, input_array), input_factor + ) + assert res_factor == input_factor + assert np.allclose(np.asarray(result_field), input_array) + + +def test_unpack_only_scalar_args(): + @to_fields(dim_sizes={}) + def multiply(f1: float, f2: float32, f3: float64, i3: int64, i2: int32, i1: int): + return f1 * f2 * f3, i1 * i2 * i3 + + f_res, i_res = multiply(0.5, 2.0, 3.0, 1, 2, 3) + assert f_res == 3.0 + assert i_res == 6 + + +@pytest.mark.parametrize("field_type", [int32, int, int64]) +def test_unpack_local_field(field_type): + ffi = cffi.FFI() + + @to_fields(dim_sizes={EdgeDim: n_edges, E2CDim: e2c_sparse_size}) + def local_field(f1: Field[[EdgeDim, E2CDim], field_type]): + return f1 + + input_field = np.arange(n_edges * e2c_sparse_size).reshape( + (n_edges, e2c_sparse_size) + ) + res_field = local_field(ffi.from_buffer("int*", input_field)) + assert np.all(res_field == input_field) + + +def test_unknown_dimension_raises_exception(): + @to_fields(dim_sizes={}) + def do_nothing(f1: Field[[VertexDim], float]): + pass + + input_array = random() + with pytest.raises(UnknownDimensionException, match=r"size of dimension "): + do_nothing(input_array) diff --git a/py2f/tests/test_code_generation.py b/tools/tests/py2f/test_code_generation.py similarity index 98% rename from py2f/tests/test_code_generation.py rename to tools/tests/py2f/test_code_generation.py index 0ae91102a..74a9326de 100644 --- a/py2f/tests/test_code_generation.py +++ b/tools/tests/py2f/test_code_generation.py @@ -14,7 +14,8 @@ import pytest from gt4py.next.type_system.type_specifications import ScalarKind -from icon4py.py2f.codegen import ( + +from icon4pytools.py2f.codegen import ( CffiPlugin, CHeaderGenerator, DimensionType, diff --git a/py2f/tests/test_parsing_wrapper.py b/tools/tests/py2f/test_parsing_wrapper.py similarity index 86% rename from py2f/tests/test_parsing_wrapper.py rename to tools/tests/py2f/test_parsing_wrapper.py index 07ccbaf6a..ea1952038 100644 --- a/py2f/tests/test_parsing_wrapper.py +++ b/tools/tests/py2f/test_parsing_wrapper.py @@ -10,12 +10,12 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -from icon4py.py2f.cffi_utils import CffiMethod -from icon4py.py2f.parsing import parse_functions_from_module +from icon4pytools.py2f.cffi_utils import CffiMethod +from icon4pytools.py2f.parsing import parse_functions_from_module def test_parse_functions(): - path = "icon4py.py2f.wrappers.diffusion_wrapper" + path = "icon4pytools.py2f.wrappers.diffusion_wrapper" plugin = parse_functions_from_module(path) assert plugin.name == "diffusion_wrapper" diff --git a/py2f/setup.py b/tools/tests/py2f/test_py2fgen.py similarity index 60% rename from py2f/setup.py rename to tools/tests/py2f/test_py2fgen.py index 9c9f7b81c..acdace136 100644 --- a/py2f/setup.py +++ b/tools/tests/py2f/test_py2fgen.py @@ -11,8 +11,16 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from setuptools import setup +import pytest +from click.testing import CliRunner +from icon4pytools.py2f.py2fgen import main -if __name__ == "__main__": - setup() + +def test_py2fgen(): + cli = CliRunner() + module = "icon4pytools.py2f.wrappers.diffusion_wrapper" + build_path = "./build" + with cli.isolated_filesystem(): + result = cli.invoke(main, [module, build_path]) + assert result.exit_code == 0 From 226657bc984f8d51fb1d4c6d28e33014c9320206 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 24 Aug 2023 22:56:11 +0200 Subject: [PATCH 240/263] update README.md files --- README.md | 22 +++++++++++++--------- model/README.md | 2 ++ model/driver/README.md | 7 ++++--- 3 files changed, 19 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 64aab8c2f..802ad9699 100644 --- a/README.md +++ b/README.md @@ -112,39 +112,43 @@ tox # Run test suite in a specific environment (use `tox -a` to see list of envs) tox -e py310 ``` + The default `tox` environment is configured to generate HTML test coverage reports in `_reports/coverage_html/`. #### data dependent tests -Some test depend on serialized data generated by a full model run. Those test are marked with -`pytest.mark.datatest` and can be skipped by specifying + +Some test depend on serialized data generated by a full model run. Those test are marked with `pytest.mark.datatest` and can be skipped by specifying + ```bash pytest -v -m "not datatest" ``` - #### testing parallel code -Test for parallel code is collected in a package's `tests/mpi_tests` folder. -All parallel tests are marked with `@pytest.mark.mpi` and are skipped if -the `--with-mpi` is not passed option is not passed to `pytest` -In order to run them, you need a MPI installation on your system: -On Debian-Linux do + +Test for parallel code is collected in a package's `tests/mpi_tests` folder. All parallel tests are marked with `@pytest.mark.mpi` and are skipped if the `--with-mpi` is not passed option is not passed to `pytest` In order to run them, you need a MPI installation on your system: On Debian-Linux do + ```bash sudo apt-get install libopenmpi-dev ``` + or + ```bash sudo apt-get install mpich ``` + on MacOs + ```bash brew install mpich ``` + Then you can run the tests with + ```bash mpirun -np 2 pytest -v -s --with-mpi path/to/test/folder/mpi_tests ``` - ### Benchmarking We use [`pytest-benchmark`](https://pytest-benchmark.readthedocs.io/en/latest/) to benchmark the execution time of stencils in icon4py. To disable benchmarking during testing you can use `--benchmark-disable` when invoking `pytest`. diff --git a/model/README.md b/model/README.md index d59191505..e0fcfa571 100644 --- a/model/README.md +++ b/model/README.md @@ -5,7 +5,9 @@ This folder contains Python implementations for multiple ICON components. It includes the following packages: - `atmosphere/dycore`: Contains implementations of the dynamical core of the ICON model +- `atmosphere/diffusion`: Contains the implementation of diffusion in the ICON model - `common`: Contains shared functionality that is required by multiple components. +- `driver`: Contains the driving code for the model ## Installation Instructions diff --git a/model/driver/README.md b/model/driver/README.md index ceeadcae1..ecee60929 100644 --- a/model/driver/README.md +++ b/model/driver/README.md @@ -14,19 +14,20 @@ The code is meant to be changed and enlarged as we port new parts of the model. It runs single node or parallel versions. For parallel runs the domain has to be decomposed previousely through a full ICON run that generates the necessary serialized data. Test data for runs with 1, 2, 4 nodes are available. +## Installation +See the general instructions in the [README.md](../../README.md) in the base folder of the repository. ## Usage ```bash export ICON4PY_ROOT= -python driver/dycore_driver.py $ICON4PY_ROOT/testdata/ser_icondata/mpitask1/mch_ch_r04b09_dsl/ser_data --n_steps=2 --run_path=/home/magdalena/temp/icon +dycore_driver $ICON4PY_ROOT/testdata/ser_icondata/mpitask1/mch_ch_r04b09_dsl/ser_data --n_steps=2 --run_path=/home/magdalena/temp/icon ``` or if running in parallel ```bash -mpirun -np 2 python driver/dycore_driver.py $ICON4PY_ROOT/testdata/ser_icondata/mpitask2/mch_ch_r04b09_dsl/ser_data --mpi=True --n_steps=2 --run_path=/home/magdalena/temp/icon - +mpirun -np 2 dycore_driver $ICON4PY_ROOT/testdata/ser_icondata/mpitask2/mch_ch_r04b09_dsl/ser_data --mpi=True --n_steps=2 --run_path=/home/magdalena/temp/icon ``` #### remarks From a8c41f3b7507c4c4b047df583480d786899c9aa2 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 24 Aug 2023 23:00:07 +0200 Subject: [PATCH 241/263] pre-commit fixes --- .../diffusion_tests/test_diffusion.py | 1 - .../model/atmosphere/dycore/__init__.py | 1 + model/common/tests/test_grid_manager.py | 1 - model/driver/README.md | 18 +++++++----------- .../src/icon4py/model/driver/io_utils.py | 1 + tools/src/icon4pytools/py2f/cffi_utils.py | 7 +++---- .../py2f/wrappers/diffusion_wrapper.py | 7 ++++++- tools/tests/py2f/test_cffi_utils.py | 10 +++------- tools/tests/py2f/test_py2fgen.py | 1 - 9 files changed, 21 insertions(+), 26 deletions(-) diff --git a/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py b/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py index 2a57d665f..20c134a24 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py @@ -275,7 +275,6 @@ def test_verify_diffusion_init_against_other_regular_savepoint( [ ("2021-06-20T12:00:10.000", "2021-06-20T12:00:10.000"), ("2021-06-20T12:00:20.000", "2021-06-20T12:00:20.000"), - ("2021-06-20T12:01:00.000", "2021-06-20T12:01:00.000"), ], ) def test_run_diffusion_single_step( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/__init__.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/__init__.py index 6cfd5ac92..dab708955 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/__init__.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/__init__.py @@ -14,6 +14,7 @@ from packaging import version as pkg_version + __all__ = [ "__author__", "__copyright__", diff --git a/model/common/tests/test_grid_manager.py b/model/common/tests/test_grid_manager.py index d87acbf5a..28b714b7d 100644 --- a/model/common/tests/test_grid_manager.py +++ b/model/common/tests/test_grid_manager.py @@ -11,7 +11,6 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import logging -from pathlib import Path from uuid import uuid4 import numpy as np diff --git a/model/driver/README.md b/model/driver/README.md index ecee60929..035ec4446 100644 --- a/model/driver/README.md +++ b/model/driver/README.md @@ -1,20 +1,17 @@ # Dummy driver for Python ICON port -`dycore_driver.py` contains a simple python program to run the experimental ICON python port. So far -code mostly draws on serialized ICON data until we increasingly can initialize and run the model by independently. +`dycore_driver.py` contains a simple python program to run the experimental ICON python port. So far code mostly draws on serialized ICON data until we increasingly can initialize and run the model by independently. -It initializes the grid from serialized data from a `mch_ch_r04b09_dsl` run and configures a -timeloop functionality based on that configuration. +It initializes the grid from serialized data from a `mch_ch_r04b09_dsl` run and configures a timeloop functionality based on that configuration. -Currently, there is does *no real timestepping*, instead it calls a dummy timestep that -serves a batch of new serialized input fields from ICON. +Currently, there is does _no real timestepping_, instead it calls a dummy timestep that serves a batch of new serialized input fields from ICON. The code is meant to be changed and enlarged as we port new parts of the model. -It runs single node or parallel versions. For parallel runs the domain has to be decomposed previousely -through a full ICON run that generates the necessary serialized data. Test data for runs with 1, 2, 4 nodes -are available. +It runs single node or parallel versions. For parallel runs the domain has to be decomposed previousely through a full ICON run that generates the necessary serialized data. Test data for runs with 1, 2, 4 nodes are available. + ## Installation + See the general instructions in the [README.md](../../README.md) in the base folder of the repository. ## Usage @@ -34,7 +31,6 @@ mpirun -np 2 dycore_driver $ICON4PY_ROOT/testdata/ser_icondata/mpitask2/mch_ch_r - First (required) arg is the folder where the serialized input data is stored. The same input data is for the unit tests is used. The path in the example is where the data is put when downloaded via the unit tests. - data for a serial (single node) run can be downloaded from https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p. -- parallel runs are possible if corresponding data is provided, which is currently available for test with 2 or 4 MPI processes: -check [fixtures.py](../common/src/icon4py/model/common/test_utils/fixtures.py) for download urls. +- parallel runs are possible if corresponding data is provided, which is currently available for test with 2 or 4 MPI processes: check [fixtures.py](../common/src/icon4py/model/common/test_utils/fixtures.py) for download urls. - The serialized data used contains only 5 timesteps so `--n_steps > 2` will throw an exception. - The code logs to file and to console. Debug logging is only going to file. The log directory can be changed with the --run_path option. diff --git a/model/driver/src/icon4py/model/driver/io_utils.py b/model/driver/src/icon4py/model/driver/io_utils.py index bb69fef9f..de08ca15a 100644 --- a/model/driver/src/icon4py/model/driver/io_utils.py +++ b/model/driver/src/icon4py/model/driver/io_utils.py @@ -32,6 +32,7 @@ from icon4py.model.common.grid.vertical import VerticalModelParams from icon4py.model.common.test_utils import serialbox_utils as sb + SB_ONLY_MSG = "Only ser_type='sb' is implemented so far." SIMULATION_START_DATE = "2021-06-20T12:00:10.000" diff --git a/tools/src/icon4pytools/py2f/cffi_utils.py b/tools/src/icon4pytools/py2f/cffi_utils.py index bcceccec2..ed12d165b 100644 --- a/tools/src/icon4pytools/py2f/cffi_utils.py +++ b/tools/src/icon4pytools/py2f/cffi_utils.py @@ -99,15 +99,14 @@ class UnknownDimensionException(Exception): pass -def to_fields( - dim_sizes: dict[Dimension, int] -): +def to_fields(dim_sizes: dict[Dimension, int]): """ Pack/Unpack Fortran 2d arrays to numpy arrays with using CFFI frombuffer. - #TODO (magdalena) handle in a better way? + Args: dim_sizes: dictionary containing the sizes of the dimension. + #TODO (magdalena) handle dimension sizes in a better way? """ ffi = cffi.FFI() dim_sizes = dim_sizes diff --git a/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py index adb517b3e..580d07dc7 100644 --- a/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py +++ b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py @@ -44,6 +44,8 @@ from icon4py.model.common.grid.vertical import VerticalGridSize, VerticalModelParams from icon4pytools.py2f.cffi_utils import CffiMethod, to_fields + + """ Wrapper module for diffusion granule. @@ -62,6 +64,7 @@ nproma = 50000 field_sizes = {EdgeDim: nproma, CellDim: nproma, VertexDim: nproma} + @to_fields(dim_sizes=field_sizes) @CffiMethod.register def diffusion_init( @@ -115,7 +118,9 @@ def diffusion_init( cells_area: Field[[CellDim], float], # dsl specific additional args mask_hdiff: Field[[CellDim, KDim], bool], - zd_vertoffset: Field[[CECDim], int32], # vertical offsets used in DSL for absolute indices zd_vertidx in mo_vertical_grid.f90 + zd_vertoffset: Field[ + [CECDim], int32 + ], # vertical offsets used in DSL for absolute indices zd_vertidx in mo_vertical_grid.f90 rbf_coeff_1: Field[[VertexDim, V2EDim], float], # -> used in rbf_vec_interpol_vertex rbf_coeff_2: Field[[VertexDim, V2EDim], float], # -> used in rbf_vec_interpol_vertex verts_edge_idx: Field[[VertexDim, V2EDim], int32], # -> mo_intp_rbf_rbf_vec_interpol_vertex diff --git a/tools/tests/py2f/test_cffi_utils.py b/tools/tests/py2f/test_cffi_utils.py index d00571682..d8fdf6196 100644 --- a/tools/tests/py2f/test_cffi_utils.py +++ b/tools/tests/py2f/test_cffi_utils.py @@ -16,8 +16,8 @@ import pytest from gt4py.next.common import Field from gt4py.next.ffront.fbuiltins import float32, float64, int32, int64 - from icon4py.model.common.dimension import E2CDim, EdgeDim, KDim, VertexDim + from icon4pytools.py2f.cffi_utils import UnknownDimensionException, to_fields @@ -42,9 +42,7 @@ def identity( ffi = cffi.FFI() input_array = random(n_vertices, levels) input_factor = 0.5 - res_factor, result_field = identity( - ffi.from_buffer(pointer_type, input_array), input_factor - ) + res_factor, result_field = identity(ffi.from_buffer(pointer_type, input_array), input_factor) assert res_factor == input_factor assert np.allclose(np.asarray(result_field), input_array) @@ -67,9 +65,7 @@ def test_unpack_local_field(field_type): def local_field(f1: Field[[EdgeDim, E2CDim], field_type]): return f1 - input_field = np.arange(n_edges * e2c_sparse_size).reshape( - (n_edges, e2c_sparse_size) - ) + input_field = np.arange(n_edges * e2c_sparse_size).reshape((n_edges, e2c_sparse_size)) res_field = local_field(ffi.from_buffer("int*", input_field)) assert np.all(res_field == input_field) diff --git a/tools/tests/py2f/test_py2fgen.py b/tools/tests/py2f/test_py2fgen.py index acdace136..8330e4a26 100644 --- a/tools/tests/py2f/test_py2fgen.py +++ b/tools/tests/py2f/test_py2fgen.py @@ -11,7 +11,6 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import pytest from click.testing import CliRunner from icon4pytools.py2f.py2fgen import main From bce21eedb4ab99f60dbf475fa64fab50126ec96f Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 25 Aug 2023 09:11:10 +0200 Subject: [PATCH 242/263] simple cleanups --- .github/workflows/icon4pytools-qa.yml | 4 ---- base-requirements-dev.txt | 6 +----- base-requirements.txt | 2 -- model/.pre-commit-config.yaml | 2 +- model/pytest.ini | 5 ----- model/tox.ini | 2 +- pytest.ini | 2 -- requirements.txt | 1 + tox.ini | 2 +- 9 files changed, 5 insertions(+), 21 deletions(-) delete mode 100644 model/pytest.ini diff --git a/.github/workflows/icon4pytools-qa.yml b/.github/workflows/icon4pytools-qa.yml index ce407e2ed..9afe586f3 100644 --- a/.github/workflows/icon4pytools-qa.yml +++ b/.github/workflows/icon4pytools-qa.yml @@ -31,10 +31,6 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v2 - - name: Install boost - run: | - sudo apt-get update - sudo apt-get install libboost-all-dev - name: Set up Python uses: actions/setup-python@v3 with: diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 021c88962..019001292 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -3,7 +3,6 @@ git+https://github.com/GridTools/serialbox#egg=serialbox&subdirectory=src/serialbox-python git+https://github.com/boeschf/GHEX.git@pipify#subdirectory=bindings/python - # PyPI bump2version>=1.0.1 coverage[toml]>=5.0 @@ -27,8 +26,5 @@ pytest-xdist[psutil]>=2.2 pytest-mpi>=0.6 setuptools>=40.8.0 wheel>=0.37.1 -tox>=3.25 +tox >= 3.25 wget>=3.2 - - - diff --git a/base-requirements.txt b/base-requirements.txt index 7b7d94fac..0720ba258 100644 --- a/base-requirements.txt +++ b/base-requirements.txt @@ -1,5 +1,3 @@ # VCS gt4py @ git+https://github.com/GridTools/gt4py.git@main pyghex @ git+https://github.com/boeschf/GHEX.git@pipify#subdirectory=bindings/python - - diff --git a/model/.pre-commit-config.yaml b/model/.pre-commit-config.yaml index 1b7302dae..f5a626ce7 100644 --- a/model/.pre-commit-config.yaml +++ b/model/.pre-commit-config.yaml @@ -76,7 +76,7 @@ repos: types_or: [json, yaml, markdown] - repo: https://github.com/asottile/yesqa - rev: v1.5.0 + rev: v1.3.0 hooks: - id: yesqa exclude: model/common/decomposition/parallel_setup.py diff --git a/model/pytest.ini b/model/pytest.ini deleted file mode 100644 index 3ed8b0a52..000000000 --- a/model/pytest.ini +++ /dev/null @@ -1,5 +0,0 @@ -[pytest] -norecursedirs = _external_src -markers = - datatest: this test uses binary data -python_functions = test_* bench_* diff --git a/model/tox.ini b/model/tox.ini index 6714d9033..9cb4ccc95 100644 --- a/model/tox.ini +++ b/model/tox.ini @@ -14,7 +14,7 @@ passenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -m "not datatest" -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src common/src driver/src + -pytest -v -m "not datatest" -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src atmosphere/diffusion/src common/src driver/src pytest -v -m "not datatest" -s -n auto --cov --cov-append --benchmark-disable commands_post = rm -rf tests/_reports/coverage_html diff --git a/pytest.ini b/pytest.ini index 3ed8b0a52..84a9d34b4 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,5 +1,3 @@ [pytest] norecursedirs = _external_src -markers = - datatest: this test uses binary data python_functions = test_* bench_* diff --git a/requirements.txt b/requirements.txt index 803c4f811..14c1cc762 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,5 +5,6 @@ ./model/atmosphere/diffusion ./model/common ./model/driver + # icon4pytools ./tools diff --git a/tox.ini b/tox.ini index 0370f167b..d1d7611fa 100644 --- a/tox.ini +++ b/tox.ini @@ -14,7 +14,7 @@ passenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -m "not datatest" -s -n auto -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/dycore/src model/common/src tools/src + -pytest -v -m "not datatest" -s -n auto -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/dycore/src model/atmosphere/diffusion/src model/driver/src model/common/src tools/src pytest -v -m "not datatest" -s -n auto --cov --cov-append --benchmark-disable commands_post = rm -rf tests/_reports/coverage_html From 2d858af8a93a235425bd779c7b501cc0171470c2 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 25 Aug 2023 19:07:37 +0200 Subject: [PATCH 243/263] (fix) move interpolation stencil to common package --- .../model/atmosphere/diffusion/diffusion.py | 8 +++--- .../model/common/interpolation/__init__.py | 12 +++++++++ .../interpolation_fields.py | 0 .../common/interpolation/stencils/__init__.py | 12 +++++++++ .../mo_intp_rbf_rbf_vec_interpol_vertex.py | 0 model/common/tests/conftest.py | 2 ++ .../common/tests/test_interpolation_fields.py | 4 +-- ...est_mo_intp_rbf_rbf_vec_interpol_vertex.py | 4 +-- .../src/icon4pytools/liskov/external/gt4py.py | 6 ++++- tools/tests/icon4pygen/test_codegen.py | 26 ++++++++++++++++--- 10 files changed, 61 insertions(+), 13 deletions(-) create mode 100644 model/common/src/icon4py/model/common/interpolation/__init__.py rename model/common/src/icon4py/model/common/{field_management => interpolation}/interpolation_fields.py (100%) create mode 100644 model/common/src/icon4py/model/common/interpolation/stencils/__init__.py rename model/{atmosphere/dycore/src/icon4py/model/atmosphere/dycore => common/src/icon4py/model/common/interpolation/stencils}/mo_intp_rbf_rbf_vec_interpol_vertex.py (100%) rename model/{atmosphere/dycore => common}/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py (95%) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py index 4a8d8a110..9aea7121d 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -67,11 +67,6 @@ from icon4py.model.atmosphere.diffusion.stencils.update_theta_and_exner import ( update_theta_and_exner, ) - -# TODO (magdalena) should go to common.math -from icon4py.model.atmosphere.dycore.mo_intp_rbf_rbf_vec_interpol_vertex import ( - mo_intp_rbf_rbf_vec_interpol_vertex, -) from icon4py.model.common.constants import ( CPD, DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO, @@ -89,6 +84,9 @@ ) from icon4py.model.common.grid.icon_grid import IconGrid from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.interpolation.stencils.mo_intp_rbf_rbf_vec_interpol_vertex import ( + mo_intp_rbf_rbf_vec_interpol_vertex, +) # flake8: noqa diff --git a/model/common/src/icon4py/model/common/interpolation/__init__.py b/model/common/src/icon4py/model/common/interpolation/__init__.py new file mode 100644 index 000000000..15dfdb009 --- /dev/null +++ b/model/common/src/icon4py/model/common/interpolation/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/model/common/src/icon4py/model/common/field_management/interpolation_fields.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py similarity index 100% rename from model/common/src/icon4py/model/common/field_management/interpolation_fields.py rename to model/common/src/icon4py/model/common/interpolation/interpolation_fields.py diff --git a/model/common/src/icon4py/model/common/interpolation/stencils/__init__.py b/model/common/src/icon4py/model/common/interpolation/stencils/__init__.py new file mode 100644 index 000000000..15dfdb009 --- /dev/null +++ b/model/common/src/icon4py/model/common/interpolation/stencils/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py similarity index 100% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_intp_rbf_rbf_vec_interpol_vertex.py rename to model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py diff --git a/model/common/tests/conftest.py b/model/common/tests/conftest.py index 684c0c381..e9024dc03 100644 --- a/model/common/tests/conftest.py +++ b/model/common/tests/conftest.py @@ -14,6 +14,7 @@ from icon4py.model.common.test_utils.data_handling import download_and_extract from icon4py.model.common.test_utils.fixtures import ( # noqa F401 + backend, base_path, damping_height, data_provider, @@ -23,6 +24,7 @@ grid_savepoint, icon_grid, interpolation_savepoint, + mesh, processor_props, ranked_data_path, ) diff --git a/model/common/tests/test_interpolation_fields.py b/model/common/tests/test_interpolation_fields.py index 7c76c446e..1d49e4ded 100644 --- a/model/common/tests/test_interpolation_fields.py +++ b/model/common/tests/test_interpolation_fields.py @@ -27,10 +27,10 @@ import pytest from icon4py.model.common.dimension import EdgeDim -from icon4py.model.common.field_management.interpolation_fields import ( +from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex +from icon4py.model.common.interpolation.interpolation_fields import ( compute_c_lin_e, ) -from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex @pytest.mark.datatest diff --git a/model/atmosphere/dycore/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py similarity index 95% rename from model/atmosphere/dycore/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py rename to model/common/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index 51fc28295..02ae28148 100644 --- a/model/atmosphere/dycore/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -15,10 +15,10 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_intp_rbf_rbf_vec_interpol_vertex import ( +from icon4py.model.common.dimension import EdgeDim, KDim, V2EDim, VertexDim +from icon4py.model.common.interpolation.stencils.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) -from icon4py.model.common.dimension import EdgeDim, KDim, V2EDim, VertexDim from icon4py.model.common.test_utils.helpers import ( StencilTest, random_field, diff --git a/tools/src/icon4pytools/liskov/external/gt4py.py b/tools/src/icon4pytools/liskov/external/gt4py.py index b84c27e51..b62c6eefd 100644 --- a/tools/src/icon4pytools/liskov/external/gt4py.py +++ b/tools/src/icon4pytools/liskov/external/gt4py.py @@ -32,7 +32,11 @@ class UpdateFieldsWithGt4PyStencils(Step): - _STENCIL_PACKAGES = ["atmosphere.dycore", "atmosphere.diffusion.stencils"] + _STENCIL_PACKAGES = [ + "atmosphere.dycore", + "atmosphere.diffusion.stencils", + "common.interpolation", + ] def __init__(self, parsed: IntegrationCodeInterface): self.parsed = parsed diff --git a/tools/tests/icon4pygen/test_codegen.py b/tools/tests/icon4pygen/test_codegen.py index 4ca1a1d1f..72c35d898 100644 --- a/tools/tests/icon4pygen/test_codegen.py +++ b/tools/tests/icon4pygen/test_codegen.py @@ -15,7 +15,9 @@ import pkgutil import re +import icon4py.model.atmosphere.diffusion.stencils as diffusion import icon4py.model.atmosphere.dycore as dycore +import icon4py.model.common.interpolation.stencils as intp import pytest from click.testing import CliRunner @@ -25,6 +27,9 @@ DYCORE_PKG = "atmosphere.dycore" +INTERPOLATION_PKG = "common.interpolation.stencils" +DIFFUSION_PKG = "atmosphere.diffusion.stencils" + LEVELS_PER_THREAD = "1" BLOCK_SIZE = "128" OUTPATH = "." @@ -36,9 +41,21 @@ def cli(): def dycore_fencils() -> list[tuple[str, str]]: - pkgpath = os.path.dirname(dycore.__file__) + return _fencils(dycore.__file__, DYCORE_PKG) + + +def interpolation_fencils() -> list[tuple[str, str]]: + return _fencils(intp.__file__, INTERPOLATION_PKG) + + +def diffusion_fencils() -> list[tuple[str, str]]: + return _fencils(diffusion.__file__, DIFFUSION_PKG) + + +def _fencils(module_name, package_name) -> list[tuple[str, str]]: + pkgpath = os.path.dirname(module_name) stencils = [name for _, name, _ in pkgutil.iter_modules([pkgpath])] - fencils = [(DYCORE_PKG, stencil) for stencil in stencils] + fencils = [(package_name, stencil) for stencil in stencils] return fencils @@ -111,7 +128,10 @@ def check_code_was_generated(stencil_name: str) -> None: check_cpp_codegen(f"{stencil_name}.cpp") -@pytest.mark.parametrize(("stencil_module", "stencil_name"), dycore_fencils()) +@pytest.mark.parametrize( + ("stencil_module", "stencil_name"), + dycore_fencils() + interpolation_fencils() + diffusion_fencils(), +) def test_codegen_dycore(cli, stencil_module, stencil_name) -> None: module_path = get_stencil_module_path(stencil_module, stencil_name) with cli.isolated_filesystem(): From d8b9090bb17fe306883ac77374b930e03b9f3051 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 29 Aug 2023 09:47:25 +0200 Subject: [PATCH 244/263] switch to GHEX main. --- base-requirements-dev.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 019001292..991b3d473 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -1,7 +1,7 @@ # VCS -e git+https://github.com/GridTools/gt4py.git@main#egg=gt4py git+https://github.com/GridTools/serialbox#egg=serialbox&subdirectory=src/serialbox-python -git+https://github.com/boeschf/GHEX.git@pipify#subdirectory=bindings/python +git+https://github.com/ghex-org/GHEX.git#subdirectory=bindings/python # PyPI bump2version>=1.0.1 From ee819759e424d5aafcacdfbcd84103b63be98f5c Mon Sep 17 00:00:00 2001 From: Magdalena Date: Thu, 31 Aug 2023 15:47:40 +0200 Subject: [PATCH 245/263] Update model/driver/README.md Co-authored-by: Enrique G. Paredes <18477+egparedes@users.noreply.github.com> --- model/driver/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/driver/README.md b/model/driver/README.md index 035ec4446..d8671299e 100644 --- a/model/driver/README.md +++ b/model/driver/README.md @@ -1,6 +1,6 @@ # Dummy driver for Python ICON port -`dycore_driver.py` contains a simple python program to run the experimental ICON python port. So far code mostly draws on serialized ICON data until we increasingly can initialize and run the model by independently. +`dycore_driver.py` contains a simple python program to run the experimental ICON python port. So far code mostly draws on serialized ICON data until we increasingly can initialize and run the model independently. It initializes the grid from serialized data from a `mch_ch_r04b09_dsl` run and configures a timeloop functionality based on that configuration. From d4e0bee36a6d58e61e5d19bf50e634962f827c7a Mon Sep 17 00:00:00 2001 From: Magdalena Date: Thu, 31 Aug 2023 16:49:46 +0200 Subject: [PATCH 246/263] Update tools/src/icon4pytools/py2f/cffi_utils.py Co-authored-by: Enrique G. Paredes <18477+egparedes@users.noreply.github.com> --- tools/src/icon4pytools/py2f/cffi_utils.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tools/src/icon4pytools/py2f/cffi_utils.py b/tools/src/icon4pytools/py2f/cffi_utils.py index ed12d165b..c25183132 100644 --- a/tools/src/icon4pytools/py2f/cffi_utils.py +++ b/tools/src/icon4pytools/py2f/cffi_utils.py @@ -35,10 +35,7 @@ class CffiMethod: @classmethod def register(cls, func): - try: - cls._registry[func.__module__].append(func.__name__) - except KeyError: - cls._registry[func.__module__] = [func.__name__] + cls._registry.setdefault(func.__module__, []).append(func.__name__) @functools.wraps(func) def wrapper(*args, **kwargs): From 3ac39f9276c793331c9b3ce2d7b63efcb5937285 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 31 Aug 2023 16:52:41 +0200 Subject: [PATCH 247/263] PR review fixes: README s move module docstring to correct location --- README.md | 34 ----------------- model/README.md | 38 +++++++++++++++++++ model/driver/README.md | 8 ++-- .../py2f/wrappers/diffusion_wrapper.py | 19 +++++----- 4 files changed, 52 insertions(+), 47 deletions(-) diff --git a/README.md b/README.md index 802ad9699..2b4738ba6 100644 --- a/README.md +++ b/README.md @@ -115,40 +115,6 @@ tox -e py310 The default `tox` environment is configured to generate HTML test coverage reports in `_reports/coverage_html/`. -#### data dependent tests - -Some test depend on serialized data generated by a full model run. Those test are marked with `pytest.mark.datatest` and can be skipped by specifying - -```bash -pytest -v -m "not datatest" -``` - -#### testing parallel code - -Test for parallel code is collected in a package's `tests/mpi_tests` folder. All parallel tests are marked with `@pytest.mark.mpi` and are skipped if the `--with-mpi` is not passed option is not passed to `pytest` In order to run them, you need a MPI installation on your system: On Debian-Linux do - -```bash -sudo apt-get install libopenmpi-dev -``` - -or - -```bash -sudo apt-get install mpich -``` - -on MacOs - -```bash -brew install mpich -``` - -Then you can run the tests with - -```bash -mpirun -np 2 pytest -v -s --with-mpi path/to/test/folder/mpi_tests -``` - ### Benchmarking We use [`pytest-benchmark`](https://pytest-benchmark.readthedocs.io/en/latest/) to benchmark the execution time of stencils in icon4py. To disable benchmarking during testing you can use `--benchmark-disable` when invoking `pytest`. diff --git a/model/README.md b/model/README.md index e0fcfa571..83bd729fd 100644 --- a/model/README.md +++ b/model/README.md @@ -24,3 +24,41 @@ pip install -r requirements-dev.txt ``` **Note**: For more information specific to each component, please refer to the README in their respective subfolders. + +### Testing + +See the repository [README](../README.md) for general information. + +#### Data dependent tests + +Some test depend on serialized data generated by a full model run. Those test are marked with `pytest.mark.datatest` and can be skipped by specifying + +```bash +pytest -v -m "not datatest" +``` + +#### Testing parallel code + +Tests for parallel codes using MPI are collected in specific subpackages of the model components test folders (e.g. `diffusion_tests/mpi_tests`). All parallel tests are marked with `@pytest.mark.mpi` and are skipped if the `--with-mpi` is not passed option is not passed to `pytest` In order to run them, you need a MPI installation on your system: On Debian-Linux do + +```bash +sudo apt-get install libopenmpi-dev +``` + +or + +```bash +sudo apt-get install mpich +``` + +on MacOs + +```bash +brew install mpich +``` + +Then you can run the tests with + +```bash +mpirun -np 2 pytest -v -s --with-mpi path/to/test/folder/mpi_tests +``` diff --git a/model/driver/README.md b/model/driver/README.md index d8671299e..aec461abb 100644 --- a/model/driver/README.md +++ b/model/driver/README.md @@ -1,10 +1,10 @@ # Dummy driver for Python ICON port -`dycore_driver.py` contains a simple python program to run the experimental ICON python port. So far code mostly draws on serialized ICON data until we increasingly can initialize and run the model independently. +`dycore_driver.py` contains a simple python program to run the experimental ICON python port. So far the code mostly draws on serialized ICON data until we increasingly can initialize and run the model independently. It initializes the grid from serialized data from a `mch_ch_r04b09_dsl` run and configures a timeloop functionality based on that configuration. -Currently, there is does _no real timestepping_, instead it calls a dummy timestep that serves a batch of new serialized input fields from ICON. +Currently, it does _no real timestepping_, instead it calls a dummy timestep that serves a batch of new serialized input fields from ICON. The code is meant to be changed and enlarged as we port new parts of the model. @@ -27,9 +27,9 @@ or if running in parallel mpirun -np 2 dycore_driver $ICON4PY_ROOT/testdata/ser_icondata/mpitask2/mch_ch_r04b09_dsl/ser_data --mpi=True --n_steps=2 --run_path=/home/magdalena/temp/icon ``` -#### remarks +#### Remarks -- First (required) arg is the folder where the serialized input data is stored. The same input data is for the unit tests is used. The path in the example is where the data is put when downloaded via the unit tests. +- First (required) arg is the folder where the serialized input data is stored. The input data is the same as is used in the unit tests. The path in the example is where the data is put when downloaded via the unit tests. - data for a serial (single node) run can be downloaded from https://polybox.ethz.ch/index.php/s/vcsCYmCFA9Qe26p. - parallel runs are possible if corresponding data is provided, which is currently available for test with 2 or 4 MPI processes: check [fixtures.py](../common/src/icon4py/model/common/test_utils/fixtures.py) for download urls. - The serialized data used contains only 5 timesteps so `--n_steps > 2` will throw an exception. diff --git a/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py index 580d07dc7..002087bbd 100644 --- a/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py +++ b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py @@ -11,6 +11,16 @@ # # SPDX-License-Identifier: GPL-3.0-or-later # flake8: noqa +""" +Wrapper module for diffusion granule. + +Module contains a diffusion_init and diffusion_run function that follow the architecture of +Fortran granule interfaces: +- all arguments needed from external sources are passed. +- passing of scalar types or fields of simple types + + +""" import numpy as np from gt4py.next.common import Field from gt4py.next.ffront.fbuiltins import int32 @@ -46,16 +56,7 @@ from icon4pytools.py2f.cffi_utils import CffiMethod, to_fields -""" -Wrapper module for diffusion granule. -Module contains a diffusion_init and diffusion_run function that follow the architecture of -Fortran granule interfaces: -- all arguments needed from external sources are passed. -- passing of scalar types or fields of simple types - - -""" # TODO (magdalena) Revise interface architecture with Fortran granules: # The module variable to match the Fortran interface: where only fields are passed. # We should rather instantiate the object init and return it. From 419f7eb9df66af8925fef364f29cd75a24be3f1c Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Mon, 4 Sep 2023 22:47:35 +0200 Subject: [PATCH 248/263] fix interpolation stencils path in liskov --- tools/src/icon4pytools/liskov/external/gt4py.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/src/icon4pytools/liskov/external/gt4py.py b/tools/src/icon4pytools/liskov/external/gt4py.py index b62c6eefd..6edcc302a 100644 --- a/tools/src/icon4pytools/liskov/external/gt4py.py +++ b/tools/src/icon4pytools/liskov/external/gt4py.py @@ -35,7 +35,7 @@ class UpdateFieldsWithGt4PyStencils(Step): _STENCIL_PACKAGES = [ "atmosphere.dycore", "atmosphere.diffusion.stencils", - "common.interpolation", + "common.interpolation.stencils", ] def __init__(self, parsed: IntegrationCodeInterface): From 9230345d72fe9203bda99c175a58f93ec9894b88 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 5 Sep 2023 09:31:12 +0200 Subject: [PATCH 249/263] add pre-commit and flake8 config from main to new packages --- model/atmosphere/diffusion/.flake8 | 42 +++++++ .../diffusion/.pre-commit-config.yaml | 114 ++++++++++++++++++ .../diffusion/diffusion_tests/conftest.py | 5 +- .../mpi_tests/test_parallel_diffusion.py | 18 +-- ...st_apply_nabla2_and_nabla4_global_to_vn.py | 8 +- .../test_apply_nabla2_and_nabla4_to_vn.py | 8 +- .../diffusion_tests/test_apply_nabla2_to_w.py | 8 +- ...st_calculate_diagnostics_for_turbulence.py | 14 +-- ...ate_horizontal_gradients_for_turbulence.py | 6 +- ...ate_nabla2_and_smag_coefficients_for_vn.py | 12 +- .../test_calculate_nabla2_for_w.py | 6 +- .../test_calculate_nabla2_of_theta.py | 8 +- .../diffusion_tests/test_calculate_nabla4.py | 12 +- .../diffusion_tests/test_diffusion.py | 54 ++------- .../diffusion_tests/test_diffusion_utils.py | 12 +- ...d_for_grid_point_cold_pools_enhancement.py | 6 +- ...fusion_nabla_of_theta_over_steep_points.py | 7 +- model/atmosphere/diffusion/pyproject.toml | 78 ++++++------ .../model/atmosphere/diffusion/diffusion.py | 102 ++++------------ .../atmosphere/diffusion/diffusion_states.py | 12 +- .../atmosphere/diffusion/diffusion_utils.py | 28 ++--- .../stencils/apply_diffusion_to_vn.py | 8 +- ...ute_horizontal_gradients_for_turbulance.py | 12 +- .../apply_nabla2_and_nabla4_global_to_vn.py | 4 +- .../apply_nabla2_to_vn_in_lateral_boundary.py | 4 +- .../calculate_diagnostics_for_turbulence.py | 4 +- ..._coefficients_for_grid_point_cold_pools.py | 4 +- ...ate_nabla2_and_smag_coefficients_for_vn.py | 9 +- .../stencils/calculate_nabla2_of_theta.py | 10 +- .../diffusion/stencils/calculate_nabla4.py | 9 +- ...n_coefficient_for_grid_point_cold_pools.py | 4 +- ...orary_fields_for_turbulence_diagnostics.py | 14 +-- ...fusion_nabla_of_theta_over_steep_points.py | 24 +--- model/driver/.flake8 | 42 +++++++ model/driver/.pre-commit-config.yaml | 114 ++++++++++++++++++ model/driver/pyproject.toml | 85 +++++++------ .../src/icon4py/model/driver/dycore_driver.py | 14 +-- .../model/driver/icon_configuration.py | 8 +- .../src/icon4py/model/driver/io_utils.py | 33 ++--- model/driver/tests/test_io_utils.py | 8 +- 40 files changed, 502 insertions(+), 468 deletions(-) create mode 100644 model/atmosphere/diffusion/.flake8 create mode 100644 model/atmosphere/diffusion/.pre-commit-config.yaml create mode 100644 model/driver/.flake8 create mode 100644 model/driver/.pre-commit-config.yaml diff --git a/model/atmosphere/diffusion/.flake8 b/model/atmosphere/diffusion/.flake8 new file mode 100644 index 000000000..31cecff5a --- /dev/null +++ b/model/atmosphere/diffusion/.flake8 @@ -0,0 +1,42 @@ +[flake8] +# Some sane defaults for the code style checker flake8 +max-line-length = 100 +max-complexity = 15 +doctests = true +extend-ignore = + # Do not perform function calls in argument defaults + B008, + # Public code object needs docstring + D1, + # Disable dargling errors by default + DAR, + # Whitespace before ':' (black formatter breaks this sometimes) + E203, + # Line too long (using Bugbear's B950 warning) + E501, + # Line break occurred before a binary operator + W503 + +exclude = + .eggs, + .gt_cache, + .ipynb_checkpoints, + .tox, + _local_, + build, + dist, + docs, + _external_src, + tests/_disabled, + setup.py + +rst-roles = + py:mod, mod, + py:func, func, + py:data, data, + py:const, const, + py:class, class, + py:meth, meth, + py:attr, attr, + py:exc, exc, + py:obj, obj, diff --git a/model/atmosphere/diffusion/.pre-commit-config.yaml b/model/atmosphere/diffusion/.pre-commit-config.yaml new file mode 100644 index 000000000..1d6ce2a03 --- /dev/null +++ b/model/atmosphere/diffusion/.pre-commit-config.yaml @@ -0,0 +1,114 @@ +# NOTE: pre-commit runs all hooks from the root folder of the repository, +# as regular git hooks do. Therefore, paths passed as arguments to the plugins +# should always be relative to the root folder. + +default_stages: [commit, push] +default_language_version: + python: python3.10 +minimum_pre_commit_version: 2.20.0 +files: "model/atmosphere/diffusion/.*" + +repos: +- repo: meta + hooks: + - id: check-hooks-apply + stages: [manual] + - id: check-useless-excludes + stages: [manual] + +- repo: https://github.com/asottile/setup-cfg-fmt + rev: v1.20.1 + hooks: + # Run only manually because it deletes comments + - id: setup-cfg-fmt + name: format setup.cfg + stages: [manual] + +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.3.0 + hooks: + - id: check-case-conflict + - id: check-merge-conflict + - id: check-shebang-scripts-are-executable + - id: check-symlinks + - id: check-yaml + - id: debug-statements + - id: destroyed-symlinks + +- repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks + rev: v2.6.0 + hooks: + - id: pretty-format-ini + args: [--autofix] + - id: pretty-format-toml + args: [--autofix] + - id: pretty-format-yaml + args: [--autofix, --preserve-quotes, --indent, "2"] + +- repo: https://github.com/pre-commit/mirrors-prettier + rev: v3.0.0-alpha.4 + hooks: + - id: prettier + types_or: [markdown, json] + +- repo: https://github.com/Lucas-C/pre-commit-hooks + rev: v1.3.0 + hooks: + - id: insert-license + name: add license for all ICON4Py Python source files + types: [python] + args: [--comment-style, "|#|", --license-filepath, model/.license_header.txt, --fuzzy-match-generates-todo] + +- repo: https://github.com/asottile/yesqa + rev: v1.3.0 + hooks: + - id: yesqa + +- repo: https://github.com/psf/black + rev: '22.3.0' + hooks: + - id: black + name: black Python formatter + args: [--config, model/atmosphere/diffusion/pyproject.toml] + +- repo: https://github.com/asottile/blacken-docs + rev: v1.12.1 + hooks: + - id: blacken-docs + name: black Python formatter for docstrings + additional_dependencies: [black==22.3.0] + +- repo: https://github.com/PyCQA/isort + rev: '5.12.0' + hooks: + - id: isort + args: [--config-root, model/atmosphere/diffusion/, --resolve-all-configs] + +- repo: https://github.com/PyCQA/flake8 + rev: '4.0.1' + hooks: + - id: flake8 + name: flake8 code style checks + additional_dependencies: + - darglint + - flake8-bugbear + - flake8-builtins + - flake8-debugger + - flake8-docstrings + - flake8-eradicate + - flake8-mutable + - pygments + args: [--config=model/atmosphere/diffusion/.flake8, model/atmosphere/diffusion/src/icon4py/] + +- repo: local + hooks: + - id: mypy + name: mypy static type checker + entry: bash -c 'echo mypy temporarily disabled' + #entry: bash -c 'cd model/atmosphere/dycore; mypy src/' -- + language: system + types_or: [python, pyi] + always_run: true + #pass_filenames: false + require_serial: true + stages: [commit] diff --git a/model/atmosphere/diffusion/diffusion_tests/conftest.py b/model/atmosphere/diffusion/diffusion_tests/conftest.py index 2ec7b0911..448cec45d 100644 --- a/model/atmosphere/diffusion/diffusion_tests/conftest.py +++ b/model/atmosphere/diffusion/diffusion_tests/conftest.py @@ -13,10 +13,7 @@ import pytest -from icon4py.model.atmosphere.diffusion.diffusion import ( - DiffusionConfig, - DiffusionType, -) +from icon4py.model.atmosphere.diffusion.diffusion import DiffusionConfig, DiffusionType from icon4py.model.common.test_utils.fixtures import ( # noqa F401 backend, damping_height, diff --git a/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py index 16edaf64c..3b68ec7a0 100644 --- a/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py +++ b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py @@ -14,14 +14,8 @@ import pytest -from icon4py.model.atmosphere.diffusion.diffusion import ( - Diffusion, - DiffusionParams, -) -from icon4py.model.common.decomposition.decomposed import ( - DecompositionInfo, - create_exchange, -) +from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.model.common.decomposition.decomposed import DecompositionInfo, create_exchange from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.model.common.grid.vertical import VerticalModelParams from icon4py.model.common.test_utils.parallel_helpers import check_comm_size @@ -68,9 +62,7 @@ def test_parallel_diffusion( metric_state = metrics_savepoint.construct_metric_state_for_diffusion() cell_geometry = grid_savepoint.construct_cell_geometry() edge_geometry = grid_savepoint.construct_edge_geometry() - interpolation_state = ( - interpolation_savepoint.construct_interpolation_state_for_diffusion() - ) + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() diffusion_params = DiffusionParams(r04b09_diffusion_config) dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") @@ -91,9 +83,7 @@ def test_parallel_diffusion( edge_params=edge_geometry, cell_params=cell_geometry, ) - print( - f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized " - ) + print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ") diagnostic_state = diffusion_savepoint_init.construct_diagnostics_for_diffusion() prognostic_state = diffusion_savepoint_init.construct_prognostics() if linit: diff --git a/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_global_to_vn.py index 1930df665..02e618fcd 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_global_to_vn.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_global_to_vn.py @@ -44,12 +44,8 @@ def input_data(self, mesh): ) @staticmethod - def reference( - mesh, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn - ): + def reference(mesh, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn): area_edge = np.expand_dims(area_edge, axis=-1) diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0) - vn = vn + area_edge * ( - kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge - ) + vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge) return dict(vn=vn) diff --git a/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_to_vn.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_to_vn.py index 092a9d945..29d3fa391 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_to_vn.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_and_nabla4_to_vn.py @@ -96,12 +96,8 @@ def input_data(self, mesh): ) @staticmethod - def reference( - mesh, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn - ): + def reference(mesh, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn): area_edge = np.expand_dims(area_edge, axis=-1) diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0) - vn = vn + area_edge * ( - kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge - ) + vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge) return dict(vn=vn) diff --git a/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w.py b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w.py index 8cc82ea2c..147fd5c59 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_apply_nabla2_to_w.py @@ -15,9 +15,7 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w import ( - apply_nabla2_to_w, -) +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w import apply_nabla2_to_w from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field @@ -38,9 +36,7 @@ def reference( ) -> np.array: geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) area = np.expand_dims(area, axis=-1) - w = w - diff_multfac_w * area * area * np.sum( - z_nabla2_c[mesh.c2e2cO] * geofac_n2s, axis=1 - ) + w = w - diff_multfac_w * area * area * np.sum(z_nabla2_c[mesh.c2e2cO] * geofac_n2s, axis=1) return dict(w=w) @pytest.fixture diff --git a/model/atmosphere/diffusion/diffusion_tests/test_calculate_diagnostics_for_turbulence.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_diagnostics_for_turbulence.py index f7e7831a2..c70fc8fa2 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_calculate_diagnostics_for_turbulence.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_diagnostics_for_turbulence.py @@ -18,11 +18,7 @@ calculate_diagnostics_for_turbulence, ) from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field class TestCalculateDiagnosticsForTurbulence(StencilTest): @@ -36,9 +32,7 @@ def reference( kc_offset_1 = np.roll(kh_c, shift=1, axis=1) div_offset_1 = np.roll(div, shift=1, axis=1) div_ic[:, 1:] = (wgtfac_c * div + (1.0 - wgtfac_c) * div_offset_1)[:, 1:] - hdef_ic[:, 1:] = ((wgtfac_c * kh_c + (1.0 - wgtfac_c) * kc_offset_1) ** 2)[ - :, 1: - ] + hdef_ic[:, 1:] = ((wgtfac_c * kh_c + (1.0 - wgtfac_c) * kc_offset_1) ** 2)[:, 1:] return dict(div_ic=div_ic, hdef_ic=hdef_ic) @pytest.fixture @@ -48,6 +42,4 @@ def input_data(self, mesh): kh_c = random_field(mesh, CellDim, KDim) div_ic = zero_field(mesh, CellDim, KDim) hdef_ic = zero_field(mesh, CellDim, KDim) - return dict( - wgtfac_c=wgtfac_c, div=div, kh_c=kh_c, div_ic=div_ic, hdef_ic=hdef_ic - ) + return dict(wgtfac_c=wgtfac_c, div=div, kh_c=kh_c, div_ic=div_ic, hdef_ic=hdef_ic) diff --git a/model/atmosphere/diffusion/diffusion_tests/test_calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_horizontal_gradients_for_turbulence.py index d8f18f442..c7a7016f5 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_calculate_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_horizontal_gradients_for_turbulence.py @@ -19,11 +19,7 @@ calculate_horizontal_gradients_for_turbulence, ) from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field class TestCalculateHorizontalGradientsForTurbulence(StencilTest): diff --git a/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index c249d23c4..b230948b5 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -17,13 +17,7 @@ from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_and_smag_coefficients_for_vn import ( calculate_nabla2_and_smag_coefficients_for_vn, ) -from icon4py.model.common.dimension import ( - E2C2VDim, - ECVDim, - EdgeDim, - KDim, - VertexDim, -) +from icon4py.model.common.dimension import E2C2VDim, ECVDim, EdgeDim, KDim, VertexDim from icon4py.model.common.test_utils.helpers import ( StencilTest, as_1D_sparse_field, @@ -117,9 +111,7 @@ def reference( + v_vert_e2c2v[:, 3] * primal_normal_vert_y[:, 3] ) - kh_smag_2 = (kh_smag_2 * inv_vert_vert_length) - ( - dvt_tang * inv_primal_edge_length - ) + kh_smag_2 = (kh_smag_2 * inv_vert_vert_length) - (dvt_tang * inv_primal_edge_length) kh_smag_2 = kh_smag_2 * kh_smag_2 diff --git a/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_w.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_w.py index 7e7e0fa21..310a91c21 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_w.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_for_w.py @@ -19,11 +19,7 @@ calculate_nabla2_for_w, ) from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field class TestCalculateNabla2ForW(StencilTest): diff --git a/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_of_theta.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_of_theta.py index 46b467d3a..0d6bd5251 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_of_theta.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla2_of_theta.py @@ -31,14 +31,10 @@ class TestCalculateNabla2OfTheta(StencilTest): OUTPUTS = ("z_temp",) @staticmethod - def reference( - mesh, z_nabla2_e: np.array, geofac_div: np.array, **kwargs - ) -> np.array: + def reference(mesh, z_nabla2_e: np.array, geofac_div: np.array, **kwargs) -> np.array: geofac_div = geofac_div.reshape(mesh.c2e.shape) geofac_div = np.expand_dims(geofac_div, axis=-1) - z_temp = np.sum( - z_nabla2_e[mesh.c2e] * geofac_div, axis=1 - ) # sum along edge dimension + z_temp = np.sum(z_nabla2_e[mesh.c2e] * geofac_div, axis=1) # sum along edge dimension return dict(z_temp=z_temp) @pytest.fixture diff --git a/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla4.py b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla4.py index 66fd13926..de26d3cd6 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla4.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_calculate_nabla4.py @@ -14,16 +14,8 @@ import numpy as np import pytest -from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import ( - calculate_nabla4, -) -from icon4py.model.common.dimension import ( - E2C2VDim, - ECVDim, - EdgeDim, - KDim, - VertexDim, -) +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import calculate_nabla4 +from icon4py.model.common.dimension import E2C2VDim, ECVDim, EdgeDim, KDim, VertexDim from icon4py.model.common.test_utils.helpers import ( StencilTest, as_1D_sparse_field, diff --git a/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py b/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py index 20c134a24..d1df49c73 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py @@ -14,10 +14,7 @@ import numpy as np import pytest -from icon4py.model.atmosphere.diffusion.diffusion import ( - Diffusion, - DiffusionParams, -) +from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.model.atmosphere.diffusion.diffusion_states import ( DiffusionDiagnosticState, PrognosticState, @@ -118,9 +115,7 @@ def test_diffusion_init( assert meta["linit"] is False assert meta["date"] == step_date_init - interpolation_state = ( - interpolation_savepoint.construct_interpolation_state_for_diffusion() - ) + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() metric_state = metrics_savepoint.construct_metric_state_for_diffusion() edge_params = grid_savepoint.construct_edge_geometry() cell_params = grid_savepoint.construct_cell_geometry() @@ -154,10 +149,7 @@ def test_diffusion_init( config.substep_as_float, ) - assert ( - diffusion.smag_offset - == 0.25 * additional_parameters.K4 * config.substep_as_float - ) + assert diffusion.smag_offset == 0.25 * additional_parameters.K4 * config.substep_as_float assert np.allclose(expected_smag_limit, diffusion.smag_limit) expected_diff_multfac_vn = diff_multfac_vn_numpy( @@ -184,15 +176,11 @@ def _verify_init_values_against_savepoint( assert savepoint.diff_multfac_w() == diffusion.diff_multfac_w # this is done in diffusion.run(...) because it depends on the dtime - scale_k( - diffusion.enh_smag_fac, dtime, diffusion.diff_multfac_smag, offset_provider={} - ) + scale_k(diffusion.enh_smag_fac, dtime, diffusion.diff_multfac_smag, offset_provider={}) assert np.allclose(savepoint.diff_multfac_smag(), diffusion.diff_multfac_smag) assert np.allclose(savepoint.smag_limit(), diffusion.smag_limit) - assert np.allclose( - savepoint.diff_multfac_n2w(), np.asarray(diffusion.diff_multfac_n2w) - ) + assert np.allclose(savepoint.diff_multfac_n2w(), np.asarray(diffusion.diff_multfac_n2w)) assert np.allclose(savepoint.diff_multfac_vn(), diffusion.diff_multfac_vn) @@ -212,9 +200,7 @@ def test_verify_diffusion_init_against_first_regular_savepoint( cell_geometry = grid_savepoint.construct_cell_geometry() edge_geometry = grid_savepoint.construct_edge_geometry() - interpolation_state = ( - interpolation_savepoint.construct_interpolation_state_for_diffusion() - ) + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() metric_state = metrics_savepoint.construct_metric_state_for_diffusion() diffusion = Diffusion() @@ -247,9 +233,7 @@ def test_verify_diffusion_init_against_other_regular_savepoint( additional_parameters = DiffusionParams(config) vertical_params = VerticalModelParams(grid_savepoint.vct_a(), damping_height) - interpolation_state = ( - interpolation_savepoint.construct_interpolation_state_for_diffusion() - ) + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() metric_state = metrics_savepoint.construct_metric_state_for_diffusion() edge_params = grid_savepoint.construct_edge_geometry() cell_params = grid_savepoint.construct_cell_geometry() @@ -290,16 +274,12 @@ def test_run_diffusion_single_step( dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() - interpolation_state = ( - interpolation_savepoint.construct_interpolation_state_for_diffusion() - ) + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() metric_state = metrics_savepoint.construct_metric_state_for_diffusion() diagnostic_state = diffusion_savepoint_init.construct_diagnostics_for_diffusion() prognostic_state = diffusion_savepoint_init.construct_prognostics() vct_a = grid_savepoint.vct_a() - vertical_params = VerticalModelParams( - vct_a=vct_a, rayleigh_damping_height=damping_height - ) + vertical_params = VerticalModelParams(vct_a=vct_a, rayleigh_damping_height=damping_height) config = r04b09_diffusion_config additional_parameters = DiffusionParams(config) @@ -314,18 +294,14 @@ def test_run_diffusion_single_step( edge_params=edge_geometry, cell_params=cell_geometry, ) - _verify_diffusion_fields( - diagnostic_state, prognostic_state, diffusion_savepoint_init - ) + _verify_diffusion_fields(diagnostic_state, prognostic_state, diffusion_savepoint_init) assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v diffusion.run( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime, ) - _verify_diffusion_fields( - diagnostic_state, prognostic_state, diffusion_savepoint_exit - ) + _verify_diffusion_fields(diagnostic_state, prognostic_state, diffusion_savepoint_exit) def _verify_diffusion_fields( @@ -375,16 +351,12 @@ def test_run_diffusion_initial_step( dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() - interpolation_state = ( - interpolation_savepoint.construct_interpolation_state_for_diffusion() - ) + interpolation_state = interpolation_savepoint.construct_interpolation_state_for_diffusion() metric_state = metrics_savepoint.construct_metric_state_for_diffusion() diagnostic_state = diffusion_savepoint_init.construct_diagnostics_for_diffusion() prognostic_state = diffusion_savepoint_init.construct_prognostics() vct_a = grid_savepoint.vct_a() - vertical_params = VerticalModelParams( - vct_a=vct_a, rayleigh_damping_height=damping_height - ) + vertical_params = VerticalModelParams(vct_a=vct_a, rayleigh_damping_height=damping_height) config = r04b09_diffusion_config additional_parameters = DiffusionParams(config) diff --git a/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py b/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py index 733243bb7..98d0516c5 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py @@ -83,9 +83,7 @@ def test_diff_multfac_vn_smag_limit_for_time_step_with_const_value(): expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps) expected_smag_limit = smag_limit_numpy(diff_multfac_vn_numpy, shape, k4, substeps) - _setup_runtime_diff_multfac_vn( - k4, efdt_ratio, out=diff_multfac_vn, offset_provider={} - ) + _setup_runtime_diff_multfac_vn(k4, efdt_ratio, out=diff_multfac_vn, offset_provider={}) _setup_smag_limit(diff_multfac_vn, out=smag_limit, offset_provider={}) assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) @@ -102,9 +100,7 @@ def test_diff_multfac_vn_smag_limit_for_loop_run_with_k4_substeps(): shape = np.asarray(diff_multfac_vn).shape expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps) expected_smag_limit = smag_limit_numpy(diff_multfac_vn_numpy, shape, k4, substeps) - _setup_runtime_diff_multfac_vn( - k4, substeps, out=diff_multfac_vn, offset_provider={} - ) + _setup_runtime_diff_multfac_vn(k4, substeps, out=diff_multfac_vn, offset_provider={}) _setup_smag_limit(diff_multfac_vn, out=smag_limit, offset_provider={}) assert np.allclose(expected_diff_multfac_vn, diff_multfac_vn) @@ -136,9 +132,7 @@ def test_init_enh_smag_fac(): enhanced_smag_fac_np = enhanced_smagorinski_factor_numpy(fac, z, np.asarray(a_vec)) - _en_smag_fac_for_zero_nshift( - a_vec, *fac, *z, out=enh_smag_fac, offset_provider={"Koff": KDim} - ) + _en_smag_fac_for_zero_nshift(a_vec, *fac, *z, out=enh_smag_fac, offset_provider={"Koff": KDim}) assert np.allclose(enhanced_smag_fac_np, np.asarray(enh_smag_fac)) diff --git a/model/atmosphere/diffusion/diffusion_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py b/model/atmosphere/diffusion/diffusion_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py index 2ab67d0d3..84d8ed516 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py @@ -18,11 +18,7 @@ temporary_field_for_grid_point_cold_pools_enhancement, ) from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field class TestTemporaryFieldForGridPointColdPoolsEnhancement(StencilTest): diff --git a/model/atmosphere/diffusion/diffusion_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/diffusion_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index 0a4f8d354..72baed394 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -56,15 +56,12 @@ def truly_horizontal_diffusion_nabla_of_theta_over_steep_points_numpy( ] sum_over = np.sum( - geofac_n2s_nbh - * (vcoef * theta_v_at_zd_vertidx + (1.0 - vcoef) * theta_v_at_zd_vertidx_p1), + geofac_n2s_nbh * (vcoef * theta_v_at_zd_vertidx + (1.0 - vcoef) * theta_v_at_zd_vertidx_p1), axis=1, ) geofac_n2s_c = np.expand_dims(geofac_n2s_c, axis=1) # add KDim - return np.where( - mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum_over), z_temp - ) + return np.where(mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum_over), z_temp) def test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points(): diff --git a/model/atmosphere/diffusion/pyproject.toml b/model/atmosphere/diffusion/pyproject.toml index a797dfae8..8f0d08a2c 100644 --- a/model/atmosphere/diffusion/pyproject.toml +++ b/model/atmosphere/diffusion/pyproject.toml @@ -1,38 +1,38 @@ [build-system] -requires = ["setuptools>=61.0", "wheel>=0.40.0"] build-backend = "setuptools.build_meta" +requires = ["setuptools>=61.0", "wheel>=0.40.0"] [project] -name = "icon4py-atmosphere-diffusion" -description = "ICON diffusion." -readme = "README.md" -requires-python = ">=3.10" -license = {file = "LICENSE"} authors = [ - {email = "gridtools@cscs.ch"}, - {name = "ETH Zurich"} + {email = "gridtools@cscs.ch"}, + {name = "ETH Zurich"} ] classifiers = [ - "Development Status :: 3 - Alpha", - "Intended Audience :: Science/Research", - "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)", - "Operating System :: POSIX", - "Programming Language :: Python", - "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: Implementation :: CPython", - "Topic :: Scientific/Engineering :: Atmospheric Science", - "Topic :: Scientific/Engineering :: Mathematics", - "Topic :: Scientific/Engineering :: Physics" + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)", + "Operating System :: POSIX", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: Implementation :: CPython", + "Topic :: Scientific/Engineering :: Atmospheric Science", + "Topic :: Scientific/Engineering :: Mathematics", + "Topic :: Scientific/Engineering :: Physics" ] dependencies = [ - "gt4py>=1.0.1", - "icon4py-common>=0.0.5", - "mpi4py<=3.1.4", - "pyghex>=0.3.0", + "gt4py>=1.0.1", + "icon4py-common>=0.0.5", + "mpi4py<=3.1.4", + "pyghex>=0.3.0" ] +description = "ICON diffusion." dynamic = ['version'] +license = {file = "LICENSE"} +name = "icon4py-atmosphere-diffusion" +readme = "README.md" +requires-python = ">=3.10" [project.urls] repository = "https://github.com/C2SM/icon4py" @@ -73,46 +73,46 @@ exclude_lines = [ ignore_errors = true [tool.coverage.run] -parallel = true branch = true +parallel = true source_pkgs = ['diffusion'] [tool.isort] +force_grid_wrap = 0 +include_trailing_comma = true +known_first_party = ['icon4py.model'] +known_third_party = ['gt4py'] lexicographical = true line_length = 100 # It should be the same as in `tool.black.line-length` above lines_after_imports = 2 +multi_line_output = 3 profile = 'black' sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'FIRSTPARTY', 'TESTS', 'LOCALFOLDER'] skip_gitignore = true skip_glob = ['*.venv/**', '_local/**'] -known_first_party = ['icon4py.model'] -known_third_party = ['gt4py'] -multi_line_output = 3 use_parentheses = true -include_trailing_comma = true -force_grid_wrap = 0 [tool.mypy] -install_types = true -non_interactive = true -exclude = [ - '^tests/*.py', -] disallow_incomplete_defs = true disallow_untyped_defs = true +exclude = [ + '^tests/*.py' +] ignore_missing_imports = false implicit_reexport = true -warn_unused_configs = true -warn_unused_ignores = true -warn_redundant_casts = true +install_types = true +non_interactive = true show_column_numbers = true show_error_codes = true +warn_redundant_casts = true +warn_unused_configs = true +warn_unused_ignores = true [tool.pytest] [tool.pytest.ini_options] -testpaths = 'diffusion_tests' markers = 'datatest: this test uses binary data' +testpaths = 'diffusion_tests' [tool.setuptools.dynamic] version = {attr = 'icon4py.model.atmosphere.diffusion.__init__.__version__'} diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py index 9aea7121d..b46ccfde3 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -43,9 +43,7 @@ setup_fields_for_initial_step, zero_field, ) -from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_vn import ( - apply_diffusion_to_vn, -) +from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_vn import apply_diffusion_to_vn from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance, ) @@ -72,16 +70,9 @@ DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO, GAS_CONSTANT_DRY_AIR, ) -from icon4py.model.common.decomposition.decomposed import ( - ExchangeRuntime, - SingleNode, -) +from icon4py.model.common.decomposition.decomposed import ExchangeRuntime, SingleNode from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim -from icon4py.model.common.grid.horizontal import ( - CellParams, - EdgeParams, - HorizontalMarkerIndex, -) +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalMarkerIndex from icon4py.model.common.grid.icon_grid import IconGrid from icon4py.model.common.grid.vertical import VerticalModelParams from icon4py.model.common.interpolation.stencils.mo_intp_rbf_rbf_vec_interpol_vertex import ( @@ -111,9 +102,7 @@ class DiffusionType(int, Enum): LINEAR_2ND_ORDER = 2 #: 2nd order linear diffusion on all vertical levels SMAGORINSKY_NO_BACKGROUND = 3 #: Smagorinsky diffusion without background diffusion LINEAR_4TH_ORDER = 4 #: 4th order linear diffusion on all vertical levels - SMAGORINSKY_4TH_ORDER = ( - 5 #: Smagorinsky diffusion with fourth-order background diffusion - ) + SMAGORINSKY_4TH_ORDER = 5 #: Smagorinsky diffusion with fourth-order background diffusion class DiffusionConfig: @@ -214,9 +203,7 @@ def __init__( #: Denominator for velocity boundary diffusion #: Called 'denom_diffu_v' in mo_gridref_nml.f90 - self.velocity_boundary_diffusion_denominator: float = ( - velocity_boundary_diffusion_denom - ) + self.velocity_boundary_diffusion_denominator: float = velocity_boundary_diffusion_denom # parameters from namelist: mo_interpol_nml.f90 @@ -250,9 +237,7 @@ def _validate(self): self.apply_to_horizontal_wind = True if not self.apply_zdiffusion_t: - raise NotImplementedError( - "zdiffu_t = False is not implemented (leaves out stencil_15)" - ) + raise NotImplementedError("zdiffu_t = False is not implemented (leaves out stencil_15)") @functools.cached_property def substep_as_float(self): @@ -276,22 +261,14 @@ def __post_init__(self, config): object.__setattr__( self, "K2", - ( - 1.0 / (config.hdiff_efdt_ratio * 8.0) - if config.hdiff_efdt_ratio > 0.0 - else 0.0 - ), + (1.0 / (config.hdiff_efdt_ratio * 8.0) if config.hdiff_efdt_ratio > 0.0 else 0.0), ) object.__setattr__(self, "K4", self.K2 / 8.0) object.__setattr__(self, "K6", self.K2 / 64.0) object.__setattr__( self, "K4W", - ( - 1.0 / (config.hdiff_w_efdt_ratio * 36.0) - if config.hdiff_w_efdt_ratio > 0 - else 0.0 - ), + (1.0 / (config.hdiff_w_efdt_ratio * 36.0) if config.hdiff_w_efdt_ratio > 0 else 0.0), ) ( @@ -425,23 +402,16 @@ def _get_start_index_for_w_diffusion() -> int32: ), ) - self.nudgezone_diff: float = 0.04 / ( - params.scaled_nudge_max_coeff + sys.float_info.epsilon - ) - self.bdy_diff: float = 0.015 / ( - params.scaled_nudge_max_coeff + sys.float_info.epsilon - ) + self.nudgezone_diff: float = 0.04 / (params.scaled_nudge_max_coeff + sys.float_info.epsilon) + self.bdy_diff: float = 0.015 / (params.scaled_nudge_max_coeff + sys.float_info.epsilon) self.fac_bdydiff_v: float = ( - math.sqrt(config.substep_as_float) - / config.velocity_boundary_diffusion_denominator + math.sqrt(config.substep_as_float) / config.velocity_boundary_diffusion_denominator if config.lhdiff_rcf else 1.0 / config.velocity_boundary_diffusion_denominator ) self.smag_offset: float = 0.25 * params.K4 * config.substep_as_float - self.diff_multfac_w: float = min( - 1.0 / 48.0, params.K4W * config.substep_as_float - ) + self.diff_multfac_w: float = min(1.0 / 48.0, params.K4W * config.substep_as_float) init_diffusion_local_fields_for_regular_timestep.with_backend(backend)( params.K4, @@ -601,12 +571,8 @@ def _do_diffusion_step( cell_start_nudging = self.grid.get_start_index( CellDim, HorizontalMarkerIndex.nudging(CellDim) ) - cell_end_local = self.grid.get_end_index( - CellDim, HorizontalMarkerIndex.local(CellDim) - ) - cell_end_halo = self.grid.get_end_index( - CellDim, HorizontalMarkerIndex.halo(CellDim) - ) + cell_end_local = self.grid.get_end_index(CellDim, HorizontalMarkerIndex.local(CellDim)) + cell_end_halo = self.grid.get_end_index(CellDim, HorizontalMarkerIndex.halo(CellDim)) edge_start_nudging_plus_one = self.grid.get_start_index( EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1 @@ -617,15 +583,11 @@ def _do_diffusion_step( edge_start_lb_plus4 = self.grid.get_start_index( EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4 ) - edge_end_local = self.grid.get_end_index( - EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - ) + edge_end_local = self.grid.get_end_index(EdgeDim, HorizontalMarkerIndex.local(EdgeDim)) edge_end_local_minus2 = self.grid.get_end_index( EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2 ) - edge_end_halo = self.grid.get_end_index( - EdgeDim, HorizontalMarkerIndex.halo(EdgeDim) - ) + edge_end_halo = self.grid.get_end_index(EdgeDim, HorizontalMarkerIndex.halo(EdgeDim)) vertex_start_lb_plus1 = self.grid.get_start_index( VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1 @@ -660,9 +622,7 @@ def _do_diffusion_step( h.wait() log.debug("communication rbf extrapolation of vn - end") - log.debug( - "running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start" - ) + log.debug("running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start") calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend)( diff_multfac_smag=self.diff_multfac_smag, tangent_orientation=self.edge_params.tangent_orientation, @@ -689,12 +649,8 @@ def _do_diffusion_step( "E2ECV": self.grid.get_e2ecv_connectivity(), }, ) - log.debug( - "running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end" - ) - log.debug( - "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start" - ) + log.debug("running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end") + log.debug("running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start") calculate_diagnostic_quantities_for_turbulence.with_backend(backend)( kh_smag_ec=self.kh_smag_ec, vn=prognostic_state.vn, @@ -714,9 +670,7 @@ def _do_diffusion_step( "Koff": KDim, }, ) - log.debug( - "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end" - ) + log.debug("running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end") # HALO EXCHANGE IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH @@ -782,12 +736,8 @@ def _do_diffusion_step( log.debug( "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" ) - copy_field.with_backend(backend)( - prognostic_state.w, self.w_tmp, offset_provider={} - ) - apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.with_backend( - backend - )( + copy_field.with_backend(backend)(prognostic_state.w, self.w_tmp, offset_provider={}) + apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.with_backend(backend)( area=self.cell_params.area, geofac_n2s=self.interpolation_state.geofac_n2s, geofac_grg_x=self.interpolation_state.geofac_grg_x, @@ -820,9 +770,7 @@ def _do_diffusion_step( log.debug( "running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" ) - calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.with_backend( - backend - )( + calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.with_backend(backend)( theta_v=prognostic_state.theta_v, theta_ref_mc=self.metric_state.theta_ref_mc, thresh_tdiff=self.thresh_tdiff, @@ -860,9 +808,7 @@ def _do_diffusion_step( log.debug( "running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start" ) - truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_backend( - backend - )( + truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_backend(backend)( mask=self.metric_state.mask_hdiff, zd_vertoffset=self.metric_state.zd_vertoffset, zd_diffcoef=self.metric_state.zd_diffcoef, diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py index 91bdf279e..ea423047c 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py @@ -68,9 +68,7 @@ class DiffusionMetricState: class DiffusionInterpolationState: """Represents the ICON interpolation state needed in diffusion.""" - e_bln_c_s: Field[ - [CEDim], float - ] # coefficent for bilinear interpolation from edge to cell () + e_bln_c_s: Field[[CEDim], float] # coefficent for bilinear interpolation from edge to cell () rbf_coeff_1: Field[ [VertexDim, V2EDim], float ] # rbf_vec_coeff_v_1(nproma, rbf_vec_dim_v, nblks_v) @@ -78,9 +76,7 @@ class DiffusionInterpolationState: [VertexDim, V2EDim], float ] # rbf_vec_coeff_v_2(nproma, rbf_vec_dim_v, nblks_v) - geofac_div: Field[ - [CEDim], float - ] # factor for divergence (nproma,cell_type,nblks_c) + geofac_div: Field[[CEDim], float] # factor for divergence (nproma,cell_type,nblks_c) geofac_n2s: Field[ [CellDim, C2E2CODim], float @@ -113,9 +109,7 @@ class PrognosticState: Corresponds to ICON t_nh_prog """ - w: Field[ - [CellDim, KDim], float - ] # vertical_wind field, w(nproma, nlevp1, nblks_c) [m/s] + w: Field[[CellDim, KDim], float] # vertical_wind field, w(nproma, nlevp1, nblks_c) [m/s] vn: Field[[EdgeDim, KDim], float] # vn(nproma, nlev, nblks_e) [m/s] exner_pressure: Field[[CellDim, KDim], float] # exner(nrpoma, nlev, nblks_c) theta_v: Field[[CellDim, KDim], float] # (nproma, nlev, nlbks_c) [K] diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py index fc17ed567..183d146cb 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py @@ -28,23 +28,17 @@ def zero_field(mesh, *dims: Dimension, dtype=float): @field_operator -def _identity_c_k( - field: Field[[CellDim, KDim], float] -) -> Field[[CellDim, KDim], float]: +def _identity_c_k(field: Field[[CellDim, KDim], float]) -> Field[[CellDim, KDim], float]: return field @program -def copy_field( - old_f: Field[[CellDim, KDim], float], new_f: Field[[CellDim, KDim], float] -): +def copy_field(old_f: Field[[CellDim, KDim], float], new_f: Field[[CellDim, KDim], float]): _identity_c_k(old_f, out=new_f) @field_operator -def _identity_e_k( - field: Field[[EdgeDim, KDim], float] -) -> Field[[EdgeDim, KDim], float]: +def _identity_e_k(field: Field[[EdgeDim, KDim], float]) -> Field[[EdgeDim, KDim], float]: return field @@ -54,9 +48,7 @@ def _scale_k(field: Field[[KDim], float], factor: float) -> Field[[KDim], float] @program -def scale_k( - field: Field[[KDim], float], factor: float, scaled_field: Field[[KDim], float] -): +def scale_k(field: Field[[KDim], float], factor: float, scaled_field: Field[[KDim], float]): _scale_k(field, factor, out=scaled_field) @@ -76,18 +68,14 @@ def _setup_smag_limit(diff_multfac_vn: Field[[KDim], float]) -> Field[[KDim], fl @field_operator -def _setup_runtime_diff_multfac_vn( - k4: float, dyn_substeps: float -) -> Field[[KDim], float]: +def _setup_runtime_diff_multfac_vn(k4: float, dyn_substeps: float) -> Field[[KDim], float]: con = 1.0 / 128.0 dyn = k4 * dyn_substeps / 3.0 return broadcast(minimum(con, dyn), (KDim,)) @field_operator -def _setup_initial_diff_multfac_vn( - k4: float, hdiff_efdt_ratio: float -) -> Field[[KDim], float]: +def _setup_initial_diff_multfac_vn(k4: float, hdiff_efdt_ratio: float) -> Field[[KDim], float]: return broadcast(k4 / 3.0 * hdiff_efdt_ratio, (KDim,)) @@ -107,9 +95,7 @@ def setup_fields_for_initial_step( diff_multfac_vn: Field[[KDim], float], smag_limit: Field[[KDim], float], ): - _setup_fields_for_initial_step( - k4, hdiff_efdt_ratio, out=(diff_multfac_vn, smag_limit) - ) + _setup_fields_for_initial_step(k4, hdiff_efdt_ratio, out=(diff_multfac_vn, smag_limit)) @field_operator diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py index 5fa0d5e50..88b74a92c 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py @@ -23,9 +23,7 @@ from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_vn_in_lateral_boundary import ( _apply_nabla2_to_vn_in_lateral_boundary, ) -from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import ( - _calculate_nabla4, -) +from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import _calculate_nabla4 from icon4py.model.common.dimension import ECVDim, EdgeDim, KDim, VertexDim @@ -73,9 +71,7 @@ def _apply_diffusion_to_vn( vn, nudgezone_diff, ), - _apply_nabla2_to_vn_in_lateral_boundary( - z_nabla2_e, area_edge, vn, fac_bdydiff_v - ), + _apply_nabla2_to_vn_in_lateral_boundary(z_nabla2_e, area_edge, vn, fac_bdydiff_v), ) if limited_area else where( diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py index 44f4ab69f..dd4f0a264 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py @@ -14,9 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where -from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w import ( - _apply_nabla2_to_w, -) +from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w import _apply_nabla2_to_w from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w_in_upper_damping_layer import ( _apply_nabla2_to_w_in_upper_damping_layer, ) @@ -54,9 +52,7 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( dwdx, dwdy = where( int32(0) < vert_idx, - _calculate_horizontal_gradients_for_turbulence( - w_old, geofac_grg_x, geofac_grg_y - ), + _calculate_horizontal_gradients_for_turbulence(w_old, geofac_grg_x, geofac_grg_y), (dwdx, dwdy), ) @@ -73,9 +69,7 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( & (vert_idx < nrdmax) & (interior_idx <= horz_idx) & (horz_idx < halo_idx), - _apply_nabla2_to_w_in_upper_damping_layer( - w, diff_multfac_n2w, area, z_nabla2_c - ), + _apply_nabla2_to_w_in_upper_damping_layer(w, diff_multfac_n2w, area, z_nabla2_c), w, ) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py index aa76bed1b..313e9ac50 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py @@ -27,9 +27,7 @@ def _apply_nabla2_and_nabla4_global_to_vn( diff_multfac_vn: Field[[KDim], float], vn: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - vn = vn + area_edge * ( - kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge - ) + vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge) return vn diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py index 1710922bd..c2f107489 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py @@ -35,6 +35,4 @@ def apply_nabla2_to_vn_in_lateral_boundary( vn: Field[[EdgeDim, KDim], float], fac_bdydiff_v: float, ): - _apply_nabla2_to_vn_in_lateral_boundary( - z_nabla2_e, area_edge, vn, fac_bdydiff_v, out=vn - ) + _apply_nabla2_to_vn_in_lateral_boundary(z_nabla2_e, area_edge, vn, fac_bdydiff_v, out=vn) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py index 60c3126a2..1c686fe58 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py @@ -38,6 +38,4 @@ def calculate_diagnostics_for_turbulence( div_ic: Field[[CellDim, KDim], float], hdef_ic: Field[[CellDim, KDim], float], ): - _calculate_diagnostics_for_turbulence( - div, kh_c, wgtfac_c, out=(div_ic[:, 1:], hdef_ic[:, 1:]) - ) + _calculate_diagnostics_for_turbulence(div, kh_c, wgtfac_c, out=(div_ic[:, 1:], hdef_ic[:, 1:])) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py index cc7eb2642..6a997cd21 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py @@ -33,9 +33,7 @@ def _calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( enh_diffu_3d = _temporary_field_for_grid_point_cold_pools_enhancement( theta_v, theta_ref_mc, thresh_tdiff ) - kh_smag_e = _enhance_diffusion_coefficient_for_grid_point_cold_pools( - kh_smag_e, enh_diffu_3d - ) + kh_smag_e = _enhance_diffusion_coefficient_for_grid_point_cold_pools(kh_smag_e, enh_diffu_3d) return kh_smag_e diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_and_smag_coefficients_for_vn.py index 8df12abc3..36ed3d10a 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_and_smag_coefficients_for_vn.py @@ -15,14 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32, maximum, minimum, sqrt -from icon4py.model.common.dimension import ( - E2C2V, - E2ECV, - ECVDim, - EdgeDim, - KDim, - VertexDim, -) +from icon4py.model.common.dimension import E2C2V, E2ECV, ECVDim, EdgeDim, KDim, VertexDim @field_operator diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py index 0420b78a0..e6517c9a1 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py @@ -15,15 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.model.common.dimension import ( - C2CE, - C2E, - C2EDim, - CEDim, - CellDim, - EdgeDim, - KDim, -) +from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim @field_operator diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py index dede3cb39..e37b3a585 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py @@ -15,14 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field -from icon4py.model.common.dimension import ( - E2C2V, - E2ECV, - ECVDim, - EdgeDim, - KDim, - VertexDim, -) +from icon4py.model.common.dimension import E2C2V, E2ECV, ECVDim, EdgeDim, KDim, VertexDim @field_operator diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py index ca00e06c5..37430b544 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py @@ -32,6 +32,4 @@ def enhance_diffusion_coefficient_for_grid_point_cold_pools( kh_smag_e: Field[[EdgeDim, KDim], float], enh_diffu_3d: Field[[CellDim, KDim], float], ): - _enhance_diffusion_coefficient_for_grid_point_cold_pools( - kh_smag_e, enh_diffu_3d, out=kh_smag_e - ) + _enhance_diffusion_coefficient_for_grid_point_cold_pools(kh_smag_e, enh_diffu_3d, out=kh_smag_e) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py index e1e38351a..cb528e6fa 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py @@ -15,15 +15,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, neighbor_sum -from icon4py.model.common.dimension import ( - C2CE, - C2E, - C2EDim, - CEDim, - CellDim, - EdgeDim, - KDim, -) +from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim @field_operator @@ -34,9 +26,7 @@ def _temporary_fields_for_turbulence_diagnostics( geofac_div: Field[[CEDim], float], diff_multfac_smag: Field[[KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - kh_c = ( - neighbor_sum(kh_smag_ec(C2E) * e_bln_c_s(C2CE), axis=C2EDim) / diff_multfac_smag - ) + kh_c = neighbor_sum(kh_smag_ec(C2E) * e_bln_c_s(C2CE), axis=C2EDim) / diff_multfac_smag div = neighbor_sum(vn(C2E) * geofac_div(C2CE), axis=C2EDim) return kh_c, div diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index 0984444ec..55eb44f2a 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -15,14 +15,7 @@ from gt4py.next.ffront.experimental import as_offset from gt4py.next.ffront.fbuiltins import Field, int32, where -from icon4py.model.common.dimension import ( - C2CEC, - C2E2C, - CECDim, - CellDim, - KDim, - Koff, -) +from icon4py.model.common.dimension import C2CEC, C2E2C, CECDim, CellDim, KDim, Koff @field_operator @@ -46,20 +39,11 @@ def _truly_horizontal_diffusion_nabla_of_theta_over_steep_points( sum_over_neighbors = ( geofac_n2s_nbh(C2CEC[0]) - * ( - vcoef(C2CEC[0]) * theta_v_0(C2E2C[0]) - + (1.0 - vcoef(C2CEC[0])) * theta_v_0_m1(C2E2C[0]) - ) + * (vcoef(C2CEC[0]) * theta_v_0(C2E2C[0]) + (1.0 - vcoef(C2CEC[0])) * theta_v_0_m1(C2E2C[0])) + geofac_n2s_nbh(C2CEC[1]) - * ( - vcoef(C2CEC[1]) * theta_v_1(C2E2C[1]) - + (1.0 - vcoef(C2CEC[1])) * theta_v_1_m1(C2E2C[1]) - ) + * (vcoef(C2CEC[1]) * theta_v_1(C2E2C[1]) + (1.0 - vcoef(C2CEC[1])) * theta_v_1_m1(C2E2C[1])) + geofac_n2s_nbh(C2CEC[2]) - * ( - vcoef(C2CEC[2]) * theta_v_2(C2E2C[2]) - + (1.0 - vcoef(C2CEC[2])) * theta_v_2_m1(C2E2C[2]) - ) + * (vcoef(C2CEC[2]) * theta_v_2(C2E2C[2]) + (1.0 - vcoef(C2CEC[2])) * theta_v_2_m1(C2E2C[2])) ) z_temp = where( diff --git a/model/driver/.flake8 b/model/driver/.flake8 new file mode 100644 index 000000000..31cecff5a --- /dev/null +++ b/model/driver/.flake8 @@ -0,0 +1,42 @@ +[flake8] +# Some sane defaults for the code style checker flake8 +max-line-length = 100 +max-complexity = 15 +doctests = true +extend-ignore = + # Do not perform function calls in argument defaults + B008, + # Public code object needs docstring + D1, + # Disable dargling errors by default + DAR, + # Whitespace before ':' (black formatter breaks this sometimes) + E203, + # Line too long (using Bugbear's B950 warning) + E501, + # Line break occurred before a binary operator + W503 + +exclude = + .eggs, + .gt_cache, + .ipynb_checkpoints, + .tox, + _local_, + build, + dist, + docs, + _external_src, + tests/_disabled, + setup.py + +rst-roles = + py:mod, mod, + py:func, func, + py:data, data, + py:const, const, + py:class, class, + py:meth, meth, + py:attr, attr, + py:exc, exc, + py:obj, obj, diff --git a/model/driver/.pre-commit-config.yaml b/model/driver/.pre-commit-config.yaml new file mode 100644 index 000000000..4ae0972f5 --- /dev/null +++ b/model/driver/.pre-commit-config.yaml @@ -0,0 +1,114 @@ +# NOTE: pre-commit runs all hooks from the root folder of the repository, +# as regular git hooks do. Therefore, paths passed as arguments to the plugins +# should always be relative to the root folder. + +default_stages: [commit, push] +default_language_version: + python: python3.10 +minimum_pre_commit_version: 2.20.0 +files: "model/driver/.*" + +repos: +- repo: meta + hooks: + - id: check-hooks-apply + stages: [manual] + - id: check-useless-excludes + stages: [manual] + +- repo: https://github.com/asottile/setup-cfg-fmt + rev: v1.20.1 + hooks: + # Run only manually because it deletes comments + - id: setup-cfg-fmt + name: format setup.cfg + stages: [manual] + +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.3.0 + hooks: + - id: check-case-conflict + - id: check-merge-conflict + - id: check-shebang-scripts-are-executable + - id: check-symlinks + - id: check-yaml + - id: debug-statements + - id: destroyed-symlinks + +- repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks + rev: v2.6.0 + hooks: + - id: pretty-format-ini + args: [--autofix] + - id: pretty-format-toml + args: [--autofix] + - id: pretty-format-yaml + args: [--autofix, --preserve-quotes, --indent, "2"] + +- repo: https://github.com/pre-commit/mirrors-prettier + rev: v3.0.0-alpha.4 + hooks: + - id: prettier + types_or: [markdown, json] + +- repo: https://github.com/Lucas-C/pre-commit-hooks + rev: v1.3.0 + hooks: + - id: insert-license + name: add license for all ICON4Py Python source files + types: [python] + args: [--comment-style, "|#|", --license-filepath, model/.license_header.txt, --fuzzy-match-generates-todo] + +- repo: https://github.com/asottile/yesqa + rev: v1.3.0 + hooks: + - id: yesqa + +- repo: https://github.com/psf/black + rev: '22.3.0' + hooks: + - id: black + name: black Python formatter + args: [--config, model/driver/pyproject.toml] + +- repo: https://github.com/asottile/blacken-docs + rev: v1.12.1 + hooks: + - id: blacken-docs + name: black Python formatter for docstrings + additional_dependencies: [black==22.3.0] + +- repo: https://github.com/PyCQA/isort + rev: '5.12.0' + hooks: + - id: isort + args: [--config-root, model/driver/, --resolve-all-configs] + +- repo: https://github.com/PyCQA/flake8 + rev: '4.0.1' + hooks: + - id: flake8 + name: flake8 code style checks + additional_dependencies: + - darglint + - flake8-bugbear + - flake8-builtins + - flake8-debugger + - flake8-docstrings + - flake8-eradicate + - flake8-mutable + - pygments + args: [--config=model/driver/.flake8, model/driver/src/icon4py/] + +- repo: local + hooks: + - id: mypy + name: mypy static type checker + entry: bash -c 'echo mypy temporarily disabled' + #entry: bash -c 'cd model/atmosphere/dycore; mypy src/' -- + language: system + types_or: [python, pyi] + always_run: true + #pass_filenames: false + require_serial: true + stages: [commit] diff --git a/model/driver/pyproject.toml b/model/driver/pyproject.toml index 188d04d2d..e9123ed8b 100644 --- a/model/driver/pyproject.toml +++ b/model/driver/pyproject.toml @@ -1,46 +1,45 @@ [build-system] -requires = ["setuptools>=61.0", "wheel>=0.40.0"] build-backend = "setuptools.build_meta" +requires = ["setuptools>=61.0", "wheel>=0.40.0"] [project] -name = "icon4py-driver" -description = "ICON model driver." -readme = "README.md" -requires-python = ">=3.10" -license = {file = "LICENSE"} authors = [ - {email = "gridtools@cscs.ch"}, - {name = "ETH Zurich"} + {email = "gridtools@cscs.ch"}, + {name = "ETH Zurich"} ] classifiers = [ - "Development Status :: 3 - Alpha", - "Intended Audience :: Science/Research", - "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)", - "Operating System :: POSIX", - "Programming Language :: Python", - "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: Implementation :: CPython", - "Topic :: Scientific/Engineering :: Atmospheric Science", - "Topic :: Scientific/Engineering :: Mathematics", - "Topic :: Scientific/Engineering :: Physics" + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)", + "Operating System :: POSIX", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: Implementation :: CPython", + "Topic :: Scientific/Engineering :: Atmospheric Science", + "Topic :: Scientific/Engineering :: Mathematics", + "Topic :: Scientific/Engineering :: Physics" ] dependencies = [ - "gt4py>=1.0.1", - "mpi4py<=3.1.4", - "pyghex>=0.3.0", - "pytz>=2023.2", - "icon4py-common>=0.0.5", - "icon4py-atmosphere-dycore>=0.0.5", - "icon4py-atmosphere-diffusion>=0.0.5", + "gt4py>=1.0.1", + "mpi4py<=3.1.4", + "pyghex>=0.3.0", + "pytz>=2023.2", + "icon4py-common>=0.0.5", + "icon4py-atmosphere-dycore>=0.0.5", + "icon4py-atmosphere-diffusion>=0.0.5" ] +description = "ICON model driver." dynamic = ['version'] +license = {file = "LICENSE"} +name = "icon4py-driver" +readme = "README.md" +requires-python = ">=3.10" [project.scripts] dycore_driver = "icon4py.model.driver.dycore_driver:main" - [project.urls] repository = "https://github.com/C2SM/icon4py" @@ -80,46 +79,46 @@ exclude_lines = [ ignore_errors = true [tool.coverage.run] -parallel = true branch = true +parallel = true source_pkgs = ['driver'] [tool.isort] +force_grid_wrap = 0 +include_trailing_comma = true +known_first_party = ['icon4py.model'] +known_third_party = ['gt4py'] lexicographical = true line_length = 100 # It should be the same as in `tool.black.line-length` above lines_after_imports = 2 +multi_line_output = 3 profile = 'black' sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'FIRSTPARTY', 'TESTS', 'LOCALFOLDER'] skip_gitignore = true skip_glob = ['*.venv/**', '_local/**'] -known_first_party = ['icon4py.model'] -known_third_party = ['gt4py'] -multi_line_output = 3 use_parentheses = true -include_trailing_comma = true -force_grid_wrap = 0 [tool.mypy] -install_types = true -non_interactive = true -exclude = [ - '^tests/*.py', -] disallow_incomplete_defs = true disallow_untyped_defs = true +exclude = [ + '^tests/*.py' +] ignore_missing_imports = false implicit_reexport = true -warn_unused_configs = true -warn_unused_ignores = true -warn_redundant_casts = true +install_types = true +non_interactive = true show_column_numbers = true show_error_codes = true +warn_redundant_casts = true +warn_unused_configs = true +warn_unused_ignores = true [tool.pytest] [tool.pytest.ini_options] -testpaths = 'tests' markers = 'datatest: this test uses binary data' +testpaths = 'tests' [tool.setuptools.dynamic] version = {attr = 'icon4py.model.driver.__init__.__version__'} diff --git a/model/driver/src/icon4py/model/driver/dycore_driver.py b/model/driver/src/icon4py/model/driver/dycore_driver.py index f9f50b2ad..9103a2443 100644 --- a/model/driver/src/icon4py/model/driver/dycore_driver.py +++ b/model/driver/src/icon4py/model/driver/dycore_driver.py @@ -21,18 +21,12 @@ from gt4py.next import Field, program from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn -from icon4py.model.atmosphere.diffusion.diffusion import ( - Diffusion, - DiffusionParams, -) +from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.model.atmosphere.diffusion.diffusion_states import ( DiffusionDiagnosticState, PrognosticState, ) -from icon4py.model.atmosphere.diffusion.diffusion_utils import ( - _identity_c_k, - _identity_e_k, -) +from icon4py.model.atmosphere.diffusion.diffusion_utils import _identity_c_k, _identity_e_k from icon4py.model.common.decomposition.decomposed import create_exchange from icon4py.model.common.decomposition.parallel_setup import ( ProcessProperties, @@ -259,9 +253,7 @@ def initialize(n_time_steps, file_path: Path, props: ProcessProperties): @click.command() @click.argument("input_path") @click.option("--run_path", default="", help="folder for output") -@click.option( - "--n_steps", default=5, help="number of time steps to run, max 5 is supported" -) +@click.option("--n_steps", default=5, help="number of time steps to run, max 5 is supported") @click.option("--mpi", default=False, help="whether or not you are running with mpi") def main(input_path, run_path, n_steps, mpi): """ diff --git a/model/driver/src/icon4py/model/driver/icon_configuration.py b/model/driver/src/icon4py/model/driver/icon_configuration.py index 54d92833f..b52f15a15 100644 --- a/model/driver/src/icon4py/model/driver/icon_configuration.py +++ b/model/driver/src/icon4py/model/driver/icon_configuration.py @@ -75,13 +75,9 @@ def _mch_ch_r04b09_config(n_steps): ) if experiment == "mch_ch_r04b09_dsl": - (model_run_config, diffusion_config, dycore_config) = _mch_ch_r04b09_config( - n_time_steps - ) + (model_run_config, diffusion_config, dycore_config) = _mch_ch_r04b09_config(n_time_steps) else: - (model_run_config, diffusion_config, dycore_config) = _default_config( - n_time_steps - ) + (model_run_config, diffusion_config, dycore_config) = _default_config(n_time_steps) return IconConfig( run_config=model_run_config, diffusion_config=diffusion_config, diff --git a/model/driver/src/icon4py/model/driver/io_utils.py b/model/driver/src/icon4py/model/driver/io_utils.py index de08ca15a..ae14933bb 100644 --- a/model/driver/src/icon4py/model/driver/io_utils.py +++ b/model/driver/src/icon4py/model/driver/io_utils.py @@ -23,10 +23,7 @@ PrognosticState, ) from icon4py.model.common.decomposition.decomposed import DecompositionInfo -from icon4py.model.common.decomposition.parallel_setup import ( - ParallelLogger, - ProcessProperties, -) +from icon4py.model.common.decomposition.parallel_setup import ParallelLogger, ProcessProperties from icon4py.model.common.grid.horizontal import CellParams, EdgeParams from icon4py.model.common.grid.icon_grid import IconGrid from icon4py.model.common.grid.vertical import VerticalModelParams @@ -58,9 +55,7 @@ def read_icon_grid( """ if ser_type == SerializationType.SB: return ( - sb.IconSerialDataProvider( - "icon_pydycore", str(path.absolute()), False, mpi_rank=rank - ) + sb.IconSerialDataProvider("icon_pydycore", str(path.absolute()), False, mpi_rank=rank) .from_savepoint_grid() .construct_icon_grid() ) @@ -112,9 +107,7 @@ def read_geometry_fields( ).from_savepoint_grid() edge_geometry = sp.construct_edge_geometry() cell_geometry = sp.construct_cell_geometry() - vertical_geometry = VerticalModelParams( - vct_a=sp.vct_a(), rayleigh_damping_height=12500 - ) + vertical_geometry = VerticalModelParams(vct_a=sp.vct_a(), rayleigh_damping_height=12500) return edge_geometry, cell_geometry, vertical_geometry else: raise NotImplementedError(SB_ONLY_MSG) @@ -157,17 +150,13 @@ def read_static_fields( interpolation_state = ( dataprovider.from_interpolation_savepoint().construct_interpolation_state_for_diffusion() ) - metric_state = ( - dataprovider.from_metrics_savepoint().construct_metric_state_for_diffusion() - ) + metric_state = dataprovider.from_metrics_savepoint().construct_metric_state_for_diffusion() return metric_state, interpolation_state else: raise NotImplementedError(SB_ONLY_MSG) -def configure_logging( - run_path: str, start_time, processor_procs: ProcessProperties = None -) -> None: +def configure_logging(run_path: str, start_time, processor_procs: ProcessProperties = None) -> None: """ Configure logging. @@ -178,13 +167,9 @@ def configure_logging( start_time: start time of the model run """ - run_dir = ( - Path(run_path).absolute() if run_path else Path(__file__).absolute().parent - ) + run_dir = Path(run_path).absolute() if run_path else Path(__file__).absolute().parent run_dir.mkdir(exist_ok=True) - logfile = run_dir.joinpath( - f"dummy_dycore_driver_{datetime.isoformat(start_time)}.log" - ) + logfile = run_dir.joinpath(f"dummy_dycore_driver_{datetime.isoformat(start_time)}.log") logfile.touch(exist_ok=True) logging.basicConfig( level=logging.DEBUG, @@ -195,9 +180,7 @@ def configure_logging( console_handler = logging.StreamHandler() console_handler.addFilter(ParallelLogger(processor_procs)) - log_format = ( - "{rank} {asctime} - {filename}: {funcName:<20}: {levelname:<7} {message}" - ) + log_format = "{rank} {asctime} - {filename}: {funcName:<20}: {levelname:<7} {message}" formatter = logging.Formatter(fmt=log_format, style="{", defaults={"rank": None}) console_handler.setFormatter(formatter) console_handler.setLevel(logging.DEBUG) diff --git a/model/driver/tests/test_io_utils.py b/model/driver/tests/test_io_utils.py index b3b7434e6..de51e2af1 100644 --- a/model/driver/tests/test_io_utils.py +++ b/model/driver/tests/test_io_utils.py @@ -25,9 +25,7 @@ @pytest.mark.datatest -@pytest.mark.parametrize( - "read_fun", (read_geometry_fields, read_static_fields, read_icon_grid) -) +@pytest.mark.parametrize("read_fun", (read_geometry_fields, read_static_fields, read_icon_grid)) def test_read_geometry_fields_not_implemented_type(read_fun, datapath): with pytest.raises(NotImplementedError, match=r"Only ser_type='sb'"): read_fun(path=datapath, ser_type=SerializationType.NC) @@ -54,9 +52,7 @@ def test_read_icon_grid_for_type_sb(datapath): @pytest.mark.datatest def test_read_static_fields_for_type_sb(datapath): - metric_state, interpolation_state = read_static_fields( - datapath, ser_type=SerializationType.SB - ) + metric_state, interpolation_state = read_static_fields(datapath, ser_type=SerializationType.SB) assert_metric_state_fields(metric_state) assert_interpolation_state_fields(interpolation_state) From ad9eb646153b525bbc3fd157eac0a27c61ec6230 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 5 Sep 2023 09:35:38 +0200 Subject: [PATCH 250/263] reformat common package --- model/common/.flake8 | 42 ++++++++++ model/common/pyproject.toml | 78 +++++++++---------- .../model/common/decomposition/decomposed.py | 42 +++------- .../common/decomposition/parallel_setup.py | 2 +- .../icon4py/model/common/grid/grid_manager.py | 22 ++---- .../icon4py/model/common/grid/horizontal.py | 19 ++--- .../icon4py/model/common/grid/icon_grid.py | 18 +---- .../model/common/test_utils/data_handling.py | 4 +- .../model/common/test_utils/fixtures.py | 4 +- .../model/common/test_utils/helpers.py | 8 +- .../common/test_utils/parallel_helpers.py | 4 +- .../common/test_utils/serialbox_utils.py | 52 ++++--------- .../model/common/test_utils/simple_mesh.py | 13 +--- model/common/tests/conftest.py | 20 ++--- .../common/tests/mpi_tests/test_decomposed.py | 20 ++--- .../tests/mpi_tests/test_parallel_setup.py | 5 +- model/common/tests/test_grid_manager.py | 36 +++------ model/common/tests/test_icon_grid.py | 4 +- .../common/tests/test_interpolation_fields.py | 4 +- ...est_mo_intp_rbf_rbf_vec_interpol_vertex.py | 6 +- model/common/tests/test_vertical.py | 9 +-- .../py2f/wrappers/diffusion_wrapper.py | 1 - 22 files changed, 153 insertions(+), 260 deletions(-) create mode 100644 model/common/.flake8 diff --git a/model/common/.flake8 b/model/common/.flake8 new file mode 100644 index 000000000..31cecff5a --- /dev/null +++ b/model/common/.flake8 @@ -0,0 +1,42 @@ +[flake8] +# Some sane defaults for the code style checker flake8 +max-line-length = 100 +max-complexity = 15 +doctests = true +extend-ignore = + # Do not perform function calls in argument defaults + B008, + # Public code object needs docstring + D1, + # Disable dargling errors by default + DAR, + # Whitespace before ':' (black formatter breaks this sometimes) + E203, + # Line too long (using Bugbear's B950 warning) + E501, + # Line break occurred before a binary operator + W503 + +exclude = + .eggs, + .gt_cache, + .ipynb_checkpoints, + .tox, + _local_, + build, + dist, + docs, + _external_src, + tests/_disabled, + setup.py + +rst-roles = + py:mod, mod, + py:func, func, + py:data, data, + py:const, const, + py:class, class, + py:meth, meth, + py:attr, attr, + py:exc, exc, + py:obj, obj, diff --git a/model/common/pyproject.toml b/model/common/pyproject.toml index f5485f947..0bc4e5f8b 100644 --- a/model/common/pyproject.toml +++ b/model/common/pyproject.toml @@ -1,38 +1,38 @@ [build-system] -requires = ["setuptools>=61.0", "wheel>=0.40.0"] build-backend = "setuptools.build_meta" +requires = ["setuptools>=61.0", "wheel>=0.40.0"] [project] -name = "icon4py-common" -description = "Shared code for the icon4py model." -readme = "README.md" -requires-python = ">=3.10" -license = {file = "LICENSE"} authors = [ - {email = "gridtools@cscs.ch"}, - {name = "ETH Zurich"} + {email = "gridtools@cscs.ch"}, + {name = "ETH Zurich"} ] classifiers = [ - "Development Status :: 3 - Alpha", - "Intended Audience :: Science/Research", - "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)", - "Operating System :: POSIX", - "Programming Language :: Python", - "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: Implementation :: CPython", - "Topic :: Scientific/Engineering :: Atmospheric Science", - "Topic :: Scientific/Engineering :: Mathematics", - "Topic :: Scientific/Engineering :: Physics" + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)", + "Operating System :: POSIX", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: Implementation :: CPython", + "Topic :: Scientific/Engineering :: Atmospheric Science", + "Topic :: Scientific/Engineering :: Mathematics", + "Topic :: Scientific/Engineering :: Physics" ] dependencies = [ - "gt4py>=1.0.1", - "mpi4py<=3.1.4", - "pyghex>=0.3.0", - "netcdf4>=1.6.0" + "gt4py>=1.0.1", + "mpi4py<=3.1.4", + "pyghex>=0.3.0", + "netcdf4>=1.6.0" ] +description = "Shared code for the icon4py model." dynamic = ['version'] +license = {file = "LICENSE"} +name = "icon4py-common" +readme = "README.md" +requires-python = ">=3.10" [project.urls] repository = "https://github.com/C2SM/icon4py" @@ -73,46 +73,46 @@ exclude_lines = [ ignore_errors = true [tool.coverage.run] -parallel = true branch = true +parallel = true source_pkgs = ['common'] [tool.isort] +force_grid_wrap = 0 +include_trailing_comma = true +known_first_party = ['icon4py.model'] +known_third_party = ['gt4py'] lexicographical = true line_length = 100 # It should be the same as in `tool.black.line-length` above lines_after_imports = 2 +multi_line_output = 3 profile = 'black' sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'FIRSTPARTY', 'TESTS', 'LOCALFOLDER'] skip_gitignore = true skip_glob = ['*.venv/**', '_local/**'] -known_first_party = ['icon4py.model'] -known_third_party = ['gt4py'] -multi_line_output = 3 use_parentheses = true -include_trailing_comma = true -force_grid_wrap = 0 [tool.mypy] -install_types = true -non_interactive = true -exclude = [ - '^tests/*.py', -] disallow_incomplete_defs = true disallow_untyped_defs = true +exclude = [ + '^tests/*.py' +] ignore_missing_imports = false implicit_reexport = true -warn_unused_configs = true -warn_unused_ignores = true -warn_redundant_casts = true +install_types = true +non_interactive = true show_column_numbers = true show_error_codes = true +warn_redundant_casts = true +warn_unused_configs = true +warn_unused_ignores = true [tool.pytest] [tool.pytest.ini_options] -testpaths = 'tests' markers = 'datatest: this test uses binary data' +testpaths = 'tests' [tool.setuptools.dynamic] version = {attr = 'icon4py.model.common.__init__.__version__'} diff --git a/model/common/src/icon4py/model/common/decomposition/decomposed.py b/model/common/src/icon4py/model/common/decomposition/decomposed.py index d46e52a15..23b278b65 100644 --- a/model/common/src/icon4py/model/common/decomposition/decomposed.py +++ b/model/common/src/icon4py/model/common/decomposition/decomposed.py @@ -22,12 +22,7 @@ from gt4py.next import Dimension from icon4py.model.common.decomposition.parallel_setup import ProcessProperties -from icon4py.model.common.dimension import ( - CellDim, - DimensionKind, - EdgeDim, - VertexDim, -) +from icon4py.model.common.dimension import CellDim, DimensionKind, EdgeDim, VertexDim from icon4py.model.common.utils import builder @@ -61,9 +56,7 @@ class EntryType(int, Enum): HALO = 2 @builder - def with_dimension( - self, dim: Dimension, global_index: np.ndarray, owner_mask: np.ndarray - ): + def with_dimension(self, dim: Dimension, global_index: np.ndarray, owner_mask: np.ndarray): masked_global_index = ma.array(global_index, mask=owner_mask) self._global_index[dim] = masked_global_index @@ -135,9 +128,7 @@ def is_ready(self) -> bool: return True -def create_exchange( - props: ProcessProperties, decomp_info: DecompositionInfo -) -> ExchangeRuntime: +def create_exchange(props: ProcessProperties, decomp_info: DecompositionInfo) -> ExchangeRuntime: """ Create an Exchange depending on the runtime size. @@ -170,9 +161,7 @@ def is_ready(self) -> bool: class GHexMultiNode: - def __init__( - self, props: ProcessProperties, domain_decomposition: DecompositionInfo - ): + def __init__(self, props: ProcessProperties, domain_decomposition: DecompositionInfo): self._context = ghex.context(ghex.mpi_comm(props.comm), True) self._domain_id_gen = DomainDescriptorIdGenerator(props) self._decomposition_info = domain_decomposition @@ -185,9 +174,7 @@ def __init__( ), EdgeDim: self._create_domain_descriptor(EdgeDim), } - log.info( - f"domain descriptors for dimensions {self._domain_descriptors.keys()} initialized" - ) + log.info(f"domain descriptors for dimensions {self._domain_descriptors.keys()} initialized") self._patterns = { CellDim: self._create_pattern(CellDim), @@ -208,12 +195,8 @@ def my_rank(self): return self._context.rank() def _create_domain_descriptor(self, dim: Dimension): - all_global = self._decomposition_info.global_index( - dim, DecompositionInfo.EntryType.ALL - ) - local_halo = self._decomposition_info.local_index( - dim, DecompositionInfo.EntryType.HALO - ) + all_global = self._decomposition_info.global_index(dim, DecompositionInfo.EntryType.ALL) + local_halo = self._decomposition_info.local_index(dim, DecompositionInfo.EntryType.HALO) # first arg is the domain ID which builds up an MPI Tag. # if those ids are not different for all domain descriptors the system might deadlock # if two parallel exchanges with the same domain id are done @@ -246,17 +229,12 @@ def exchange(self, dim: Dimension, *fields: tuple): pattern = self._patterns[dim] assert pattern is not None, f"pattern for {dim.value} not found" domain_descriptor = self._domain_descriptors[dim] - assert ( - domain_descriptor is not None - ), f"domain descriptor for {dim.value} not found" + assert domain_descriptor is not None, f"domain descriptor for {dim.value} not found" applied_patterns = [ - pattern(unstructured.field_descriptor(domain_descriptor, np.asarray(f))) - for f in fields + pattern(unstructured.field_descriptor(domain_descriptor, np.asarray(f))) for f in fields ] handle = self._comm.exchange(applied_patterns) - log.info( - f"exchange for {len(fields)} fields of dimension ='{dim.value}' initiated." - ) + log.info(f"exchange for {len(fields)} fields of dimension ='{dim.value}' initiated.") return MultiNodeResult(handle, applied_patterns) diff --git a/model/common/src/icon4py/model/common/decomposition/parallel_setup.py b/model/common/src/icon4py/model/common/decomposition/parallel_setup.py index 877885002..fd7c68816 100644 --- a/model/common/src/icon4py/model/common/decomposition/parallel_setup.py +++ b/model/common/src/icon4py/model/common/decomposition/parallel_setup.py @@ -91,6 +91,6 @@ def __init__(self, process_properties: ProcessProperties = None): if process_properties and process_properties.comm_size > 1: self._rank_info = f"rank={process_properties.rank}/{process_properties.comm_size} [{process_properties.comm_name}] " - def filter(self, record: logging.LogRecord) -> bool: # noqa: A003 + def filter(self, record: logging.LogRecord) -> bool: record.rank = self._rank_info return True diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py index 56bc1dcaf..79382c354 100644 --- a/model/common/src/icon4py/model/common/grid/grid_manager.py +++ b/model/common/src/icon4py/model/common/grid/grid_manager.py @@ -36,11 +36,7 @@ VertexDim, ) from icon4py.model.common.grid.horizontal import HorizontalGridSize -from icon4py.model.common.grid.icon_grid import ( - GridConfig, - IconGrid, - VerticalGridSize, -) +from icon4py.model.common.grid.icon_grid import GridConfig, IconGrid, VerticalGridSize class GridFileName(str, Enum): @@ -169,9 +165,7 @@ def __init__(self, dataset: Dataset): def dimension(self, name: GridFileName) -> int: return self._dataset.dimensions[name].size - def int_field( - self, name: GridFileName, transpose=True, dtype=np.int32 - ) -> np.ndarray: + def int_field(self, name: GridFileName, transpose=True, dtype=np.int32) -> np.ndarray: try: nc_variable = self._dataset.variables[name] @@ -392,15 +386,9 @@ def _from_grid_dataset(self, dataset: Dataset) -> IconGrid: V2E2VDim: v2e2v, } ) - .with_start_end_indices( - CellDim, start_indices[CellDim], end_indices[CellDim] - ) - .with_start_end_indices( - EdgeDim, start_indices[EdgeDim], end_indices[EdgeDim] - ) - .with_start_end_indices( - VertexDim, start_indices[VertexDim], end_indices[VertexDim] - ) + .with_start_end_indices(CellDim, start_indices[CellDim], end_indices[CellDim]) + .with_start_end_indices(EdgeDim, start_indices[EdgeDim], end_indices[EdgeDim]) + .with_start_end_indices(VertexDim, start_indices[VertexDim], end_indices[VertexDim]) ) return icon_grid diff --git a/model/common/src/icon4py/model/common/grid/horizontal.py b/model/common/src/icon4py/model/common/grid/horizontal.py index 8aa9155d5..fc05b90bd 100644 --- a/model/common/src/icon4py/model/common/grid/horizontal.py +++ b/model/common/src/icon4py/model/common/grid/horizontal.py @@ -195,9 +195,7 @@ def __init__( defined int ICON in mo_model_domain.f90:t_grid_edges%primal_edge_length """ - self.inverse_primal_edge_lengths: Field[ - [EdgeDim], float - ] = inverse_primal_edge_lengths + self.inverse_primal_edge_lengths: Field[[EdgeDim], float] = inverse_primal_edge_lengths """ Inverse of the triangle edge length: 1.0/primal_edge_length. @@ -211,18 +209,14 @@ def __init__( defined int ICON in mo_model_domain.f90:t_grid_edges%dual_edge_length """ - self.inverse_dual_edge_lengths: Field[ - [EdgeDim], float - ] = inverse_dual_edge_lengths + self.inverse_dual_edge_lengths: Field[[EdgeDim], float] = inverse_dual_edge_lengths """ Inverse of hexagon/pentagon edge length: 1.0/dual_edge_length. defined int ICON in mo_model_domain.f90:t_grid_edges%inv_dual_edge_length """ - self.inverse_vertex_vertex_lengths: Field[ - [EdgeDim], float - ] = inverse_vertex_vertex_lengths + self.inverse_vertex_vertex_lengths: Field[[EdgeDim], float] = inverse_vertex_vertex_lengths r""" Inverse distance between outer vertices of adjacent cells. @@ -239,9 +233,10 @@ def __init__( defined int ICON in mo_model_domain.f90:t_grid_edges%inv_vert_vert_length """ - self.primal_normal_vert: tuple[ - Field[[ECVDim], float], Field[[ECVDim], float] - ] = (primal_normal_vert_x, primal_normal_vert_y) + self.primal_normal_vert: tuple[Field[[ECVDim], float], Field[[ECVDim], float]] = ( + primal_normal_vert_x, + primal_normal_vert_y, + ) """ Normal of the triangle edge, projected onto the location of the vertices diff --git a/model/common/src/icon4py/model/common/grid/icon_grid.py b/model/common/src/icon4py/model/common/grid/icon_grid.py index 3332ef405..d7e8585d2 100644 --- a/model/common/src/icon4py/model/common/grid/icon_grid.py +++ b/model/common/src/icon4py/model/common/grid/icon_grid.py @@ -19,15 +19,7 @@ from gt4py.next.iterator.embedded import NeighborTableOffsetProvider from typing_extensions import deprecated -from icon4py.model.common.dimension import ( - CECDim, - CEDim, - CellDim, - ECVDim, - EdgeDim, - KDim, - VertexDim, -) +from icon4py.model.common.dimension import CECDim, CEDim, CellDim, ECVDim, EdgeDim, KDim, VertexDim from icon4py.model.common.grid.horizontal import HorizontalGridSize from icon4py.model.common.grid.vertical import VerticalGridSize from icon4py.model.common.utils import builder @@ -125,9 +117,7 @@ def get_indices_from_to( start_marker to the end of the region given by the end_marker """ if dim.kind != DimensionKind.HORIZONTAL: - raise ValueError( - "only defined for {} dimension kind ", DimensionKind.HORIZONTAL - ) + raise ValueError("only defined for {} dimension kind ", DimensionKind.HORIZONTAL) return self.start_indices[dim][start_marker], self.end_indices[dim][end_marker] def get_start_index(self, dim: Dimension, marker: int) -> int32: @@ -196,9 +186,7 @@ def _neighbortable_offset_provider_for_1d_sparse_fields( neighbor_axis: Dimension, ): table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) - return NeighborTableOffsetProvider( - table, origin_axis, neighbor_axis, table.shape[1] - ) + return NeighborTableOffsetProvider(table, origin_axis, neighbor_axis, table.shape[1]) def get_c2cec_connectivity(self): return self._neighbortable_offset_provider_for_1d_sparse_fields( diff --git a/model/common/src/icon4py/model/common/test_utils/data_handling.py b/model/common/src/icon4py/model/common/test_utils/data_handling.py index e069b24f6..b00bb2d32 100644 --- a/model/common/src/icon4py/model/common/test_utils/data_handling.py +++ b/model/common/src/icon4py/model/common/test_utils/data_handling.py @@ -20,9 +20,7 @@ def download_and_extract(uri: str, local_path: Path, data_file: str): local_path.mkdir(parents=True, exist_ok=True) if not any(local_path.iterdir()): - print( - f"directory {local_path} is empty: downloading data from {uri} and extracting" - ) + print(f"directory {local_path} is empty: downloading data from {uri} and extracting") wget.download(uri, out=data_file) # extract downloaded file if not tarfile.is_tarfile(data_file): diff --git a/model/common/src/icon4py/model/common/test_utils/fixtures.py b/model/common/src/icon4py/model/common/test_utils/fixtures.py index 66921a308..bf575b446 100644 --- a/model/common/src/icon4py/model/common/test_utils/fixtures.py +++ b/model/common/src/icon4py/model/common/test_utils/fixtures.py @@ -76,9 +76,7 @@ def download_ser_data(request, processor_props, ranked_data_path): @pytest.fixture(scope="session") -def data_provider( - download_ser_data, datapath, processor_props -) -> IconSerialDataProvider: +def data_provider(download_ser_data, datapath, processor_props) -> IconSerialDataProvider: return IconSerialDataProvider( fname_prefix="icon_pydycore", path=str(datapath), diff --git a/model/common/src/icon4py/model/common/test_utils/helpers.py b/model/common/src/icon4py/model/common/test_utils/helpers.py index 9060bb89c..682cc974c 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -66,9 +66,7 @@ def random_field( extend: Optional[dict[gt_common.Dimension, int]] = None, ) -> it_embedded.MutableLocatedField: return it_embedded.np_as_located_field(*dims)( - np.random.default_rng().uniform( - low=low, high=high, size=_shape(mesh, *dims, extend=extend) - ) + np.random.default_rng().uniform(low=low, high=high, size=_shape(mesh, *dims, extend=extend)) ) @@ -115,9 +113,7 @@ def flatten_first_two_dims( def _test_validation(self, mesh, backend, input_data): - reference_outputs = self.reference( - mesh, **{k: np.array(v) for k, v in input_data.items()} - ) + reference_outputs = self.reference(mesh, **{k: np.array(v) for k, v in input_data.items()}) self.PROGRAM.with_backend(backend)( **input_data, offset_provider=mesh.get_offset_provider(), diff --git a/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py b/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py index 7fce8ba16..6dbe56db6 100644 --- a/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/parallel_helpers.py @@ -19,6 +19,4 @@ def check_comm_size(props: ProcessProperties, sizes=(1, 2, 4)): if props.comm_size not in sizes: - pytest.xfail( - f"wrong comm size: {props.comm_size}: test only works for comm-sizes: {sizes}" - ) + pytest.xfail(f"wrong comm size: {props.comm_size}: test only works for comm-sizes: {sizes}") 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 b4e50b4f6..f4550ccb5 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 @@ -43,16 +43,8 @@ V2EDim, VertexDim, ) -from icon4py.model.common.grid.horizontal import ( - CellParams, - EdgeParams, - HorizontalGridSize, -) -from icon4py.model.common.grid.icon_grid import ( - GridConfig, - IconGrid, - VerticalGridSize, -) +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalGridSize +from icon4py.model.common.grid.icon_grid import GridConfig, IconGrid, VerticalGridSize from icon4py.model.common.test_utils.helpers import as_1D_sparse_field @@ -287,9 +279,7 @@ def construct_icon_grid(self) -> IconGrid: C2E2CODim: c2e2c0, } ) - .with_connectivities( - {E2VDim: self.e2v(), V2EDim: self.v2e(), E2C2VDim: self.e2c2v()} - ) + .with_connectivities({E2VDim: self.e2v(), V2EDim: self.v2e(), E2C2VDim: self.e2c2v()}) ) return grid @@ -322,9 +312,9 @@ class InterpolationSavepoint(IconSavepoint): def geofac_grg(self): grg = np.squeeze(self.serializer.read("geofac_grg", self.savepoint)) num_cells = self.sizes[CellDim] - return np_as_located_field(CellDim, C2E2CODim)( - grg[:num_cells, :, 0] - ), np_as_located_field(CellDim, C2E2CODim)(grg[:num_cells, :, 1]) + return np_as_located_field(CellDim, C2E2CODim)(grg[:num_cells, :, 0]), np_as_located_field( + CellDim, C2E2CODim + )(grg[:num_cells, :, 1]) def zd_intcoef(self): return self._get_field("vcoef", CellDim, C2E2CDim, KDim) @@ -383,9 +373,7 @@ def zd_diffcoef(self): def zd_intcoef(self): return self._read_and_reorder_sparse_field("vcoef", CellDim) - def _read_and_reorder_sparse_field( - self, name: str, horizontal_dim: Dimension, sparse_size=3 - ): + def _read_and_reorder_sparse_field(self, name: str, horizontal_dim: Dimension, sparse_size=3): ser_input = np.squeeze(self.serializer.read(name, self.savepoint))[ : self.sizes[horizontal_dim], :, : ] @@ -538,9 +526,7 @@ def __init__(self, fname_prefix, path=".", do_print=False, mpi_rank=0): def _init_serializer(self, do_print: bool): if not self.fname: self.log.warning(" WARNING: no filename! closing serializer") - self.serializer = ser.Serializer( - ser.OpenModeKind.Read, self.file_path, self.fname - ) + self.serializer = ser.Serializer(ser.OpenModeKind.Read, self.file_path, self.fname) if do_print: self.print_info() @@ -572,14 +558,9 @@ def from_savepoint_diffusion_init( date: str, ) -> IconDiffusionInitSavepoint: savepoint = ( - self.serializer.savepoint["call-diffusion-init"] - .linit[linit] - .date[date] - .as_savepoint() - ) - return IconDiffusionInitSavepoint( - savepoint, self.serializer, size=self.grid_size + self.serializer.savepoint["call-diffusion-init"].linit[linit].date[date].as_savepoint() ) + return IconDiffusionInitSavepoint(savepoint, self.serializer, size=self.grid_size) def from_interpolation_savepoint(self) -> InterpolationSavepoint: savepoint = self.serializer.savepoint["interpolation_state"].as_savepoint() @@ -589,15 +570,8 @@ def from_metrics_savepoint(self) -> MetricSavepoint: savepoint = self.serializer.savepoint["metric_state"].as_savepoint() return MetricSavepoint(savepoint, self.serializer, size=self.grid_size) - def from_savepoint_diffusion_exit( - self, linit: bool, date: str - ) -> IconDiffusionExitSavepoint: + def from_savepoint_diffusion_exit(self, linit: bool, date: str) -> IconDiffusionExitSavepoint: savepoint = ( - self.serializer.savepoint["call-diffusion-exit"] - .linit[linit] - .date[date] - .as_savepoint() - ) - return IconDiffusionExitSavepoint( - savepoint, self.serializer, size=self.grid_size + self.serializer.savepoint["call-diffusion-exit"].linit[linit].date[date].as_savepoint() ) + return IconDiffusionExitSavepoint(savepoint, self.serializer, size=self.grid_size) diff --git a/model/common/src/icon4py/model/common/test_utils/simple_mesh.py b/model/common/src/icon4py/model/common/test_utils/simple_mesh.py index 430c42e20..4715be0fb 100644 --- a/model/common/src/icon4py/model/common/test_utils/simple_mesh.py +++ b/model/common/src/icon4py/model/common/test_utils/simple_mesh.py @@ -14,10 +14,7 @@ from dataclasses import dataclass import numpy as np -from gt4py.next.iterator.embedded import ( - NeighborTableOffsetProvider, - StridedNeighborOffsetProvider, -) +from gt4py.next.iterator.embedded import NeighborTableOffsetProvider, StridedNeighborOffsetProvider from icon4py.model.common.dimension import ( C2E2CDim, @@ -461,16 +458,12 @@ def get_e2c2v_offset_provider(self) -> NeighborTableOffsetProvider: def get_e2ecv_offset_provider(self): old_shape = self.e2c2v.shape e2ecv_table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) - return NeighborTableOffsetProvider( - e2ecv_table, EdgeDim, ECVDim, e2ecv_table.shape[1] - ) + return NeighborTableOffsetProvider(e2ecv_table, EdgeDim, ECVDim, e2ecv_table.shape[1]) def get_c2ce_offset_provider(self): old_shape = self.c2e.shape c2ce_table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) - return NeighborTableOffsetProvider( - c2ce_table, CellDim, CEDim, c2ce_table.shape[1] - ) + return NeighborTableOffsetProvider(c2ce_table, CellDim, CEDim, c2ce_table.shape[1]) def get_offset_provider(self): return { diff --git a/model/common/tests/conftest.py b/model/common/tests/conftest.py index e9024dc03..1824c1b89 100644 --- a/model/common/tests/conftest.py +++ b/model/common/tests/conftest.py @@ -32,18 +32,12 @@ grids_path = base_path.joinpath("grids") r04b09_dsl_grid_path = grids_path.joinpath("mch_ch_r04b09_dsl") -r04b09_dsl_data_file = r04b09_dsl_grid_path.joinpath( - "mch_ch_r04b09_dsl_grids_v1.tar.gz" -).name +r04b09_dsl_data_file = r04b09_dsl_grid_path.joinpath("mch_ch_r04b09_dsl_grids_v1.tar.gz").name r02b04_global_grid_path = grids_path.joinpath("r02b04_global") -r02b04_global_data_file = r02b04_global_grid_path.joinpath( - "icon_grid_0013_R02B04_G.tar.gz" -).name +r02b04_global_data_file = r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_G.tar.gz").name -mch_ch_r04b09_dsl_grid_uri = ( - "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download" -) +mch_ch_r04b09_dsl_grid_uri = "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download" r02b04_global_grid_uri = "https://polybox.ethz.ch/index.php/s/0EM8O8U53GKGsst/download" @@ -59,9 +53,5 @@ def get_grid_files(): Session scoped fixture which is a prerequisite of all the other fixtures in this file. """ - download_and_extract( - mch_ch_r04b09_dsl_grid_uri, r04b09_dsl_grid_path, r04b09_dsl_data_file - ) - download_and_extract( - r02b04_global_grid_uri, r02b04_global_grid_path, r02b04_global_data_file - ) + download_and_extract(mch_ch_r04b09_dsl_grid_uri, r04b09_dsl_grid_path, r04b09_dsl_data_file) + download_and_extract(r02b04_global_grid_uri, r02b04_global_grid_path, r02b04_global_data_file) diff --git a/model/common/tests/mpi_tests/test_decomposed.py b/model/common/tests/mpi_tests/test_decomposed.py index 27c5be973..b363461f0 100644 --- a/model/common/tests/mpi_tests/test_decomposed.py +++ b/model/common/tests/mpi_tests/test_decomposed.py @@ -67,14 +67,10 @@ def test_decomposition_info_masked( my_owned = owned[my_rank] assert all_indices.shape[0] == my_total - owned_indices = decomposition_info.global_index( - dim, DecompositionInfo.EntryType.OWNED - ) + owned_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.OWNED) assert owned_indices.shape[0] == my_owned - halo_indices = decomposition_info.global_index( - dim, DecompositionInfo.EntryType.HALO - ) + halo_indices = decomposition_info.global_index(dim, DecompositionInfo.EntryType.HALO) assert halo_indices.shape[0] == my_total - my_owned _assert_index_partitioning(all_indices, halo_indices, owned_indices) @@ -122,9 +118,7 @@ def test_decomposition_info_local_index( assert halo_indices.shape[0] < all_indices.shape[0] assert np.alltrue(halo_indices <= np.max(all_indices)) - owned_indices = decomposition_info.local_index( - dim, DecompositionInfo.EntryType.OWNED - ) + owned_indices = decomposition_info.local_index(dim, DecompositionInfo.EntryType.OWNED) assert owned_indices.shape[0] == my_owned assert owned_indices.shape[0] <= all_indices.shape[0] assert np.alltrue(owned_indices <= np.max(all_indices)) @@ -172,15 +166,11 @@ def test_decomposition_info_matches_gridsize( == icon_grid.num_cells() ) assert ( - decomposition_info.global_index( - VertexDim, DecompositionInfo.EntryType.ALL - ).shape[0] + decomposition_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape[0] == icon_grid.num_vertices() ) assert ( - decomposition_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape[ - 0 - ] + decomposition_info.global_index(EdgeDim, DecompositionInfo.EntryType.ALL).shape[0] == icon_grid.num_edges() ) diff --git a/model/common/tests/mpi_tests/test_parallel_setup.py b/model/common/tests/mpi_tests/test_parallel_setup.py index 1cfd6856c..bcf714005 100644 --- a/model/common/tests/mpi_tests/test_parallel_setup.py +++ b/model/common/tests/mpi_tests/test_parallel_setup.py @@ -15,10 +15,7 @@ import pytest from mpi4py import MPI -from icon4py.model.common.decomposition.parallel_setup import ( - get_processor_properties, - init_mpi, -) +from icon4py.model.common.decomposition.parallel_setup import get_processor_properties, init_mpi @pytest.mark.mpi diff --git a/model/common/tests/test_grid_manager.py b/model/common/tests/test_grid_manager.py index 28b714b7d..6f5f6c47a 100644 --- a/model/common/tests/test_grid_manager.py +++ b/model/common/tests/test_grid_manager.py @@ -43,9 +43,7 @@ def simple_mesh_gridfile(tmp_path): dataset.createDimension(GridFile.DimensionName.EDGE_NAME, size=mesh.n_edges) dataset.createDimension(GridFile.DimensionName.CELL_NAME, size=mesh.n_cells) - dataset.createDimension( - GridFile.DimensionName.NEIGHBORS_TO_EDGE_SIZE, size=mesh.n_e2v - ) + dataset.createDimension(GridFile.DimensionName.NEIGHBORS_TO_EDGE_SIZE, size=mesh.n_e2v) dataset.createDimension(GridFile.DimensionName.DIAMOND_EDGE_SIZE, size=mesh.n_e2c2e) dataset.createDimension(GridFile.DimensionName.MAX_CHILD_DOMAINS, size=1) # add dummy values for the grf dimensions @@ -72,12 +70,8 @@ def simple_mesh_gridfile(tmp_path): (GridFile.DimensionName.VERTEX_NAME,), ) - dataset.createDimension( - GridFile.DimensionName.NEIGHBORS_TO_CELL_SIZE, size=mesh.n_c2e - ) - dataset.createDimension( - GridFile.DimensionName.NEIGHBORS_TO_VERTEX_SIZE, size=mesh.n_v2c - ) + dataset.createDimension(GridFile.DimensionName.NEIGHBORS_TO_CELL_SIZE, size=mesh.n_c2e) + dataset.createDimension(GridFile.DimensionName.NEIGHBORS_TO_VERTEX_SIZE, size=mesh.n_v2c) _add_to_dataset( dataset, @@ -217,15 +211,9 @@ def test_gridfile_vertex_cell_edge_dimensions(grid_savepoint, r04b09_dsl_gridfil data = Dataset(r04b09_dsl_gridfile, "r") grid_file = GridFile(data) - assert grid_file.dimension(GridFile.DimensionName.CELL_NAME) == grid_savepoint.num( - CellDim - ) - assert grid_file.dimension(GridFile.DimensionName.EDGE_NAME) == grid_savepoint.num( - EdgeDim - ) - assert grid_file.dimension( - GridFile.DimensionName.VERTEX_NAME - ) == grid_savepoint.num(VertexDim) + assert grid_file.dimension(GridFile.DimensionName.CELL_NAME) == grid_savepoint.num(CellDim) + assert grid_file.dimension(GridFile.DimensionName.EDGE_NAME) == grid_savepoint.num(EdgeDim) + assert grid_file.dimension(GridFile.DimensionName.VERTEX_NAME) == grid_savepoint.num(VertexDim) def test_grid_parser_index_fields(simple_mesh_gridfile, caplog): @@ -404,9 +392,7 @@ def init_grid_manager(fname): @pytest.mark.parametrize("dim, size", [(CellDim, 18), (EdgeDim, 27), (VertexDim, 9)]) def test_grid_manager_getsize(simple_mesh_gridfile, dim, size, caplog): caplog.set_level(logging.DEBUG) - gm = GridManager( - IndexTransformation(), simple_mesh_gridfile, VerticalGridSize(num_lev=80) - ) + gm = GridManager(IndexTransformation(), simple_mesh_gridfile, VerticalGridSize(num_lev=80)) gm() assert size == gm.get_size(dim) @@ -484,9 +470,7 @@ def test_gt4py_transform_offset_by_1_where_valid(size): def test_get_start_index(r04b09_dsl_gridfile, icon_grid, dim, marker, index): grid_from_manager = init_grid_manager(r04b09_dsl_gridfile).get_grid() assert grid_from_manager.get_start_index(dim, marker) == index - assert grid_from_manager.get_start_index(dim, marker) == icon_grid.get_start_index( - dim, marker - ) + assert grid_from_manager.get_start_index(dim, marker) == icon_grid.get_start_index(dim, marker) @pytest.mark.datatest @@ -534,6 +518,4 @@ def test_get_start_index(r04b09_dsl_gridfile, icon_grid, dim, marker, index): def test_get_end_index(r04b09_dsl_gridfile, icon_grid, dim, marker, index): grid_from_manager = init_grid_manager(r04b09_dsl_gridfile).get_grid() assert grid_from_manager.get_end_index(dim, marker) == index - assert grid_from_manager.get_end_index(dim, marker) == icon_grid.get_end_index( - dim, marker - ) + assert grid_from_manager.get_end_index(dim, marker) == icon_grid.get_end_index(dim, marker) diff --git a/model/common/tests/test_icon_grid.py b/model/common/tests/test_icon_grid.py index 0de42f6ab..e9b5f39ce 100644 --- a/model/common/tests/test_icon_grid.py +++ b/model/common/tests/test_icon_grid.py @@ -338,9 +338,7 @@ def test_horizontal_edge_markers(icon_grid, start_marker, end_marker, expected_b ), ], ) -def test_horizontal_vertex_markers( - icon_grid, start_marker, end_marker, expected_bounds -): +def test_horizontal_vertex_markers(icon_grid, start_marker, end_marker, expected_bounds): assert ( icon_grid.get_indices_from_to( VertexDim, diff --git a/model/common/tests/test_interpolation_fields.py b/model/common/tests/test_interpolation_fields.py index 1d49e4ded..e4c8eb5b9 100644 --- a/model/common/tests/test_interpolation_fields.py +++ b/model/common/tests/test_interpolation_fields.py @@ -28,9 +28,7 @@ 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.interpolation.interpolation_fields import compute_c_lin_e @pytest.mark.datatest diff --git a/model/common/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index 02ae28148..189beae08 100644 --- a/model/common/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -19,11 +19,7 @@ from icon4py.model.common.interpolation.stencils.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) -from icon4py.model.common.test_utils.helpers import ( - StencilTest, - random_field, - zero_field, -) +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field class TestMoIntpRbfRbfVecInterpolVertex(StencilTest): diff --git a/model/common/tests/test_vertical.py b/model/common/tests/test_vertical.py index 42b2d4039..f205941b9 100644 --- a/model/common/tests/test_vertical.py +++ b/model/common/tests/test_vertical.py @@ -28,19 +28,14 @@ def test_nrdmax_calculation(max_h, damping, delta): vct_a = np.arange(0, max_h, delta) vct_a = vct_a[::-1] vertical_params = VerticalModelParams(rayleigh_damping_height=damping, vct_a=vct_a) - assert ( - vertical_params.index_of_damping_layer - == vct_a.shape[0] - math.ceil(damping / delta) - 1 - ) + assert vertical_params.index_of_damping_layer == vct_a.shape[0] - math.ceil(damping / delta) - 1 @pytest.mark.datatest def test_nrdmax_calculation_from_icon_input(grid_savepoint, damping_height): a = grid_savepoint.vct_a() nrdmax = grid_savepoint.nrdmax() - vertical_params = VerticalModelParams( - rayleigh_damping_height=damping_height, vct_a=a - ) + vertical_params = VerticalModelParams(rayleigh_damping_height=damping_height, vct_a=a) assert nrdmax == vertical_params.index_of_damping_layer a_array = np.asarray(a) assert a_array[nrdmax] > damping_height diff --git a/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py index 002087bbd..89c2dbb59 100644 --- a/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py +++ b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py @@ -56,7 +56,6 @@ from icon4pytools.py2f.cffi_utils import CffiMethod, to_fields - # TODO (magdalena) Revise interface architecture with Fortran granules: # The module variable to match the Fortran interface: where only fields are passed. # We should rather instantiate the object init and return it. From 75e976603b530dd7f9d47276f5b377357ca46fc8 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 5 Sep 2023 09:44:21 +0200 Subject: [PATCH 251/263] pre-commit on common --- model/common/.pre-commit-config.yaml | 1 + .../src/icon4py/model/common/decomposition/parallel_setup.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/model/common/.pre-commit-config.yaml b/model/common/.pre-commit-config.yaml index 3d45afc46..13dc05548 100644 --- a/model/common/.pre-commit-config.yaml +++ b/model/common/.pre-commit-config.yaml @@ -63,6 +63,7 @@ repos: rev: v1.3.0 hooks: - id: yesqa + exclude: 'decomposition/parallel_setup.py' - repo: https://github.com/psf/black rev: '22.3.0' diff --git a/model/common/src/icon4py/model/common/decomposition/parallel_setup.py b/model/common/src/icon4py/model/common/decomposition/parallel_setup.py index fd7c68816..877885002 100644 --- a/model/common/src/icon4py/model/common/decomposition/parallel_setup.py +++ b/model/common/src/icon4py/model/common/decomposition/parallel_setup.py @@ -91,6 +91,6 @@ def __init__(self, process_properties: ProcessProperties = None): if process_properties and process_properties.comm_size > 1: self._rank_info = f"rank={process_properties.rank}/{process_properties.comm_size} [{process_properties.comm_name}] " - def filter(self, record: logging.LogRecord) -> bool: + def filter(self, record: logging.LogRecord) -> bool: # noqa: A003 record.rank = self._rank_info return True From 9d0a358830900ba80447de741a4b5f8d86d97527 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 5 Sep 2023 09:47:57 +0200 Subject: [PATCH 252/263] pre-commit on dycore --- model/atmosphere/dycore/pyproject.toml | 6 +++--- .../mo_advection_traj_btraj_compute_o1_dsl.py | 6 ++---- ...nterpolation_scalar_cells2verts_scalar_ri_dsl.py | 4 +--- .../dycore/mo_solve_nonhydro_stencil_02.py | 4 +--- .../dycore/mo_solve_nonhydro_stencil_10.py | 4 +--- ...solve_nonhydro_stencil_16_fused_btraj_traj_o1.py | 6 ++---- .../dycore/mo_solve_nonhydro_stencil_17.py | 5 +---- .../dycore/mo_solve_nonhydro_stencil_18.py | 4 +--- .../dycore/mo_solve_nonhydro_stencil_20.py | 10 ++-------- .../dycore/mo_solve_nonhydro_stencil_22.py | 4 +--- .../dycore/mo_solve_nonhydro_stencil_24.py | 4 +--- .../dycore/mo_solve_nonhydro_stencil_39.py | 4 +--- .../dycore/mo_solve_nonhydro_stencil_40.py | 12 +++--------- .../dycore/mo_solve_nonhydro_stencil_42.py | 4 +--- .../dycore/mo_solve_nonhydro_stencil_55.py | 13 ++----------- .../dycore/mo_solve_nonhydro_stencil_56_63.py | 4 +--- .../dycore/mo_solve_nonhydro_stencil_58.py | 4 +--- .../dycore/mo_solve_nonhydro_stencil_60.py | 4 +--- .../dycore/mo_solve_nonhydro_stencil_65.py | 4 +--- .../dycore/mo_solve_nonhydro_stencil_67.py | 4 +--- .../dycore/mo_solve_nonhydro_stencil_68.py | 5 +---- .../dycore/mo_velocity_advection_stencil_04.py | 4 +--- .../dycore/mo_velocity_advection_stencil_07.py | 4 +--- .../dycore/mo_velocity_advection_stencil_10.py | 4 +--- .../dycore/mo_velocity_advection_stencil_14.py | 8 ++------ .../dycore/mo_velocity_advection_stencil_16.py | 8 ++------ .../dycore/mo_velocity_advection_stencil_19.py | 5 +---- .../test_mo_advection_traj_btraj_compute_o1_dsl.py | 6 ++---- .../tests/test_mo_solve_nonhydro_stencil_02.py | 4 +--- .../tests/test_mo_solve_nonhydro_stencil_04.py | 9 +++------ .../tests/test_mo_solve_nonhydro_stencil_05.py | 8 ++------ .../tests/test_mo_solve_nonhydro_stencil_06.py | 4 +--- .../tests/test_mo_solve_nonhydro_stencil_10.py | 8 ++------ .../tests/test_mo_solve_nonhydro_stencil_15.py | 4 +--- ...solve_nonhydro_stencil_16_fused_btraj_traj_o1.py | 6 ++---- .../tests/test_mo_solve_nonhydro_stencil_19.py | 5 ++--- .../tests/test_mo_solve_nonhydro_stencil_20.py | 8 ++------ .../tests/test_mo_solve_nonhydro_stencil_21.py | 4 +--- .../tests/test_mo_solve_nonhydro_stencil_25.py | 4 +--- .../tests/test_mo_solve_nonhydro_stencil_26.py | 4 +--- .../tests/test_mo_solve_nonhydro_stencil_28.py | 4 +--- .../tests/test_mo_solve_nonhydro_stencil_29.py | 4 +--- .../tests/test_mo_solve_nonhydro_stencil_39.py | 8 ++------ .../tests/test_mo_solve_nonhydro_stencil_40.py | 12 +++--------- .../tests/test_mo_solve_nonhydro_stencil_42.py | 4 +--- .../tests/test_mo_solve_nonhydro_stencil_44.py | 4 +--- .../tests/test_mo_solve_nonhydro_stencil_47.py | 4 +--- .../tests/test_mo_solve_nonhydro_stencil_51.py | 4 +--- .../tests/test_mo_solve_nonhydro_stencil_52.py | 8 ++------ .../tests/test_mo_solve_nonhydro_stencil_54.py | 4 +--- .../tests/test_mo_solve_nonhydro_stencil_55.py | 5 +---- .../tests/test_mo_solve_nonhydro_stencil_58.py | 4 +--- .../tests/test_mo_solve_nonhydro_stencil_60.py | 4 +--- .../tests/test_mo_solve_nonhydro_stencil_62.py | 4 +--- .../tests/test_mo_solve_nonhydro_stencil_65.py | 4 +--- .../tests/test_mo_solve_nonhydro_stencil_68.py | 5 +---- .../tests/test_mo_velocity_advection_stencil_02.py | 12 +++--------- .../tests/test_mo_velocity_advection_stencil_08.py | 4 +--- .../tests/test_mo_velocity_advection_stencil_09.py | 4 +--- .../tests/test_mo_velocity_advection_stencil_10.py | 8 ++------ .../tests/test_mo_velocity_advection_stencil_13.py | 4 +--- .../tests/test_mo_velocity_advection_stencil_19.py | 5 +---- 62 files changed, 87 insertions(+), 254 deletions(-) diff --git a/model/atmosphere/dycore/pyproject.toml b/model/atmosphere/dycore/pyproject.toml index 8721ffa6f..842b53ceb 100644 --- a/model/atmosphere/dycore/pyproject.toml +++ b/model/atmosphere/dycore/pyproject.toml @@ -22,8 +22,8 @@ classifiers = [ "Topic :: Scientific/Engineering :: Physics" ] dependencies = [ - "gt4py>=1.0.1", - "icon4py-common>=0.0.5", + "gt4py>=1.0.1", + "icon4py-common>=0.0.5" ] description = "ICON dynamical core." dynamic = ['version'] @@ -110,8 +110,8 @@ warn_unused_ignores = true [tool.pytest] [tool.pytest.ini_options] -testpaths = 'tests' markers = 'datatest: this test uses binary data' +testpaths = 'tests' [tool.setuptools.dynamic] version = {attr = 'icon4py.model.atmosphere.dycore.__init__.__version__'} diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py index 06a94e5ec..cd8a2e9bf 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py @@ -43,13 +43,11 @@ def _mo_advection_traj_btraj_compute_o1_dsl( p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) z_ntdistv_bary_1 = -( - p_vn * p_dthalf - + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) + p_vn * p_dthalf + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf - + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) + p_vt * p_dthalf + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) ) p_distv_bary_1 = where( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py index 06a4125d8..295be5cab 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py @@ -33,6 +33,4 @@ def mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl( c_intp: Field[[VertexDim, V2CDim], float], p_vert_out: Field[[VertexDim, KDim], float], ): - _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl( - p_cell_in, c_intp, out=p_vert_out - ) + _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl(p_cell_in, c_intp, out=p_vert_out) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py index 509a9e615..ca293f5a9 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_02.py @@ -24,9 +24,7 @@ def _mo_solve_nonhydro_stencil_02( exner_ref_mc: Field[[CellDim, KDim], float], exner_pr: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - z_exner_ex_pr = (1.0 + exner_exfac) * ( - exner - exner_ref_mc - ) - exner_exfac * exner_pr + z_exner_ex_pr = (1.0 + exner_exfac) * (exner - exner_ref_mc) - exner_exfac * exner_pr exner_pr = exner - exner_ref_mc return z_exner_ex_pr, exner_pr diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py index bfb6b25f6..b70912057 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_10.py @@ -43,9 +43,7 @@ def _mo_solve_nonhydro_stencil_10( ]: z_w_backtraj = -(w - w_concorr_c) * dtime * 0.5 / ddqz_z_half z_rho_tavg_m1 = wgt_nnow_rth * rho_now(Koff[-1]) + wgt_nnew_rth * rho_var(Koff[-1]) - z_theta_tavg_m1 = wgt_nnow_rth * theta_now(Koff[-1]) + wgt_nnew_rth * theta_var( - Koff[-1] - ) + z_theta_tavg_m1 = wgt_nnow_rth * theta_now(Koff[-1]) + wgt_nnew_rth * theta_var(Koff[-1]) z_rho_tavg = wgt_nnow_rth * rho_now + wgt_nnew_rth * rho_var z_theta_tavg = wgt_nnow_rth * theta_now + wgt_nnew_rth * theta_var rho_ic = ( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py index 0818b0cc5..124776266 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py @@ -33,13 +33,11 @@ def _compute_btraj( lvn_pos = where(p_vn > 0.0, True, False) z_ntdistv_bary_1 = -( - p_vn * p_dthalf - + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) + p_vn * p_dthalf + where(lvn_pos, pos_on_tplane_e_1(E2EC[0]), pos_on_tplane_e_1(E2EC[1])) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf - + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) + p_vt * p_dthalf + where(lvn_pos, pos_on_tplane_e_2(E2EC[0]), pos_on_tplane_e_2(E2EC[1])) ) p_distv_bary_1 = where( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py index f7841c6d0..4c628b66c 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_17.py @@ -28,10 +28,7 @@ def _mo_solve_nonhydro_stencil_17( ) -> Field[[EdgeDim, KDim], float]: scalfac_dd3d = broadcast(scalfac_dd3d, (EdgeDim, KDim)) z_graddiv_vn = z_graddiv_vn + ( - hmask_dd3d - * scalfac_dd3d - * inv_dual_edge_length - * (z_dwdz_dd(E2C[1]) - z_dwdz_dd(E2C[0])) + hmask_dd3d * scalfac_dd3d * inv_dual_edge_length * (z_dwdz_dd(E2C[1]) - z_dwdz_dd(E2C[0])) ) return z_graddiv_vn diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py index 9d99e03ba..5046f22bf 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_18.py @@ -23,9 +23,7 @@ def _mo_solve_nonhydro_stencil_18( inv_dual_edge_length: Field[[EdgeDim], float], z_exner_ex_pr: Field[[CellDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - z_gradh_exner = inv_dual_edge_length * ( - z_exner_ex_pr(E2C[1]) - z_exner_ex_pr(E2C[0]) - ) + z_gradh_exner = inv_dual_edge_length * (z_exner_ex_pr(E2C[1]) - z_exner_ex_pr(E2C[0])) return z_gradh_exner diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py index 737d0f57d..8971bb810 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_20.py @@ -41,18 +41,12 @@ def _mo_solve_nonhydro_stencil_20( ( z_exner_ex_pr_1(E2C[1]) + zdiff_gradp(E2EC[1]) - * ( - z_dexner_dz_c1_1(E2C[1]) - + zdiff_gradp(E2EC[1]) * z_dexner_dz_c2_1(E2C[1]) - ) + * (z_dexner_dz_c1_1(E2C[1]) + zdiff_gradp(E2EC[1]) * z_dexner_dz_c2_1(E2C[1])) ) - ( z_exner_ex_pr_0(E2C[0]) + zdiff_gradp(E2EC[0]) - * ( - z_dexner_dz_c1_0(E2C[0]) - + zdiff_gradp(E2EC[0]) * z_dexner_dz_c2_0(E2C[0]) - ) + * (z_dexner_dz_c1_0(E2C[0]) + zdiff_gradp(E2EC[0]) * z_dexner_dz_c2_0(E2C[0])) ) ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py index f1919b193..84fb050f6 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_22.py @@ -25,9 +25,7 @@ def _mo_solve_nonhydro_stencil_22( z_hydro_corr: Field[[EdgeDim], float], z_gradh_exner: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - z_gradh_exner = where( - ipeidx_dsl, z_gradh_exner + z_hydro_corr * pg_exdist, z_gradh_exner - ) + z_gradh_exner = where(ipeidx_dsl, z_gradh_exner + z_hydro_corr * pg_exdist, z_gradh_exner) return z_gradh_exner diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py index 66f3c2234..a23544d44 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_24.py @@ -28,9 +28,7 @@ def _mo_solve_nonhydro_stencil_24( dtime: float, cpd: float, ) -> Field[[EdgeDim, KDim], float]: - vn_nnew = vn_nnow + dtime * ( - ddt_vn_adv_ntl1 + ddt_vn_phy - cpd * z_theta_v_e * z_gradh_exner - ) + vn_nnew = vn_nnow + dtime * (ddt_vn_adv_ntl1 + ddt_vn_phy - cpd * z_theta_v_e * z_gradh_exner) return vn_nnew diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py index f50801788..4b192e07b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_39.py @@ -25,9 +25,7 @@ def _mo_solve_nonhydro_stencil_39( wgtfac_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: z_w_concorr_me_offset_1 = z_w_concorr_me(Koff[-1]) - z_w_concorr_mc_m1 = neighbor_sum( - e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim - ) + z_w_concorr_mc_m1 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim) z_w_concorr_mc_m0 = neighbor_sum(e_bln_c_s * z_w_concorr_me(C2E), axis=C2EDim) w_concorr_c = wgtfac_c * z_w_concorr_mc_m0 + (1.0 - wgtfac_c) * z_w_concorr_mc_m1 return w_concorr_c diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py index 8c1f88676..5b57d8f19 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_40.py @@ -28,15 +28,9 @@ def _mo_solve_nonhydro_stencil_40( z_w_concorr_me_offset_2 = z_w_concorr_me(Koff[-2]) z_w_concorr_me_offset_3 = z_w_concorr_me(Koff[-3]) - z_w_concorr_mc_m1 = neighbor_sum( - e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim - ) - z_w_concorr_mc_m2 = neighbor_sum( - e_bln_c_s * z_w_concorr_me_offset_2(C2E), axis=C2EDim - ) - z_w_concorr_mc_m3 = neighbor_sum( - e_bln_c_s * z_w_concorr_me_offset_3(C2E), axis=C2EDim - ) + z_w_concorr_mc_m1 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_1(C2E), axis=C2EDim) + z_w_concorr_mc_m2 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_2(C2E), axis=C2EDim) + z_w_concorr_mc_m3 = neighbor_sum(e_bln_c_s * z_w_concorr_me_offset_3(C2E), axis=C2EDim) return ( wgtfacq_c(Koff[-1]) * z_w_concorr_mc_m1 diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py index 8606aa719..5b6b0d51e 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_42.py @@ -33,9 +33,7 @@ def _mo_solve_nonhydro_stencil_42( cpd: float, ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: z_w_expl = w_nnow + dtime * ( - wgt_nnow_vel * ddt_w_adv_ntl1 - + wgt_nnew_vel * ddt_w_adv_ntl2 - - cpd * z_th_ddz_exner_c + wgt_nnow_vel * ddt_w_adv_ntl1 + wgt_nnew_vel * ddt_w_adv_ntl2 - cpd * z_th_ddz_exner_c ) z_contr_w_fl_l = rho_ic * (-w_concorr_c + vwind_expl_wgt * w_nnow) return z_w_expl, z_contr_w_fl_l diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py index d55620484..a354730e7 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py @@ -42,17 +42,8 @@ def _mo_solve_nonhydro_stencil_55( rho_new = z_rho_expl - vwind_impl_wgt * dtime * inv_ddqz_z_full * ( rho_ic * w - rho_ic(Koff[1]) * w(Koff[1]) ) - exner_new = ( - z_exner_expl - + exner_ref_mc - - z_beta * (z_alpha * w - z_alpha(Koff[1]) * w(Koff[1])) - ) - theta_v_new = ( - rho_now - * theta_v_now - * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) - / rho_new - ) + exner_new = z_exner_expl + exner_ref_mc - z_beta * (z_alpha * w - z_alpha(Koff[1]) * w(Koff[1])) + theta_v_new = rho_now * theta_v_now * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) / rho_new return rho_new, exner_new, theta_v_new diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py index b1f718731..e5e70e7ba 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py @@ -24,9 +24,7 @@ def _mo_solve_nonhydro_stencil_56_63( w: Field[[CellDim, KDim], float], w_concorr_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - z_dwdz_dd = inv_ddqz_z_full * ( - (w - w(Koff[1])) - (w_concorr_c - w_concorr_c(Koff[1])) - ) + z_dwdz_dd = inv_ddqz_z_full * ((w - w(Koff[1])) - (w_concorr_c - w_concorr_c(Koff[1]))) return z_dwdz_dd diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py index 8b5347ada..17879f2db 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py @@ -27,9 +27,7 @@ def _mo_solve_nonhydro_stencil_58( mass_flx_ic: Field[[CellDim, KDim], float], r_nsubsteps: float, ) -> Field[[CellDim, KDim], float]: - mass_flx_ic = mass_flx_ic + ( - r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w) - ) + mass_flx_ic = mass_flx_ic + (r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w)) return mass_flx_ic diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py index d0435fdde..995322d72 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py @@ -26,9 +26,7 @@ def _mo_solve_nonhydro_stencil_60( ndyn_substeps_var: float, dtime: float, ) -> Field[[CellDim, KDim], float]: - exner_dyn_incr = exner - ( - exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy - ) + exner_dyn_incr = exner - (exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy) return exner_dyn_incr diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py index 3ff46a988..eea826bad 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py @@ -30,9 +30,7 @@ def _mo_solve_nonhydro_stencil_65( r_nsubsteps: float, ) -> Field[[CellDim, KDim], float]: mass_flx_ic = mass_flx_ic + ( - r_nsubsteps - * rho_ic - * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) + r_nsubsteps * rho_ic * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) ) return mass_flx_ic diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py index 0c949bd69..6f46f314d 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py @@ -39,6 +39,4 @@ def mo_solve_nonhydro_stencil_67( rd_o_cvd: float, rd_o_p0ref: float, ): - _mo_solve_nonhydro_stencil_67( - rho, theta_v, exner, rd_o_cvd, rd_o_p0ref, out=(theta_v, exner) - ) + _mo_solve_nonhydro_stencil_67(rho, theta_v, exner, rd_o_cvd, rd_o_p0ref, out=(theta_v, exner)) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py index 8725c196c..15230b8af 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py @@ -31,10 +31,7 @@ def _mo_solve_nonhydro_stencil_68( ) -> Field[[CellDim, KDim], float]: theta_v_new = where( mask_prog_halo_c, - rho_now - * theta_v_now - * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) - / rho_new, + rho_now * theta_v_now * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) / rho_new, theta_v_new, ) return theta_v_new diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py index 72ad7e547..c7c051da0 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py @@ -37,6 +37,4 @@ def mo_velocity_advection_stencil_04( vt: Field[[EdgeDim, KDim], float], z_w_concorr_me: Field[[EdgeDim, KDim], float], ): - _mo_velocity_advection_stencil_04( - vn, ddxn_z_full, ddxt_z_full, vt, out=z_w_concorr_me - ) + _mo_velocity_advection_stencil_04(vn, ddxn_z_full, ddxt_z_full, vt, out=z_w_concorr_me) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py index 31e916167..b4b4c697f 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_07.py @@ -30,9 +30,7 @@ def _mo_velocity_advection_stencil_07( ) -> Field[[EdgeDim, KDim], float]: return vn_ie * inv_dual_edge_length * ( w(E2C[0]) - w(E2C[1]) - ) + z_vt_ie * inv_primal_edge_length * tangent_orientation * ( - z_w_v(E2V[0]) - z_w_v(E2V[1]) - ) + ) + z_vt_ie * inv_primal_edge_length * tangent_orientation * (z_w_v(E2V[0]) - z_w_v(E2V[1])) @program(grid_type=GridType.UNSTRUCTURED) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py index 4cdbc7d73..4dd324603 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_10.py @@ -23,9 +23,7 @@ def _mo_velocity_advection_stencil_10( z_w_concorr_mc: Field[[CellDim, KDim], float], wgtfac_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc( - Koff[-1] - ) + w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc(Koff[-1]) return w_concorr_c diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py index d18b7afb7..b61a42a54 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py @@ -39,13 +39,9 @@ def _mo_velocity_advection_stencil_14( # should reduce the vertical cfl to a scalar and the levmask to a per level boolean field (see Fortran dycore). vcfl = where(cfl_clipping, z_w_con_c * dtime / ddqz_z_half, 0.0) - z_w_con_c = where( - (cfl_clipping) & (vcfl < -0.85), -0.85 * ddqz_z_half / dtime, z_w_con_c - ) + z_w_con_c = where((cfl_clipping) & (vcfl < -0.85), -0.85 * ddqz_z_half / dtime, z_w_con_c) - z_w_con_c = where( - (cfl_clipping) & (vcfl > 0.85), 0.85 * ddqz_z_half / dtime, z_w_con_c - ) + z_w_con_c = where((cfl_clipping) & (vcfl > 0.85), 0.85 * ddqz_z_half / dtime, z_w_con_c) return cfl_clipping, vcfl, z_w_con_c diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py index 38b60dd96..a1823ac6b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_16.py @@ -26,9 +26,7 @@ def _mo_velocity_advection_stencil_16( coeff2_dwdz: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: ddt_w_adv = -z_w_con_c * ( - w(Koff[-1]) * coeff1_dwdz - - w(Koff[1]) * coeff2_dwdz - + w * (coeff2_dwdz - coeff1_dwdz) + w(Koff[-1]) * coeff1_dwdz - w(Koff[1]) * coeff2_dwdz + w * (coeff2_dwdz - coeff1_dwdz) ) return ddt_w_adv @@ -41,6 +39,4 @@ def mo_velocity_advection_stencil_16( coeff2_dwdz: Field[[CellDim, KDim], float], ddt_w_adv: Field[[CellDim, KDim], float], ): - _mo_velocity_advection_stencil_16( - z_w_con_c, w, coeff1_dwdz, coeff2_dwdz, out=ddt_w_adv[:, 1:] - ) + _mo_velocity_advection_stencil_16(z_w_con_c, w, coeff1_dwdz, coeff2_dwdz, out=ddt_w_adv[:, 1:]) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py index 6476e2a57..621e6befe 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_19.py @@ -45,10 +45,7 @@ def _mo_velocity_advection_stencil_19( ): ddt_vn_adv = -( (coeff_gradekin(E2EC[0]) - coeff_gradekin(E2EC[1])) * z_kin_hor_e - + ( - -coeff_gradekin(E2EC[0]) * z_ekinh(E2C[0]) - + coeff_gradekin(E2EC[1]) * z_ekinh(E2C[1]) - ) + + (-coeff_gradekin(E2EC[0]) * z_ekinh(E2C[0]) + coeff_gradekin(E2EC[1]) * z_ekinh(E2C[1])) + vt * (f_e + 0.5 * neighbor_sum(zeta(E2V), axis=E2VDim)) + neighbor_sum(z_w_con_c_full(E2C) * c_lin_e, axis=E2CDim) * (vn_ie - vn_ie(Koff[1])) diff --git a/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py index 371732361..becc4c26b 100644 --- a/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py @@ -70,12 +70,10 @@ def reference( p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) z_ntdistv_bary_1 = -( - p_vn * p_dthalf - + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) + p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf - + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) + p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) ) p_distv_bary_1 = np.where( diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py index c04d712ff..05975778c 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_02.py @@ -34,9 +34,7 @@ def reference( exner_exfac: np.array, **kwargs, ) -> dict: - z_exner_ex_pr = (1 + exner_exfac) * ( - exner - exner_ref_mc - ) - exner_exfac * exner_pr + z_exner_ex_pr = (1 + exner_exfac) * (exner - exner_ref_mc) - exner_exfac * exner_pr exner_pr = exner - exner_ref_mc return dict(z_exner_ex_pr=z_exner_ex_pr, exner_pr=exner_pr) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py index 5b1cf61da..e247e9f7a 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_04.py @@ -33,12 +33,9 @@ def reference( z_exner_ic: np.array, ) -> np.array: z_exner_ic[:, 3:] = ( - np.roll(wgtfacq_c, shift=1, axis=1) - * np.roll(z_exner_ex_pr, shift=1, axis=1) - + np.roll(wgtfacq_c, shift=2, axis=1) - * np.roll(z_exner_ex_pr, shift=2, axis=1) - + np.roll(wgtfacq_c, shift=3, axis=1) - * np.roll(z_exner_ex_pr, shift=3, axis=1) + np.roll(wgtfacq_c, shift=1, axis=1) * np.roll(z_exner_ex_pr, shift=1, axis=1) + + np.roll(wgtfacq_c, shift=2, axis=1) * np.roll(z_exner_ex_pr, shift=2, axis=1) + + np.roll(wgtfacq_c, shift=3, axis=1) * np.roll(z_exner_ex_pr, shift=3, axis=1) )[:, 3:] return {"z_exner_ic": z_exner_ic} diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py index 7a5289af1..6d54d033a 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_05.py @@ -26,13 +26,9 @@ class TestMoSolveNonhydroStencil05(StencilTest): OUTPUTS = ("z_exner_ic",) @staticmethod - def reference( - mesh, wgtfac_c: np.array, z_exner_ex_pr: np.array, **kwargs - ) -> np.array: + def reference(mesh, wgtfac_c: np.array, z_exner_ex_pr: np.array, **kwargs) -> np.array: z_exner_ex_pr_offset_1 = np.roll(z_exner_ex_pr, shift=1, axis=1) - z_exner_ic = ( - wgtfac_c * z_exner_ex_pr + (1.0 - wgtfac_c) * z_exner_ex_pr_offset_1 - ) + z_exner_ic = wgtfac_c * z_exner_ex_pr + (1.0 - wgtfac_c) * z_exner_ex_pr_offset_1 return dict(z_exner_ic=z_exner_ic) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py index 12ef9a3cb..e6086dc44 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_06.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil06(StencilTest): OUTPUTS = ("z_dexner_dz_c_1",) @staticmethod - def reference( - mesh, z_exner_ic: np.array, inv_ddqz_z_full: np.array, **kwargs - ) -> np.array: + def reference(mesh, z_exner_ic: np.array, inv_ddqz_z_full: np.array, **kwargs) -> np.array: z_dexner_dz_c_1 = (z_exner_ic[:, :-1] - z_exner_ic[:, 1:]) * inv_ddqz_z_full return dict(z_dexner_dz_c_1=z_dexner_dz_c_1) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py index c2d1d80e2..d66d399fa 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_10.py @@ -55,9 +55,7 @@ def reference( z_w_backtraj = -(w - w_concorr_c) * dtime * 0.5 / ddqz_z_half z_rho_tavg_m1 = wgt_nnow_rth * rho_now_offset + wgt_nnew_rth * rho_var_offset - z_theta_tavg_m1 = ( - wgt_nnow_rth * theta_now_offset + wgt_nnew_rth * theta_var_offset - ) + z_theta_tavg_m1 = wgt_nnow_rth * theta_now_offset + wgt_nnew_rth * theta_var_offset z_rho_tavg = wgt_nnow_rth * rho_now + wgt_nnew_rth * rho_var z_theta_tavg = wgt_nnow_rth * theta_now + wgt_nnew_rth * theta_var rho_ic = ( @@ -67,9 +65,7 @@ def reference( ) z_theta_v_pr_mc_m1 = z_theta_tavg_m1 - theta_ref_mc_offset z_theta_v_pr_mc = z_theta_tavg - theta_ref_mc - z_theta_v_pr_ic = ( - wgtfac_c * z_theta_v_pr_mc + (1 - wgtfac_c) * z_theta_v_pr_mc_m1 - ) + z_theta_v_pr_ic = wgtfac_c * z_theta_v_pr_mc + (1 - wgtfac_c) * z_theta_v_pr_mc_m1 theta_v_ic = ( wgtfac_c * z_theta_tavg + (1 - wgtfac_c) * z_theta_tavg_m1 diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py index 40f10d203..bb92f1fbe 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_15.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil15(StencilTest): OUTPUTS = ("z_rho_e", "z_theta_v_e") @staticmethod - def reference( - mesh, z_rho_e: np.array, z_theta_v_e: np.array, **kwargs - ) -> tuple[np.array]: + def reference(mesh, z_rho_e: np.array, z_theta_v_e: np.array, **kwargs) -> tuple[np.array]: z_rho_e = np.zeros_like(z_rho_e) z_theta_v_e = np.zeros_like(z_theta_v_e) return dict(z_rho_e=z_rho_e, z_theta_v_e=z_theta_v_e) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py index 8dcbd9b88..461c0a654 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py @@ -48,12 +48,10 @@ def compute_btraj_numpy( dual_normal_cell_2 = np.expand_dims(dual_normal_cell_2, axis=-1) z_ntdistv_bary_1 = -( - p_vn * p_dthalf - + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) + p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf - + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) + p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) ) p_distv_bary_1 = np.where( diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py index 56d05c7b9..5a782f7c5 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_19.py @@ -41,9 +41,8 @@ def reference( z_exner_ex_pr_e2c = z_exner_ex_pr[mesh.e2c] z_exner_ex_weighted = z_exner_ex_pr_e2c[:, 1] - z_exner_ex_pr_e2c[:, 0] - z_gradh_exner = ( - inv_dual_edge_length * z_exner_ex_weighted - - ddxn_z_full * np.sum(c_lin_e * z_dexner_dz_c_1[mesh.e2c], axis=1) + z_gradh_exner = inv_dual_edge_length * z_exner_ex_weighted - ddxn_z_full * np.sum( + c_lin_e * z_dexner_dz_c_1[mesh.e2c], axis=1 ) return dict(z_gradh_exner=z_gradh_exner) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py index 6fa501ec8..320fb7dc4 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_20.py @@ -47,12 +47,8 @@ def _apply_index_field(shape, to_index, neighbor_table, offset_field): inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, -1) z_exner_ex_pr_at_kidx = _apply_index_field(full_shape, z_exner_ex_pr, e2c, ikoffset) - z_dexner_dz_c_1_at_kidx = _apply_index_field( - full_shape, z_dexner_dz_c_1, e2c, ikoffset - ) - z_dexner_dz_c_2_at_kidx = _apply_index_field( - full_shape, z_dexner_dz_c_2, e2c, ikoffset - ) + z_dexner_dz_c_1_at_kidx = _apply_index_field(full_shape, z_dexner_dz_c_1, e2c, ikoffset) + z_dexner_dz_c_2_at_kidx = _apply_index_field(full_shape, z_dexner_dz_c_2, e2c, ikoffset) def at_neighbor(i): return z_exner_ex_pr_at_kidx[:, i, :] + zdiff_gradp[:, i, :] * ( diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py index c5d316662..a601630e6 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_21.py @@ -57,9 +57,7 @@ def _apply_index_field(shape, to_index, neighbor_table, offset_field): full_shape, theta_v_ic, e2c, ikoffset ) - inv_ddqz_z_full_at_kidx, _ = _apply_index_field( - full_shape, inv_ddqz_z_full, e2c, ikoffset - ) + inv_ddqz_z_full_at_kidx, _ = _apply_index_field(full_shape, inv_ddqz_z_full, e2c, ikoffset) z_theta1 = ( theta_v_at_kidx[:, 0, :] diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py index fd5765f80..85fdb0788 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_25.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil25(StencilTest): OUTPUTS = ("z_graddiv2_vn",) @staticmethod - def reference( - mesh, geofac_grdiv: np.array, z_graddiv_vn: np.array, **kwargs - ) -> np.array: + def reference(mesh, geofac_grdiv: np.array, z_graddiv_vn: np.array, **kwargs) -> np.array: geofac_grdiv = np.expand_dims(geofac_grdiv, axis=-1) z_graddiv2_vn = np.sum(z_graddiv_vn[mesh.e2c2eO] * geofac_grdiv, axis=1) return dict(z_graddiv2_vn=z_graddiv2_vn) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py index dded703d1..ec2389a8d 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_26.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil26(StencilTest): OUTPUTS = ("vn",) @staticmethod - def reference( - mesh, z_graddiv_vn: np.array, vn: np.array, scal_divdamp_o2, **kwargs - ) -> dict: + def reference(mesh, z_graddiv_vn: np.array, vn: np.array, scal_divdamp_o2, **kwargs) -> dict: vn = vn + (scal_divdamp_o2 * z_graddiv_vn) return dict(vn=vn) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py index 2eeba98f4..17234d448 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_28.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil28(StencilTest): OUTPUTS = ("vn",) @staticmethod - def reference( - mesh, vn_incr: np.array, vn: np.array, iau_wgt_dyn, **kwargs - ) -> np.array: + def reference(mesh, vn_incr: np.array, vn: np.array, iau_wgt_dyn, **kwargs) -> np.array: vn = vn + (iau_wgt_dyn * vn_incr) return dict(vn=vn) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py index 1520b78be..cc48d754d 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_29.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil29(StencilTest): OUTPUTS = ("vn_new",) @staticmethod - def reference( - mesh, grf_tend_vn: np.array, vn_now: np.array, dtime, **kwargs - ) -> dict: + def reference(mesh, grf_tend_vn: np.array, vn_now: np.array, dtime, **kwargs) -> dict: vn_new = vn_now + dtime * grf_tend_vn return dict(vn_new=vn_new) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py index 46062545d..053013639 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_39.py @@ -36,12 +36,8 @@ def reference( e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) z_w_concorr_me_offset_1 = np.roll(z_w_concorr_me, shift=1, axis=1) z_w_concorr_mc_m0 = np.sum(e_bln_c_s * z_w_concorr_me[mesh.c2e], axis=1) - z_w_concorr_mc_m1 = np.sum( - e_bln_c_s * z_w_concorr_me_offset_1[mesh.c2e], axis=1 - ) - w_concorr_c = ( - wgtfac_c * z_w_concorr_mc_m0 + (1.0 - wgtfac_c) * z_w_concorr_mc_m1 - ) + z_w_concorr_mc_m1 = np.sum(e_bln_c_s * z_w_concorr_me_offset_1[mesh.c2e], axis=1) + w_concorr_c = wgtfac_c * z_w_concorr_mc_m0 + (1.0 - wgtfac_c) * z_w_concorr_mc_m1 return dict(w_concorr_c=w_concorr_c) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py index 896b38a04..6774dc946 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_40.py @@ -37,15 +37,9 @@ def reference( z_w_concorr_me_offset_1 = np.roll(z_w_concorr_me, shift=1, axis=1) z_w_concorr_me_offset_2 = np.roll(z_w_concorr_me, shift=2, axis=1) z_w_concorr_me_offset_3 = np.roll(z_w_concorr_me, shift=3, axis=1) - z_w_concorr_mc_m1 = np.sum( - e_bln_c_s * z_w_concorr_me_offset_1[mesh.c2e], axis=1 - ) - z_w_concorr_mc_m2 = np.sum( - e_bln_c_s * z_w_concorr_me_offset_2[mesh.c2e], axis=1 - ) - z_w_concorr_mc_m3 = np.sum( - e_bln_c_s * z_w_concorr_me_offset_3[mesh.c2e], axis=1 - ) + z_w_concorr_mc_m1 = np.sum(e_bln_c_s * z_w_concorr_me_offset_1[mesh.c2e], axis=1) + z_w_concorr_mc_m2 = np.sum(e_bln_c_s * z_w_concorr_me_offset_2[mesh.c2e], axis=1) + z_w_concorr_mc_m3 = np.sum(e_bln_c_s * z_w_concorr_me_offset_3[mesh.c2e], axis=1) w_concorr_c = ( np.roll(wgtfacq_c, shift=1, axis=1) * z_w_concorr_mc_m1 + np.roll(wgtfacq_c, shift=2, axis=1) * z_w_concorr_mc_m2 diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py index 6835c5fbd..2e0882c12 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_42.py @@ -42,9 +42,7 @@ def reference( **kwargs, ) -> tuple[np.array]: z_w_expl = w_nnow + dtime * ( - wgt_nnow_vel * ddt_w_adv_ntl1 - + wgt_nnew_vel * ddt_w_adv_ntl2 - - cpd * z_th_ddz_exner_c + wgt_nnow_vel * ddt_w_adv_ntl1 + wgt_nnew_vel * ddt_w_adv_ntl2 - cpd * z_th_ddz_exner_c ) vwind_expl_wgt = np.expand_dims(vwind_expl_wgt, axis=-1) z_contr_w_fl_l = rho_ic * (-w_concorr_c + vwind_expl_wgt * w_nnow) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py index ea0e2aed0..6008665b5 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_44.py @@ -40,9 +40,7 @@ def reference( cvd, **kwargs, ) -> dict: - z_beta = ( - dtime * rd * exner_nnow / (cvd * rho_nnow * theta_v_nnow) * inv_ddqz_z_full - ) + z_beta = dtime * rd * exner_nnow / (cvd * rho_nnow * theta_v_nnow) * inv_ddqz_z_full vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) z_alpha = vwind_impl_wgt * theta_v_ic * rho_ic diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py index 5892ce5b6..66aa9b6b0 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_47.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil47(StencilTest): OUTPUTS = ("w_nnew", "z_contr_w_fl_l") @staticmethod - def reference( - mesh, w_concorr_c: np.array, z_contr_w_fl_l: np.array, **kwargs - ) -> dict: + def reference(mesh, w_concorr_c: np.array, z_contr_w_fl_l: np.array, **kwargs) -> dict: w_nnew = w_concorr_c z_contr_w_fl_l = np.zeros_like(z_contr_w_fl_l) return dict(w_nnew=w_nnew, z_contr_w_fl_l=z_contr_w_fl_l) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py index d1cb07a66..f7e12c96f 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_51.py @@ -57,9 +57,7 @@ def mo_solve_nonhydro_stencil_51_numpy( z_c = -z_gamma * z_beta * z_alpha_k_plus_1 z_b = 1.0 + z_gamma * z_alpha[:, :-1] * (z_beta_k_minus_1 + z_beta) z_q = mo_solve_nonhydro_stencil_51_z_q_numpy(z_c, z_b) - w_nnew = mo_solve_nonhydro_stencil_51_w_nnew_numpy( - z_gamma, z_b, z_w_expl, z_exner_expl - ) + w_nnew = mo_solve_nonhydro_stencil_51_w_nnew_numpy(z_gamma, z_b, z_w_expl, z_exner_expl) return z_q, w_nnew diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py index 85b153189..6c0cc28e0 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_52.py @@ -49,15 +49,11 @@ def mo_solve_nonhydro_stencil_52_numpy( for k in range(1, k_size): z_a[:, k] = -z_gamma[:, k] * z_beta[:, k - 1] * z_alpha[:, k - 1] z_c[:, k] = -z_gamma[:, k] * z_beta[:, k] * z_alpha[:, k + 1] - z_b[:, k] = 1.0 + z_gamma[:, k] * z_alpha[:, k] * ( - z_beta[:, k - 1] + z_beta[:, k] - ) + z_b[:, k] = 1.0 + z_gamma[:, k] * z_alpha[:, k] * (z_beta[:, k - 1] + z_beta[:, k]) z_g[:, k] = 1.0 / (z_b[:, k] + z_a[:, k] * z_q[:, k - 1]) z_q[:, k] = -z_c[:, k] * z_g[:, k] - w[:, k] = z_w_expl[:, k] - z_gamma[:, k] * ( - z_exner_expl[:, k - 1] - z_exner_expl[:, k] - ) + w[:, k] = z_w_expl[:, k] - z_gamma[:, k] * (z_exner_expl[:, k - 1] - z_exner_expl[:, k]) w[:, k] = (w[:, k] - z_a[:, k] * w[:, k - 1]) * z_g[:, k] return z_q, w diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py index bc2e501a2..eeb17c2cc 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_54.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil54(StencilTest): OUTPUTS = ("w",) @staticmethod - def reference( - mesh, z_raylfac: np.array, w_1: np.array, w: np.array, **kwargs - ) -> np.array: + def reference(mesh, z_raylfac: np.array, w_1: np.array, w: np.array, **kwargs) -> np.array: z_raylfac = np.expand_dims(z_raylfac, axis=0) w_1 = np.expand_dims(w_1, axis=-1) w = z_raylfac * w + (1.0 - z_raylfac) * w_1 diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py index 1ba408306..b8e216639 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_55.py @@ -58,10 +58,7 @@ def reference( - z_beta * (z_alpha[:, :-1] * w_offset_0 - z_alpha_offset_1 * w_offset_1) ) theta_v_new = ( - rho_now - * theta_v_now - * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) - / rho_new + rho_now * theta_v_now * ((exner_new / exner_now - 1.0) * cvd_o_rd + 1.0) / rho_new ) return dict(rho_new=rho_new, exner_new=exner_new, theta_v_new=theta_v_new) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py index 93e8c6a95..dff001052 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_58.py @@ -37,9 +37,7 @@ def reference( **kwargs, ) -> dict: vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) - mass_flx_ic = mass_flx_ic + ( - r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w) - ) + mass_flx_ic = mass_flx_ic + (r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w)) return dict(mass_flx_ic=mass_flx_ic) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py index 407dc06e9..634446cb3 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_60.py @@ -35,9 +35,7 @@ def reference( dtime: float, **kwargs, ) -> np.array: - exner_dyn_incr = exner - ( - exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy - ) + exner_dyn_incr = exner - (exner_dyn_incr + ndyn_substeps_var * dtime * ddt_exner_phy) return dict(exner_dyn_incr=exner_dyn_incr) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py index d31fffed0..9c3d3442c 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_62.py @@ -26,9 +26,7 @@ class TestMoSolveNonhydroStencil62(StencilTest): OUTPUTS = ("w_new",) @staticmethod - def reference( - mesh, w_now: np.array, grf_tend_w: np.array, dtime: float, **kwargs - ) -> np.array: + def reference(mesh, w_now: np.array, grf_tend_w: np.array, dtime: float, **kwargs) -> np.array: w_new = w_now + dtime * grf_tend_w return dict(w_new=w_new) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py index 77e46f592..b5b5b9c0c 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_65.py @@ -41,9 +41,7 @@ def reference( vwind_expl_wgt = np.expand_dims(vwind_expl_wgt, axis=-1) vwind_impl_wgt = np.expand_dims(vwind_impl_wgt, axis=-1) mass_flx_ic = mass_flx_ic + ( - r_nsubsteps - * rho_ic - * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) + r_nsubsteps * rho_ic * (vwind_expl_wgt * w_now + vwind_impl_wgt * w_new - w_concorr_c) ) return dict(mass_flx_ic=mass_flx_ic) diff --git a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py index 081dbb08d..3a53d5270 100644 --- a/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py +++ b/model/atmosphere/dycore/tests/test_mo_solve_nonhydro_stencil_68.py @@ -42,10 +42,7 @@ def reference( theta_v_new = np.where( mask_prog_halo_c, - rho_now - * theta_v_now - * ((exner_new / exner_now - 1) * cvd_o_rd + 1.0) - / rho_new, + rho_now * theta_v_now * ((exner_new / exner_now - 1) * cvd_o_rd + 1.0) / rho_new, theta_v_new, ) return dict(theta_v_new=theta_v_new) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py index 60fabe060..142b7e56c 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_02.py @@ -26,24 +26,18 @@ class TestMoVelocityAdvectionStencil02VnIe(StencilTest): OUTPUTS = ("vn_ie", "z_kin_hor_e") @staticmethod - def mo_velocity_advection_stencil_02_vn_ie_numpy( - wgtfac_e: np.array, vn: np.array - ) -> np.array: + def mo_velocity_advection_stencil_02_vn_ie_numpy(wgtfac_e: np.array, vn: np.array) -> np.array: vn_ie_k_minus_1 = np.roll(vn, shift=1, axis=1) vn_ie = wgtfac_e * vn + (1.0 - wgtfac_e) * vn_ie_k_minus_1 return vn_ie @staticmethod - def mo_velocity_advection_stencil_02_z_kin_hor_e_numpy( - vn: np.array, vt: np.array - ) -> np.array: + def mo_velocity_advection_stencil_02_z_kin_hor_e_numpy(vn: np.array, vt: np.array) -> np.array: z_kin_hor_e = 0.5 * (vn * vn + vt * vt) return z_kin_hor_e @classmethod - def reference( - cls, mesh, wgtfac_e: np.array, vn: np.array, vt: np.array, **kwargs - ) -> dict: + def reference(cls, mesh, wgtfac_e: np.array, vn: np.array, vt: np.array, **kwargs) -> dict: vn_ie = cls.mo_velocity_advection_stencil_02_vn_ie_numpy(wgtfac_e, vn) z_kin_hor_e = cls.mo_velocity_advection_stencil_02_z_kin_hor_e_numpy(vn, vt) return dict(vn_ie=vn_ie, z_kin_hor_e=z_kin_hor_e) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py index 0a8bf8d0f..562129855 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_08.py @@ -26,9 +26,7 @@ class TestMoVelocityAdvectionStencil08(StencilTest): OUTPUTS = ("z_ekinh",) @staticmethod - def reference( - mesh, z_kin_hor_e: np.array, e_bln_c_s: np.array, **kwargs - ) -> np.array: + def reference(mesh, z_kin_hor_e: np.array, e_bln_c_s: np.array, **kwargs) -> np.array: e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) z_ekinh = np.sum(z_kin_hor_e[mesh.c2e] * e_bln_c_s, axis=1) return dict(z_ekinh=z_ekinh) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py index 1b0d8d4f1..e12005423 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_09.py @@ -26,9 +26,7 @@ class TestMoVelocityAdvectionStencil09(StencilTest): OUTPUTS = ("z_w_concorr_mc",) @staticmethod - def reference( - mesh, z_w_concorr_me: np.array, e_bln_c_s: np.array, **kwargs - ) -> np.array: + def reference(mesh, z_w_concorr_me: np.array, e_bln_c_s: np.array, **kwargs) -> np.array: e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) z_w_concorr_mc = np.sum(z_w_concorr_me[mesh.c2e] * e_bln_c_s, axis=1) return dict(z_w_concorr_mc=z_w_concorr_mc) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py index c561da5a8..d7d7e3d2c 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_10.py @@ -26,13 +26,9 @@ class TestMoVelocityAdvectionStencil10(StencilTest): OUTPUTS = ("w_concorr_c",) @staticmethod - def reference( - mesh, wgtfac_c: np.array, z_w_concorr_mc: np.array, **kwargs - ) -> np.array: + def reference(mesh, wgtfac_c: np.array, z_w_concorr_mc: np.array, **kwargs) -> np.array: z_w_concorr_mc_k_minus_1 = np.roll(z_w_concorr_mc, shift=1, axis=1) - w_concorr_c = ( - wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc_k_minus_1 - ) + w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc_k_minus_1 return dict(w_concorr_c=w_concorr_c) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py index 336ce24d3..eb7fb3fc4 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_13.py @@ -26,9 +26,7 @@ class TestMoVelocityAdvectionStencil13(StencilTest): OUTPUTS = ("z_w_con_c",) @staticmethod - def reference( - mesh, w_concorr_c: np.array, z_w_con_c: np.array, **kwargs - ) -> np.array: + def reference(mesh, w_concorr_c: np.array, z_w_con_c: np.array, **kwargs) -> np.array: z_w_con_c = z_w_con_c - w_concorr_c return dict(z_w_con_c=z_w_con_c) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py index 330b4c9b0..88063c317 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py +++ b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_19.py @@ -55,10 +55,7 @@ def mo_velocity_advection_stencil_19_numpy( ddt_vn_adv = -( (coeff_gradekin[:, 0] - coeff_gradekin[:, 1]) * z_kin_hor_e - + ( - -coeff_gradekin[:, 0] * z_ekinh_e2c[:, 0] - + coeff_gradekin[:, 1] * z_ekinh_e2c[:, 1] - ) + + (-coeff_gradekin[:, 0] * z_ekinh_e2c[:, 0] + coeff_gradekin[:, 1] * z_ekinh_e2c[:, 1]) + vt * (f_e + 0.5 * np.sum(zeta[e2v], axis=1)) + np.sum(z_w_con_c_full[e2c] * c_lin_e, axis=1) * (vn_ie[:, :-1] - vn_ie[:, 1:]) From 64a9154c927a1e111e422a06c2f9b3b90451b2cb Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 5 Sep 2023 09:51:20 +0200 Subject: [PATCH 253/263] pre-commit on tools --- tools/pyproject.toml | 33 +++++++++++++++------------------ 1 file changed, 15 insertions(+), 18 deletions(-) diff --git a/tools/pyproject.toml b/tools/pyproject.toml index 8a0a8c100..a06bef8db 100644 --- a/tools/pyproject.toml +++ b/tools/pyproject.toml @@ -19,13 +19,13 @@ classifiers = [ 'Topic :: Scientific/Engineering :: Physics' ] dependencies = [ - "icon4py-atmosphere-diffusion", - "icon4py-atmosphere-dycore", - "gt4py>=1.0.1", - "icon4py-common", - "tabulate>=0.8.9", - "fprettify>=0.3.7", - "cffi>=1.5", + "icon4py-atmosphere-diffusion", + "icon4py-atmosphere-dycore", + "gt4py>=1.0.1", + "icon4py-common", + "tabulate>=0.8.9", + "fprettify>=0.3.7", + "cffi>=1.5" ] description = 'Tools and utilities for integrating icon4py code into the ICON model.' dynamic = ['version'] @@ -35,9 +35,9 @@ readme = 'README.md' requires-python = '>=3.10' [project.scripts] -icon_liskov = "icon4pytools.liskov.cli:main" -icon4pygen = "icon4pytools.icon4pygen.cli:main" f2ser = "icon4pytools.f2ser.cli:main" +icon4pygen = "icon4pytools.icon4pygen.cli:main" +icon_liskov = "icon4pytools.liskov.cli:main" py2fgen = "icon4pytools.py2f.py2fgen:main" [project.urls] @@ -109,17 +109,14 @@ skip_glob = ['*.venv/**', '_local/**'] use_parentheses = true [tool.mypy] -install_types = true -non_interactive = true -exclude = [ - '^tests/f2ser/*.py', - '^tests/icon4pygen/*.py', - '^tests/liskov/*.py', - '^tests/py2f/*.py', -] disallow_incomplete_defs = true disallow_untyped_defs = true -exclude = ['^tests/f2ser/*.py', '^tests/icon4pygen/*.py', '^tests/liskov/*.py'] +exclude = [ + '^tests/f2ser/*.py', + '^tests/icon4pygen/*.py', + '^tests/liskov/*.py', + '^tests/py2f/*.py' +] ignore_missing_imports = false implicit_reexport = true install_types = true From 6951c73f9b6342d191b845c3dcd42b945c936a35 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 6 Sep 2023 09:12:41 +0200 Subject: [PATCH 254/263] fix order of arguments in apply_diffusion_to_vn.py since it creates problems with icon4pygen. --- .../atmosphere/diffusion/stencils/apply_diffusion_to_vn.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py index 88b74a92c..36fded768 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - +from gt4py.next import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32, where @@ -91,7 +91,7 @@ def _apply_diffusion_to_vn( return vn -@program +@program(grid_type=GridType.UNSTRUCTURED) def apply_diffusion_to_vn( u_vert: Field[[VertexDim, KDim], float], v_vert: Field[[VertexDim, KDim], float], @@ -109,11 +109,11 @@ def apply_diffusion_to_vn( nudgezone_diff: float, fac_bdydiff_v: float, start_2nd_nudge_line_idx_e: int32, + limited_area: bool, horizontal_start: int32, horizontal_end: int32, vertical_start: int32, vertical_end: int32, - limited_area: bool, ): _apply_diffusion_to_vn( u_vert, From d39748a673323034617e360e3820412bb936f6c5 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 6 Sep 2023 10:08:02 +0200 Subject: [PATCH 255/263] review fixes --- .../diffusion/diffusion_tests/conftest.py | 18 ++++- .../mpi_tests/test_parallel_diffusion.py | 4 +- .../diffusion_tests/test_diffusion.py | 50 ++---------- .../diffusion_tests/test_diffusion_utils.py | 25 +----- .../diffusion/diffusion_tests/utils.py | 77 +++++++++++++++++++ .../src/icon4py/model/common/dimension.py | 2 +- .../icon4py/model/common/grid/horizontal.py | 26 +++---- .../icon4py/model/common/grid/icon_grid.py | 4 +- .../model/common/test_utils/data_handling.py | 2 +- tools/pyproject.toml | 22 +++--- 10 files changed, 126 insertions(+), 104 deletions(-) create mode 100644 model/atmosphere/diffusion/diffusion_tests/utils.py diff --git a/model/atmosphere/diffusion/diffusion_tests/conftest.py b/model/atmosphere/diffusion/diffusion_tests/conftest.py index 448cec45d..d42920444 100644 --- a/model/atmosphere/diffusion/diffusion_tests/conftest.py +++ b/model/atmosphere/diffusion/diffusion_tests/conftest.py @@ -14,7 +14,7 @@ import pytest from icon4py.model.atmosphere.diffusion.diffusion import DiffusionConfig, DiffusionType -from icon4py.model.common.test_utils.fixtures import ( # noqa F401 +from icon4py.model.common.test_utils.fixtures import ( # noqa: F401 # import fixtures from test_utils package backend, damping_height, data_provider, @@ -36,7 +36,9 @@ @pytest.fixture -def r04b09_diffusion_config(ndyn_substeps) -> DiffusionConfig: # noqa F811 +def r04b09_diffusion_config( + ndyn_substeps, # noqa: F811 # imported `ndyn_substeps` fxiture +) -> DiffusionConfig: """ Create DiffusionConfig matching MCH_CH_r04b09_dsl. @@ -60,7 +62,11 @@ def r04b09_diffusion_config(ndyn_substeps) -> DiffusionConfig: # noqa F811 @pytest.fixture -def diffusion_savepoint_init(data_provider, linit, step_date_init): # noqa F811 +def diffusion_savepoint_init( + data_provider, # noqa: F811 # imported fixtures data_provider + linit, # noqa: F811 # imported fixtures linit + step_date_init, # noqa: F811 # imported fixtures data_provider +): """ Load data from ICON savepoint at start of diffusion module. @@ -73,7 +79,11 @@ def diffusion_savepoint_init(data_provider, linit, step_date_init): # noqa F811 @pytest.fixture -def diffusion_savepoint_exit(data_provider, linit, step_date_exit): # noqa F811 +def diffusion_savepoint_exit( + data_provider, # noqa: F811 # imported fixtures data_provider` + linit, # noqa: F811 # imported fixtures linit` + step_date_exit, # noqa: F811 # imported fixtures step_date_exit` +): """ Load data from ICON savepoint at exist of diffusion module. diff --git a/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py index 3b68ec7a0..8a6602799 100644 --- a/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py +++ b/model/atmosphere/diffusion/diffusion_tests/mpi_tests/test_parallel_diffusion.py @@ -20,7 +20,7 @@ from icon4py.model.common.grid.vertical import VerticalModelParams from icon4py.model.common.test_utils.parallel_helpers import check_comm_size -from ..test_diffusion import _verify_diffusion_fields +from ..utils import verify_diffusion_fields @pytest.mark.mpi @@ -100,7 +100,7 @@ def test_parallel_diffusion( ) print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ") - _verify_diffusion_fields( + verify_diffusion_fields( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, diffusion_savepoint=diffusion_savepoint_exit, diff --git a/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py b/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py index d1df49c73..383d10f34 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py @@ -15,22 +15,16 @@ import pytest from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams -from icon4py.model.atmosphere.diffusion.diffusion_states import ( - DiffusionDiagnosticState, - PrognosticState, -) from icon4py.model.atmosphere.diffusion.diffusion_utils import scale_k from icon4py.model.common.grid.horizontal import CellParams, EdgeParams from icon4py.model.common.grid.vertical import VerticalModelParams -from icon4py.model.common.test_utils.serialbox_utils import ( - IconDiffusionExitSavepoint, - IconDiffusionInitSavepoint, -) +from icon4py.model.common.test_utils.serialbox_utils import IconDiffusionInitSavepoint -from .test_diffusion_utils import ( +from .utils import ( diff_multfac_vn_numpy, enhanced_smagorinski_factor_numpy, smag_limit_numpy, + verify_diffusion_fields, ) @@ -294,46 +288,14 @@ def test_run_diffusion_single_step( edge_params=edge_geometry, cell_params=cell_geometry, ) - _verify_diffusion_fields(diagnostic_state, prognostic_state, diffusion_savepoint_init) + verify_diffusion_fields(diagnostic_state, prognostic_state, diffusion_savepoint_init) assert diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v diffusion.run( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, dtime=dtime, ) - _verify_diffusion_fields(diagnostic_state, prognostic_state, diffusion_savepoint_exit) - - -def _verify_diffusion_fields( - diagnostic_state: DiffusionDiagnosticState, - prognostic_state: PrognosticState, - diffusion_savepoint: IconDiffusionExitSavepoint, -): - ref_div_ic = np.asarray(diffusion_savepoint.div_ic()) - val_div_ic = np.asarray(diagnostic_state.div_ic) - ref_hdef_ic = np.asarray(diffusion_savepoint.hdef_ic()) - val_hdef_ic = np.asarray(diagnostic_state.hdef_ic) - assert np.allclose(ref_div_ic, val_div_ic) - assert np.allclose(ref_hdef_ic, val_hdef_ic) - ref_w = np.asarray(diffusion_savepoint.w()) - val_w = np.asarray(prognostic_state.w) - ref_dwdx = np.asarray(diffusion_savepoint.dwdx()) - val_dwdx = np.asarray(diagnostic_state.dwdx) - ref_dwdy = np.asarray(diffusion_savepoint.dwdy()) - val_dwdy = np.asarray(diagnostic_state.dwdy) - assert np.allclose(ref_dwdx, val_dwdx) - assert np.allclose(ref_dwdy, val_dwdy) - - ref_vn = np.asarray(diffusion_savepoint.vn()) - val_vn = np.asarray(prognostic_state.vn) - assert np.allclose(ref_vn, val_vn) - assert np.allclose(ref_w, val_w) - ref_exner = np.asarray(diffusion_savepoint.exner()) - ref_theta_v = np.asarray(diffusion_savepoint.theta_v()) - val_theta_v = np.asarray(prognostic_state.theta_v) - val_exner = np.asarray(prognostic_state.exner_pressure) - assert np.allclose(ref_theta_v, val_theta_v) - assert np.allclose(ref_exner, val_exner) + verify_diffusion_fields(diagnostic_state, prognostic_state, diffusion_savepoint_exit) @pytest.mark.datatest @@ -379,7 +341,7 @@ def test_run_diffusion_initial_step( dtime=dtime, ) - _verify_diffusion_fields( + verify_diffusion_fields( diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, diffusion_savepoint=diffusion_savepoint_exit, diff --git a/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py b/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py index 98d0516c5..0ba7ce353 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_diffusion_utils.py @@ -27,20 +27,13 @@ from icon4py.model.common.test_utils.helpers import random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh - -def diff_multfac_vn_numpy(shape, k4, substeps): - factor = min(1.0 / 128.0, k4 * substeps / 3.0) - return factor * np.ones(shape) +from .utils import diff_multfac_vn_numpy, enhanced_smagorinski_factor_numpy, smag_limit_numpy def initial_diff_multfac_vn_numpy(shape, k4, hdiff_efdt_ratio): return k4 * hdiff_efdt_ratio / 3.0 * np.ones(shape) -def smag_limit_numpy(func, *args): - return 0.125 - 4.0 * func(*args) - - def test_scale_k(): mesh = SimpleMesh() field = random_field(mesh, KDim) @@ -107,22 +100,6 @@ def test_diff_multfac_vn_smag_limit_for_loop_run_with_k4_substeps(): assert np.allclose(expected_smag_limit, smag_limit) -def enhanced_smagorinski_factor_numpy(factor_in, heigths_in, a_vec): - alin = (factor_in[1] - factor_in[0]) / (heigths_in[1] - heigths_in[0]) - df32 = factor_in[2] - factor_in[1] - df42 = factor_in[3] - factor_in[1] - dz32 = heigths_in[2] - heigths_in[1] - dz42 = heigths_in[3] - heigths_in[1] - bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32)) - aqdr = df32 / dz32 - bqdr * dz32 - zf = 0.5 * (a_vec[:-1] + a_vec[1:]) - max0 = np.maximum(0.0, zf - heigths_in[0]) - dzlin = np.minimum(heigths_in[1] - heigths_in[0], max0) - max1 = np.maximum(0.0, zf - heigths_in[1]) - dzqdr = np.minimum(heigths_in[3] - heigths_in[1], max1) - return factor_in[0] + dzlin * alin + dzqdr * (aqdr + dzqdr * bqdr) - - def test_init_enh_smag_fac(): mesh = SimpleMesh() enh_smag_fac = zero_field(mesh, KDim) diff --git a/model/atmosphere/diffusion/diffusion_tests/utils.py b/model/atmosphere/diffusion/diffusion_tests/utils.py new file mode 100644 index 000000000..b2fbf987d --- /dev/null +++ b/model/atmosphere/diffusion/diffusion_tests/utils.py @@ -0,0 +1,77 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + PrognosticState, +) +from icon4py.model.common.test_utils.serialbox_utils import IconDiffusionExitSavepoint + + +def verify_diffusion_fields( + diagnostic_state: DiffusionDiagnosticState, + prognostic_state: PrognosticState, + diffusion_savepoint: IconDiffusionExitSavepoint, +): + ref_div_ic = np.asarray(diffusion_savepoint.div_ic()) + val_div_ic = np.asarray(diagnostic_state.div_ic) + ref_hdef_ic = np.asarray(diffusion_savepoint.hdef_ic()) + val_hdef_ic = np.asarray(diagnostic_state.hdef_ic) + assert np.allclose(ref_div_ic, val_div_ic) + assert np.allclose(ref_hdef_ic, val_hdef_ic) + ref_w = np.asarray(diffusion_savepoint.w()) + val_w = np.asarray(prognostic_state.w) + ref_dwdx = np.asarray(diffusion_savepoint.dwdx()) + val_dwdx = np.asarray(diagnostic_state.dwdx) + ref_dwdy = np.asarray(diffusion_savepoint.dwdy()) + val_dwdy = np.asarray(diagnostic_state.dwdy) + assert np.allclose(ref_dwdx, val_dwdx) + assert np.allclose(ref_dwdy, val_dwdy) + + ref_vn = np.asarray(diffusion_savepoint.vn()) + val_vn = np.asarray(prognostic_state.vn) + assert np.allclose(ref_vn, val_vn) + assert np.allclose(ref_w, val_w) + ref_exner = np.asarray(diffusion_savepoint.exner()) + ref_theta_v = np.asarray(diffusion_savepoint.theta_v()) + val_theta_v = np.asarray(prognostic_state.theta_v) + val_exner = np.asarray(prognostic_state.exner_pressure) + assert np.allclose(ref_theta_v, val_theta_v) + assert np.allclose(ref_exner, val_exner) + + +def smag_limit_numpy(func, *args): + return 0.125 - 4.0 * func(*args) + + +def diff_multfac_vn_numpy(shape, k4, substeps): + factor = min(1.0 / 128.0, k4 * substeps / 3.0) + return factor * np.ones(shape) + + +def enhanced_smagorinski_factor_numpy(factor_in, heigths_in, a_vec): + alin = (factor_in[1] - factor_in[0]) / (heigths_in[1] - heigths_in[0]) + df32 = factor_in[2] - factor_in[1] + df42 = factor_in[3] - factor_in[1] + dz32 = heigths_in[2] - heigths_in[1] + dz42 = heigths_in[3] - heigths_in[1] + bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32)) + aqdr = df32 / dz32 - bqdr * dz32 + zf = 0.5 * (a_vec[:-1] + a_vec[1:]) + max0 = np.maximum(0.0, zf - heigths_in[0]) + dzlin = np.minimum(heigths_in[1] - heigths_in[0], max0) + max1 = np.maximum(0.0, zf - heigths_in[1]) + dzqdr = np.minimum(heigths_in[3] - heigths_in[1], max1) + return factor_in[0] + dzlin * alin + dzqdr * (aqdr + dzqdr * bqdr) diff --git a/model/common/src/icon4py/model/common/dimension.py b/model/common/src/icon4py/model/common/dimension.py index ca4ec442b..fd6bb6111 100644 --- a/model/common/src/icon4py/model/common/dimension.py +++ b/model/common/src/icon4py/model/common/dimension.py @@ -18,7 +18,7 @@ KDim = Dimension("K", kind=DimensionKind.VERTICAL) KHalfDim = Dimension("KHalf", kind=DimensionKind.VERTICAL) EdgeDim = Dimension("Edge") -CellDim: Dimension = Dimension("Cell") +CellDim = Dimension("Cell") VertexDim = Dimension("Vertex") CEDim = Dimension("CE") CECDim = Dimension("CEC") diff --git a/model/common/src/icon4py/model/common/grid/horizontal.py b/model/common/src/icon4py/model/common/grid/horizontal.py index fc05b90bd..a68c95d1b 100644 --- a/model/common/src/icon4py/model/common/grid/horizontal.py +++ b/model/common/src/icon4py/model/common/grid/horizontal.py @@ -22,20 +22,18 @@ class HorizontalMarkerIndex: """ Handles constants indexing into the start_index and end_index fields. - ICON uses a double indexing scheme for field indices marking the start and end of special - grid zone: The constants defined here (from mo_impl_constants.f90 and mo_impl_constants_grf.f90) - are the indices that are used to index into the start_idx and end_idx arrays - provided by the grid file where for each dimension the start index of the horizontal - "zones" are defined: - f.ex. an inlined access of the field F: Field[[CellDim], double] at the starting point of the lateral boundary zone would be - - F[start_idx_c[_LATERAL_BOUNDARY_CELLS] - - - ICON uses a custom index range from [ICON_INDEX_OFFSET... ] such that the index 0 marks the - internal entities for _all_ dimensions (Cell, Edge, Vertex) that is why we define these - additional INDEX_OFFSETs here in order to swap back to a 0 base python array. - + ICON uses a double indexing scheme for field indices marking the start and end of special + grid zone: The constants defined here (from mo_impl_constants.f90 and mo_impl_constants_grf.f90) + are the indices that are used to index into the start_idx and end_idx arrays + provided by the grid file where for each dimension the start index of the horizontal + "zones" are defined: + f.ex. an inlined access of the field F: Field[[CellDim], double] at the starting point of the lateral boundary zone would be + + F[start_idx_c[_LATERAL_BOUNDARY_CELLS] + + ICON uses a custom index range from [ICON_INDEX_OFFSET... ] such that the index 0 marks the + internal entities for _all_ dimensions (Cell, Edge, Vertex) that is why we define these + additional INDEX_OFFSETs here in order to swap back to a 0 base python array. """ NUM_GHOST_ROWS: Final[int] = 2 diff --git a/model/common/src/icon4py/model/common/grid/icon_grid.py b/model/common/src/icon4py/model/common/grid/icon_grid.py index d7e8585d2..f514ca128 100644 --- a/model/common/src/icon4py/model/common/grid/icon_grid.py +++ b/model/common/src/icon4py/model/common/grid/icon_grid.py @@ -25,9 +25,7 @@ from icon4py.model.common.utils import builder -@dataclass( - frozen=True, -) +@dataclass(frozen=True) class GridConfig: horizontal_config: HorizontalGridSize vertical_config: VerticalGridSize diff --git a/model/common/src/icon4py/model/common/test_utils/data_handling.py b/model/common/src/icon4py/model/common/test_utils/data_handling.py index b00bb2d32..7a8a49c34 100644 --- a/model/common/src/icon4py/model/common/test_utils/data_handling.py +++ b/model/common/src/icon4py/model/common/test_utils/data_handling.py @@ -17,7 +17,7 @@ import wget -def download_and_extract(uri: str, local_path: Path, data_file: str): +def download_and_extract(uri: str, local_path: Path, data_file: str) -> None: local_path.mkdir(parents=True, exist_ok=True) if not any(local_path.iterdir()): print(f"directory {local_path} is empty: downloading data from {uri} and extracting") diff --git a/tools/pyproject.toml b/tools/pyproject.toml index a06bef8db..9df4270d6 100644 --- a/tools/pyproject.toml +++ b/tools/pyproject.toml @@ -19,13 +19,13 @@ classifiers = [ 'Topic :: Scientific/Engineering :: Physics' ] dependencies = [ - "icon4py-atmosphere-diffusion", - "icon4py-atmosphere-dycore", - "gt4py>=1.0.1", - "icon4py-common", - "tabulate>=0.8.9", - "fprettify>=0.3.7", - "cffi>=1.5" + 'icon4py-atmosphere-diffusion', + 'icon4py-atmosphere-dycore', + 'gt4py>=1.0.1', + 'icon4py-common', + 'tabulate>=0.8.9', + 'fprettify>=0.3.7', + 'cffi>=1.5' ] description = 'Tools and utilities for integrating icon4py code into the ICON model.' dynamic = ['version'] @@ -35,10 +35,10 @@ readme = 'README.md' requires-python = '>=3.10' [project.scripts] -f2ser = "icon4pytools.f2ser.cli:main" -icon4pygen = "icon4pytools.icon4pygen.cli:main" -icon_liskov = "icon4pytools.liskov.cli:main" -py2fgen = "icon4pytools.py2f.py2fgen:main" +f2ser = 'icon4pytools.f2ser.cli:main' +icon4pygen = 'icon4pytools.icon4pygen.cli:main' +icon_liskov = 'icon4pytools.liskov.cli:main' +py2fgen = 'icon4pytools.py2f.py2fgen:main' [project.urls] repository = 'https://github.com/C2SM/icon4py' From 23eb9ff28c2a158c4bb58cd3f1208579644ac14d Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 6 Sep 2023 10:32:20 +0200 Subject: [PATCH 256/263] fix yesqa configuration to work with flake8 for A003 --- model/common/.pre-commit-config.yaml | 11 ++++++++++- .../model/common/decomposition/parallel_setup.py | 4 +++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/model/common/.pre-commit-config.yaml b/model/common/.pre-commit-config.yaml index 13dc05548..0a905ba43 100644 --- a/model/common/.pre-commit-config.yaml +++ b/model/common/.pre-commit-config.yaml @@ -63,7 +63,16 @@ repos: rev: v1.3.0 hooks: - id: yesqa - exclude: 'decomposition/parallel_setup.py' + additional_dependencies: + - flake8==4.0.1 + - darglint + - flake8-bugbear + - flake8-builtins + - flake8-debugger + - flake8-docstrings + - flake8-eradicate + - flake8-mutable + - pygments - repo: https://github.com/psf/black rev: '22.3.0' diff --git a/model/common/src/icon4py/model/common/decomposition/parallel_setup.py b/model/common/src/icon4py/model/common/decomposition/parallel_setup.py index 877885002..11d5eaeb6 100644 --- a/model/common/src/icon4py/model/common/decomposition/parallel_setup.py +++ b/model/common/src/icon4py/model/common/decomposition/parallel_setup.py @@ -91,6 +91,8 @@ def __init__(self, process_properties: ProcessProperties = None): if process_properties and process_properties.comm_size > 1: self._rank_info = f"rank={process_properties.rank}/{process_properties.comm_size} [{process_properties.comm_name}] " - def filter(self, record: logging.LogRecord) -> bool: # noqa: A003 + def filter( # noqa: A003 # overwriting logging.Filter.filter() + self, record: logging.LogRecord + ) -> bool: record.rank = self._rank_info return True From 9f8d5a733378664ea84a46fcfc69c89ad69f02eb Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 6 Sep 2023 11:03:16 +0200 Subject: [PATCH 257/263] add icon4py.atmosphere.diffusion and icon4py.driver to the icon4py-qa.yml --- .github/workflows/icon4py-qa.yml | 48 ++++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/.github/workflows/icon4py-qa.yml b/.github/workflows/icon4py-qa.yml index 18be464b9..fd4c6178e 100644 --- a/.github/workflows/icon4py-qa.yml +++ b/.github/workflows/icon4py-qa.yml @@ -46,6 +46,30 @@ jobs: run: | pre-commit run --config model/common/.pre-commit-config.yaml --all-files + pre-commit-icon4py-model-driver: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + cache: 'pip' + cache-dependency-path: | + **/pyproject.toml + **/base-requirements.txt + **/base-requirements-dev.txt + **/requirements.txt + **/requirements-dev.txt + - name: Install icon4py-model-driver + working-directory: model/driver + run: | + python -m pip install --upgrade pip setuptools wheel + python -m pip install -r ./requirements-dev.txt + - name: Run checks icon4py-model-driver + run: | + pre-commit run --config model/driver/.pre-commit-config.yaml --all-files + pre-commit-icon4py-model-atmosphere-dycore: runs-on: ubuntu-latest steps: @@ -69,3 +93,27 @@ jobs: - name: Run checks icon4py-model-atmosphere-dycore run: | pre-commit run --config model/atmosphere/dycore/.pre-commit-config.yaml --all-files + + pre-commit-icon4py-model-atmosphere-diffusion: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + cache: 'pip' + cache-dependency-path: | + **/pyproject.toml + **/base-requirements.txt + **/base-requirements-dev.txt + **/requirements.txt + **/requirements-dev.txt + - name: Install icon4py-model-atmosphere-diffusion + working-directory: model/atmosphere/diffusion + run: | + python -m pip install --upgrade pip setuptools wheel + python -m pip install -r ./requirements-dev.txt + - name: Run checks icon4py-model-atmosphere-diffusion + run: | + pre-commit run --config model/atmosphere/diffusion/.pre-commit-config.yaml --all-files From d60e4f5c4a795c1053293cd64e129386b8890790 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 6 Sep 2023 12:44:22 +0200 Subject: [PATCH 258/263] icon4py-qa.yml install boost --- .github/workflows/icon4py-qa.yml | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/.github/workflows/icon4py-qa.yml b/.github/workflows/icon4py-qa.yml index fd4c6178e..76f790b59 100644 --- a/.github/workflows/icon4py-qa.yml +++ b/.github/workflows/icon4py-qa.yml @@ -25,6 +25,10 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 + - name: Install system dependencies + run: | + sudo apt-get update + sudo apt-get install libboost-all-dev - name: Set up Python uses: actions/setup-python@v4 with: @@ -50,6 +54,10 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 + - name: Install system dependencies + run: | + sudo apt-get update + sudo apt-get install libboost-all-dev - name: Set up Python uses: actions/setup-python@v4 with: @@ -74,6 +82,10 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 + - name: Install system dependencies + run: | + sudo apt-get update + sudo apt-get install libboost-all-dev - name: Set up Python uses: actions/setup-python@v4 with: @@ -98,6 +110,10 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 + - name: Install system dependencies + run: | + sudo apt-get update + sudo apt-get install libboost-all-dev - name: Set up Python uses: actions/setup-python@v4 with: From acae8aa8d6d2ad04843761f8e81c2808ae8f4bf6 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 6 Sep 2023 12:58:46 +0200 Subject: [PATCH 259/263] fix wrong path in requirements-dev.txt of model/driver install boost in icon4pytools-qa.yml --- .github/workflows/icon4pytools-qa.yml | 4 ++++ model/driver/requirements-dev.txt | 3 ++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/.github/workflows/icon4pytools-qa.yml b/.github/workflows/icon4pytools-qa.yml index 0f550ce84..af1138fed 100644 --- a/.github/workflows/icon4pytools-qa.yml +++ b/.github/workflows/icon4pytools-qa.yml @@ -31,6 +31,10 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 + - name: Install boost + run: | + sudo apt-get update + sudo apt-get install libboost-all-dev - name: Set up Python uses: actions/setup-python@v4 with: diff --git a/model/driver/requirements-dev.txt b/model/driver/requirements-dev.txt index 3f54c69af..68e2b1610 100644 --- a/model/driver/requirements-dev.txt +++ b/model/driver/requirements-dev.txt @@ -1,4 +1,5 @@ -r ../../base-requirements-dev.txt -e ../common --e ../atmosphere +-e ../atmosphere/diffusion +-e ../atmosphere/dycore -e . From 9ac623fbc0385b56635ad5e78ba5687fb4fc853d Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 6 Sep 2023 17:14:36 +0200 Subject: [PATCH 260/263] combine pre-commit runs of all model packages in one job --- .github/workflows/icon4py-qa.yml | 82 +------------------ .../model/driver/icon_configuration.py | 4 +- 2 files changed, 6 insertions(+), 80 deletions(-) diff --git a/.github/workflows/icon4py-qa.yml b/.github/workflows/icon4py-qa.yml index 76f790b59..55b260cad 100644 --- a/.github/workflows/icon4py-qa.yml +++ b/.github/workflows/icon4py-qa.yml @@ -21,7 +21,7 @@ on: types: [submitted] jobs: - pre-commit-icon4py-model-common: + pre-commit-icon4py-model: runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 @@ -40,8 +40,8 @@ jobs: **/base-requirements-dev.txt **/requirements.txt **/requirements-dev.txt - - name: Install icon4py-common - working-directory: model/common + - name: Install icon4py-model packages + working-directory: model run: | python -m pip install --upgrade pip setuptools wheel python -m pip install -r ./requirements-dev.txt @@ -49,87 +49,13 @@ jobs: - name: Run checks icon4py-model-common run: | pre-commit run --config model/common/.pre-commit-config.yaml --all-files - - pre-commit-icon4py-model-driver: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v3 - - name: Install system dependencies - run: | - sudo apt-get update - sudo apt-get install libboost-all-dev - - name: Set up Python - uses: actions/setup-python@v4 - with: - python-version: "3.10" - cache: 'pip' - cache-dependency-path: | - **/pyproject.toml - **/base-requirements.txt - **/base-requirements-dev.txt - **/requirements.txt - **/requirements-dev.txt - - name: Install icon4py-model-driver - working-directory: model/driver - run: | - python -m pip install --upgrade pip setuptools wheel - python -m pip install -r ./requirements-dev.txt - name: Run checks icon4py-model-driver run: | pre-commit run --config model/driver/.pre-commit-config.yaml --all-files - - pre-commit-icon4py-model-atmosphere-dycore: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v3 - - name: Install system dependencies - run: | - sudo apt-get update - sudo apt-get install libboost-all-dev - - name: Set up Python - uses: actions/setup-python@v4 - with: - python-version: "3.10" - cache: 'pip' - cache-dependency-path: | - **/pyproject.toml - **/base-requirements.txt - **/base-requirements-dev.txt - **/requirements.txt - **/requirements-dev.txt - - name: Install icon4py-model-atmosphere-dycore - working-directory: model/atmosphere/dycore - run: | - python -m pip install --upgrade pip setuptools wheel - python -m pip install -r ./requirements-dev.txt - name: Run checks icon4py-model-atmosphere-dycore run: | pre-commit run --config model/atmosphere/dycore/.pre-commit-config.yaml --all-files - - pre-commit-icon4py-model-atmosphere-diffusion: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v3 - - name: Install system dependencies - run: | - sudo apt-get update - sudo apt-get install libboost-all-dev - - name: Set up Python - uses: actions/setup-python@v4 - with: - python-version: "3.10" - cache: 'pip' - cache-dependency-path: | - **/pyproject.toml - **/base-requirements.txt - **/base-requirements-dev.txt - **/requirements.txt - **/requirements-dev.txt - - name: Install icon4py-model-atmosphere-diffusion - working-directory: model/atmosphere/diffusion - run: | - python -m pip install --upgrade pip setuptools wheel - python -m pip install -r ./requirements-dev.txt - name: Run checks icon4py-model-atmosphere-diffusion run: | pre-commit run --config model/atmosphere/diffusion/.pre-commit-config.yaml --all-files + diff --git a/model/driver/src/icon4py/model/driver/icon_configuration.py b/model/driver/src/icon4py/model/driver/icon_configuration.py index b52f15a15..187c4f4f4 100644 --- a/model/driver/src/icon4py/model/driver/icon_configuration.py +++ b/model/driver/src/icon4py/model/driver/icon_configuration.py @@ -14,7 +14,7 @@ from dataclasses import dataclass from typing import Optional -from icon4py.model.atmosphere.diffusion.diffusion import DiffusionConfig +from icon4py.model.atmosphere.diffusion.diffusion import DiffusionConfig, DiffusionType n_substeps_reduced = 2 @@ -46,7 +46,7 @@ def _default_run_config(n_steps: int): def mch_ch_r04b09_diffusion_config(): return DiffusionConfig( - diffusion_type=5, + diffusion_type=DiffusionType.SMAGORINSKY_4TH_ORDER, hdiff_w=True, n_substeps=n_substeps_reduced, hdiff_vn=True, From 393b7c2ddc6e6d2fd99948b4a4875d9e70845e62 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 8 Sep 2023 15:52:15 +0200 Subject: [PATCH 261/263] remove unnecessary test file --- .../model/atmosphere/diffusion/diffusion.py | 1 + model/driver/tests/test_dycore_driver.py | 68 ------------------- 2 files changed, 1 insertion(+), 68 deletions(-) delete mode 100644 model/driver/tests/test_dycore_driver.py diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py index b46ccfde3..d857a4f27 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -736,6 +736,7 @@ def _do_diffusion_step( log.debug( "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" ) + # TODO (magdalena) get rid of this copying. So far passing an empty buffer instead did not verify? copy_field.with_backend(backend)(prognostic_state.w, self.w_tmp, offset_provider={}) apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.with_backend(backend)( area=self.cell_params.area, diff --git a/model/driver/tests/test_dycore_driver.py b/model/driver/tests/test_dycore_driver.py deleted file mode 100644 index 88c4fb95d..000000000 --- a/model/driver/tests/test_dycore_driver.py +++ /dev/null @@ -1,68 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np - -from icon4py.model.common.dimension import CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import random_field, zero_field -from icon4py.model.common.test_utils.simple_mesh import SimpleMesh -from icon4py.model.driver.dycore_driver import _copy_diagnostic_and_prognostics - - -def test_copy_diagnostic_and_prognostics(): - mesh = SimpleMesh() - f1_in = random_field(mesh, CellDim, KDim) - f1_out = zero_field(mesh, CellDim, KDim) - f2_in = random_field(mesh, CellDim, KDim) - f2_out = zero_field(mesh, CellDim, KDim) - f3_in = random_field(mesh, CellDim, KDim) - f3_out = zero_field(mesh, CellDim, KDim) - f4_in = random_field(mesh, CellDim, KDim) - f4_out = zero_field(mesh, CellDim, KDim) - f5_in = random_field(mesh, EdgeDim, KDim) - f5_out = zero_field(mesh, EdgeDim, KDim) - f6_in = random_field(mesh, CellDim, KDim) - f6_out = zero_field(mesh, CellDim, KDim) - f7_in = random_field(mesh, CellDim, KDim) - f7_out = zero_field(mesh, CellDim, KDim) - f8_in = random_field(mesh, CellDim, KDim) - f8_out = zero_field(mesh, CellDim, KDim) - - _copy_diagnostic_and_prognostics( - f1_in, - f1_out, - f2_in, - f2_out, - f3_in, - f3_out, - f4_in, - f4_out, - f5_in, - f5_out, - f6_in, - f6_out, - f7_in, - f7_out, - f8_in, - f8_out, - offset_provider={}, - ) - - assert np.allclose(f1_in, f1_out) - assert np.allclose(f2_in, f2_out) - assert np.allclose(f3_in, f3_out) - assert np.allclose(f4_in, f4_out) - assert np.allclose(f5_in, f5_out) - assert np.allclose(f6_in, f6_out) - assert np.allclose(f7_in, f7_out) - assert np.allclose(f8_in, f8_out) From c6d8bfbf501716a162e66a2d79a6d0edbe9cd042 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 8 Sep 2023 17:49:23 +0200 Subject: [PATCH 262/263] fix: register custom marker in tox.ini fix: explicitly skip tests that access data fix: add missing datatest marker --- .../common/src/icon4py/model/common/test_utils/fixtures.py | 6 ++++-- model/common/tests/conftest.py | 4 +++- model/common/tests/test_icon_grid.py | 2 +- model/common/tests/test_vertical.py | 2 +- model/tox.ini | 3 +++ 5 files changed, 12 insertions(+), 5 deletions(-) diff --git a/model/common/src/icon4py/model/common/test_utils/fixtures.py b/model/common/src/icon4py/model/common/test_utils/fixtures.py index bf575b446..d28608d0b 100644 --- a/model/common/src/icon4py/model/common/test_utils/fixtures.py +++ b/model/common/src/icon4py/model/common/test_utils/fixtures.py @@ -51,14 +51,16 @@ def ranked_data_path(processor_props): def datapath(ranked_data_path): return ranked_data_path.joinpath("mch_ch_r04b09_dsl/ser_data") - @pytest.fixture(scope="session") -def download_ser_data(request, processor_props, ranked_data_path): +def download_ser_data(request, processor_props, ranked_data_path, pytestconfig): """ Get the binary ICON data from a remote server. Session scoped fixture which is a prerequisite of all the other fixtures in this file. """ + if 'not datatest' in pytestconfig.option.markexpr: + pytest.skip('not running datatest marked tests') + try: uri = data_uris[processor_props.comm_size] diff --git a/model/common/tests/conftest.py b/model/common/tests/conftest.py index 1824c1b89..0c4d19c49 100644 --- a/model/common/tests/conftest.py +++ b/model/common/tests/conftest.py @@ -47,11 +47,13 @@ def r04b09_dsl_gridfile(get_grid_files): @pytest.fixture(scope="session") -def get_grid_files(): +def get_grid_files(pytestconfig): """ Get the grid files used for testing. Session scoped fixture which is a prerequisite of all the other fixtures in this file. """ + if "not datatest" in pytestconfig.option.markexpr: + pytest.skip("not running datatest marked tests") download_and_extract(mch_ch_r04b09_dsl_grid_uri, r04b09_dsl_grid_path, r04b09_dsl_data_file) download_and_extract(r02b04_global_grid_uri, r02b04_global_grid_path, r02b04_global_data_file) diff --git a/model/common/tests/test_icon_grid.py b/model/common/tests/test_icon_grid.py index e9b5f39ce..bdd399942 100644 --- a/model/common/tests/test_icon_grid.py +++ b/model/common/tests/test_icon_grid.py @@ -392,7 +392,7 @@ def test_cross_check_marker_equivalences(icon_grid): HorizontalMarkerIndex.nudging(EdgeDim), ) - +@pytest.mark.datatest def test_grid_size(grid_savepoint): assert 10663 == grid_savepoint.num(VertexDim) assert 20896 == grid_savepoint.num(CellDim) diff --git a/model/common/tests/test_vertical.py b/model/common/tests/test_vertical.py index f205941b9..f40d32062 100644 --- a/model/common/tests/test_vertical.py +++ b/model/common/tests/test_vertical.py @@ -41,6 +41,6 @@ def test_nrdmax_calculation_from_icon_input(grid_savepoint, damping_height): assert a_array[nrdmax] > damping_height assert a_array[nrdmax + 1] < damping_height - +@pytest.mark.datatest def test_grid_size(grid_savepoint): assert 65 == grid_savepoint.num(KDim) diff --git a/model/tox.ini b/model/tox.ini index 9cb4ccc95..c9fd7a989 100644 --- a/model/tox.ini +++ b/model/tox.ini @@ -31,3 +31,6 @@ setenv = skip_install = true commands = commands_post = + +[pytest] +markers = datatest: this test uses serialized data From 9713364c2a55be42776b9e16e50a1a7cbc7ef695 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Sun, 10 Sep 2023 15:46:11 +0200 Subject: [PATCH 263/263] add command line option for running datatest --- model/README.md | 10 ++++- model/atmosphere/diffusion/pyproject.toml | 3 +- model/atmosphere/diffusion/tests/conftest.py | 24 +++++++++++ model/atmosphere/dycore/pyproject.toml | 1 - model/atmosphere/dycore/tests/conftest.py | 5 +++ model/common/pyproject.toml | 1 - .../model/common/test_utils/fixtures.py | 5 ++- .../model/common/test_utils/pytest_config.py | 40 +++++++++++++++++++ model/common/tests/conftest.py | 9 ++++- model/common/tests/test_icon_grid.py | 1 + model/common/tests/test_vertical.py | 1 + model/driver/pyproject.toml | 1 - model/driver/tests/conftest.py | 7 +++- model/tox.ini | 6 +-- tox.ini | 4 +- 15 files changed, 100 insertions(+), 18 deletions(-) create mode 100644 model/atmosphere/diffusion/tests/conftest.py create mode 100644 model/common/src/icon4py/model/common/test_utils/pytest_config.py diff --git a/model/README.md b/model/README.md index 83bd729fd..0f076a932 100644 --- a/model/README.md +++ b/model/README.md @@ -31,10 +31,16 @@ See the repository [README](../README.md) for general information. #### Data dependent tests -Some test depend on serialized data generated by a full model run. Those test are marked with `pytest.mark.datatest` and can be skipped by specifying +Some test depend on serialized data generated by a full model run. +Those test are marked with `pytest.mark.datatest` and are only run when the `--datatest` +option is specified. Note that due to `pytests` configuration discovery +you need to specify the base a package directory i.e. one that contains a `pyproject.toml` for the +commandline option to work. ```bash -pytest -v -m "not datatest" + +pytest -v --datatest model/common +pytest -v --datatest model/atmosphere/diffusion ``` #### Testing parallel code diff --git a/model/atmosphere/diffusion/pyproject.toml b/model/atmosphere/diffusion/pyproject.toml index 8f0d08a2c..d6ab9aa52 100644 --- a/model/atmosphere/diffusion/pyproject.toml +++ b/model/atmosphere/diffusion/pyproject.toml @@ -111,8 +111,7 @@ warn_unused_ignores = true [tool.pytest] [tool.pytest.ini_options] -markers = 'datatest: this test uses binary data' -testpaths = 'diffusion_tests' +testpaths = ['tests', 'diffusion_tests'] [tool.setuptools.dynamic] version = {attr = 'icon4py.model.atmosphere.diffusion.__init__.__version__'} diff --git a/model/atmosphere/diffusion/tests/conftest.py b/model/atmosphere/diffusion/tests/conftest.py new file mode 100644 index 000000000..7ba065d2c --- /dev/null +++ b/model/atmosphere/diffusion/tests/conftest.py @@ -0,0 +1,24 @@ +# 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 + +""" +Initialize pytest. + +Workaround for pytest not discovering those configuration function, when they are added to the +diffusion_test/conftest.py folder +""" +from icon4py.model.common.test_utils.pytest_config import ( # noqa: F401 + pytest_addoption, + pytest_configure, + pytest_runtest_setup, +) diff --git a/model/atmosphere/dycore/pyproject.toml b/model/atmosphere/dycore/pyproject.toml index 842b53ceb..a29c25ea6 100644 --- a/model/atmosphere/dycore/pyproject.toml +++ b/model/atmosphere/dycore/pyproject.toml @@ -110,7 +110,6 @@ warn_unused_ignores = true [tool.pytest] [tool.pytest.ini_options] -markers = 'datatest: this test uses binary data' testpaths = 'tests' [tool.setuptools.dynamic] diff --git a/model/atmosphere/dycore/tests/conftest.py b/model/atmosphere/dycore/tests/conftest.py index 8e6147e6d..9e9274b36 100644 --- a/model/atmosphere/dycore/tests/conftest.py +++ b/model/atmosphere/dycore/tests/conftest.py @@ -25,3 +25,8 @@ step_date_exit, step_date_init, ) +from icon4py.model.common.test_utils.pytest_config import ( # noqa: F401 + pytest_addoption, + pytest_configure, + pytest_runtest_setup, +) diff --git a/model/common/pyproject.toml b/model/common/pyproject.toml index 6904b9941..22a32d8ac 100644 --- a/model/common/pyproject.toml +++ b/model/common/pyproject.toml @@ -112,7 +112,6 @@ warn_unused_ignores = true [tool.pytest] [tool.pytest.ini_options] -markers = 'datatest: this test uses binary data' testpaths = 'tests' [tool.setuptools.dynamic] diff --git a/model/common/src/icon4py/model/common/test_utils/fixtures.py b/model/common/src/icon4py/model/common/test_utils/fixtures.py index d28608d0b..c6a52f0c9 100644 --- a/model/common/src/icon4py/model/common/test_utils/fixtures.py +++ b/model/common/src/icon4py/model/common/test_utils/fixtures.py @@ -51,6 +51,7 @@ def ranked_data_path(processor_props): def datapath(ranked_data_path): return ranked_data_path.joinpath("mch_ch_r04b09_dsl/ser_data") + @pytest.fixture(scope="session") def download_ser_data(request, processor_props, ranked_data_path, pytestconfig): """ @@ -58,8 +59,8 @@ def download_ser_data(request, processor_props, ranked_data_path, pytestconfig): Session scoped fixture which is a prerequisite of all the other fixtures in this file. """ - if 'not datatest' in pytestconfig.option.markexpr: - pytest.skip('not running datatest marked tests') + if not pytestconfig.getoption("datatest"): + pytest.skip("not running datatest marked tests") try: uri = data_uris[processor_props.comm_size] diff --git a/model/common/src/icon4py/model/common/test_utils/pytest_config.py b/model/common/src/icon4py/model/common/test_utils/pytest_config.py new file mode 100644 index 000000000..de4425c8a --- /dev/null +++ b/model/common/src/icon4py/model/common/test_utils/pytest_config.py @@ -0,0 +1,40 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import pytest + + +def pytest_configure(config): + config.addinivalue_line("markers", "datatest: this test uses binary data") + + +def pytest_addoption(parser): + """Add --datatest commandline option for pytest. + + Makes sure the option is set only once even when running tests of several model packages in one session. + """ + try: + parser.addoption( + "--datatest", + action="store_true", + help="running tests that use serialized data, can be slow since data might be downloaded from online storage", + default=False, + ) + except ValueError: + pass + + +def pytest_runtest_setup(item): + for _ in item.iter_markers(name="datatest"): + if not item.config.getoption("--datatest"): + pytest.skip("need '--datatest' option to run") diff --git a/model/common/tests/conftest.py b/model/common/tests/conftest.py index 0c4d19c49..63571d684 100644 --- a/model/common/tests/conftest.py +++ b/model/common/tests/conftest.py @@ -13,7 +13,7 @@ import pytest from icon4py.model.common.test_utils.data_handling import download_and_extract -from icon4py.model.common.test_utils.fixtures import ( # noqa F401 +from icon4py.model.common.test_utils.fixtures import ( # noqa: F401 backend, base_path, damping_height, @@ -28,6 +28,11 @@ processor_props, ranked_data_path, ) +from icon4py.model.common.test_utils.pytest_config import ( # noqa: F401 + pytest_addoption, + pytest_configure, + pytest_runtest_setup, +) grids_path = base_path.joinpath("grids") @@ -53,7 +58,7 @@ def get_grid_files(pytestconfig): Session scoped fixture which is a prerequisite of all the other fixtures in this file. """ - if "not datatest" in pytestconfig.option.markexpr: + if not pytestconfig.getoption("datatest"): pytest.skip("not running datatest marked tests") download_and_extract(mch_ch_r04b09_dsl_grid_uri, r04b09_dsl_grid_path, r04b09_dsl_data_file) download_and_extract(r02b04_global_grid_uri, r02b04_global_grid_path, r02b04_global_data_file) diff --git a/model/common/tests/test_icon_grid.py b/model/common/tests/test_icon_grid.py index bdd399942..83f20ee41 100644 --- a/model/common/tests/test_icon_grid.py +++ b/model/common/tests/test_icon_grid.py @@ -392,6 +392,7 @@ def test_cross_check_marker_equivalences(icon_grid): HorizontalMarkerIndex.nudging(EdgeDim), ) + @pytest.mark.datatest def test_grid_size(grid_savepoint): assert 10663 == grid_savepoint.num(VertexDim) diff --git a/model/common/tests/test_vertical.py b/model/common/tests/test_vertical.py index f40d32062..4d6d71b9a 100644 --- a/model/common/tests/test_vertical.py +++ b/model/common/tests/test_vertical.py @@ -41,6 +41,7 @@ def test_nrdmax_calculation_from_icon_input(grid_savepoint, damping_height): assert a_array[nrdmax] > damping_height assert a_array[nrdmax + 1] < damping_height + @pytest.mark.datatest def test_grid_size(grid_savepoint): assert 65 == grid_savepoint.num(KDim) diff --git a/model/driver/pyproject.toml b/model/driver/pyproject.toml index e9123ed8b..b011e5f5b 100644 --- a/model/driver/pyproject.toml +++ b/model/driver/pyproject.toml @@ -117,7 +117,6 @@ warn_unused_ignores = true [tool.pytest] [tool.pytest.ini_options] -markers = 'datatest: this test uses binary data' testpaths = 'tests' [tool.setuptools.dynamic] diff --git a/model/driver/tests/conftest.py b/model/driver/tests/conftest.py index 777fbf547..d9ad439dc 100644 --- a/model/driver/tests/conftest.py +++ b/model/driver/tests/conftest.py @@ -11,9 +11,14 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from icon4py.model.common.test_utils.fixtures import ( # noqa F401 +from icon4py.model.common.test_utils.fixtures import ( # noqa: F401 datapath, download_ser_data, processor_props, ranked_data_path, ) +from icon4py.model.common.test_utils.pytest_config import ( # noqa: F401 + pytest_addoption, + pytest_configure, + pytest_runtest_setup, +) diff --git a/model/tox.ini b/model/tox.ini index c9fd7a989..9ef46bd3c 100644 --- a/model/tox.ini +++ b/model/tox.ini @@ -14,8 +14,8 @@ passenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -m "not datatest" -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src atmosphere/diffusion/src common/src driver/src - pytest -v -m "not datatest" -s -n auto --cov --cov-append --benchmark-disable + -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src atmosphere/diffusion/src common/src driver/src + pytest -v -s -n auto --cov --cov-append --benchmark-disable commands_post = rm -rf tests/_reports/coverage_html -coverage html @@ -32,5 +32,3 @@ skip_install = true commands = commands_post = -[pytest] -markers = datatest: this test uses serialized data diff --git a/tox.ini b/tox.ini index d1d7611fa..4b8680a57 100644 --- a/tox.ini +++ b/tox.ini @@ -14,8 +14,8 @@ passenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -m "not datatest" -s -n auto -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/dycore/src model/atmosphere/diffusion/src model/driver/src model/common/src tools/src - pytest -v -m "not datatest" -s -n auto --cov --cov-append --benchmark-disable + -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/dycore/src model/atmosphere/diffusion/src model/driver/src model/common/src tools/src + pytest -v -s -n auto --cov --cov-append --benchmark-disable commands_post = rm -rf tests/_reports/coverage_html -coverage html