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/compute_contravariant_correction.py similarity index 89% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_04.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_contravariant_correction.py index b412a98b5f..b751cc3bc8 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/compute_contravariant_correction.py @@ -20,12 +20,13 @@ @field_operator -def _mo_velocity_advection_stencil_04( +def _compute_contravariant_correction( vn: Field[[EdgeDim, KDim], wpfloat], ddxn_z_full: Field[[EdgeDim, KDim], vpfloat], ddxt_z_full: Field[[EdgeDim, KDim], vpfloat], vt: Field[[EdgeDim, KDim], vpfloat], ) -> Field[[EdgeDim, KDim], vpfloat]: + """Formerly known as _mo_solve_nonhydro_stencil_35 or mo_velocity_advection_stencil_04.""" ddxn_z_full_wp = astype(ddxn_z_full, wpfloat) z_w_concorr_me_wp = vn * ddxn_z_full_wp + astype(vt * ddxt_z_full, wpfloat) @@ -33,7 +34,7 @@ def _mo_velocity_advection_stencil_04( @program(grid_type=GridType.UNSTRUCTURED) -def mo_velocity_advection_stencil_04( +def compute_contravariant_correction( vn: Field[[EdgeDim, KDim], wpfloat], ddxn_z_full: Field[[EdgeDim, KDim], vpfloat], ddxt_z_full: Field[[EdgeDim, KDim], vpfloat], @@ -44,7 +45,7 @@ def mo_velocity_advection_stencil_04( vertical_start: int32, vertical_end: int32, ): - _mo_velocity_advection_stencil_04( + _compute_contravariant_correction( vn, ddxn_z_full, ddxt_z_full, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_07.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_pertubation_of_rho_and_theta.py similarity index 90% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_07.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_pertubation_of_rho_and_theta.py index 25e392a08e..4f01570b00 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_07.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_pertubation_of_rho_and_theta.py @@ -20,12 +20,13 @@ @field_operator -def _mo_solve_nonhydro_stencil_07( +def _compute_pertubation_of_rho_and_theta( rho: Field[[CellDim, KDim], wpfloat], rho_ref_mc: Field[[CellDim, KDim], vpfloat], theta_v: Field[[CellDim, KDim], wpfloat], theta_ref_mc: Field[[CellDim, KDim], vpfloat], ) -> tuple[Field[[CellDim, KDim], vpfloat], Field[[CellDim, KDim], vpfloat]]: + """Formerly known as mo_solve_nonhydro_stencil_07 or mo_solve_nonhydro_stencil_13.""" rho_ref_mc_wp, theta_ref_mc_wp = astype((rho_ref_mc, theta_ref_mc), wpfloat) z_rth_pr_1_wp = rho - rho_ref_mc_wp @@ -34,7 +35,7 @@ def _mo_solve_nonhydro_stencil_07( @program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_07( +def compute_pertubation_of_rho_and_theta( rho: Field[[CellDim, KDim], wpfloat], rho_ref_mc: Field[[CellDim, KDim], vpfloat], theta_v: Field[[CellDim, KDim], wpfloat], @@ -46,7 +47,7 @@ def mo_solve_nonhydro_stencil_07( vertical_start: int32, vertical_end: int32, ): - _mo_solve_nonhydro_stencil_07( + _compute_pertubation_of_rho_and_theta( rho, rho_ref_mc, theta_v, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_06.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/extrapolate_at_top.py similarity index 91% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_06.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/extrapolate_at_top.py index f9f26485c6..cf3c226110 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_06.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/extrapolate_at_top.py @@ -20,10 +20,11 @@ @field_operator -def _mo_velocity_advection_stencil_06( +def _extrapolate_at_top( wgtfacq_e: Field[[EdgeDim, KDim], vpfloat], vn: Field[[EdgeDim, KDim], wpfloat], ) -> Field[[EdgeDim, KDim], vpfloat]: + """Formerly known as mo_velocity_advection_stencil_06 or mo_solve_nonhydro_stencil_38.""" wgtfacq_e_wp = astype(wgtfacq_e, wpfloat) vn_ie_wp = ( @@ -36,7 +37,7 @@ def _mo_velocity_advection_stencil_06( @program(grid_type=GridType.UNSTRUCTURED) -def mo_velocity_advection_stencil_06( +def extrapolate_at_top( wgtfacq_e: Field[[EdgeDim, KDim], vpfloat], vn: Field[[EdgeDim, KDim], wpfloat], vn_ie: Field[[EdgeDim, KDim], vpfloat], @@ -45,7 +46,7 @@ def mo_velocity_advection_stencil_06( vertical_start: int32, vertical_end: int32, ): - _mo_velocity_advection_stencil_06( + _extrapolate_at_top( wgtfacq_e, vn, out=vn_ie, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_1_to_7.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_1_to_7.py index f20d94b01e..b4b55aa847 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_1_to_7.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_1_to_7.py @@ -14,6 +14,10 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import broadcast, int32, where +from icon4py.model.atmosphere.dycore.compute_contravariant_correction import ( + _compute_contravariant_correction, +) +from icon4py.model.atmosphere.dycore.extrapolate_at_top import _extrapolate_at_top from icon4py.model.atmosphere.dycore.mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl import ( _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl, ) @@ -26,15 +30,9 @@ from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_03 import ( _mo_velocity_advection_stencil_03, ) -from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_04 import ( - _mo_velocity_advection_stencil_04, -) from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_05 import ( _mo_velocity_advection_stencil_05, ) -from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_06 import ( - _mo_velocity_advection_stencil_06, -) from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_07 import ( _mo_velocity_advection_stencil_07, ) @@ -93,11 +91,11 @@ def _fused_velocity_advection_stencil_1_to_6( (vn_ie, z_vt_ie, z_kin_hor_e), ) - vn_ie = where(k == nlevp1, _mo_velocity_advection_stencil_06(wgtfacq_e_dsl, vn), vn_ie) + vn_ie = where(k == nlevp1, _extrapolate_at_top(wgtfacq_e_dsl, vn), vn_ie) z_w_concorr_me = where( nflatlev < k < nlevp1, - _mo_velocity_advection_stencil_04(vn, ddxn_z_full, ddxt_z_full, vt), + _compute_contravariant_correction(vn, ddxn_z_full, ddxt_z_full, vt), z_w_concorr_me, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_8_to_14.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_8_to_14.py index 7226b5542d..515fc5e996 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_8_to_14.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_8_to_14.py @@ -14,12 +14,7 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, maximum, where -from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_08 import ( - _mo_velocity_advection_stencil_08, -) -from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_09 import ( - _mo_velocity_advection_stencil_09, -) +from icon4py.model.atmosphere.dycore.interpolate_to_cell_center import _interpolate_to_cell_center from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_10 import ( _mo_velocity_advection_stencil_10, ) @@ -70,14 +65,14 @@ def _fused_velocity_advection_stencil_8_to_14( ]: z_ekinh = where( k < nlev, - _mo_velocity_advection_stencil_08(z_kin_hor_e, e_bln_c_s), + _interpolate_to_cell_center(z_kin_hor_e, e_bln_c_s), z_ekinh, ) z_w_concorr_mc = ( where( nflatlev < k < nlev, - _mo_velocity_advection_stencil_09(z_w_concorr_me, e_bln_c_s), + _interpolate_to_cell_center(z_w_concorr_me, e_bln_c_s), z_w_concorr_mc, ) if istep == 1 diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_08.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_cell_center.py similarity index 69% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_08.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_cell_center.py index ec4374085a..79fe0ed310 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_08.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_cell_center.py @@ -20,30 +20,30 @@ @field_operator -def _mo_velocity_advection_stencil_08( - z_kin_hor_e: Field[[EdgeDim, KDim], vpfloat], +def _interpolate_to_cell_center( + interpolant: Field[[EdgeDim, KDim], vpfloat], e_bln_c_s: Field[[CEDim], wpfloat], ) -> Field[[CellDim, KDim], vpfloat]: - z_kin_hor_e_wp = astype(z_kin_hor_e, wpfloat) - - z_ekinh_wp = neighbor_sum(z_kin_hor_e_wp(C2E) * e_bln_c_s(C2CE), axis=C2EDim) - return astype(z_ekinh_wp, vpfloat) + """Formerly known as mo_velocity_advection_stencil_08 or mo_velocity_advection_stencil_09.""" + interpolant_wp = astype(interpolant, wpfloat) + interpolation_wp = neighbor_sum(e_bln_c_s(C2CE) * interpolant_wp(C2E), axis=C2EDim) + return astype(interpolation_wp, vpfloat) @program(grid_type=GridType.UNSTRUCTURED) -def mo_velocity_advection_stencil_08( - z_kin_hor_e: Field[[EdgeDim, KDim], vpfloat], +def interpolate_to_cell_center( + interpolant: Field[[EdgeDim, KDim], vpfloat], e_bln_c_s: Field[[CEDim], wpfloat], - z_ekinh: Field[[CellDim, KDim], vpfloat], + interpolation: Field[[CellDim, KDim], vpfloat], horizontal_start: int32, horizontal_end: int32, vertical_start: int32, vertical_end: int32, ): - _mo_velocity_advection_stencil_08( - z_kin_hor_e, + _interpolate_to_cell_center( + interpolant, e_bln_c_s, - out=z_ekinh, + out=interpolation, domain={ CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end), diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_13.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_13.py deleted file mode 100644 index 6126bd6b8b..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_13.py +++ /dev/null @@ -1,59 +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.common import GridType -from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, astype, int32 - -from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.type_alias import vpfloat, wpfloat - - -@field_operator -def _mo_solve_nonhydro_stencil_13( - rho: Field[[CellDim, KDim], wpfloat], - rho_ref_mc: Field[[CellDim, KDim], vpfloat], - theta_v: Field[[CellDim, KDim], wpfloat], - theta_ref_mc: Field[[CellDim, KDim], vpfloat], -) -> tuple[Field[[CellDim, KDim], vpfloat], Field[[CellDim, KDim], vpfloat]]: - rho_ref_mc_wp, theta_ref_mc_wp = astype((rho_ref_mc, theta_ref_mc), wpfloat) - - z_rth_pr_1_wp = rho - rho_ref_mc_wp - z_rth_pr_2_wp = theta_v - theta_ref_mc_wp - return astype((z_rth_pr_1_wp, z_rth_pr_2_wp), vpfloat) - - -@program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_13( - rho: Field[[CellDim, KDim], wpfloat], - rho_ref_mc: Field[[CellDim, KDim], vpfloat], - theta_v: Field[[CellDim, KDim], wpfloat], - theta_ref_mc: Field[[CellDim, KDim], vpfloat], - z_rth_pr_1: Field[[CellDim, KDim], vpfloat], - z_rth_pr_2: Field[[CellDim, KDim], vpfloat], - horizontal_start: int32, - horizontal_end: int32, - vertical_start: int32, - vertical_end: int32, -): - _mo_solve_nonhydro_stencil_13( - rho, - rho_ref_mc, - theta_v, - theta_ref_mc, - out=(z_rth_pr_1, z_rth_pr_2), - domain={ - CellDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_35.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_35.py deleted file mode 100644 index 571ad4659a..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_35.py +++ /dev/null @@ -1,57 +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.common import GridType -from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, astype, int32 - -from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.type_alias import vpfloat, wpfloat - - -@field_operator -def _mo_solve_nonhydro_stencil_35( - vn: Field[[EdgeDim, KDim], wpfloat], - ddxn_z_full: Field[[EdgeDim, KDim], vpfloat], - ddxt_z_full: Field[[EdgeDim, KDim], vpfloat], - vt: Field[[EdgeDim, KDim], vpfloat], -) -> Field[[EdgeDim, KDim], vpfloat]: - ddxn_z_full_wp = astype(ddxn_z_full, wpfloat) - - z_w_concorr_me_wp = vn * ddxn_z_full_wp + astype(vt * ddxt_z_full, wpfloat) - return astype(z_w_concorr_me_wp, vpfloat) - - -@program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_35( - vn: Field[[EdgeDim, KDim], wpfloat], - ddxn_z_full: Field[[EdgeDim, KDim], vpfloat], - ddxt_z_full: Field[[EdgeDim, KDim], vpfloat], - vt: Field[[EdgeDim, KDim], vpfloat], - z_w_concorr_me: Field[[EdgeDim, KDim], vpfloat], - horizontal_start: int32, - horizontal_end: int32, - vertical_start: int32, - vertical_end: int32, -): - _mo_solve_nonhydro_stencil_35( - vn, - ddxn_z_full, - ddxt_z_full, - vt, - out=z_w_concorr_me, - domain={ - EdgeDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_38.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_38.py deleted file mode 100644 index 5cf39f3564..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_38.py +++ /dev/null @@ -1,55 +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.common import GridType -from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, astype, int32 - -from icon4py.model.common.dimension import EdgeDim, KDim, Koff -from icon4py.model.common.type_alias import vpfloat, wpfloat - - -@field_operator -def _mo_solve_nonhydro_stencil_38( - vn: Field[[EdgeDim, KDim], wpfloat], - wgtfacq_e: Field[[EdgeDim, KDim], vpfloat], -) -> Field[[EdgeDim, KDim], vpfloat]: - wgtfacq_e_wp = astype(wgtfacq_e, wpfloat) - - vn_ie_wp = ( - wgtfacq_e_wp(Koff[-1]) * vn(Koff[-1]) - + wgtfacq_e_wp(Koff[-2]) * vn(Koff[-2]) - + wgtfacq_e_wp(Koff[-3]) * vn(Koff[-3]) - ) - return astype(vn_ie_wp, vpfloat) - - -@program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_38( - vn: Field[[EdgeDim, KDim], wpfloat], - wgtfacq_e: Field[[EdgeDim, KDim], vpfloat], - vn_ie: Field[[EdgeDim, KDim], vpfloat], - horizontal_start: int32, - horizontal_end: int32, - vertical_start: int32, - vertical_end: int32, -): - _mo_solve_nonhydro_stencil_38( - vn, - wgtfacq_e, - out=vn_ie, - domain={ - EdgeDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_09.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_09.py deleted file mode 100644 index 250e396136..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_09.py +++ /dev/null @@ -1,50 +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.common import GridType -from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum - -from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.type_alias import vpfloat, wpfloat - - -@field_operator -def _mo_velocity_advection_stencil_09( - z_w_concorr_me: Field[[EdgeDim, KDim], vpfloat], - e_bln_c_s: Field[[CEDim], wpfloat], -) -> Field[[CellDim, KDim], vpfloat]: - z_w_concorr_me_wp = astype(z_w_concorr_me, wpfloat) - z_w_concorr_mc_wp = neighbor_sum(e_bln_c_s(C2CE) * z_w_concorr_me_wp(C2E), axis=C2EDim) - return astype(z_w_concorr_mc_wp, vpfloat) - - -@program(grid_type=GridType.UNSTRUCTURED) -def mo_velocity_advection_stencil_09( - z_w_concorr_me: Field[[EdgeDim, KDim], vpfloat], - e_bln_c_s: Field[[CEDim], wpfloat], - z_w_concorr_mc: Field[[CellDim, KDim], vpfloat], - horizontal_start: int32, - horizontal_end: int32, - vertical_start: int32, - vertical_end: int32, -): - _mo_velocity_advection_stencil_09( - z_w_concorr_me, - e_bln_c_s, - out=z_w_concorr_mc, - domain={ - CellDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py index 45002fbbed..fa2d1c681c 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py @@ -20,6 +20,9 @@ import icon4py.model.atmosphere.dycore.nh_solve.solve_nonhydro_program as nhsolve_prog import icon4py.model.common.constants as constants +from icon4py.model.atmosphere.dycore.compute_pertubation_of_rho_and_theta import ( + compute_pertubation_of_rho_and_theta, +) from icon4py.model.atmosphere.dycore.mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl import ( mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl, ) @@ -38,9 +41,6 @@ from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_12 import ( mo_solve_nonhydro_stencil_12, ) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_13 import ( - mo_solve_nonhydro_stencil_13, -) from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_17 import ( mo_solve_nonhydro_stencil_17, ) @@ -783,7 +783,7 @@ def run_predictor_step( # Add computation of z_grad_rth (perturbation density and virtual potential temperature at main levels) # at outer halo points: needed for correct calculation of the upwind gradients for Miura scheme - mo_solve_nonhydro_stencil_13.with_backend(backend)( + compute_pertubation_of_rho_and_theta.with_backend(backend)( rho=prognostic_state[nnow].rho, rho_ref_mc=self.metric_state_nonhydro.rho_ref_mc, theta_v=prognostic_state[nnow].theta_v, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro_program.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro_program.py index 9c92a8b004..778777a319 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro_program.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro_program.py @@ -14,6 +14,13 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, where +from icon4py.model.atmosphere.dycore.compute_contravariant_correction import ( + _compute_contravariant_correction, +) +from icon4py.model.atmosphere.dycore.compute_pertubation_of_rho_and_theta import ( + _compute_pertubation_of_rho_and_theta, +) +from icon4py.model.atmosphere.dycore.extrapolate_at_top import _extrapolate_at_top from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_02 import ( _mo_solve_nonhydro_stencil_02, ) @@ -26,9 +33,6 @@ from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_06 import ( _mo_solve_nonhydro_stencil_06, ) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_07 import ( - _mo_solve_nonhydro_stencil_07, -) from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_08 import ( _mo_solve_nonhydro_stencil_08, ) @@ -44,18 +48,12 @@ from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1 import ( _mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1, ) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_35 import ( - _mo_solve_nonhydro_stencil_35, -) from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_36 import ( _mo_solve_nonhydro_stencil_36, ) from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_37 import ( _mo_solve_nonhydro_stencil_37, ) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_38 import ( - _mo_solve_nonhydro_stencil_38, -) from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_39 import ( _mo_solve_nonhydro_stencil_39, ) @@ -276,7 +274,7 @@ def _predictor_stencils_7_8_9( ]: (z_rth_pr_1, z_rth_pr_2) = where( k_field == int32(0), - _mo_solve_nonhydro_stencil_07(rho, rho_ref_mc, theta_v, theta_ref_mc), + _compute_pertubation_of_rho_and_theta(rho, rho_ref_mc, theta_v, theta_ref_mc), (z_rth_pr_1, z_rth_pr_2), ) @@ -486,7 +484,7 @@ def _predictor_stencils_35_36( ]: z_w_concorr_me = where( k_field >= nflatlev_startindex, - _mo_solve_nonhydro_stencil_35(vn, ddxn_z_full, ddxt_z_full, vt), + _compute_contravariant_correction(vn, ddxn_z_full, ddxt_z_full, vt), z_w_concorr_me, ) (vn_ie, z_vt_ie, z_kin_hor_e) = where( @@ -557,7 +555,7 @@ def _predictor_stencils_37_38( _mo_solve_nonhydro_stencil_37(vn, vt), (vn_ie, z_vt_ie, z_kin_hor_e), ) - vn_ie = where(k_field == nlev, _mo_solve_nonhydro_stencil_38(vn, wgtfacq_e_dsl), vn_ie) + vn_ie = where(k_field == nlev, _extrapolate_at_top(vn, wgtfacq_e_dsl), vn_ie) return vn_ie, z_vt_ie, z_kin_hor_e diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection.py index 7888e09356..aa1f9944b4 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection.py @@ -22,6 +22,7 @@ ) import icon4py.model.atmosphere.dycore.velocity.velocity_advection_program as velocity_prog +from icon4py.model.atmosphere.dycore.interpolate_to_cell_center import interpolate_to_cell_center from icon4py.model.atmosphere.dycore.mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl import ( mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl, ) @@ -40,9 +41,6 @@ from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_07 import ( mo_velocity_advection_stencil_07, ) -from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_08 import ( - mo_velocity_advection_stencil_08, -) from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_15 import ( mo_velocity_advection_stencil_15, ) @@ -250,7 +248,7 @@ def run_predictor_step( vertical_end=self.grid.num_levels, offset_provider={}, ) - velocity_prog.mo_velocity_advection_stencil_06.with_backend(backend)( + velocity_prog.extrapolate_at_top.with_backend(backend)( wgtfacq_e=self.metric_state.wgtfacq_e, vn=prognostic_state.vn, vn_ie=diagnostic_state.vn_ie, @@ -281,7 +279,7 @@ def run_predictor_step( }, ) - mo_velocity_advection_stencil_08.with_backend(backend)( + interpolate_to_cell_center.with_backend(backend)( z_kin_hor_e=z_kin_hor_e, e_bln_c_s=self.interpolation_state.e_bln_c_s, z_ekinh=self.z_ekinh, @@ -549,7 +547,7 @@ def run_corrector_step( }, ) - mo_velocity_advection_stencil_08.with_backend(backend)( + interpolate_to_cell_center.with_backend(backend)( z_kin_hor_e=z_kin_hor_e, e_bln_c_s=self.interpolation_state.e_bln_c_s, z_ekinh=self.z_ekinh, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection_program.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection_program.py index babed9ee2d..d96748f7c4 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection_program.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection_program.py @@ -14,18 +14,14 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, where -from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_04 import ( - _mo_velocity_advection_stencil_04, +from icon4py.model.atmosphere.dycore.compute_contravariant_correction import ( + _compute_contravariant_correction, ) +from icon4py.model.atmosphere.dycore.extrapolate_at_top import _extrapolate_at_top +from icon4py.model.atmosphere.dycore.interpolate_to_cell_center import _interpolate_to_cell_center from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_05 import ( _mo_velocity_advection_stencil_05, ) -from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_06 import ( - _mo_velocity_advection_stencil_06, -) -from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_09 import ( - _mo_velocity_advection_stencil_09, -) from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_10 import ( _mo_velocity_advection_stencil_10, ) @@ -71,7 +67,7 @@ def _fused_stencils_4_5( ]: z_w_concorr_me = where( (k_field >= nflatlev_startindex) & (k_field < nlev), - _mo_velocity_advection_stencil_04(vn, ddxn_z_full, ddxt_z_full, vt), + _compute_contravariant_correction(vn, ddxn_z_full, ddxt_z_full, vt), z_w_concorr_me, ) @@ -123,7 +119,7 @@ def fused_stencils_4_5( @program -def mo_velocity_advection_stencil_06( +def extrapolate_at_top( wgtfacq_e: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], vn_ie: Field[[EdgeDim, KDim], float], @@ -132,7 +128,7 @@ def mo_velocity_advection_stencil_06( vertical_start: int32, vertical_end: int32, ): - _mo_velocity_advection_stencil_06( + _extrapolate_at_top( wgtfacq_e, vn, out=vn_ie, @@ -156,7 +152,7 @@ def _fused_stencils_9_10( ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: local_z_w_concorr_mc = where( (k_field >= nflatlev_startindex) & (k_field < nlev), - _mo_velocity_advection_stencil_09(z_w_concorr_me, e_bln_c_s), + _interpolate_to_cell_center(z_w_concorr_me, e_bln_c_s), local_z_w_concorr_mc, ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_04.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_contravariant_correction.py similarity index 86% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_04.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_contravariant_correction.py index 2dc2ccb763..d9fb9a0c04 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_04.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_contravariant_correction.py @@ -15,23 +15,23 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_04 import ( - mo_velocity_advection_stencil_04, +from icon4py.model.atmosphere.dycore.compute_contravariant_correction import ( + compute_contravariant_correction, ) from icon4py.model.common.dimension import EdgeDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field from icon4py.model.common.type_alias import vpfloat, wpfloat -def mo_velocity_advection_stencil_04_numpy( +def compute_contravariant_correction_numpy( vn: np.array, ddxn_z_full: np.array, ddxt_z_full: np.array, vt: np.array ) -> np.array: z_w_concorr_me = vn * ddxn_z_full + vt * ddxt_z_full return z_w_concorr_me -class TestMoVelocityAdvectionStencil04(StencilTest): - PROGRAM = mo_velocity_advection_stencil_04 +class TestMoSolveNonhydroStencil35(StencilTest): + PROGRAM = compute_contravariant_correction OUTPUTS = ("z_w_concorr_me",) @staticmethod @@ -43,7 +43,7 @@ def reference( vt: np.array, **kwargs, ) -> dict: - z_w_concorr_me = mo_velocity_advection_stencil_04_numpy(vn, ddxn_z_full, ddxt_z_full, vt) + z_w_concorr_me = compute_contravariant_correction_numpy(vn, ddxn_z_full, ddxt_z_full, vt) return dict(z_w_concorr_me=z_w_concorr_me) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_07.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_pertubation_of_rho_and_theta.py similarity index 92% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_07.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_pertubation_of_rho_and_theta.py index dae2c6daea..189851e6c1 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_07.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_pertubation_of_rho_and_theta.py @@ -15,8 +15,8 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_07 import ( - mo_solve_nonhydro_stencil_07, +from icon4py.model.atmosphere.dycore.compute_pertubation_of_rho_and_theta import ( + compute_pertubation_of_rho_and_theta, ) from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field @@ -24,7 +24,7 @@ class TestMoSolveNonhydroStencil07(StencilTest): - PROGRAM = mo_solve_nonhydro_stencil_07 + PROGRAM = compute_pertubation_of_rho_and_theta OUTPUTS = ("z_rth_pr_1", "z_rth_pr_2") @staticmethod diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_06.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_extrapolate_at_top.py similarity index 86% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_06.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_extrapolate_at_top.py index ba8890ce19..d5001d2194 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_06.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_extrapolate_at_top.py @@ -15,15 +15,13 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_06 import ( - mo_velocity_advection_stencil_06, -) +from icon4py.model.atmosphere.dycore.extrapolate_at_top import extrapolate_at_top from icon4py.model.common.dimension import EdgeDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field from icon4py.model.common.type_alias import vpfloat, wpfloat -def mo_velocity_advection_stencil_06_numpy(wgtfacq_e: np.array, vn: np.array) -> np.array: +def extrapolate_at_top_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) @@ -40,12 +38,12 @@ def mo_velocity_advection_stencil_06_numpy(wgtfacq_e: np.array, vn: np.array) -> class TestMoVelocityAdvectionStencil06(StencilTest): - PROGRAM = mo_velocity_advection_stencil_06 + PROGRAM = extrapolate_at_top OUTPUTS = ("vn_ie",) @staticmethod def reference(grid, wgtfacq_e: np.array, vn: np.array, **kwargs) -> dict: - vn_ie = mo_velocity_advection_stencil_06_numpy(wgtfacq_e, vn) + vn_ie = extrapolate_at_top_numpy(wgtfacq_e, vn) return dict(vn_ie=vn_ie) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_1_to_7.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_1_to_7.py index 0b4f40aab7..7d1c9c800f 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_1_to_7.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_1_to_7.py @@ -22,15 +22,15 @@ from icon4py.model.common.dimension import CellDim, E2C2EDim, EdgeDim, KDim, V2CDim, VertexDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field +from .test_compute_contravariant_correction import compute_contravariant_correction_numpy +from .test_extrapolate_at_top import extrapolate_at_top_numpy from .test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl import ( mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl_numpy, ) from .test_mo_velocity_advection_stencil_01 import mo_velocity_advection_stencil_01_numpy from .test_mo_velocity_advection_stencil_02 import mo_velocity_advection_stencil_02_numpy from .test_mo_velocity_advection_stencil_03 import mo_velocity_advection_stencil_03_numpy -from .test_mo_velocity_advection_stencil_04 import mo_velocity_advection_stencil_04_numpy from .test_mo_velocity_advection_stencil_05 import mo_velocity_advection_stencil_05_numpy -from .test_mo_velocity_advection_stencil_06 import mo_velocity_advection_stencil_06_numpy from .test_mo_velocity_advection_stencil_07 import mo_velocity_advection_stencil_07_numpy @@ -91,14 +91,12 @@ def _fused_velocity_advection_stencil_1_to_6_numpy( ) condition4 = k == nlevp1 - vn_ie = np.where( - condition4, mo_velocity_advection_stencil_06_numpy(wgtfacq_e_dsl, vn), vn_ie - ) + vn_ie = np.where(condition4, extrapolate_at_top_numpy(wgtfacq_e_dsl, vn), vn_ie) condition5 = (nflatlev < k) & (k < nlevp1) z_w_concorr_me = np.where( condition5, - mo_velocity_advection_stencil_04_numpy(vn, ddxn_z_full, ddxt_z_full, vt), + compute_contravariant_correction_numpy(vn, ddxn_z_full, ddxt_z_full, vt), z_w_concorr_me, ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_8_to_14.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_8_to_14.py index 8b81cf1f54..b2ad189db8 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_8_to_14.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_8_to_14.py @@ -27,8 +27,7 @@ zero_field, ) -from .test_mo_velocity_advection_stencil_08 import mo_velocity_advection_stencil_08_numpy -from .test_mo_velocity_advection_stencil_09 import mo_velocity_advection_stencil_09_numpy +from .test_interpolate_to_cell_center import interpolate_to_cell_center_numpy from .test_mo_velocity_advection_stencil_10 import mo_velocity_advection_stencil_10_numpy from .test_mo_velocity_advection_stencil_11 import mo_velocity_advection_stencil_11_numpy from .test_mo_velocity_advection_stencil_12 import mo_velocity_advection_stencil_12_numpy @@ -75,14 +74,14 @@ def reference( z_ekinh = np.where( k < nlev, - mo_velocity_advection_stencil_08_numpy(grid, z_kin_hor_e, e_bln_c_s), + interpolate_to_cell_center_numpy(grid, z_kin_hor_e, e_bln_c_s), z_ekinh, ) if istep == 1: z_w_concorr_mc = np.where( (nflatlev < k) & (k < nlev), - mo_velocity_advection_stencil_09_numpy(grid, z_w_concorr_me, e_bln_c_s), + interpolate_to_cell_center_numpy(grid, z_w_concorr_me, e_bln_c_s), z_w_concorr_mc, ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_08.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_to_cell_center.py similarity index 63% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_08.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_to_cell_center.py index 22729ebb86..683d9dc77b 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_08.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_to_cell_center.py @@ -15,9 +15,7 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_08 import ( - mo_velocity_advection_stencil_08, -) +from icon4py.model.atmosphere.dycore.interpolate_to_cell_center import interpolate_to_cell_center from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim from icon4py.model.common.test_utils.helpers import ( StencilTest, @@ -28,37 +26,37 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat -def mo_velocity_advection_stencil_08_numpy( - grid, z_kin_hor_e: np.array, e_bln_c_s: np.array, **kwargs +def interpolate_to_cell_center_numpy( + grid, interpolant: 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[grid.connectivities[C2EDim]] + interpolation = np.sum( + interpolant[grid.connectivities[C2EDim]] * e_bln_c_s[grid.get_offset_provider("C2CE").table], axis=1, ) - return z_ekinh + return interpolation -class TestMoVelocityAdvectionStencil08(StencilTest): - PROGRAM = mo_velocity_advection_stencil_08 - OUTPUTS = ("z_ekinh",) +class TestMoVelocityAdvectionStencil09(StencilTest): + PROGRAM = interpolate_to_cell_center + OUTPUTS = ("interpolation",) @staticmethod - def reference(grid, z_kin_hor_e: np.array, e_bln_c_s: np.array, **kwargs) -> dict: - z_ekinh = mo_velocity_advection_stencil_08_numpy(grid, z_kin_hor_e, e_bln_c_s) - return dict(z_ekinh=z_ekinh) + def reference(grid, interpolant: np.array, e_bln_c_s: np.array, **kwargs) -> dict: + interpolation = interpolate_to_cell_center_numpy(grid, interpolant, e_bln_c_s) + return dict(interpolation=interpolation) @pytest.fixture def input_data(self, grid): - z_kin_hor_e = random_field(grid, EdgeDim, KDim, dtype=vpfloat) + interpolant = random_field(grid, EdgeDim, KDim, dtype=vpfloat) e_bln_c_s = random_field(grid, CellDim, C2EDim, dtype=wpfloat) - z_ekinh = zero_field(grid, CellDim, KDim, dtype=vpfloat) + interpolation = zero_field(grid, CellDim, KDim, dtype=vpfloat) return dict( - z_kin_hor_e=z_kin_hor_e, + interpolant=interpolant, e_bln_c_s=as_1D_sparse_field(e_bln_c_s, CEDim), - z_ekinh=z_ekinh, + interpolation=interpolation, horizontal_start=int32(0), horizontal_end=int32(grid.num_cells), vertical_start=int32(0), diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_13.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_13.py deleted file mode 100644 index e763ab81b4..0000000000 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_13.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 - -import numpy as np -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_13 import ( - mo_solve_nonhydro_stencil_13, -) -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.type_alias import vpfloat, wpfloat - - -class TestMoSolveNonhydroStencil13(StencilTest): - PROGRAM = mo_solve_nonhydro_stencil_13 - OUTPUTS = ("z_rth_pr_1", "z_rth_pr_2") - - @staticmethod - def reference( - grid, - rho: np.array, - rho_ref_mc: np.array, - theta_v: np.array, - theta_ref_mc: np.array, - **kwargs, - ) -> dict: - z_rth_pr_1 = rho - rho_ref_mc - z_rth_pr_2 = theta_v - theta_ref_mc - return dict(z_rth_pr_1=z_rth_pr_1, z_rth_pr_2=z_rth_pr_2) - - @pytest.fixture - def input_data(self, grid): - rho = random_field(grid, CellDim, KDim, dtype=wpfloat) - rho_ref_mc = random_field(grid, CellDim, KDim, dtype=vpfloat) - theta_v = random_field(grid, CellDim, KDim, dtype=wpfloat) - theta_ref_mc = random_field(grid, CellDim, KDim, dtype=vpfloat) - z_rth_pr_1 = zero_field(grid, CellDim, KDim, dtype=vpfloat) - z_rth_pr_2 = zero_field(grid, CellDim, KDim, dtype=vpfloat) - - return dict( - rho=rho, - rho_ref_mc=rho_ref_mc, - theta_v=theta_v, - theta_ref_mc=theta_ref_mc, - z_rth_pr_1=z_rth_pr_1, - z_rth_pr_2=z_rth_pr_2, - horizontal_start=int32(0), - horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), - vertical_end=int32(grid.num_levels), - ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_35.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_35.py deleted file mode 100644 index 9c80b08654..0000000000 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_35.py +++ /dev/null @@ -1,60 +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 -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_35 import ( - mo_solve_nonhydro_stencil_35, -) -from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field -from icon4py.model.common.type_alias import vpfloat, wpfloat - - -class TestMoSolveNonhydroStencil35(StencilTest): - PROGRAM = mo_solve_nonhydro_stencil_35 - OUTPUTS = ("z_w_concorr_me",) - - @staticmethod - def reference( - grid, - vn: np.array, - ddxn_z_full: np.array, - ddxt_z_full: np.array, - vt: np.array, - **kwargs, - ) -> dict: - z_w_concorr_me = vn * ddxn_z_full + vt * ddxt_z_full - return dict(z_w_concorr_me=z_w_concorr_me) - - @pytest.fixture - def input_data(self, grid): - vn = random_field(grid, EdgeDim, KDim, dtype=wpfloat) - ddxn_z_full = random_field(grid, EdgeDim, KDim, dtype=vpfloat) - ddxt_z_full = random_field(grid, EdgeDim, KDim, dtype=vpfloat) - vt = random_field(grid, EdgeDim, KDim, dtype=vpfloat) - z_w_concorr_me = zero_field(grid, EdgeDim, KDim, dtype=vpfloat) - - return dict( - vn=vn, - ddxn_z_full=ddxn_z_full, - ddxt_z_full=ddxt_z_full, - vt=vt, - z_w_concorr_me=z_w_concorr_me, - horizontal_start=int32(0), - horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), - vertical_end=int32(grid.num_levels), - ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_38.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_38.py deleted file mode 100644 index dd28114282..0000000000 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_38.py +++ /dev/null @@ -1,54 +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 -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_38 import ( - mo_solve_nonhydro_stencil_38, -) -from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field -from icon4py.model.common.type_alias import vpfloat, wpfloat - - -class TestMoSolveNonhydroStencil38(StencilTest): - PROGRAM = mo_solve_nonhydro_stencil_38 - OUTPUTS = ("vn_ie",) - - @staticmethod - def reference(grid, vn: np.array, wgtfacq_e: np.array, **kwargs) -> dict: - vn_ie = np.zeros_like(vn) - vn_ie[:, -1] = ( - np.roll(wgtfacq_e, shift=1, axis=1) * np.roll(vn, shift=1, axis=1) - + np.roll(wgtfacq_e, shift=2, axis=1) * np.roll(vn, shift=2, axis=1) - + np.roll(wgtfacq_e, shift=3, axis=1) * np.roll(vn, shift=3, axis=1) - )[:, -1] - return dict(vn_ie=vn_ie) - - @pytest.fixture - def input_data(self, grid): - wgtfacq_e = zero_field(grid, EdgeDim, KDim, dtype=vpfloat) - vn = random_field(grid, EdgeDim, KDim, dtype=wpfloat) - vn_ie = zero_field(grid, EdgeDim, KDim, dtype=vpfloat) - - return dict( - vn=vn, - wgtfacq_e=wgtfacq_e, - vn_ie=vn_ie, - horizontal_start=int32(0), - horizontal_end=int32(grid.num_edges), - vertical_start=int32(grid.num_levels - 1), - vertical_end=int32(grid.num_levels), - ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_09.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_09.py deleted file mode 100644 index de6d09fc72..0000000000 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_09.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 numpy as np -import pytest -from gt4py.next.ffront.fbuiltins import int32 - -from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_09 import ( - mo_velocity_advection_stencil_09, -) -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, -) -from icon4py.model.common.type_alias import vpfloat, wpfloat - - -def mo_velocity_advection_stencil_09_numpy( - grid, 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[grid.connectivities[C2EDim]] - * e_bln_c_s[grid.get_offset_provider("C2CE").table], - axis=1, - ) - return z_w_concorr_mc - - -class TestMoVelocityAdvectionStencil09(StencilTest): - PROGRAM = mo_velocity_advection_stencil_09 - OUTPUTS = ("z_w_concorr_mc",) - - @staticmethod - def reference(grid, z_w_concorr_me: np.array, e_bln_c_s: np.array, **kwargs) -> dict: - z_w_concorr_mc = mo_velocity_advection_stencil_09_numpy(grid, z_w_concorr_me, e_bln_c_s) - return dict(z_w_concorr_mc=z_w_concorr_mc) - - @pytest.fixture - def input_data(self, grid): - z_w_concorr_me = random_field(grid, EdgeDim, KDim, dtype=vpfloat) - e_bln_c_s = random_field(grid, CellDim, C2EDim, dtype=wpfloat) - z_w_concorr_mc = zero_field(grid, CellDim, KDim, dtype=vpfloat) - - return dict( - z_w_concorr_me=z_w_concorr_me, - e_bln_c_s=as_1D_sparse_field(e_bln_c_s, CEDim), - z_w_concorr_mc=z_w_concorr_mc, - horizontal_start=int32(0), - horizontal_end=int32(grid.num_cells), - vertical_start=int32(0), - vertical_end=int32(grid.num_levels), - )