diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.py new file mode 100644 index 0000000000..c9c359621f --- /dev/null +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.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 + +from functional.common import Field +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import int32, maximum, minimum, where + +from icon4py.common.dimension import CellDim, KDim + + +@field_operator +def _hflx_limiter_mo_stencil_02( + refin_ctrl: Field[[CellDim], int32], + p_cc: Field[[CellDim, KDim], float], + z_tracer_new_low: Field[[CellDim, KDim], float], + lo_bound: int32, + hi_bound: int32, +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + condition = (refin_ctrl == lo_bound) | (refin_ctrl == hi_bound) + z_tracer_new_tmp = where( + condition, + minimum(1.1 * p_cc, maximum(0.9 * p_cc, z_tracer_new_low)), + z_tracer_new_low, + ) + return where( + condition, + (maximum(p_cc, z_tracer_new_tmp), minimum(p_cc, z_tracer_new_tmp)), + (z_tracer_new_low, z_tracer_new_low), + ) + + +@program +def hflx_limiter_mo_stencil_02( + refin_ctrl: Field[[CellDim], int32], + p_cc: Field[[CellDim, KDim], float], + z_tracer_new_low: Field[[CellDim, KDim], float], + lo_bound: int32, + hi_bound: int32, + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], +): + _hflx_limiter_mo_stencil_02( + refin_ctrl, + p_cc, + z_tracer_new_low, + lo_bound, + hi_bound, + out=(z_tracer_max, z_tracer_min), + ) diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py new file mode 100644 index 0000000000..5969a171fd --- /dev/null +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.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.common import Field +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import max_over, maximum, min_over, minimum + +from icon4py.common.dimension import C2E2C, C2E2CDim, CellDim, KDim + + +@field_operator +def _hflx_limiter_mo_stencil_03_min_max( + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], + beta_fct: float, + r_beta_fct: float, +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + z_max = beta_fct * maximum( + max_over(z_tracer_max(C2E2C), axis=C2E2CDim), z_tracer_max + ) + z_min = r_beta_fct * minimum( + min_over(z_tracer_min(C2E2C), axis=C2E2CDim), z_tracer_min + ) + return z_max, z_min + + +@program +def hflx_limiter_mo_stencil_03_min_max( + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], + beta_fct: float, + r_beta_fct: float, + z_max: Field[[CellDim, KDim], float], + z_min: Field[[CellDim, KDim], float], +): + _hflx_limiter_mo_stencil_03_min_max( + z_tracer_max, z_tracer_min, beta_fct, r_beta_fct, out=(z_max, z_min) + ) + + +@field_operator +def _hflx_limiter_mo_stencil_03( + z_mflx_anti_in: Field[[CellDim, KDim], float], + z_mflx_anti_out: Field[[CellDim, KDim], float], + z_tracer_new_low: Field[[CellDim, KDim], float], + z_max: Field[[CellDim, KDim], float], + z_min: Field[[CellDim, KDim], float], + dbl_eps: float, +) -> tuple[[Field[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in + dbl_eps) + r_m = (z_tracer_new_low - z_min) / (z_mflx_anti_out + dbl_eps) + return r_p, r_m + + +@program +def hflx_limiter_mo_stencil_03( + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], + beta_fct: float, + r_beta_fct: float, + z_max: Field[[CellDim, KDim], float], + z_min: Field[[CellDim, KDim], float], + z_mflx_anti_in: Field[[CellDim, KDim], float], + z_mflx_anti_out: Field[[CellDim, KDim], float], + z_tracer_new_low: Field[[CellDim, KDim], float], + dbl_eps: float, + r_p: Field[[CellDim, KDim], float], + r_m: Field[[CellDim, KDim], float], +): + + _hflx_limiter_mo_stencil_03_min_max( + z_tracer_max, z_tracer_min, beta_fct, r_beta_fct, out=(z_min, z_max) + ) + _hflx_limiter_mo_stencil_03( + z_mflx_anti_in, + z_mflx_anti_out, + z_tracer_new_low, + z_max, + z_min, + dbl_eps, + out=(r_p, r_m), + ) diff --git a/advection/src/icon4py/advection/hflx_limiter_mo_stencil_04.py b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_04.py new file mode 100644 index 0000000000..8a97740918 --- /dev/null +++ b/advection/src/icon4py/advection/hflx_limiter_mo_stencil_04.py @@ -0,0 +1,44 @@ +# 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 minimum, where + +from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim + + +@field_operator +def _hflx_limiter_mo_stencil_04( + z_anti: Field[[EdgeDim, KDim], float], + r_m: Field[[CellDim, KDim], float], + r_p: Field[[CellDim, KDim], float], + z_mflx_low: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + r_frac = where( + z_anti >= 0.0, + minimum(r_m(E2C[0]), r_p(E2C[1])), + minimum(r_m(E2C[1]), r_p(E2C[0])), + ) + return z_mflx_low + minimum(1.0, r_frac) * z_anti + + +@program +def hflx_limiter_mo_stencil_04( + z_anti: Field[[EdgeDim, KDim], float], + r_m: Field[[CellDim, KDim], float], + r_p: Field[[CellDim, KDim], float], + z_mflx_low: Field[[EdgeDim, KDim], float], + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], +): + _hflx_limiter_mo_stencil_04(z_anti, r_m, r_p, z_mflx_low, out=p_mflx_tracer_h) diff --git a/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py b/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py index 857ee7584d..5d4d239f08 100644 --- a/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py +++ b/advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py @@ -12,23 +12,23 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, where +from functional.ffront.fbuiltins import Field, int32, where from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim @field_operator def _hflx_limiter_pd_stencil_02( - refin_ctrl: Field[[EdgeDim], float], + refin_ctrl: Field[[EdgeDim], int32], r_m: Field[[CellDim, KDim], float], p_mflx_tracer_h_in: Field[[EdgeDim, KDim], float], - bound: float, + bound: int32, ) -> Field[[EdgeDim, KDim], float]: p_mflx_tracer_h_out = where( refin_ctrl == bound, p_mflx_tracer_h_in, where( - (p_mflx_tracer_h_in > 0.0) | (p_mflx_tracer_h_in == 0.0), + p_mflx_tracer_h_in >= 0.0, p_mflx_tracer_h_in * r_m(E2C[0]), p_mflx_tracer_h_in * r_m(E2C[1]), ), @@ -38,10 +38,10 @@ def _hflx_limiter_pd_stencil_02( @program def hflx_limiter_pd_stencil_02( - refin_ctrl: Field[[EdgeDim], float], + refin_ctrl: Field[[EdgeDim], int32], r_m: Field[[CellDim, KDim], float], p_mflx_tracer_h: Field[[EdgeDim, KDim], float], - bound: float, + bound: int32, ): _hflx_limiter_pd_stencil_02( refin_ctrl, diff --git a/advection/src/icon4py/advection/step_advection_stencil_01.py b/advection/src/icon4py/advection/step_advection_stencil_01.py new file mode 100644 index 0000000000..0a2ac7e3c8 --- /dev/null +++ b/advection/src/icon4py/advection/step_advection_stencil_01.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 + +from functional.common import Field +from functional.ffront.decorator import field_operator, program + +from icon4py.common.dimension import CellDim, KDim, Koff + + +@field_operator +def _step_advection_stencil_01( + rhodz_ast: Field[[CellDim, KDim], float], + p_mflx_contra_v: Field[[CellDim, KDim], float], + deepatmo_divzl: Field[[KDim], float], + deepatmo_divzu: Field[[KDim], float], + p_dtime: float, +) -> Field[[CellDim, KDim], float]: + k_offset_up_low = p_dtime * ( + p_mflx_contra_v(Koff[1]) * deepatmo_divzl - p_mflx_contra_v * deepatmo_divzu + ) + return rhodz_ast + k_offset_up_low + + +@program +def step_advection_stencil_01( + rhodz_ast: Field[[CellDim, KDim], float], + p_mflx_contra_v: Field[[CellDim, KDim], float], + deepatmo_divzl: Field[[KDim], float], + deepatmo_divzu: Field[[KDim], float], + p_dtime: float, + rhodz_ast2: Field[[CellDim, KDim], float], +): + _step_advection_stencil_01( + rhodz_ast, + p_mflx_contra_v, + deepatmo_divzl, + deepatmo_divzu, + p_dtime, + out=rhodz_ast2, + ) diff --git a/advection/src/icon4py/advection/step_advection_stencil_02.py b/advection/src/icon4py/advection/step_advection_stencil_02.py new file mode 100644 index 0000000000..e004903f5a --- /dev/null +++ b/advection/src/icon4py/advection/step_advection_stencil_02.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 + +from functional.common import Field +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import maximum + +from icon4py.common.dimension import CellDim, KDim, Koff + + +@field_operator +def _step_advection_stencil_02( + p_rhodz_new: Field[[CellDim, KDim], float], + p_mflx_contra_v: Field[[CellDim, KDim], float], + deepatmo_divzl: Field[[KDim], float], + deepatmo_divzu: Field[[KDim], float], + p_dtime: float, +) -> Field[[CellDim, KDim], float]: + return maximum(0.1 * p_rhodz_new, p_rhodz_new) - p_dtime * ( + p_mflx_contra_v(Koff[1]) * deepatmo_divzl - p_mflx_contra_v * deepatmo_divzu + ) + + +@program +def step_advection_stencil_02( + p_rhodz_new: Field[[CellDim, KDim], float], + p_mflx_contra_v: Field[[CellDim, KDim], float], + deepatmo_divzl: Field[[KDim], float], + deepatmo_divzu: Field[[KDim], float], + p_dtime: float, + rhodz_ast2: Field[[CellDim, KDim], float], +): + _step_advection_stencil_02( + p_rhodz_new, + p_mflx_contra_v, + deepatmo_divzl, + deepatmo_divzu, + p_dtime, + out=rhodz_ast2, + ) diff --git a/advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py b/advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py new file mode 100644 index 0000000000..f1c8abca13 --- /dev/null +++ b/advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py @@ -0,0 +1,49 @@ +# 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 icon4py.common.dimension import C2CEC, C2E2C, CECDim, CellDim, KDim + + +@field_operator +def _upwind_hflux_miura_stencil_02( + p_cc: Field[[CellDim, KDim], float], + lsq_pseudoinv_1: Field[[CECDim], float], + lsq_pseudoinv_2: Field[[CECDim], float], +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + p_coeff_1 = ( + lsq_pseudoinv_1(C2CEC[0]) * (p_cc(C2E2C[0]) - p_cc) + + lsq_pseudoinv_1(C2CEC[1]) * (p_cc(C2E2C[1]) - p_cc) + + lsq_pseudoinv_1(C2CEC[2]) * (p_cc(C2E2C[2]) - p_cc) + ) + p_coeff_2 = ( + lsq_pseudoinv_2(C2CEC[0]) * (p_cc(C2E2C[0]) - p_cc) + + lsq_pseudoinv_2(C2CEC[1]) * (p_cc(C2E2C[1]) - p_cc) + + lsq_pseudoinv_2(C2CEC[2]) * (p_cc(C2E2C[2]) - p_cc) + ) + return p_coeff_1, p_coeff_2 + + +@program +def upwind_hflux_miura_stencil_02( + p_cc: Field[[CellDim, KDim], float], + lsq_pseudoinv_1: Field[[CECDim], float], + lsq_pseudoinv_2: Field[[CECDim], float], + p_coeff_1: Field[[CellDim, KDim], float], + p_coeff_2: Field[[CellDim, KDim], float], +): + _upwind_hflux_miura_stencil_02( + p_cc, lsq_pseudoinv_1, lsq_pseudoinv_2, out=(p_coeff_1, p_coeff_2) + ) diff --git a/advection/src/icon4py/advection/upwind_vflux_ppm_stencil_01.py b/advection/src/icon4py/advection/upwind_vflux_ppm_stencil_01.py new file mode 100644 index 0000000000..6a027a7482 --- /dev/null +++ b/advection/src/icon4py/advection/upwind_vflux_ppm_stencil_01.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.common import Field +from functional.ffront.decorator import field_operator, program + +from icon4py.common.dimension import CellDim, KDim + + +@field_operator +def _upwind_vflux_ppm_stencil_01( + z_face_up: Field[[CellDim, KDim], float], + z_face_low: Field[[CellDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + z_delta_q = 0.5 * (z_face_up - z_face_low) + z_a1 = p_cc - 0.5 * (z_face_up + z_face_low) + + return z_delta_q, z_a1 + + +@program +def upwind_vflux_ppm_stencil_01( + z_face_up: Field[[CellDim, KDim], float], + z_face_low: Field[[CellDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], + z_delta_q: Field[[CellDim, KDim], float], + z_a1: Field[[CellDim, KDim], float], +): + _upwind_vflux_ppm_stencil_01(z_face_up, z_face_low, p_cc, out=(z_delta_q, z_a1)) diff --git a/advection/tests/test_hflx_limiter_mo_stencil_02.py b/advection/tests/test_hflx_limiter_mo_stencil_02.py new file mode 100644 index 0000000000..13ed590a59 --- /dev/null +++ b/advection/tests/test_hflx_limiter_mo_stencil_02.py @@ -0,0 +1,103 @@ +# 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 numpy import int32 + +from icon4py.advection.hflx_limiter_mo_stencil_02 import ( + hflx_limiter_mo_stencil_02, +) +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import constant_field, random_field, zero_field + + +def hflx_limiter_mo_stencil_02_numpy( + refin_ctrl: np.ndarray, + p_cc: np.ndarray, + z_tracer_new_low: np.ndarray, + lo_bound: float, + hi_bound: float, +): + refin_ctrl = np.expand_dims(refin_ctrl, axis=1) + condition = np.logical_or( + np.equal(refin_ctrl, lo_bound * np.ones(refin_ctrl.shape, dtype=int32)), + np.equal(refin_ctrl, hi_bound * np.ones(refin_ctrl.shape, dtype=int32)), + ) + z_tracer_new_tmp = np.where( + condition, + np.minimum(1.1 * p_cc, np.maximum(0.9 * p_cc, z_tracer_new_low)), + z_tracer_new_low, + ) + z_tracer_max = np.where( + condition, np.maximum(p_cc, z_tracer_new_tmp), z_tracer_new_low + ) + z_tracer_min = np.where( + condition, np.minimum(p_cc, z_tracer_new_tmp), z_tracer_new_low + ) + return z_tracer_max, z_tracer_min + + +def test_hflx_limiter_mo_stencil_02_some_matching_condition(): + mesh = SimpleMesh() + hi_bound = np.int32(1) + lo_bound = np.int32(5) + refin_ctrl = constant_field(mesh, hi_bound, CellDim, dtype=int32) + refin_ctrl[0:2] = np.int32(3) + p_cc = random_field(mesh, CellDim, KDim) + z_tracer_new_low = random_field(mesh, CellDim, KDim) + z_tracer_max = zero_field(mesh, CellDim, KDim) + z_tracer_min = zero_field(mesh, CellDim, KDim) + ref_max, ref_min = hflx_limiter_mo_stencil_02_numpy( + np.asarray(refin_ctrl), + np.asarray(p_cc), + np.asarray(z_tracer_new_low), + lo_bound, + hi_bound, + ) + + hflx_limiter_mo_stencil_02( + refin_ctrl, + p_cc, + z_tracer_new_low, + lo_bound, + hi_bound, + z_tracer_max, + z_tracer_min, + offset_provider={}, + ) + assert np.allclose(z_tracer_max, ref_max) + assert np.allclose(z_tracer_min, ref_min) + + +def test_hflx_limiter_mo_stencil_02_none_matching_condition(): + mesh = SimpleMesh() + hi_bound = np.int32(3) + lo_bound = np.int32(1) + refin_ctrl = constant_field(mesh, 2, CellDim, dtype=int32) + p_cc = random_field(mesh, CellDim, KDim) + z_tracer_new_low = random_field(mesh, CellDim, KDim) + z_tracer_max = zero_field(mesh, CellDim, KDim) + z_tracer_min = zero_field(mesh, CellDim, KDim) + hflx_limiter_mo_stencil_02( + refin_ctrl, + p_cc, + z_tracer_new_low, + lo_bound, + hi_bound, + z_tracer_max, + z_tracer_min, + offset_provider={}, + ) + assert np.allclose(z_tracer_min, z_tracer_new_low) + assert np.allclose(z_tracer_max, z_tracer_new_low) diff --git a/advection/tests/test_hflx_limiter_mo_stencil_03.py b/advection/tests/test_hflx_limiter_mo_stencil_03.py new file mode 100644 index 0000000000..f696b358b3 --- /dev/null +++ b/advection/tests/test_hflx_limiter_mo_stencil_03.py @@ -0,0 +1,126 @@ +# 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.advection.hflx_limiter_mo_stencil_03 import ( + hflx_limiter_mo_stencil_03, + hflx_limiter_mo_stencil_03_min_max, +) +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + + +def hflx_limiter_mo_stencil_03_numpy( + c2e2c: np.ndarray, + z_tracer_max: np.ndarray, + z_tracer_min: np.ndarray, + beta_fct: float, + r_beta_fct: float, + z_mflx_anti_in: np.ndarray, + z_mflx_anti_out: np.ndarray, + z_tracer_new_low: np.ndarray, + dbl_eps: float, +): + z_max, z_min = hflx_limiter_mo_stencil_03_min_max_numpy( + c2e2c, z_tracer_max, z_tracer_min, beta_fct, r_beta_fct + ) + r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in + dbl_eps) + r_m = (z_tracer_new_low - z_min) / (z_mflx_anti_out * dbl_eps) + return r_p, r_m + + +def hflx_limiter_mo_stencil_03_min_max_numpy( + c2e2c: np.array, + z_tracer_max: np.ndarray, + z_tracer_min: np.ndarray, + beta_fct: float, + r_beta_fct: float, +) -> tuple[np.ndarray]: + z_max = beta_fct * np.maximum(np.max(z_tracer_max[c2e2c], axis=1), z_tracer_max) + z_min = r_beta_fct * np.minimum(np.min(z_tracer_min[c2e2c], axis=1), z_tracer_min) + return z_max, z_min + + +def test_hflx_diffusion_mo_stencil_03_min_max(): + mesh = SimpleMesh() + z_tracer_max = random_field(mesh, CellDim, KDim) + z_tracer_min = random_field(mesh, CellDim, KDim) + z_max = zero_field(mesh, CellDim, KDim) + z_min = zero_field(mesh, CellDim, KDim) + beta_fct = 0.9 + r_beta_fct = 0.3 + z_max_ref, z_min_ref = hflx_limiter_mo_stencil_03_min_max_numpy( + mesh.c2e2c, + np.asarray(z_tracer_max), + np.asarray(z_tracer_min), + beta_fct, + r_beta_fct, + ) + hflx_limiter_mo_stencil_03_min_max( + z_tracer_max, + z_tracer_min, + beta_fct, + r_beta_fct, + z_max, + z_min, + offset_provider={"C2E2C": mesh.get_c2e2c_offset_provider()}, + ) + assert np.allclose(z_max, z_max_ref) + assert np.allclose(z_min, z_min_ref) + + +def test_hflx_diffusion_mo_stencil_03(): + mesh = SimpleMesh() + z_tracer_max = random_field(mesh, CellDim, KDim) + z_tracer_min = random_field(mesh, CellDim, KDim) + z_max = zero_field(mesh, CellDim, KDim) + z_min = zero_field(mesh, CellDim, KDim) + beta_fct = 0.4 + r_beta_fct = 0.6 + z_mflx_anti_in = random_field(mesh, CellDim, KDim) + z_mflx_anti_out = random_field(mesh, CellDim, KDim) + z_tracer_new_low = random_field(mesh, CellDim, KDim) + dbl_eps = 1e-5 + r_p = zero_field(mesh, CellDim, KDim) + r_m = zero_field(mesh, CellDim, KDim) + + r_p_ref, r_m_ref = hflx_limiter_mo_stencil_03_numpy( + mesh.c2e2c, + np.asarray(z_tracer_max), + np.asarray(z_tracer_min), + beta_fct, + r_beta_fct, + np.asarray(z_mflx_anti_in), + np.asarray(z_mflx_anti_out), + np.asarray(z_tracer_new_low), + dbl_eps, + ) + + hflx_limiter_mo_stencil_03( + z_tracer_max, + z_tracer_min, + beta_fct, + r_beta_fct, + z_max, + z_min, + z_mflx_anti_in, + z_mflx_anti_out, + z_tracer_new_low, + dbl_eps, + r_p, + r_m, + offset_provider={"C2E2C": mesh.get_c2e2c_offset_provider()}, + ) + np.allclose(r_p_ref, r_p) + np.allclose(r_m_ref, r_m) diff --git a/advection/tests/test_hflx_limiter_mo_stencil_04.py b/advection/tests/test_hflx_limiter_mo_stencil_04.py new file mode 100644 index 0000000000..bdb7211b88 --- /dev/null +++ b/advection/tests/test_hflx_limiter_mo_stencil_04.py @@ -0,0 +1,61 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.advection.hflx_limiter_mo_stencil_04 import ( + hflx_limiter_mo_stencil_04, +) +from icon4py.common.dimension import CellDim, EdgeDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + + +def hflx_limiter_mo_stencil_04_numpy( + e2c: np.ndarray, + z_anti: np.ndarray, + r_m: np.ndarray, + r_p: np.ndarray, + z_mflx_low: np.ndarray, +): + r_frac = np.where( + z_anti >= 0, + np.minimum(r_m[e2c[:, 0]], r_p[e2c[:, 1]]), + np.minimum(r_m[e2c[:, 1]], r_p[e2c[:, 0]]), + ) + return z_mflx_low + np.minimum(1.0, r_frac) * z_anti + + +def test_hflx_limiter_mo_stencil_04(): + mesh = SimpleMesh() + z_anti = random_field(mesh, EdgeDim, KDim, low=-2.0, high=2.0) + r_m = random_field(mesh, CellDim, KDim) + r_p = random_field(mesh, CellDim, KDim) + z_mflx_low = random_field(mesh, EdgeDim, KDim) + p_mflx_tracer_h = zero_field(mesh, EdgeDim, KDim) + ref = hflx_limiter_mo_stencil_04_numpy( + mesh.e2c, + np.asarray(z_anti), + np.asarray(r_m), + np.asarray(r_p), + np.asarray(z_mflx_low), + ) + hflx_limiter_mo_stencil_04( + z_anti, + r_m, + r_p, + z_mflx_low, + p_mflx_tracer_h, + offset_provider={"E2C": mesh.get_e2c_offset_provider()}, + ) + assert np.allclose(p_mflx_tracer_h, ref) diff --git a/advection/tests/test_hflx_limiter_pd_stencil_02.py b/advection/tests/test_hflx_limiter_pd_stencil_02.py index fdb93587cb..57eb18aa59 100644 --- a/advection/tests/test_hflx_limiter_pd_stencil_02.py +++ b/advection/tests/test_hflx_limiter_pd_stencil_02.py @@ -18,7 +18,7 @@ ) from icon4py.common.dimension import CellDim, EdgeDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh -from icon4py.testutils.utils import random_field +from icon4py.testutils.utils import constant_field, random_field def hflx_limiter_pd_stencil_02_numpy( @@ -42,14 +42,12 @@ def hflx_limiter_pd_stencil_02_numpy( return p_mflx_tracer_h_out -def test_hflx_limiter_pd_stencil_02(): +def test_hflx_limiter_pd_stencil_02_nowhere_matching_refin_ctl(): mesh = SimpleMesh() - - refin_ctrl = random_field(mesh, EdgeDim) + bound = np.int32(7) + refin_ctrl = constant_field(mesh, 4, EdgeDim, dtype=np.int32) r_m = random_field(mesh, CellDim, KDim) p_mflx_tracer_h_in = random_field(mesh, EdgeDim, KDim) - p_mflx_tracer_h_out = random_field(mesh, EdgeDim, KDim) - bound = np.float64(7.0) ref = hflx_limiter_pd_stencil_02_numpy( mesh.e2c, @@ -63,10 +61,67 @@ def test_hflx_limiter_pd_stencil_02(): refin_ctrl, r_m, p_mflx_tracer_h_in, - p_mflx_tracer_h_out, bound, offset_provider={ "E2C": mesh.get_e2c_offset_provider(), }, ) - assert np.allclose(p_mflx_tracer_h_out, ref) + assert np.allclose(p_mflx_tracer_h_in, ref) + + +def test_hflx_limiter_pd_stencil_02_everywhere_matching_refin_ctl(): + mesh = SimpleMesh() + bound = np.int32(7) + refin_ctrl = constant_field(mesh, bound, EdgeDim, dtype=np.int32) + r_m = random_field(mesh, CellDim, KDim) + p_mflx_tracer_h_in = random_field(mesh, EdgeDim, KDim) + + hflx_limiter_pd_stencil_02( + refin_ctrl, + r_m, + p_mflx_tracer_h_in, + bound, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + assert np.allclose(p_mflx_tracer_h_in, p_mflx_tracer_h_in) + + +def test_hflx_limiter_pd_stencil_02_partly_matching_refin_ctl(): + mesh = SimpleMesh() + bound = np.int32(4) + refin_ctrl = constant_field(mesh, 5, EdgeDim, dtype=np.int32) + refin_ctrl[2:6] = bound + r_m = random_field(mesh, CellDim, KDim) + p_mflx_tracer_h_in = random_field(mesh, EdgeDim, KDim) + + hflx_limiter_pd_stencil_02( + refin_ctrl, + r_m, + p_mflx_tracer_h_in, + bound, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + assert np.allclose(p_mflx_tracer_h_in, p_mflx_tracer_h_in) + + +def test_hflx_limiter_pd_stencil_02_everywhere_matching_refin_ctl_does_not_change_inout_arg(): + mesh = SimpleMesh() + bound = np.int32(7) + refin_ctrl = constant_field(mesh, bound, EdgeDim, dtype=np.int32) + r_m = random_field(mesh, CellDim, KDim) + p_mflx_tracer_h_in = random_field(mesh, EdgeDim, KDim) + + hflx_limiter_pd_stencil_02( + refin_ctrl, + r_m, + p_mflx_tracer_h_in, + bound, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + assert np.allclose(p_mflx_tracer_h_in, p_mflx_tracer_h_in) diff --git a/advection/tests/test_step_advection_stencil_01.py b/advection/tests/test_step_advection_stencil_01.py new file mode 100644 index 0000000000..fc2290de19 --- /dev/null +++ b/advection/tests/test_step_advection_stencil_01.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 numpy as np + +from icon4py.advection.step_advection_stencil_01 import step_advection_stencil_01 +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + + +def step_advection_stencil_01_numpy( + rhodz_ast: np.ndarray, + p_mflx_contra_v: np.ndarray, + deepatmo_divzl: np.ndarray, + deepatmo_divzu: np.ndarray, + pd_time: float, +) -> np.ndarray: + p_mlfx_contra_v1 = np.roll(p_mflx_contra_v, axis=1, shift=-1) + tmp = pd_time * ( + p_mlfx_contra_v1 * deepatmo_divzl - p_mflx_contra_v * deepatmo_divzu + ) + return rhodz_ast + tmp + + +def test_step_advection_stencil_01(): + mesh = SimpleMesh() + rhodz_ast = random_field(mesh, CellDim, KDim) + p_mflx_contra = random_field(mesh, CellDim, KDim) + deepatmo_divzl = random_field(mesh, KDim) + deepatmo_divzu = random_field(mesh, KDim) + result = zero_field(mesh, CellDim, KDim) + p_dtime = 0.1 + + ref = step_advection_stencil_01_numpy( + np.asarray(rhodz_ast), + np.asarray(p_mflx_contra), + np.asarray(deepatmo_divzl), + np.asarray(deepatmo_divzu), + p_dtime, + ) + + step_advection_stencil_01( + rhodz_ast, + p_mflx_contra, + deepatmo_divzl, + deepatmo_divzu, + p_dtime, + result, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(ref[:, :-1], result[:, :-1]) diff --git a/advection/tests/test_step_advection_stencil_02.py b/advection/tests/test_step_advection_stencil_02.py new file mode 100644 index 0000000000..a9a5e99ca9 --- /dev/null +++ b/advection/tests/test_step_advection_stencil_02.py @@ -0,0 +1,64 @@ +# 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.advection.step_advection_stencil_02 import step_advection_stencil_02 +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + + +def step_advection_stencil_02_numpy( + rhodz_new: np.ndarray, + p_mflx_contra_v: np.ndarray, + deepatmo_divzl: np.ndarray, + deepatmo_divzu: np.ndarray, + pd_time: float, +) -> np.ndarray: + + tmp = ( + np.roll(p_mflx_contra_v, axis=1, shift=-1) * deepatmo_divzl + - p_mflx_contra_v * deepatmo_divzu + ) + return np.maximum(0.1 * rhodz_new, rhodz_new) - pd_time * tmp + + +def test_step_advection_stencil_02(): + mesh = SimpleMesh() + rhodz_ast = random_field(mesh, CellDim, KDim) + p_mflx_contra = random_field(mesh, CellDim, KDim) + deepatmo_divzl = random_field(mesh, KDim) + deepatmo_divzu = random_field(mesh, KDim) + result = zero_field(mesh, CellDim, KDim) + p_dtime = 0.1 + + ref = step_advection_stencil_02_numpy( + np.asarray(rhodz_ast), + np.asarray(p_mflx_contra), + np.asarray(deepatmo_divzl), + np.asarray(deepatmo_divzu), + p_dtime, + ) + + step_advection_stencil_02( + rhodz_ast, + p_mflx_contra, + deepatmo_divzl, + deepatmo_divzu, + p_dtime, + result, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(ref[:, :-1], result[:, :-1]) diff --git a/advection/tests/test_upwind_hflux_miura_stencil_02.py b/advection/tests/test_upwind_hflux_miura_stencil_02.py new file mode 100644 index 0000000000..3dae540683 --- /dev/null +++ b/advection/tests/test_upwind_hflux_miura_stencil_02.py @@ -0,0 +1,72 @@ +# 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.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.advection.upwind_hflux_miura_stencil_02 import ( + upwind_hflux_miura_stencil_02, +) +from icon4py.common.dimension import C2E2CDim, CECDim, CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import as_1D_sparse_field, random_field, zero_field + + +def upwind_hflux_miura_stencil_02_numpy( + c2e2c: np.ndarray, + p_cc: np.ndarray, + lsq_pseudoinv_1: np.ndarray, + lsq_pseudoinv_2: np.ndarray, +) -> tuple[np.ndarray]: + p_cc_e = np.expand_dims(p_cc, axis=1) + n_diff = p_cc[c2e2c] - p_cc_e + lsq_pseudoinv_2 = np.expand_dims(lsq_pseudoinv_2, axis=-1) + lsq_pseudoinv_1 = np.expand_dims(lsq_pseudoinv_1, axis=-1) + p_coeff_1 = np.sum(lsq_pseudoinv_1 * n_diff, axis=1) + p_coeff_2 = np.sum(lsq_pseudoinv_2 * n_diff, axis=1) + return p_coeff_1, p_coeff_2 + + +def test_upwind_hflux_miura_stencil_02(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + lsq_pseudoinv_1 = random_field(mesh, CellDim, C2E2CDim) + lsq_pseudoinv_1_field = as_1D_sparse_field(lsq_pseudoinv_1, CECDim) + + lsq_pseudoinv_2 = random_field(mesh, CellDim, C2E2CDim) + lsq_pseudoinv_2_field = as_1D_sparse_field(lsq_pseudoinv_2, CECDim) + p_coeff_1 = zero_field(mesh, CellDim, KDim) + p_coeff_2 = zero_field(mesh, CellDim, KDim) + + ref_1, ref_2 = upwind_hflux_miura_stencil_02_numpy( + mesh.c2e2c, + np.asarray(p_cc), + np.asarray(lsq_pseudoinv_1), + np.asarray(lsq_pseudoinv_2), + ) + + upwind_hflux_miura_stencil_02( + p_cc, + lsq_pseudoinv_1_field, + lsq_pseudoinv_2_field, + p_coeff_1, + p_coeff_2, + offset_provider={ + "C2E2C": mesh.get_c2e2c_offset_provider(), + "C2CEC": StridedNeighborOffsetProvider(CellDim, CECDim, mesh.n_c2e2c), + }, + ) + co1 = np.asarray(p_coeff_1) + co2 = np.asarray(p_coeff_2) + assert np.allclose(ref_1, co1) + assert np.allclose(ref_2, co2) diff --git a/advection/tests/test_upwind_vflux_ppm_stencil_01.py b/advection/tests/test_upwind_vflux_ppm_stencil_01.py new file mode 100644 index 0000000000..36f0d8246e --- /dev/null +++ b/advection/tests/test_upwind_vflux_ppm_stencil_01.py @@ -0,0 +1,49 @@ +# 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.advection.upwind_vflux_ppm_stencil_01 import ( + upwind_vflux_ppm_stencil_01, +) +from icon4py.common.dimension import CellDim, KDim +from icon4py.testutils.simple_mesh import SimpleMesh +from icon4py.testutils.utils import random_field, zero_field + + +def upwind_vflux_ppm_stencil_01_numpy( + z_face_up: np.ndarray, z_face_low: np.ndarray, p_cc: np.ndarray +) -> tuple[np.ndarray]: + z_delta_q = 0.5 * (z_face_up - z_face_low) + z_a1 = p_cc - 0.5 * (z_face_up + z_face_low) + return z_delta_q, z_a1 + + +def test_upwind_vflux_ppm_stencil_01(): + mesh = SimpleMesh() + z_face_up = random_field(mesh, CellDim, KDim) + z_face_down = random_field(mesh, CellDim, KDim) + p_cc = random_field(mesh, CellDim, KDim) + z_delta_q = zero_field(mesh, CellDim, KDim) + z_a1 = zero_field(mesh, CellDim, KDim) + + ref_z_delta_q, ref_z_a1 = upwind_vflux_ppm_stencil_01_numpy( + np.asarray(z_face_up), np.asarray(z_face_down), np.asarray(p_cc) + ) + + upwind_vflux_ppm_stencil_01( + z_face_up, z_face_down, p_cc, z_delta_q, z_a1, offset_provider={} + ) + + assert np.allclose(ref_z_delta_q, z_delta_q) + assert np.allclose(ref_z_a1, z_a1) diff --git a/common/src/icon4py/common/dimension.py b/common/src/icon4py/common/dimension.py index 499bd5ac0d..53f7904193 100644 --- a/common/src/icon4py/common/dimension.py +++ b/common/src/icon4py/common/dimension.py @@ -22,6 +22,7 @@ CEDim = Dimension("CE") ECDim = Dimension("EC") ECVDim = Dimension("ECV") +CECDim = Dimension("CEC") E2CDim = Dimension("E2C", DimensionKind.LOCAL) E2VDim = Dimension("E2V", DimensionKind.LOCAL) C2EDim = Dimension("C2E", DimensionKind.LOCAL) @@ -45,4 +46,6 @@ E2C2EO = FieldOffset("E2C2EO", source=EdgeDim, target=(EdgeDim, E2C2EODim)) E2C2E = FieldOffset("E2C2E", source=EdgeDim, target=(EdgeDim, E2C2EDim)) C2E2C = FieldOffset("C2E2C", source=CellDim, target=(CellDim, C2E2CDim)) +C2CEC = FieldOffset("C2CEC", source=CECDim, target=(CellDim, C2E2CDim)) + Koff = FieldOffset("Koff", source=KDim, target=(KDim,))