From 7696d1a08b2f6ae9d700e7a43516f816ae0034d0 Mon Sep 17 00:00:00 2001 From: ninaburg <83002751+ninaburg@users.noreply.github.com> Date: Thu, 25 Jan 2024 11:25:14 +0100 Subject: [PATCH] Refactor advection stencils to use `StencilTest` - Part 2 (#351) --- .../divide_flux_area_list_stencil_01.py | 3 +- .../divide_flux_area_list_stencil_02.py | 3 +- .../hflux_ffsl_hybrid_stencil_01a.py | 6 +- .../advection/hflux_ffsl_hybrid_stencil_02.py | 3 +- .../prep_gauss_quadrature_c_list_stencil.py | 3 +- .../prep_gauss_quadrature_c_stencil.py | 3 +- .../advection/recon_lsq_cell_c_svd_stencil.py | 4 +- .../advection/recon_lsq_cell_l_svd_stencil.py | 4 +- .../advection/step_advection_stencil_01.py | 4 +- .../advection/step_advection_stencil_02.py | 4 +- .../advection/step_advection_stencil_03.py | 3 +- .../advection/step_advection_stencil_04.py | 3 +- .../upwind_hflux_miura3_stencil_01.py | 4 +- .../upwind_hflux_miura_cycl_stencil_01.py | 4 +- .../upwind_hflux_miura_cycl_stencil_02.py | 3 +- .../upwind_hflux_miura_cycl_stencil_03a.py | 4 +- .../upwind_hflux_miura_cycl_stencil_03b.py | 4 +- .../upwind_hflux_miura_stencil_01.py | 4 +- .../test_divide_flux_area_list_stencil_01.py | 2036 +++++++++++------ .../test_divide_flux_area_list_stencil_02.py | 562 +++-- .../test_hflux_ffsl_hybrid_stencil_01a.py | 336 ++- .../test_hflux_ffsl_hybrid_stencil_02.py | 64 +- ...st_prep_gauss_quadrature_c_list_stencil.py | 1028 +++++---- .../test_prep_gauss_quadrature_c_stencil.py | 944 ++++---- .../test_recon_lsq_cell_c_svd_stencil.py | 586 +++-- .../test_recon_lsq_cell_l_svd_stencil.py | 107 +- .../test_step_advection_stencil_01.py | 78 +- .../test_step_advection_stencil_02.py | 79 +- .../test_step_advection_stencil_03.py | 62 +- .../test_step_advection_stencil_04.py | 58 +- .../test_upwind_hflux_miura3_stencil_01.py | 356 ++- ...test_upwind_hflux_miura_cycl_stencil_01.py | 138 +- ...test_upwind_hflux_miura_cycl_stencil_02.py | 139 +- ...est_upwind_hflux_miura_cycl_stencil_03a.py | 58 +- ...est_upwind_hflux_miura_cycl_stencil_03b.py | 65 +- .../test_upwind_hflux_miura_stencil_01.py | 138 +- .../src/icon4py/model/common/grid/icon.py | 8 + .../src/icon4py/model/common/grid/simple.py | 7 + 38 files changed, 3906 insertions(+), 3009 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py index c18d8eac86..2fa1051c61 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py @@ -13,6 +13,7 @@ import sys +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where @@ -665,7 +666,7 @@ def _divide_flux_area_list_stencil_01( ) -@program +@program(grid_type=GridType.UNSTRUCTURED) def divide_flux_area_list_stencil_01( famask_int: Field[[EdgeDim, KDim], int32], p_vn: Field[[EdgeDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py index 08650dade8..ec0d9e0ad4 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py @@ -13,6 +13,7 @@ import sys +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where @@ -174,7 +175,7 @@ def _divide_flux_area_list_stencil_02( ) -@program +@program(grid_type=GridType.UNSTRUCTURED) def divide_flux_area_list_stencil_02( famask_int: Field[[EdgeDim, KDim], int32], p_vn: Field[[EdgeDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py index 6047efeaa7..154bb669e5 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py @@ -11,9 +11,9 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import Field +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import int32, where +from gt4py.next.ffront.fbuiltins import Field, int32, where from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim @@ -108,7 +108,7 @@ def _hflux_ffsl_hybrid_stencil_01a( return p_out_e_hybrid_1a -@program +@program(grid_type=GridType.UNSTRUCTURED) def hflux_ffsl_hybrid_stencil_01a( z_lsq_coeff_1: Field[[CellDim, KDim], float], z_lsq_coeff_2: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py index 6f6477cd84..d6c7851ca5 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py @@ -11,6 +11,7 @@ # # 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 @@ -28,7 +29,7 @@ def _hflux_ffsl_hybrid_stencil_02( return p_out_e_hybrid_2 -@program +@program(grid_type=GridType.UNSTRUCTURED) def hflux_ffsl_hybrid_stencil_02( p_out_e_hybrid_2: Field[[EdgeDim, KDim], float], p_mass_flx_e: Field[[EdgeDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py index 55898089a9..7bd6a261ba 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py @@ -11,6 +11,7 @@ # # 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, where @@ -323,7 +324,7 @@ def _prep_gauss_quadrature_c_list_stencil( ) -@program +@program(grid_type=GridType.UNSTRUCTURED) def prep_gauss_quadrature_c_list_stencil( famask_int: Field[[EdgeDim, KDim], int32], p_coords_dreg_v_1_x: Field[[EdgeDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py index 0f249751cb..1b4d69061f 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py @@ -11,6 +11,7 @@ # # 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 ( # noqa: A004 # import gt4py builtin Field, @@ -293,7 +294,7 @@ def _prep_gauss_quadrature_c_stencil( ) -@program +@program(grid_type=GridType.UNSTRUCTURED) def prep_gauss_quadrature_c_stencil( p_coords_dreg_v_1_x: Field[[EdgeDim, KDim], float], p_coords_dreg_v_2_x: Field[[EdgeDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py index 6dfb719abe..7fba83fd5b 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import Field +from gt4py.next.common import Field, GridType from gt4py.next.ffront.decorator import field_operator, program from icon4py.model.common.dimension import C2CECEC, C2E2C2E2C, CECECDim, CellDim, KDim @@ -176,7 +176,7 @@ def _recon_lsq_cell_c_svd_stencil( ) -@program +@program(grid_type=GridType.UNSTRUCTURED) def recon_lsq_cell_c_svd_stencil( p_cc: Field[[CellDim, KDim], float], lsq_pseudoinv_1: Field[[CECECDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py index 5e85447556..db468a2515 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import Field +from gt4py.next.common import Field, GridType from gt4py.next.ffront.decorator import field_operator, program from icon4py.model.common.dimension import C2CEC, C2E2C, CECDim, CellDim, KDim @@ -41,7 +41,7 @@ def _recon_lsq_cell_l_svd_stencil( return p_coeff_1_dsl, p_coeff_2_dsl, p_coeff_3_dsl -@program +@program(grid_type=GridType.UNSTRUCTURED) def recon_lsq_cell_l_svd_stencil( p_cc: Field[[CellDim, KDim], float], lsq_pseudoinv_1: Field[[CECDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_01.py index 0ac10cc6a8..2949505bda 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_01.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import Field +from gt4py.next.common import Field, GridType from gt4py.next.ffront.decorator import field_operator, program from icon4py.model.common.dimension import CellDim, KDim, Koff @@ -31,7 +31,7 @@ def _step_advection_stencil_01( return rhodz_ast + k_offset_up_low -@program +@program(grid_type=GridType.UNSTRUCTURED) def step_advection_stencil_01( rhodz_ast: Field[[CellDim, KDim], float], p_mflx_contra_v: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_02.py index af201f9550..9d7f10e5c8 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_02.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import Field +from gt4py.next.common import Field, GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import maximum @@ -31,7 +31,7 @@ def _step_advection_stencil_02( ) -@program +@program(grid_type=GridType.UNSTRUCTURED) def step_advection_stencil_02( p_rhodz_new: Field[[CellDim, KDim], float], p_mflx_contra_v: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py index b1e775c62d..a865e013bf 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py @@ -11,6 +11,7 @@ # # 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, maximum @@ -27,7 +28,7 @@ def _step_advection_stencil_03( return p_tracer_new -@program +@program(grid_type=GridType.UNSTRUCTURED) def step_advection_stencil_03( p_tracer_now: Field[[CellDim, KDim], float], p_grf_tend_tracer: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py index 52e0c79402..4eebab8cde 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py @@ -11,6 +11,7 @@ # # 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 @@ -27,7 +28,7 @@ def _step_advection_stencil_04( return opt_ddt_tracer_adv -@program +@program(grid_type=GridType.UNSTRUCTURED) def step_advection_stencil_04( p_tracer_now: Field[[CellDim, KDim], float], p_tracer_new: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py index 1d2e822514..2df727e075 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import Field +from gt4py.next.common import Field, GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, where @@ -114,7 +114,7 @@ def _upwind_hflux_miura3_stencil_01( return p_out_e_miura3 -@program +@program(grid_type=GridType.UNSTRUCTURED) def upwind_hflux_miura3_stencil_01( z_lsq_coeff_1: Field[[CellDim, KDim], float], z_lsq_coeff_2: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py index 098617d113..45823abb00 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import Field +from gt4py.next.common import Field, GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, where @@ -51,7 +51,7 @@ def _upwind_hflux_miura_cycl_stencil_01( return z_tracer_mflx_dsl -@program +@program(grid_type=GridType.UNSTRUCTURED) def upwind_hflux_miura_cycl_stencil_01( z_lsq_coeff_1_dsl: Field[[CellDim, KDim], float], z_lsq_coeff_2_dsl: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py index 502fc5f573..1c9e1779de 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum @@ -48,7 +49,7 @@ def _upwind_hflux_miura_cycl_stencil_02( return (z_rhofluxdiv_c_out, z_fluxdiv_c_dsl, z_rho_new_dsl, z_tracer_new_dsl) -@program +@program(grid_type=GridType.UNSTRUCTURED) def upwind_hflux_miura_cycl_stencil_02( nsub: int32, p_mass_flx_e: Field[[EdgeDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py index 53cdc339d1..17a7416b31 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import Field +from gt4py.next.common import Field, GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import float64 @@ -27,7 +27,7 @@ def _upwind_hflux_miura_cycl_stencil_03a( return p_out_e -@program +@program(grid_type=GridType.UNSTRUCTURED) def upwind_hflux_miura_cycl_stencil_03a( z_tracer_mflx_1_dsl: Field[[EdgeDim, KDim], float], z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py index 4b5c73e401..51b8a8d117 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import Field +from gt4py.next.common import Field, GridType from gt4py.next.ffront.decorator import field_operator, program from icon4py.model.common.dimension import EdgeDim, KDim @@ -27,7 +27,7 @@ def _upwind_hflux_miura_cycl_stencil_03b( return p_out_e -@program +@program(grid_type=GridType.UNSTRUCTURED) def upwind_hflux_miura_cycl_stencil_03b( z_tracer_mflx_1_dsl: Field[[EdgeDim, KDim], float], z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py index 5e828de057..3edf4fc43a 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import Field +from gt4py.next.common import Field, GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, where @@ -39,7 +39,7 @@ def _upwind_hflux_miura_stencil_01( return p_out_e -@program +@program(grid_type=GridType.UNSTRUCTURED) def upwind_hflux_miura_stencil_01( z_lsq_coeff_1: Field[[CellDim, KDim], float], z_lsq_coeff_2: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py index 2ca14fd280..4a2258286c 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py @@ -19,25 +19,25 @@ divide_flux_area_list_stencil_01, ) from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid from icon4py.model.common.test_utils.helpers import ( + StencilTest, as_1D_sparse_field, random_field, random_mask, + reshape, zero_field, ) -# FUNCTIONS -# Checking turn when travelling along three points, used to check whether lines inters. -def ccw( - p0_lon: np.array, - p0_lat: np.array, - p1_lon: np.array, - p1_lat: np.array, - p2_lon: np.array, - p2_lat: np.array, -) -> np.array: +# Check whether lines inters. +def _ccw_numpy( + p0_lon, + p0_lat, + p1_lon, + p1_lat, + p2_lon, + p2_lat, +): dx1 = p1_lon - p0_lon dy1 = p1_lat - p0_lat @@ -52,25 +52,25 @@ def ccw( return ccw_out -# Checks whether two lines intersect -def lintersect( - line1_p1_lon: np.array, - line1_p1_lat: np.array, - line1_p2_lon: np.array, - line1_p2_lat: np.array, - line2_p1_lon: np.array, - line2_p1_lat: np.array, - line2_p2_lon: np.array, - line2_p2_lat: np.array, -) -> np.array: - intersect1 = ccw( +# Check whether two lines intersect +def _lintersect_numpy( + line1_p1_lon, + line1_p1_lat, + line1_p2_lon, + line1_p2_lat, + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, +): + intersect1 = _ccw_numpy( line1_p1_lon, line1_p1_lat, line1_p2_lon, line1_p2_lat, line2_p1_lon, line2_p1_lat, - ) * ccw( + ) * _ccw_numpy( line1_p1_lon, line1_p1_lat, line1_p2_lon, @@ -78,14 +78,14 @@ def lintersect( line2_p2_lon, line2_p2_lat, ) - intersect2 = ccw( + intersect2 = _ccw_numpy( line2_p1_lon, line2_p1_lat, line2_p2_lon, line2_p2_lat, line1_p1_lon, line1_p1_lat, - ) * ccw( + ) * _ccw_numpy( line2_p1_lon, line2_p1_lat, line2_p2_lon, @@ -99,16 +99,16 @@ def lintersect( # Compute intersection point of two lines in 2D -def line_intersect( - line1_p1_lon: np.array, - line1_p1_lat: np.array, - line1_p2_lon: np.array, - line1_p2_lat: np.array, - line2_p1_lon: np.array, - line2_p1_lat: np.array, - line2_p2_lon: np.array, - line2_p2_lat: np.array, -) -> tuple[np.array]: +def _line_intersect_numpy( + line1_p1_lon, + line1_p1_lat, + line1_p2_lon, + line1_p2_lat, + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, +): m1 = (line1_p2_lat - line1_p1_lat) / (line1_p2_lon - line1_p1_lon) m2 = (line2_p2_lat - line2_p1_lat) / (line2_p2_lon - line2_p1_lon) @@ -118,517 +118,378 @@ def line_intersect( return intersect_1, intersect_2 -def divide_flux_area_list_stencil_01_numpy( - e2c: np.array, - famask_int: np.array, - p_vn: np.array, - ptr_v3_lon: np.array, - ptr_v3_lat: np.array, - tangent_orientation_dsl: np.array, - dreg_patch0_1_lon_dsl: np.array, - dreg_patch0_1_lat_dsl: np.array, - dreg_patch0_2_lon_dsl: np.array, - dreg_patch0_2_lat_dsl: np.array, - dreg_patch0_3_lon_dsl: np.array, - dreg_patch0_3_lat_dsl: np.array, - dreg_patch0_4_lon_dsl: np.array, - dreg_patch0_4_lat_dsl: np.array, -): - ptr_v3_lon_e = np.expand_dims(ptr_v3_lon, axis=-1) - ptr_v3_lat_e = np.expand_dims(ptr_v3_lat, axis=-1) - ptr_v3_lon_e = np.expand_dims(ptr_v3_lon, axis=-1) - ptr_v3_lat_e = np.expand_dims(ptr_v3_lat, axis=-1) - tangent_orientation_dsl = np.expand_dims(tangent_orientation_dsl, axis=-1) - - arrival_pts_1_lon_dsl = dreg_patch0_1_lon_dsl - arrival_pts_1_lat_dsl = dreg_patch0_1_lat_dsl - arrival_pts_2_lon_dsl = dreg_patch0_2_lon_dsl - arrival_pts_2_lat_dsl = dreg_patch0_2_lat_dsl - depart_pts_1_lon_dsl = dreg_patch0_4_lon_dsl # indices have to be switched so that dep 1 belongs to arr 1 (and d2->a2) - depart_pts_1_lat_dsl = dreg_patch0_4_lat_dsl - depart_pts_2_lon_dsl = dreg_patch0_3_lon_dsl - depart_pts_2_lat_dsl = dreg_patch0_3_lat_dsl - - lvn_pos = np.where(p_vn >= 0.0, True, False) - - # get flux area departure-line segment - fl_line_p1_lon = depart_pts_1_lon_dsl - fl_line_p1_lat = depart_pts_1_lat_dsl - fl_line_p2_lon = depart_pts_2_lon_dsl - fl_line_p2_lat = depart_pts_2_lat_dsl - - # get triangle edge 1 (A1V3) - tri_line1_p1_lon = arrival_pts_1_lon_dsl - tri_line1_p1_lat = arrival_pts_1_lat_dsl - tri_line1_p2_lon = np.where( - lvn_pos, - np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), - np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), - ) - tri_line1_p2_lat = np.where( - lvn_pos, - np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), - np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), +class TestDivideFluxAreaListStencil01(StencilTest): + PROGRAM = divide_flux_area_list_stencil_01 + OUTPUTS = ( + "dreg_patch0_1_lon_dsl", + "dreg_patch0_1_lat_dsl", + "dreg_patch0_2_lon_dsl", + "dreg_patch0_2_lat_dsl", + "dreg_patch0_3_lon_dsl", + "dreg_patch0_3_lat_dsl", + "dreg_patch0_4_lon_dsl", + "dreg_patch0_4_lat_dsl", + "dreg_patch1_1_lon_vmask", + "dreg_patch1_1_lat_vmask", + "dreg_patch1_2_lon_vmask", + "dreg_patch1_2_lat_vmask", + "dreg_patch1_3_lon_vmask", + "dreg_patch1_3_lat_vmask", + "dreg_patch1_4_lon_vmask", + "dreg_patch1_4_lat_vmask", + "dreg_patch2_1_lon_vmask", + "dreg_patch2_1_lat_vmask", + "dreg_patch2_2_lon_vmask", + "dreg_patch2_2_lat_vmask", + "dreg_patch2_3_lon_vmask", + "dreg_patch2_3_lat_vmask", + "dreg_patch2_4_lon_vmask", + "dreg_patch2_4_lat_vmask", ) - # get triangle edge 2 (A2V3) - tri_line2_p1_lon = arrival_pts_2_lon_dsl - tri_line2_p1_lat = arrival_pts_2_lat_dsl - tri_line2_p2_lon = np.where( - lvn_pos, - np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), - np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), - ) - tri_line2_p2_lat = np.where( - lvn_pos, - np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), - np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), - ) + @staticmethod + def _generate_flux_area_geometry( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + p_vn, + ptr_v3_lon_e, + ptr_v3_lat_e, + ): + arrival_pts_1_lon_dsl = dreg_patch0_1_lon_dsl + arrival_pts_1_lat_dsl = dreg_patch0_1_lat_dsl + arrival_pts_2_lon_dsl = dreg_patch0_2_lon_dsl + arrival_pts_2_lat_dsl = dreg_patch0_2_lat_dsl + depart_pts_1_lon_dsl = dreg_patch0_4_lon_dsl + depart_pts_1_lat_dsl = dreg_patch0_4_lat_dsl + depart_pts_2_lon_dsl = dreg_patch0_3_lon_dsl + depart_pts_2_lat_dsl = dreg_patch0_3_lat_dsl - # Create first mask does departure-line segment intersects with A1V3 - lintersect_line1 = lintersect( - fl_line_p1_lon, - fl_line_p1_lat, - fl_line_p2_lon, - fl_line_p2_lat, - tri_line1_p1_lon, - tri_line1_p1_lat, - tri_line1_p2_lon, - tri_line1_p2_lat, - ) - # Create first mask does departure-line segment intersects with A2V3 - lintersect_line2 = lintersect( - fl_line_p1_lon, - fl_line_p1_lat, - fl_line_p2_lon, - fl_line_p2_lat, - tri_line2_p1_lon, - tri_line2_p1_lat, - tri_line2_p2_lon, - tri_line2_p2_lat, - ) + lvn_pos = np.where(p_vn >= 0.0, True, False) - lvn_sys_pos = np.where( - (p_vn * np.broadcast_to(tangent_orientation_dsl, p_vn.shape)) >= 0.0, - True, - False, - ) - famask_bool = np.where(famask_int == int32(1), True, False) - # ------------------------------------------------- Case 1 - mask_case1 = np.logical_and.reduce([lintersect_line1, lintersect_line2, famask_bool]) - ps1_x, ps1_y = line_intersect( - fl_line_p1_lon, - fl_line_p1_lat, - fl_line_p2_lon, - fl_line_p2_lat, - tri_line1_p1_lon, - tri_line1_p1_lat, - tri_line1_p2_lon, - tri_line1_p2_lat, - ) - ps2_x, ps2_y = line_intersect( - fl_line_p1_lon, - fl_line_p1_lat, - fl_line_p2_lon, - fl_line_p2_lat, - tri_line2_p1_lon, - tri_line2_p1_lat, - tri_line2_p2_lon, - tri_line2_p2_lat, - ) + fl_line_p1_lon = depart_pts_1_lon_dsl + fl_line_p1_lat = depart_pts_1_lat_dsl + fl_line_p2_lon = depart_pts_2_lon_dsl + fl_line_p2_lat = depart_pts_2_lat_dsl - # Case 1 - patch 0 - dreg_patch0_1_lon_dsl = np.where(mask_case1, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = np.where(mask_case1, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_2_lon_dsl = np.where( - mask_case1, - np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), - dreg_patch0_2_lon_dsl, - ) - dreg_patch0_2_lat_dsl = np.where( + tri_line1_p1_lon = arrival_pts_1_lon_dsl + tri_line1_p1_lat = arrival_pts_1_lat_dsl + tri_line1_p2_lon = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), + ) + tri_line1_p2_lat = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), + ) + + tri_line2_p1_lon = arrival_pts_2_lon_dsl + tri_line2_p1_lat = arrival_pts_2_lat_dsl + tri_line2_p2_lon = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), + ) + tri_line2_p2_lat = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), + ) + + return ( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + ) + + @staticmethod + def _apply_case1_patch0( mask_case1, - np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), - dreg_patch0_2_lat_dsl, - ) - dreg_patch0_3_lon_dsl = np.where(mask_case1, ps2_x, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = np.where(mask_case1, ps2_y, dreg_patch0_3_lat_dsl) - dreg_patch0_4_lon_dsl = np.where( + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + ps1_x, + ps1_y, + ps2_x, + ps2_y, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + ): + dreg_patch0_1_lon_dsl = arrival_pts_1_lon_dsl + dreg_patch0_1_lat_dsl = arrival_pts_1_lat_dsl + dreg_patch0_2_lon_dsl = np.where( + mask_case1, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + arrival_pts_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case1, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + arrival_pts_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case1, ps2_x, depart_pts_2_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case1, ps2_y, depart_pts_2_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case1, + np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + depart_pts_1_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case1, + np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), + depart_pts_1_lat_dsl, + ) + + return ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) + + @staticmethod + def _apply_case1_patch1( mask_case1, - np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), - dreg_patch0_4_lon_dsl, - ) - dreg_patch0_4_lat_dsl = np.where( + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + ps1_x, + ps1_y, + ): + dreg_patch1_1_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) + dreg_patch1_1_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) + dreg_patch1_4_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) + dreg_patch1_4_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) + dreg_patch1_2_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0 + ) + dreg_patch1_2_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0 + ) + dreg_patch1_3_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0 + ) + dreg_patch1_3_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0 + ) + + return ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ) + + @staticmethod + def _apply_case1_patch2( mask_case1, - np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), - dreg_patch0_4_lat_dsl, - ) - # Case 1 - patch 1 - dreg_patch1_1_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) - dreg_patch1_1_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) - dreg_patch1_4_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) - dreg_patch1_4_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) - dreg_patch1_2_lon_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0 - ) - dreg_patch1_2_lat_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0 - ) - dreg_patch1_3_lon_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0 - ) - dreg_patch1_3_lat_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0 - ) - # Case 1 - patch 2 - dreg_patch2_1_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) - dreg_patch2_1_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) - dreg_patch2_4_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) - dreg_patch2_4_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) - dreg_patch2_2_lon_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), 0.0 - ) - dreg_patch2_2_lat_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), 0.0 - ) - dreg_patch2_3_lon_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), 0.0 - ) - dreg_patch2_3_lat_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), 0.0 - ) + lvn_sys_pos, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + ps2_x, + ps2_y, + ): + # Case 1 - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) + dreg_patch2_1_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) + dreg_patch2_4_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) + dreg_patch2_4_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) + dreg_patch2_2_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), 0.0 + ) + dreg_patch2_2_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), 0.0 + ) + dreg_patch2_3_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), 0.0 + ) + dreg_patch2_3_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), 0.0 + ) - # ------------------------------------------------- Case 2a - mask_case2a = np.logical_and.reduce( - [lintersect_line1, np.logical_not(lintersect_line2), famask_bool] - ) - # Case 2a - patch 0 - dreg_patch0_1_lon_dsl = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_2_lon_dsl = np.where( + return ( + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + ) + + @staticmethod + def _apply_case2a_patch0( mask_case2a, - np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + ps1_x, + ps1_y, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, dreg_patch0_2_lon_dsl, - ) - dreg_patch0_2_lat_dsl = np.where( - mask_case2a, - np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), dreg_patch0_2_lat_dsl, - ) - dreg_patch0_3_lon_dsl = np.where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = np.where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) - dreg_patch0_4_lon_dsl = np.where( - mask_case2a, - np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, dreg_patch0_4_lon_dsl, - ) - dreg_patch0_4_lat_dsl = np.where( - mask_case2a, - np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), dreg_patch0_4_lat_dsl, - ) - # Case 2a - patch 1 - dreg_patch1_1_lon_vmask = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) - dreg_patch1_4_lon_vmask = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask) - dreg_patch1_4_lat_vmask = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask) - dreg_patch1_2_lon_vmask = np.where( + ): + dreg_patch0_1_lon_dsl = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case2a, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case2a, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case2a, + np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case2a, + np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + + return ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) + + @staticmethod + def _apply_case2a_patch1( mask_case2a, - np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + ps1_x, + ps1_y, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, dreg_patch1_2_lon_vmask, - ) - dreg_patch1_2_lat_vmask = np.where( - mask_case2a, - np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), dreg_patch1_2_lat_vmask, - ) - dreg_patch1_3_lon_vmask = np.where( - mask_case2a, - np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), dreg_patch1_3_lon_vmask, - ) - dreg_patch1_3_lat_vmask = np.where( - mask_case2a, - np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), dreg_patch1_3_lat_vmask, - ) - # Case 2a - patch 2 - dreg_patch2_1_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_1_lat_vmask) - dreg_patch2_2_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_2_lon_vmask) - dreg_patch2_2_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_2_lat_vmask) - dreg_patch2_3_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_3_lon_vmask) - dreg_patch2_3_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_3_lat_vmask) - dreg_patch2_4_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_4_lon_vmask) - dreg_patch2_4_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_4_lat_vmask) - - # -------------------------------------------------- Case 2b - mask_case2b = np.logical_and.reduce( - [lintersect_line2, np.logical_not(lintersect_line1), famask_bool] - ) - # Case 2b - patch 0 - dreg_patch0_1_lon_dsl = np.where(mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = np.where(mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_2_lon_dsl = np.where( - mask_case2b, - np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), - dreg_patch0_2_lon_dsl, - ) - dreg_patch0_2_lat_dsl = np.where( - mask_case2b, - np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), - dreg_patch0_2_lat_dsl, - ) - dreg_patch0_3_lon_dsl = np.where(mask_case2b, ps2_x, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = np.where(mask_case2b, ps2_y, dreg_patch0_3_lat_dsl) - dreg_patch0_4_lon_dsl = np.where( - mask_case2b, - np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), - dreg_patch0_4_lon_dsl, - ) - dreg_patch0_4_lat_dsl = np.where( - mask_case2b, - np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), - dreg_patch0_4_lat_dsl, - ) - # Case 2b - patch 1 - dreg_patch1_1_lon_vmask = np.where(mask_case2b, 0.0, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = np.where(mask_case2b, 0.0, dreg_patch1_1_lat_vmask) - dreg_patch1_2_lon_vmask = np.where(mask_case2b, 0.0, dreg_patch1_2_lon_vmask) - dreg_patch1_2_lat_vmask = np.where(mask_case2b, 0.0, dreg_patch1_2_lat_vmask) - dreg_patch1_3_lon_vmask = np.where(mask_case2b, 0.0, dreg_patch1_3_lon_vmask) - dreg_patch1_3_lat_vmask = np.where(mask_case2b, 0.0, dreg_patch1_3_lat_vmask) - dreg_patch1_4_lon_vmask = np.where(mask_case2b, 0.0, dreg_patch1_4_lon_vmask) - dreg_patch1_4_lat_vmask = np.where(mask_case2b, 0.0, dreg_patch1_4_lat_vmask) - # Case 2b - patch 2 - dreg_patch2_1_lon_vmask = np.where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = np.where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) - dreg_patch2_4_lon_vmask = np.where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask) - dreg_patch2_4_lat_vmask = np.where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask) - dreg_patch2_2_lon_vmask = np.where( - mask_case2b, - np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), - dreg_patch2_2_lon_vmask, - ) - dreg_patch2_2_lat_vmask = np.where( - mask_case2b, - np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), - dreg_patch2_2_lat_vmask, - ) - dreg_patch2_3_lon_vmask = np.where( - mask_case2b, - np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), - dreg_patch2_3_lon_vmask, - ) - dreg_patch2_3_lat_vmask = np.where( - mask_case2b, - np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), - dreg_patch2_3_lat_vmask, - ) + ): + dreg_patch1_1_lon_vmask = np.where( + mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask + ) + dreg_patch1_1_lat_vmask = np.where( + mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask + ) + dreg_patch1_4_lon_vmask = np.where( + mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask + ) + dreg_patch1_4_lat_vmask = np.where( + mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask + ) + dreg_patch1_2_lon_vmask = np.where( + mask_case2a, + np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), + dreg_patch1_2_lon_vmask, + ) + dreg_patch1_2_lat_vmask = np.where( + mask_case2a, + np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), + dreg_patch1_2_lat_vmask, + ) + dreg_patch1_3_lon_vmask = np.where( + mask_case2a, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), + dreg_patch1_3_lon_vmask, + ) + dreg_patch1_3_lat_vmask = np.where( + mask_case2a, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), + dreg_patch1_3_lat_vmask, + ) - # flux area edge 1 and 2 - fl_e1_p1_lon = arrival_pts_1_lon_dsl - fl_e1_p1_lat = arrival_pts_1_lat_dsl - fl_e1_p2_lon = depart_pts_1_lon_dsl - fl_e1_p2_lat = depart_pts_1_lat_dsl - fl_e2_p1_lon = arrival_pts_2_lon_dsl - fl_e2_p1_lat = arrival_pts_2_lat_dsl - fl_e2_p2_lon = depart_pts_2_lon_dsl - fl_e2_p2_lat = depart_pts_2_lat_dsl - - # ----------------------------------------------- Case 3a - # Check whether flux area edge 2 intersects with triangle edge 1 - lintersect_e2_line1 = lintersect( - fl_e2_p1_lon, - fl_e2_p1_lat, - fl_e2_p2_lon, - fl_e2_p2_lat, - tri_line1_p1_lon, - tri_line1_p1_lat, - tri_line1_p2_lon, - tri_line1_p2_lat, - ) - mask_case3a = np.logical_and(lintersect_e2_line1, famask_bool) - pi1_x, pi1_y = line_intersect( - fl_e2_p1_lon, - fl_e2_p1_lat, - fl_e2_p2_lon, - fl_e2_p2_lat, - tri_line1_p1_lon, - tri_line1_p1_lat, - tri_line1_p2_lon, - tri_line1_p2_lat, - ) - # Case 3a - patch 0 - dreg_patch0_1_lon_dsl = np.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = np.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_2_lon_dsl = np.where( - mask_case3a, - np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), - dreg_patch0_2_lon_dsl, - ) - dreg_patch0_2_lat_dsl = np.where( - mask_case3a, - np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), - dreg_patch0_2_lat_dsl, - ) - dreg_patch0_3_lon_dsl = np.where(mask_case3a, ps2_x, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = np.where(mask_case3a, ps2_y, dreg_patch0_3_lat_dsl) - dreg_patch0_4_lon_dsl = np.where( - mask_case3a, - np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), - dreg_patch0_4_lon_dsl, - ) - dreg_patch0_4_lat_dsl = np.where( - mask_case3a, - np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), - dreg_patch0_4_lat_dsl, - ) - # Case 3a - patch 1 - dreg_patch1_1_lon_vmask = np.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = np.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) - dreg_patch1_2_lon_vmask = np.where( - mask_case3a, - np.where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), - dreg_patch1_2_lon_vmask, - ) - dreg_patch1_2_lat_vmask = np.where( - mask_case3a, - np.where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), - dreg_patch1_2_lat_vmask, - ) - dreg_patch1_3_lon_vmask = np.where(mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask) - dreg_patch1_3_lat_vmask = np.where(mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask) - dreg_patch1_4_lon_vmask = np.where( - mask_case3a, - np.where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), - dreg_patch1_4_lon_vmask, - ) - dreg_patch1_4_lat_vmask = np.where( - mask_case3a, - np.where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), - dreg_patch1_4_lat_vmask, - ) - # Case 3a - patch 2 - dreg_patch2_1_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_1_lat_vmask) - dreg_patch2_2_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_2_lon_vmask) - dreg_patch2_2_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_2_lat_vmask) - dreg_patch2_3_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_3_lon_vmask) - dreg_patch2_3_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_3_lat_vmask) - dreg_patch2_4_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_4_lon_vmask) - dreg_patch2_4_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_4_lat_vmask) - - # ------------------------------------------------ Case 3b - # Check whether flux area edge 1 intersects with triangle edge 2 - lintersect_e1_line2 = lintersect( - fl_e1_p1_lon, - fl_e1_p1_lat, - fl_e1_p2_lon, - fl_e1_p2_lat, - tri_line2_p1_lon, - tri_line2_p1_lat, - tri_line2_p2_lon, - tri_line2_p2_lat, - ) - mask_case3b = lintersect_e1_line2 & famask_bool - pi2_x, pi2_y = line_intersect( - fl_e1_p1_lon, - fl_e1_p1_lat, - fl_e1_p2_lon, - fl_e1_p2_lat, - tri_line2_p1_lon, - tri_line2_p1_lat, - tri_line2_p2_lon, - tri_line2_p2_lat, - ) - # Case 3b - patch 0 - dreg_patch0_1_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_4_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) - dreg_patch0_4_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) - dreg_patch0_2_lon_dsl = np.where( - mask_case3b, - np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), - dreg_patch0_2_lon_dsl, - ) - dreg_patch0_2_lat_dsl = np.where( - mask_case3b, - np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), - dreg_patch0_2_lat_dsl, - ) - dreg_patch0_3_lon_dsl = np.where( - mask_case3b, - np.where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), - dreg_patch0_3_lon_dsl, - ) - dreg_patch0_3_lat_dsl = np.where( - mask_case3b, - np.where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), - dreg_patch0_3_lat_dsl, - ) - # Case 3b - patch 1 - dreg_patch1_1_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_1_lat_vmask) - dreg_patch1_2_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_2_lon_vmask) - dreg_patch1_2_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_2_lat_vmask) - dreg_patch1_3_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_3_lon_vmask) - dreg_patch1_3_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_3_lat_vmask) - dreg_patch1_4_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lon_vmask) - dreg_patch1_4_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lat_vmask) - # Case 3b - patch 2 - dreg_patch2_1_lon_vmask = np.where(mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = np.where(mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) - dreg_patch2_2_lon_vmask = np.where( - mask_case3b, - np.where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), - dreg_patch2_2_lon_vmask, - ) - dreg_patch2_2_lat_vmask = np.where( - mask_case3b, - np.where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), - dreg_patch2_2_lat_vmask, - ) - dreg_patch2_3_lon_vmask = np.where(mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask) - dreg_patch2_3_lat_vmask = np.where(mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask) - dreg_patch2_4_lon_vmask = np.where( - mask_case3b, - np.where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), - dreg_patch2_4_lon_vmask, - ) - dreg_patch2_4_lat_vmask = np.where( - mask_case3b, - np.where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), - dreg_patch2_4_lat_vmask, - ) + return ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ) - # --------------------------------------------- Case 4 - # NB: Next line acts as the "ELSE IF", indices that already previously matched one of the above conditions - # can't be overwritten by this new condition. - indices_previously_matched = np.logical_or.reduce( - [mask_case3b, mask_case3a, mask_case2b, mask_case2a, mask_case1] - ) - # mask_case4 = (abs(p_vn) < 0.1) & famask_bool & (not indices_previously_matched) we insert also the error indices - mask_case4 = np.logical_and.reduce([famask_bool, np.logical_not(indices_previously_matched)]) - # Case 4 - patch 0 - no change - # Case 4 - patch 1 - dreg_patch1_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lat_vmask) - dreg_patch1_2_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_2_lon_vmask) - dreg_patch1_2_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_2_lat_vmask) - dreg_patch1_3_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_3_lon_vmask) - dreg_patch1_3_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_3_lat_vmask) - dreg_patch1_4_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_4_lon_vmask) - dreg_patch1_4_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_4_lat_vmask) - # Case 4 - patch 2 - dreg_patch2_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_1_lat_vmask) - dreg_patch2_2_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_2_lon_vmask) - dreg_patch2_2_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_2_lat_vmask) - dreg_patch2_3_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_3_lon_vmask) - dreg_patch2_3_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_3_lat_vmask) - dreg_patch2_4_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_4_lon_vmask) - - return ( + @staticmethod + def _apply_case2b_patch0( + mask_case2b, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + ps2_x, + ps2_y, dreg_patch0_1_lon_dsl, dreg_patch0_1_lat_dsl, dreg_patch0_2_lon_dsl, @@ -637,6 +498,46 @@ def divide_flux_area_list_stencil_01_numpy( dreg_patch0_3_lat_dsl, dreg_patch0_4_lon_dsl, dreg_patch0_4_lat_dsl, + ): + dreg_patch0_1_lon_dsl = np.where(mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case2b, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case2b, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + + return ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) + + @staticmethod + def _apply_case2b_patch1( + mask_case2b, dreg_patch1_1_lon_vmask, dreg_patch1_1_lat_vmask, dreg_patch1_2_lon_vmask, @@ -645,101 +546,105 @@ def divide_flux_area_list_stencil_01_numpy( dreg_patch1_3_lat_vmask, dreg_patch1_4_lon_vmask, dreg_patch1_4_lat_vmask, + ): + zeros_array = np.zeros_like(mask_case2b) + + dreg_patch1_1_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_4_lat_vmask) + + return ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + ) + + @staticmethod + def _apply_case2b_patch2( + mask_case2b, + lvn_sys_pos, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + ps2_x, + ps2_y, dreg_patch2_1_lon_vmask, dreg_patch2_1_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, dreg_patch2_2_lon_vmask, dreg_patch2_2_lat_vmask, dreg_patch2_3_lon_vmask, dreg_patch2_3_lat_vmask, - dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask, - ) + ): + dreg_patch2_1_lon_vmask = np.where( + mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask + ) + dreg_patch2_1_lat_vmask = np.where( + mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask + ) + dreg_patch2_4_lon_vmask = np.where( + mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask + ) + dreg_patch2_4_lat_vmask = np.where( + mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask + ) + dreg_patch2_2_lon_vmask = np.where( + mask_case2b, + np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), + dreg_patch2_2_lon_vmask, + ) + dreg_patch2_2_lat_vmask = np.where( + mask_case2b, + np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), + dreg_patch2_2_lat_vmask, + ) + dreg_patch2_3_lon_vmask = np.where( + mask_case2b, + np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), + dreg_patch2_3_lon_vmask, + ) + dreg_patch2_3_lat_vmask = np.where( + mask_case2b, + np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), + dreg_patch2_3_lat_vmask, + ) -@pytest.mark.slow_tests -def test_divide_flux_area_list_stencil_01(backend): - grid = SimpleGrid() - - famask_int = random_mask(grid, EdgeDim, KDim, dtype=int32) - p_vn = random_field(grid, EdgeDim, KDim) - ptr_v3_lon = random_field(grid, EdgeDim, E2CDim) - ptr_v3_lon_field = as_1D_sparse_field(ptr_v3_lon, ECDim) - ptr_v3_lat = random_field(grid, EdgeDim, E2CDim) - ptr_v3_lat_field = as_1D_sparse_field(ptr_v3_lat, ECDim) - tangent_orientation_dsl = random_field(grid, EdgeDim) - dreg_patch0_1_lon_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_1_lat_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_2_lon_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_2_lat_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_3_lon_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_3_lat_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_4_lon_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch0_4_lat_dsl = random_field(grid, EdgeDim, KDim) - dreg_patch1_1_lon_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch1_1_lat_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch1_2_lon_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch1_2_lat_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch1_3_lon_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch1_3_lat_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch1_4_lon_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch1_4_lat_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch2_1_lon_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch2_1_lat_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch2_2_lon_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch2_2_lat_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch2_3_lon_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch2_3_lat_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch2_4_lon_vmask = zero_field(grid, EdgeDim, KDim) - dreg_patch2_4_lat_vmask = zero_field(grid, EdgeDim, KDim) - - ( - ref_1, - ref_2, - ref_3, - ref_4, - ref_5, - ref_6, - ref_7, - ref_8, - ref_9, - ref_10, - ref_11, - ref_12, - ref_13, - ref_14, - ref_15, - ref_16, - ref_17, - ref_18, - ref_19, - ref_20, - ref_21, - ref_22, - ref_23, - ref_24, - ) = divide_flux_area_list_stencil_01_numpy( - grid.connectivities[E2CDim], - np.asarray(famask_int), - np.asarray(p_vn), - np.asarray(ptr_v3_lon), - np.asarray(ptr_v3_lat), - np.asarray(tangent_orientation_dsl), - np.asarray(dreg_patch0_1_lon_dsl), - np.asarray(dreg_patch0_1_lat_dsl), - np.asarray(dreg_patch0_2_lon_dsl), - np.asarray(dreg_patch0_2_lat_dsl), - np.asarray(dreg_patch0_3_lon_dsl), - np.asarray(dreg_patch0_3_lat_dsl), - np.asarray(dreg_patch0_4_lon_dsl), - np.asarray(dreg_patch0_4_lat_dsl), - ) + return ( + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + ) - divide_flux_area_list_stencil_01.with_backend(backend)( - famask_int, - p_vn, - ptr_v3_lon_field, - ptr_v3_lat_field, - tangent_orientation_dsl, + @staticmethod + def _apply_case3a_patch0( + mask_case3a, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + lvn_sys_pos, + ps2_x, + ps2_y, dreg_patch0_1_lon_dsl, dreg_patch0_1_lat_dsl, dreg_patch0_2_lon_dsl, @@ -748,14 +653,175 @@ def test_divide_flux_area_list_stencil_01(backend): dreg_patch0_3_lat_dsl, dreg_patch0_4_lon_dsl, dreg_patch0_4_lat_dsl, + ): + dreg_patch0_1_lon_dsl = np.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case3a, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case3a, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + + return ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) + + @staticmethod + def _apply_case3a_patch1( + mask_case3a, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + pi1_x, + pi1_y, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, dreg_patch1_1_lon_vmask, dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, dreg_patch1_2_lon_vmask, dreg_patch1_2_lat_vmask, dreg_patch1_3_lon_vmask, dreg_patch1_3_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, + ): + dreg_patch1_1_lon_vmask = np.where( + mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask + ) + dreg_patch1_1_lat_vmask = np.where( + mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask + ) + dreg_patch1_2_lon_vmask = np.where( + mask_case3a, + np.where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), + dreg_patch1_2_lon_vmask, + ) + dreg_patch1_2_lat_vmask = np.where( + mask_case3a, + np.where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), + dreg_patch1_2_lat_vmask, + ) + dreg_patch1_3_lon_vmask = np.where( + mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask + ) + dreg_patch1_3_lat_vmask = np.where( + mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask + ) + dreg_patch1_4_lon_vmask = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), + dreg_patch1_4_lon_vmask, + ) + dreg_patch1_4_lat_vmask = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), + dreg_patch1_4_lat_vmask, + ) + + return ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ) + + @staticmethod + def _apply_case3b_patch0( + mask_case3b, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + pi2_x, + pi2_y, + lvn_sys_pos, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + ): + dreg_patch0_1_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_4_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) + dreg_patch0_4_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case3b, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case3b, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where( + mask_case3b, + np.where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), + dreg_patch0_3_lon_dsl, + ) + dreg_patch0_3_lat_dsl = np.where( + mask_case3b, + np.where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), + dreg_patch0_3_lat_dsl, + ) + + return ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + ) + + @staticmethod + def _apply_case3b_patch2( + mask_case3b, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + pi2_x, + pi2_y, + lvn_sys_pos, dreg_patch2_1_lon_vmask, dreg_patch2_1_lat_vmask, dreg_patch2_2_lon_vmask, @@ -764,32 +830,716 @@ def test_divide_flux_area_list_stencil_01(backend): dreg_patch2_3_lat_vmask, dreg_patch2_4_lon_vmask, dreg_patch2_4_lat_vmask, - offset_provider={ - "E2C": grid.get_offset_provider("E2C"), - "E2EC": grid.get_offset_provider("E2EC"), - }, - ) - assert np.allclose(dreg_patch0_1_lon_dsl, ref_1) - assert np.allclose(dreg_patch0_1_lat_dsl, ref_2) - assert np.allclose(dreg_patch0_2_lon_dsl, ref_3) - assert np.allclose(dreg_patch0_2_lat_dsl, ref_4) - assert np.allclose(dreg_patch0_3_lon_dsl, ref_5) - assert np.allclose(dreg_patch0_3_lat_dsl, ref_6) - assert np.allclose(dreg_patch0_4_lon_dsl, ref_7) - assert np.allclose(dreg_patch0_4_lat_dsl, ref_8) - assert np.allclose(dreg_patch1_1_lon_vmask, ref_9) - assert np.allclose(dreg_patch1_1_lat_vmask, ref_10) - assert np.allclose(dreg_patch1_2_lon_vmask, ref_11) - assert np.allclose(dreg_patch1_2_lat_vmask, ref_12) - assert np.allclose(dreg_patch1_3_lon_vmask, ref_13) - assert np.allclose(dreg_patch1_3_lat_vmask, ref_14) - assert np.allclose(dreg_patch1_4_lon_vmask, ref_15) - assert np.allclose(dreg_patch1_4_lat_vmask, ref_16) - assert np.allclose(dreg_patch2_1_lon_vmask, ref_17) - assert np.allclose(dreg_patch2_1_lat_vmask, ref_18) - assert np.allclose(dreg_patch2_2_lon_vmask, ref_19) - assert np.allclose(dreg_patch2_2_lat_vmask, ref_20) - assert np.allclose(dreg_patch2_3_lon_vmask, ref_21) - assert np.allclose(dreg_patch2_3_lat_vmask, ref_22) - assert np.allclose(dreg_patch2_4_lon_vmask, ref_23) - assert np.allclose(dreg_patch2_4_lat_vmask, ref_24) + ): + dreg_patch2_1_lon_vmask = np.where( + mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask + ) + dreg_patch2_1_lat_vmask = np.where( + mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask + ) + dreg_patch2_2_lon_vmask = np.where( + mask_case3b, + np.where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), + dreg_patch2_2_lon_vmask, + ) + dreg_patch2_2_lat_vmask = np.where( + mask_case3b, + np.where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), + dreg_patch2_2_lat_vmask, + ) + dreg_patch2_3_lon_vmask = np.where( + mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask + ) + dreg_patch2_3_lat_vmask = np.where( + mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask + ) + dreg_patch2_4_lon_vmask = np.where( + mask_case3b, + np.where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), + dreg_patch2_4_lon_vmask, + ) + dreg_patch2_4_lat_vmask = np.where( + mask_case3b, + np.where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), + dreg_patch2_4_lat_vmask, + ) + + return ( + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + ) + + @classmethod + def reference( + cls, + grid, + famask_int, + p_vn, + ptr_v3_lon, + ptr_v3_lat, + tangent_orientation_dsl, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + **kwargs, + ): + e2c = grid.connectivities[E2CDim] + ptr_v3_lon = reshape(ptr_v3_lon, e2c.shape) + ptr_v3_lon_e = np.expand_dims(ptr_v3_lon, axis=-1) + ptr_v3_lat = reshape(ptr_v3_lat, e2c.shape) + ptr_v3_lat_e = np.expand_dims(ptr_v3_lat, axis=-1) + tangent_orientation_dsl = np.expand_dims(tangent_orientation_dsl, axis=-1) + + result_tuple = cls._generate_flux_area_geometry( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + p_vn, + ptr_v3_lon_e, + ptr_v3_lat_e, + ) + + ( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + ) = result_tuple + + # Create first mask does departure-line segment intersects with A1V3 + lintersect_line1 = _lintersect_numpy( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + # Create first mask does departure-line segment intersects with A2V3 + lintersect_line2 = _lintersect_numpy( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + + lvn_sys_pos = np.where( + (p_vn * np.broadcast_to(tangent_orientation_dsl, p_vn.shape)) >= 0.0, + True, + False, + ) + famask_bool = np.where(famask_int == int32(1), True, False) + mask_case1 = np.logical_and.reduce([lintersect_line1, lintersect_line2, famask_bool]) + ps1_x, ps1_y = _line_intersect_numpy( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + ps2_x, ps2_y = _line_intersect_numpy( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + + # ------------------------------------------------- Case 1 + # Case 1 - patch 0 + ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) = cls._apply_case1_patch0( + mask_case1, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + ps1_x, + ps1_y, + ps2_x, + ps2_y, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + ) + # Case 1 - patch 1 + ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ) = cls._apply_case1_patch1( + mask_case1, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + ps1_x, + ps1_y, + ) + # Case 1 - patch 2 + ( + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + ) = cls._apply_case1_patch2( + mask_case1, + lvn_sys_pos, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + ps2_x, + ps2_y, + ) + + # ------------------------------------------------- Case 2a + mask_case2a = np.logical_and.reduce( + [lintersect_line1, np.logical_not(lintersect_line2), famask_bool] + ) + # Case 2a - patch 0 + result_tuple_patch0 = cls._apply_case2a_patch0( + mask_case2a, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + ps1_x, + ps1_y, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) + + ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) = result_tuple_patch0 + # Case 2a - patch 1 + result_tuple_patch1 = cls._apply_case2a_patch1( + mask_case2a, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + ps1_x, + ps1_y, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ) + + ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ) = result_tuple_patch1 + # Case 2a - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_4_lat_vmask) + + # -------------------------------------------------- Case 2b + mask_case2b = np.logical_and.reduce( + [lintersect_line2, np.logical_not(lintersect_line1), famask_bool] + ) + # Case 2b - patch 0 + result_tuple_patch0_case2b = cls._apply_case2b_patch0( + mask_case2b, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + ps2_x, + ps2_y, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) + + ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) = result_tuple_patch0_case2b + + # Case 2b - patch 1 + result_tuple_patch1_case2b = cls._apply_case2b_patch1( + mask_case2b, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + ) + + ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + ) = result_tuple_patch1_case2b + + # Case 2b - patch 2 + result_tuple_patch2_case2b = cls._apply_case2b_patch2( + mask_case2b, + lvn_sys_pos, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + ps2_x, + ps2_y, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + ) + + ( + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + ) = result_tuple_patch2_case2b + + # flux area edge 1 and 2 + fl_e1_p1_lon = arrival_pts_1_lon_dsl + fl_e1_p1_lat = arrival_pts_1_lat_dsl + fl_e1_p2_lon = depart_pts_1_lon_dsl + fl_e1_p2_lat = depart_pts_1_lat_dsl + fl_e2_p1_lon = arrival_pts_2_lon_dsl + fl_e2_p1_lat = arrival_pts_2_lat_dsl + fl_e2_p2_lon = depart_pts_2_lon_dsl + fl_e2_p2_lat = depart_pts_2_lat_dsl + + # ----------------------------------------------- Case 3a + # Check whether flux area edge 2 intersects with triangle edge 1 + lintersect_e2_line1 = _lintersect_numpy( + fl_e2_p1_lon, + fl_e2_p1_lat, + fl_e2_p2_lon, + fl_e2_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + mask_case3a = np.logical_and(lintersect_e2_line1, famask_bool) + pi1_x, pi1_y = _line_intersect_numpy( + fl_e2_p1_lon, + fl_e2_p1_lat, + fl_e2_p2_lon, + fl_e2_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + # Case 3a - patch 0 + result = cls._apply_case3a_patch0( + mask_case3a, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + lvn_sys_pos, + ps2_x, + ps2_y, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) + + ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) = result + + # Case 3a - patch 1 + ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ) = cls._apply_case3a_patch1( + mask_case3a, + lvn_sys_pos, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + pi1_x, + pi1_y, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ) + # Case 3a - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_4_lat_vmask) + + # ------------------------------------------------ Case 3b + # Check whether flux area edge 1 intersects with triangle edge 2 + lintersect_e1_line2 = _lintersect_numpy( + fl_e1_p1_lon, + fl_e1_p1_lat, + fl_e1_p2_lon, + fl_e1_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + mask_case3b = lintersect_e1_line2 & famask_bool + pi2_x, pi2_y = _line_intersect_numpy( + fl_e1_p1_lon, + fl_e1_p1_lat, + fl_e1_p2_lon, + fl_e1_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + # Case 3b - patch 0 + ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + ) = cls._apply_case3b_patch0( + mask_case3b, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + pi2_x, + pi2_y, + lvn_sys_pos, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + ) + # Case 3b - patch 1 + dreg_patch1_1_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lat_vmask) + + # Case 3b - patch 2 + ( + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + ) = cls._apply_case3b_patch2( + mask_case3b, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + depart_pts_1_lon_dsl, + depart_pts_1_lat_dsl, + depart_pts_2_lon_dsl, + depart_pts_2_lat_dsl, + pi2_x, + pi2_y, + lvn_sys_pos, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + ) + # --------------------------------------------- Case 4 + # NB: Next line acts as the "ELSE IF", indices that already previously matched one of the above conditions + # can't be overwritten by this new condition. + indices_previously_matched = np.logical_or.reduce( + [mask_case3b, mask_case3a, mask_case2b, mask_case2a, mask_case1] + ) + # mask_case4 = (abs(p_vn) < 0.1) & famask_bool & (not indices_previously_matched) we insert also the error indices + mask_case4 = np.logical_and.reduce( + [famask_bool, np.logical_not(indices_previously_matched)] + ) + # Case 4 - patch 0 - no change + # Case 4 - patch 1 + dreg_patch1_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_4_lat_vmask) + # Case 4 - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_4_lon_vmask) + + return dict( + dreg_patch0_1_lon_dsl=dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl=dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl=dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl=dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl=dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl=dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl=dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl=dreg_patch0_4_lat_dsl, + dreg_patch1_1_lon_vmask=dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask=dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask=dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask=dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask=dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask=dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask=dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask=dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask=dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask=dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask=dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask=dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask=dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask=dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask=dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask=dreg_patch2_4_lat_vmask, + ) + + @pytest.mark.slow_tests + @pytest.fixture + def input_data(self, grid): + famask_int = random_mask(grid, EdgeDim, KDim, dtype=int32) + p_vn = random_field(grid, EdgeDim, KDim) + ptr_v3_lon = random_field(grid, EdgeDim, E2CDim) + ptr_v3_lon_field = as_1D_sparse_field(ptr_v3_lon, ECDim) + ptr_v3_lat = random_field(grid, EdgeDim, E2CDim) + ptr_v3_lat_field = as_1D_sparse_field(ptr_v3_lat, ECDim) + tangent_orientation_dsl = random_field(grid, EdgeDim) + dreg_patch0_1_lon_dsl = random_field(grid, EdgeDim, KDim) + dreg_patch0_1_lat_dsl = random_field(grid, EdgeDim, KDim) + dreg_patch0_2_lon_dsl = random_field(grid, EdgeDim, KDim) + dreg_patch0_2_lat_dsl = random_field(grid, EdgeDim, KDim) + dreg_patch0_3_lon_dsl = random_field(grid, EdgeDim, KDim) + dreg_patch0_3_lat_dsl = random_field(grid, EdgeDim, KDim) + dreg_patch0_4_lon_dsl = random_field(grid, EdgeDim, KDim) + dreg_patch0_4_lat_dsl = random_field(grid, EdgeDim, KDim) + dreg_patch1_1_lon_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch1_1_lat_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch1_2_lon_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch1_2_lat_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch1_3_lon_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch1_3_lat_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch1_4_lon_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch1_4_lat_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch2_1_lon_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch2_1_lat_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch2_2_lon_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch2_2_lat_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch2_3_lon_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch2_3_lat_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch2_4_lon_vmask = zero_field(grid, EdgeDim, KDim) + dreg_patch2_4_lat_vmask = zero_field(grid, EdgeDim, KDim) + return dict( + famask_int=famask_int, + p_vn=p_vn, + ptr_v3_lon=ptr_v3_lon_field, + ptr_v3_lat=ptr_v3_lat_field, + tangent_orientation_dsl=tangent_orientation_dsl, + dreg_patch0_1_lon_dsl=dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl=dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl=dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl=dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl=dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl=dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl=dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl=dreg_patch0_4_lat_dsl, + dreg_patch1_1_lon_vmask=dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask=dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask=dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask=dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask=dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask=dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask=dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask=dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask=dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask=dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask=dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask=dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask=dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask=dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask=dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask=dreg_patch2_4_lat_vmask, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_02.py index 1d0d499c37..6b1daabc13 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_02.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_02.py @@ -12,321 +12,279 @@ # 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.advection.divide_flux_area_list_stencil_02 import ( divide_flux_area_list_stencil_02, ) from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, random_mask - +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + random_mask, + reshape, +) -def divide_flux_area_list_stencil_02_numpy( - e2c: np.array, - famask_int: np.array, - p_vn: np.array, - bf_cc_patch1_lon: np.array, - bf_cc_patch1_lat: np.array, - bf_cc_patch2_lon: np.array, - bf_cc_patch2_lat: np.array, - butterfly_idx_patch1_vnpos: np.array, - butterfly_idx_patch1_vnneg: np.array, - butterfly_blk_patch1_vnpos: np.array, - butterfly_blk_patch1_vnneg: np.array, - butterfly_idx_patch2_vnpos: np.array, - butterfly_idx_patch2_vnneg: np.array, - butterfly_blk_patch2_vnpos: np.array, - butterfly_blk_patch2_vnneg: np.array, - dreg_patch1_1_lon_vmask: np.array, - dreg_patch1_1_lat_vmask: np.array, - dreg_patch1_2_lon_vmask: np.array, - dreg_patch1_2_lat_vmask: np.array, - dreg_patch1_3_lon_vmask: np.array, - dreg_patch1_3_lat_vmask: np.array, - dreg_patch1_4_lon_vmask: np.array, - dreg_patch1_4_lat_vmask: np.array, - dreg_patch2_1_lon_vmask: np.array, - dreg_patch2_1_lat_vmask: np.array, - dreg_patch2_2_lon_vmask: np.array, - dreg_patch2_2_lat_vmask: np.array, - dreg_patch2_3_lon_vmask: np.array, - dreg_patch2_3_lat_vmask: np.array, - dreg_patch2_4_lon_vmask: np.array, - dreg_patch2_4_lat_vmask: np.array, -): - famask_bool = np.where(famask_int == int32(1), True, False) - lvn_pos = np.where(p_vn >= np.broadcast_to(0.0, p_vn.shape), True, False) - # Translation of patch 1 and patch 2 in system relative to respective cell - bf_cc_patch1_lon_e = np.expand_dims(bf_cc_patch1_lon, axis=-1) - bf_cc_patch1_lat_e = np.expand_dims(bf_cc_patch1_lat, axis=-1) - bf_cc_patch2_lon_e = np.expand_dims(bf_cc_patch2_lon, axis=-1) - bf_cc_patch2_lat_e = np.expand_dims(bf_cc_patch2_lat, axis=-1) - bf_cc_patch1_lon = np.where( - famask_bool, - np.where(lvn_pos, bf_cc_patch1_lon_e[:, 0], bf_cc_patch1_lon_e[:, 1]), - 0.0, - ) - bf_cc_patch1_lat = np.where( - famask_bool, - np.where(lvn_pos, bf_cc_patch1_lat_e[:, 0], bf_cc_patch1_lat_e[:, 1]), - 0.0, - ) - bf_cc_patch2_lon = np.where( - famask_bool, - np.where(lvn_pos, bf_cc_patch2_lon_e[:, 0], bf_cc_patch2_lon_e[:, 1]), - 0.0, - ) - bf_cc_patch2_lat = np.where( - famask_bool, - np.where(lvn_pos, bf_cc_patch2_lat_e[:, 0], bf_cc_patch2_lat_e[:, 1]), - 0.0, +class TestDivideFluxAreaListStencil02(StencilTest): + PROGRAM = divide_flux_area_list_stencil_02 + OUTPUTS = ( + "dreg_patch1_1_lon_vmask", + "dreg_patch1_1_lat_vmask", + "dreg_patch1_2_lon_vmask", + "dreg_patch1_2_lat_vmask", + "dreg_patch1_3_lon_vmask", + "dreg_patch1_3_lat_vmask", + "dreg_patch1_4_lon_vmask", + "dreg_patch1_4_lat_vmask", + "dreg_patch2_1_lon_vmask", + "dreg_patch2_1_lat_vmask", + "dreg_patch2_2_lon_vmask", + "dreg_patch2_2_lat_vmask", + "dreg_patch2_3_lon_vmask", + "dreg_patch2_3_lat_vmask", + "dreg_patch2_4_lon_vmask", + "dreg_patch2_4_lat_vmask", + "patch1_cell_idx_vmask", + "patch1_cell_blk_vmask", + "patch2_cell_idx_vmask", + "patch2_cell_blk_vmask", ) - # patch1 in translated system - dreg_patch1_1_lon_vmask = dreg_patch1_1_lon_vmask - bf_cc_patch1_lon - dreg_patch1_1_lat_vmask = dreg_patch1_1_lat_vmask - bf_cc_patch1_lat - dreg_patch1_2_lon_vmask = dreg_patch1_2_lon_vmask - bf_cc_patch1_lon - dreg_patch1_2_lat_vmask = dreg_patch1_2_lat_vmask - bf_cc_patch1_lat - dreg_patch1_3_lon_vmask = dreg_patch1_3_lon_vmask - bf_cc_patch1_lon - dreg_patch1_3_lat_vmask = dreg_patch1_3_lat_vmask - bf_cc_patch1_lat - dreg_patch1_4_lon_vmask = dreg_patch1_4_lon_vmask - bf_cc_patch1_lon - dreg_patch1_4_lat_vmask = dreg_patch1_4_lat_vmask - bf_cc_patch1_lat - # patch2 in translated system - dreg_patch2_1_lon_vmask = dreg_patch2_1_lon_vmask - bf_cc_patch2_lon - dreg_patch2_1_lat_vmask = dreg_patch2_1_lat_vmask - bf_cc_patch2_lat - dreg_patch2_2_lon_vmask = dreg_patch2_2_lon_vmask - bf_cc_patch2_lon - dreg_patch2_2_lat_vmask = dreg_patch2_2_lat_vmask - bf_cc_patch2_lat - dreg_patch2_3_lon_vmask = dreg_patch2_3_lon_vmask - bf_cc_patch2_lon - dreg_patch2_3_lat_vmask = dreg_patch2_3_lat_vmask - bf_cc_patch2_lat - dreg_patch2_4_lon_vmask = dreg_patch2_4_lon_vmask - bf_cc_patch2_lon - dreg_patch2_4_lat_vmask = dreg_patch2_4_lat_vmask - bf_cc_patch2_lat + @staticmethod + def reference( + grid, + famask_int: np.array, + p_vn: np.array, + bf_cc_patch1_lon: np.array, + bf_cc_patch1_lat: np.array, + bf_cc_patch2_lon: np.array, + bf_cc_patch2_lat: np.array, + butterfly_idx_patch1_vnpos: np.array, + butterfly_idx_patch1_vnneg: np.array, + butterfly_blk_patch1_vnpos: np.array, + butterfly_blk_patch1_vnneg: np.array, + butterfly_idx_patch2_vnpos: np.array, + butterfly_idx_patch2_vnneg: np.array, + butterfly_blk_patch2_vnpos: np.array, + butterfly_blk_patch2_vnneg: np.array, + dreg_patch1_1_lon_vmask: np.array, + dreg_patch1_1_lat_vmask: np.array, + dreg_patch1_2_lon_vmask: np.array, + dreg_patch1_2_lat_vmask: np.array, + dreg_patch1_3_lon_vmask: np.array, + dreg_patch1_3_lat_vmask: np.array, + dreg_patch1_4_lon_vmask: np.array, + dreg_patch1_4_lat_vmask: np.array, + dreg_patch2_1_lon_vmask: np.array, + dreg_patch2_1_lat_vmask: np.array, + dreg_patch2_2_lon_vmask: np.array, + dreg_patch2_2_lat_vmask: np.array, + dreg_patch2_3_lon_vmask: np.array, + dreg_patch2_3_lat_vmask: np.array, + dreg_patch2_4_lon_vmask: np.array, + dreg_patch2_4_lat_vmask: np.array, + **kwargs, + ): + e2c = grid.connectivities[E2CDim] + famask_bool = np.where(famask_int == int32(1), True, False) + lvn_pos = np.where(p_vn >= np.broadcast_to(0.0, p_vn.shape), True, False) + # Translation of patch 1 and patch 2 in system relative to respective cell + bf_cc_patch1_lon = reshape(bf_cc_patch1_lon, e2c.shape) + bf_cc_patch1_lon_e = np.expand_dims(bf_cc_patch1_lon, axis=-1) + bf_cc_patch1_lat = reshape(bf_cc_patch1_lat, e2c.shape) + bf_cc_patch1_lat_e = np.expand_dims(bf_cc_patch1_lat, axis=-1) + bf_cc_patch2_lon = reshape(bf_cc_patch2_lon, e2c.shape) + bf_cc_patch2_lon_e = np.expand_dims(bf_cc_patch2_lon, axis=-1) + bf_cc_patch2_lat = reshape(bf_cc_patch2_lat, e2c.shape) + bf_cc_patch2_lat_e = np.expand_dims(bf_cc_patch2_lat, axis=-1) - # Store global index of the underlying grid cell - # Adapt dimensions to fit ofr multiple levels - butterfly_idx_patch1_vnpos_3d = np.broadcast_to( - np.expand_dims(butterfly_idx_patch1_vnpos, axis=-1), p_vn.shape - ) - butterfly_idx_patch1_vnneg_3d = np.broadcast_to( - np.expand_dims(butterfly_idx_patch1_vnneg, axis=-1), p_vn.shape - ) - butterfly_idx_patch2_vnpos_3d = np.broadcast_to( - np.expand_dims(butterfly_idx_patch2_vnpos, axis=-1), p_vn.shape - ) - butterfly_idx_patch2_vnneg_3d = np.broadcast_to( - np.expand_dims(butterfly_idx_patch2_vnneg, axis=-1), p_vn.shape - ) - butterfly_blk_patch1_vnpos_3d = np.broadcast_to( - np.expand_dims(butterfly_blk_patch1_vnpos, axis=-1), p_vn.shape - ) - butterfly_blk_patch1_vnneg_3d = np.broadcast_to( - np.expand_dims(butterfly_blk_patch1_vnneg, axis=-1), p_vn.shape - ) - butterfly_blk_patch2_vnpos_3d = np.broadcast_to( - np.expand_dims(butterfly_blk_patch2_vnpos, axis=-1), p_vn.shape - ) - butterfly_blk_patch2_vnneg_3d = np.broadcast_to( - np.expand_dims(butterfly_blk_patch2_vnneg, axis=-1), p_vn.shape - ) - patch1_cell_idx_vmask = np.where( - famask_bool, - np.where(lvn_pos, butterfly_idx_patch1_vnpos_3d, butterfly_idx_patch1_vnneg_3d), - int32(0), - ) - patch2_cell_idx_vmask = np.where( - famask_bool, - np.where(lvn_pos, butterfly_idx_patch2_vnpos_3d, butterfly_idx_patch2_vnneg_3d), - int32(0), - ) - patch1_cell_blk_vmask = np.where( - famask_bool, - np.where(lvn_pos, butterfly_blk_patch1_vnpos_3d, butterfly_blk_patch1_vnneg_3d), - int32(0), - ) - patch2_cell_blk_vmask = np.where( - famask_bool, - np.where(lvn_pos, butterfly_blk_patch2_vnpos_3d, butterfly_blk_patch2_vnneg_3d), - int32(0), - ) - return ( - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask, - dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask, - dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask, - patch1_cell_idx_vmask, - patch1_cell_blk_vmask, - patch2_cell_idx_vmask, - patch2_cell_blk_vmask, - ) + bf_cc_patch1_lon = np.where( + famask_bool, + np.where(lvn_pos, bf_cc_patch1_lon_e[:, 0], bf_cc_patch1_lon_e[:, 1]), + 0.0, + ) + bf_cc_patch1_lat = np.where( + famask_bool, + np.where(lvn_pos, bf_cc_patch1_lat_e[:, 0], bf_cc_patch1_lat_e[:, 1]), + 0.0, + ) + bf_cc_patch2_lon = np.where( + famask_bool, + np.where(lvn_pos, bf_cc_patch2_lon_e[:, 0], bf_cc_patch2_lon_e[:, 1]), + 0.0, + ) + bf_cc_patch2_lat = np.where( + famask_bool, + np.where(lvn_pos, bf_cc_patch2_lat_e[:, 0], bf_cc_patch2_lat_e[:, 1]), + 0.0, + ) + # patch1 in translated system + dreg_patch1_1_lon_vmask = dreg_patch1_1_lon_vmask - bf_cc_patch1_lon + dreg_patch1_1_lat_vmask = dreg_patch1_1_lat_vmask - bf_cc_patch1_lat + dreg_patch1_2_lon_vmask = dreg_patch1_2_lon_vmask - bf_cc_patch1_lon + dreg_patch1_2_lat_vmask = dreg_patch1_2_lat_vmask - bf_cc_patch1_lat + dreg_patch1_3_lon_vmask = dreg_patch1_3_lon_vmask - bf_cc_patch1_lon + dreg_patch1_3_lat_vmask = dreg_patch1_3_lat_vmask - bf_cc_patch1_lat + dreg_patch1_4_lon_vmask = dreg_patch1_4_lon_vmask - bf_cc_patch1_lon + dreg_patch1_4_lat_vmask = dreg_patch1_4_lat_vmask - bf_cc_patch1_lat + # patch2 in translated system + dreg_patch2_1_lon_vmask = dreg_patch2_1_lon_vmask - bf_cc_patch2_lon + dreg_patch2_1_lat_vmask = dreg_patch2_1_lat_vmask - bf_cc_patch2_lat + dreg_patch2_2_lon_vmask = dreg_patch2_2_lon_vmask - bf_cc_patch2_lon + dreg_patch2_2_lat_vmask = dreg_patch2_2_lat_vmask - bf_cc_patch2_lat + dreg_patch2_3_lon_vmask = dreg_patch2_3_lon_vmask - bf_cc_patch2_lon + dreg_patch2_3_lat_vmask = dreg_patch2_3_lat_vmask - bf_cc_patch2_lat + dreg_patch2_4_lon_vmask = dreg_patch2_4_lon_vmask - bf_cc_patch2_lon + dreg_patch2_4_lat_vmask = dreg_patch2_4_lat_vmask - bf_cc_patch2_lat -def test_divide_flux_area_list_stencil_02(backend): - grid = SimpleGrid() + # Store global index of the underlying grid cell + # Adapt dimensions to fit ofr multiple levels + butterfly_idx_patch1_vnpos_3d = np.broadcast_to( + np.expand_dims(butterfly_idx_patch1_vnpos, axis=-1), p_vn.shape + ) + butterfly_idx_patch1_vnneg_3d = np.broadcast_to( + np.expand_dims(butterfly_idx_patch1_vnneg, axis=-1), p_vn.shape + ) + butterfly_idx_patch2_vnpos_3d = np.broadcast_to( + np.expand_dims(butterfly_idx_patch2_vnpos, axis=-1), p_vn.shape + ) + butterfly_idx_patch2_vnneg_3d = np.broadcast_to( + np.expand_dims(butterfly_idx_patch2_vnneg, axis=-1), p_vn.shape + ) + butterfly_blk_patch1_vnpos_3d = np.broadcast_to( + np.expand_dims(butterfly_blk_patch1_vnpos, axis=-1), p_vn.shape + ) + butterfly_blk_patch1_vnneg_3d = np.broadcast_to( + np.expand_dims(butterfly_blk_patch1_vnneg, axis=-1), p_vn.shape + ) + butterfly_blk_patch2_vnpos_3d = np.broadcast_to( + np.expand_dims(butterfly_blk_patch2_vnpos, axis=-1), p_vn.shape + ) + butterfly_blk_patch2_vnneg_3d = np.broadcast_to( + np.expand_dims(butterfly_blk_patch2_vnneg, axis=-1), p_vn.shape + ) + patch1_cell_idx_vmask = np.where( + famask_bool, + np.where(lvn_pos, butterfly_idx_patch1_vnpos_3d, butterfly_idx_patch1_vnneg_3d), + int32(0), + ) + patch2_cell_idx_vmask = np.where( + famask_bool, + np.where(lvn_pos, butterfly_idx_patch2_vnpos_3d, butterfly_idx_patch2_vnneg_3d), + int32(0), + ) + patch1_cell_blk_vmask = np.where( + famask_bool, + np.where(lvn_pos, butterfly_blk_patch1_vnpos_3d, butterfly_blk_patch1_vnneg_3d), + int32(0), + ) + patch2_cell_blk_vmask = np.where( + famask_bool, + np.where(lvn_pos, butterfly_blk_patch2_vnpos_3d, butterfly_blk_patch2_vnneg_3d), + int32(0), + ) - famask_int = random_mask(grid, EdgeDim, KDim, dtype=int32) - p_vn = random_field(grid, EdgeDim, KDim) - bf_cc_patch1_lon = random_field(grid, EdgeDim, E2CDim) - bf_cc_patch1_lon_field = as_1D_sparse_field(bf_cc_patch1_lon, ECDim) - bf_cc_patch1_lat = random_field(grid, EdgeDim, E2CDim) - bf_cc_patch1_lat_field = as_1D_sparse_field(bf_cc_patch1_lat, ECDim) - bf_cc_patch2_lon = random_field(grid, EdgeDim, E2CDim) - bf_cc_patch2_lon_field = as_1D_sparse_field(bf_cc_patch2_lon, ECDim) - bf_cc_patch2_lat = random_field(grid, EdgeDim, E2CDim) - bf_cc_patch2_lat_field = as_1D_sparse_field(bf_cc_patch2_lat, ECDim) - butterfly_idx_patch1_vnpos = random_mask(grid, EdgeDim, dtype=int32) - butterfly_idx_patch1_vnneg = random_mask(grid, EdgeDim, dtype=int32) - butterfly_blk_patch1_vnpos = random_mask(grid, EdgeDim, dtype=int32) - butterfly_blk_patch1_vnneg = random_mask(grid, EdgeDim, dtype=int32) - butterfly_idx_patch2_vnpos = random_mask(grid, EdgeDim, dtype=int32) - butterfly_idx_patch2_vnneg = random_mask(grid, EdgeDim, dtype=int32) - butterfly_blk_patch2_vnpos = random_mask(grid, EdgeDim, dtype=int32) - butterfly_blk_patch2_vnneg = random_mask(grid, EdgeDim, dtype=int32) - dreg_patch1_1_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_1_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_2_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_2_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_3_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_3_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_4_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch1_4_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_1_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_1_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_2_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_2_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_3_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_3_lat_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_4_lon_vmask = random_field(grid, EdgeDim, KDim) - dreg_patch2_4_lat_vmask = random_field(grid, EdgeDim, KDim) - patch1_cell_idx_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) - patch1_cell_blk_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) - patch2_cell_idx_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) - patch2_cell_blk_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) + return dict( + dreg_patch1_1_lon_vmask=dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask=dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask=dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask=dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask=dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask=dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask=dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask=dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask=dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask=dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask=dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask=dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask=dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask=dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask=dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask=dreg_patch2_4_lat_vmask, + patch1_cell_idx_vmask=patch1_cell_idx_vmask, + patch1_cell_blk_vmask=patch1_cell_blk_vmask, + patch2_cell_idx_vmask=patch2_cell_idx_vmask, + patch2_cell_blk_vmask=patch2_cell_blk_vmask, + ) - ( - ref_1, - ref_2, - ref_3, - ref_4, - ref_5, - ref_6, - ref_7, - ref_8, - ref_9, - ref_10, - ref_11, - ref_12, - ref_13, - ref_14, - ref_15, - ref_16, - ref_17, - ref_18, - ref_19, - ref_20, - ) = divide_flux_area_list_stencil_02_numpy( - grid.connectivities[E2CDim], - famask_int.asnumpy(), - p_vn.asnumpy(), - bf_cc_patch1_lon.asnumpy(), - bf_cc_patch1_lat.asnumpy(), - bf_cc_patch2_lon.asnumpy(), - bf_cc_patch2_lat.asnumpy(), - butterfly_idx_patch1_vnpos.asnumpy(), - butterfly_idx_patch1_vnneg.asnumpy(), - butterfly_blk_patch1_vnpos.asnumpy(), - butterfly_blk_patch1_vnneg.asnumpy(), - butterfly_idx_patch2_vnpos.asnumpy(), - butterfly_idx_patch2_vnneg.asnumpy(), - butterfly_blk_patch2_vnpos.asnumpy(), - butterfly_blk_patch2_vnneg.asnumpy(), - dreg_patch1_1_lon_vmask.asnumpy(), - dreg_patch1_1_lat_vmask.asnumpy(), - dreg_patch1_2_lon_vmask.asnumpy(), - dreg_patch1_2_lat_vmask.asnumpy(), - dreg_patch1_3_lon_vmask.asnumpy(), - dreg_patch1_3_lat_vmask.asnumpy(), - dreg_patch1_4_lon_vmask.asnumpy(), - dreg_patch1_4_lat_vmask.asnumpy(), - dreg_patch2_1_lon_vmask.asnumpy(), - dreg_patch2_1_lat_vmask.asnumpy(), - dreg_patch2_2_lon_vmask.asnumpy(), - dreg_patch2_2_lat_vmask.asnumpy(), - dreg_patch2_3_lon_vmask.asnumpy(), - dreg_patch2_3_lat_vmask.asnumpy(), - dreg_patch2_4_lon_vmask.asnumpy(), - dreg_patch2_4_lat_vmask.asnumpy(), - ) + @pytest.fixture + def input_data(self, grid): + famask_int = random_mask(grid, EdgeDim, KDim, dtype=int32) + p_vn = random_field(grid, EdgeDim, KDim) + bf_cc_patch1_lon = random_field(grid, EdgeDim, E2CDim) + bf_cc_patch1_lon_field = as_1D_sparse_field(bf_cc_patch1_lon, ECDim) + bf_cc_patch1_lat = random_field(grid, EdgeDim, E2CDim) + bf_cc_patch1_lat_field = as_1D_sparse_field(bf_cc_patch1_lat, ECDim) + bf_cc_patch2_lon = random_field(grid, EdgeDim, E2CDim) + bf_cc_patch2_lon_field = as_1D_sparse_field(bf_cc_patch2_lon, ECDim) + bf_cc_patch2_lat = random_field(grid, EdgeDim, E2CDim) + bf_cc_patch2_lat_field = as_1D_sparse_field(bf_cc_patch2_lat, ECDim) + butterfly_idx_patch1_vnpos = random_mask(grid, EdgeDim, dtype=int32) + butterfly_idx_patch1_vnneg = random_mask(grid, EdgeDim, dtype=int32) + butterfly_blk_patch1_vnpos = random_mask(grid, EdgeDim, dtype=int32) + butterfly_blk_patch1_vnneg = random_mask(grid, EdgeDim, dtype=int32) + butterfly_idx_patch2_vnpos = random_mask(grid, EdgeDim, dtype=int32) + butterfly_idx_patch2_vnneg = random_mask(grid, EdgeDim, dtype=int32) + butterfly_blk_patch2_vnpos = random_mask(grid, EdgeDim, dtype=int32) + butterfly_blk_patch2_vnneg = random_mask(grid, EdgeDim, dtype=int32) + dreg_patch1_1_lon_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch1_1_lat_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch1_2_lon_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch1_2_lat_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch1_3_lon_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch1_3_lat_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch1_4_lon_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch1_4_lat_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch2_1_lon_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch2_1_lat_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch2_2_lon_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch2_2_lat_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch2_3_lon_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch2_3_lat_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch2_4_lon_vmask = random_field(grid, EdgeDim, KDim) + dreg_patch2_4_lat_vmask = random_field(grid, EdgeDim, KDim) + patch1_cell_idx_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) + patch1_cell_blk_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) + patch2_cell_idx_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) + patch2_cell_blk_vmask = random_mask(grid, EdgeDim, KDim, dtype=int32) - divide_flux_area_list_stencil_02.with_backend(backend)( - famask_int, - p_vn, - bf_cc_patch1_lon_field, - bf_cc_patch1_lat_field, - bf_cc_patch2_lon_field, - bf_cc_patch2_lat_field, - butterfly_idx_patch1_vnpos, - butterfly_idx_patch1_vnneg, - butterfly_blk_patch1_vnpos, - butterfly_blk_patch1_vnneg, - butterfly_idx_patch2_vnpos, - butterfly_idx_patch2_vnneg, - butterfly_blk_patch2_vnpos, - butterfly_blk_patch2_vnneg, - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask, - dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask, - dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask, - patch1_cell_idx_vmask, - patch1_cell_blk_vmask, - patch2_cell_idx_vmask, - patch2_cell_blk_vmask, - offset_provider={ - "E2C": grid.get_offset_provider("E2C"), - "E2EC": grid.get_offset_provider("E2EC"), - }, - ) - assert np.allclose(dreg_patch1_1_lon_vmask.asnumpy(), ref_1) - assert np.allclose(dreg_patch1_1_lat_vmask.asnumpy(), ref_2) - assert np.allclose(dreg_patch1_2_lon_vmask.asnumpy(), ref_3) - assert np.allclose(dreg_patch1_2_lat_vmask.asnumpy(), ref_4) - assert np.allclose(dreg_patch1_3_lon_vmask.asnumpy(), ref_5) - assert np.allclose(dreg_patch1_3_lat_vmask.asnumpy(), ref_6) - assert np.allclose(dreg_patch1_4_lon_vmask.asnumpy(), ref_7) - assert np.allclose(dreg_patch1_4_lat_vmask.asnumpy(), ref_8) - assert np.allclose(dreg_patch2_1_lon_vmask.asnumpy(), ref_9) - assert np.allclose(dreg_patch2_1_lat_vmask.asnumpy(), ref_10) - assert np.allclose(dreg_patch2_2_lon_vmask.asnumpy(), ref_11) - assert np.allclose(dreg_patch2_2_lat_vmask.asnumpy(), ref_12) - assert np.allclose(dreg_patch2_3_lon_vmask.asnumpy(), ref_13) - assert np.allclose(dreg_patch2_3_lat_vmask.asnumpy(), ref_14) - assert np.allclose(dreg_patch2_4_lon_vmask.asnumpy(), ref_15) - assert np.allclose(dreg_patch2_4_lat_vmask.asnumpy(), ref_16) - assert np.allclose(patch1_cell_idx_vmask.asnumpy(), ref_17) - assert np.allclose(patch1_cell_blk_vmask.asnumpy(), ref_18) - assert np.allclose(patch2_cell_idx_vmask.asnumpy(), ref_19) - assert np.allclose(patch2_cell_blk_vmask.asnumpy(), ref_20) + return dict( + famask_int=famask_int, + p_vn=p_vn, + bf_cc_patch1_lon=bf_cc_patch1_lon_field, + bf_cc_patch1_lat=bf_cc_patch1_lat_field, + bf_cc_patch2_lon=bf_cc_patch2_lon_field, + bf_cc_patch2_lat=bf_cc_patch2_lat_field, + butterfly_idx_patch1_vnpos=butterfly_idx_patch1_vnpos, + butterfly_idx_patch1_vnneg=butterfly_idx_patch1_vnneg, + butterfly_blk_patch1_vnpos=butterfly_blk_patch1_vnpos, + butterfly_blk_patch1_vnneg=butterfly_blk_patch1_vnneg, + butterfly_idx_patch2_vnpos=butterfly_idx_patch2_vnpos, + butterfly_idx_patch2_vnneg=butterfly_idx_patch2_vnneg, + butterfly_blk_patch2_vnpos=butterfly_blk_patch2_vnpos, + butterfly_blk_patch2_vnneg=butterfly_blk_patch2_vnneg, + dreg_patch1_1_lon_vmask=dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask=dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask=dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask=dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask=dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask=dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask=dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask=dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask=dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask=dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask=dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask=dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask=dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask=dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask=dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask=dreg_patch2_4_lat_vmask, + patch1_cell_idx_vmask=patch1_cell_idx_vmask, + patch1_cell_blk_vmask=patch1_cell_blk_vmask, + patch2_cell_idx_vmask=patch2_cell_idx_vmask, + patch2_cell_blk_vmask=patch2_cell_blk_vmask, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py index cc404c707c..acfa579ff3 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_01a.py @@ -12,193 +12,173 @@ # 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.advection.hflux_ffsl_hybrid_stencil_01a import ( hflux_ffsl_hybrid_stencil_01a, ) from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import constant_field, random_field, zero_field - - -def hflux_ffsl_hybrid_stencil_01a_numpy( - e2c: np.array, - z_lsq_coeff_1: np.ndarray, - z_lsq_coeff_2: np.ndarray, - z_lsq_coeff_3: np.ndarray, - z_lsq_coeff_4: np.ndarray, - z_lsq_coeff_5: np.ndarray, - z_lsq_coeff_6: np.ndarray, - z_lsq_coeff_7: np.ndarray, - z_lsq_coeff_8: np.ndarray, - z_lsq_coeff_9: np.ndarray, - z_lsq_coeff_10: np.ndarray, - z_quad_vector_sum0_1: np.ndarray, - z_quad_vector_sum0_2: np.ndarray, - z_quad_vector_sum0_3: np.ndarray, - z_quad_vector_sum0_4: np.ndarray, - z_quad_vector_sum0_5: np.ndarray, - z_quad_vector_sum0_6: np.ndarray, - z_quad_vector_sum0_7: np.ndarray, - z_quad_vector_sum0_8: np.ndarray, - z_quad_vector_sum0_9: np.ndarray, - z_quad_vector_sum0_10: np.ndarray, - patch0_cell_rel_idx_dsl: np.ndarray, -): - z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] - z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] - z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] - z_lsq_coeff_4_e2c = z_lsq_coeff_4[e2c] - z_lsq_coeff_5_e2c = z_lsq_coeff_5[e2c] - z_lsq_coeff_6_e2c = z_lsq_coeff_6[e2c] - z_lsq_coeff_7_e2c = z_lsq_coeff_7[e2c] - z_lsq_coeff_8_e2c = z_lsq_coeff_8[e2c] - z_lsq_coeff_9_e2c = z_lsq_coeff_9[e2c] - z_lsq_coeff_10_e2c = z_lsq_coeff_10[e2c] - - p_out_e_hybrid_1a = ( - np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_1_e2c[:, 1], - z_lsq_coeff_1_e2c[:, 0], - ) - * z_quad_vector_sum0_1 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_2_e2c[:, 1], - z_lsq_coeff_2_e2c[:, 0], - ) - * z_quad_vector_sum0_2 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_3_e2c[:, 1], - z_lsq_coeff_3_e2c[:, 0], - ) - * z_quad_vector_sum0_3 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_4_e2c[:, 1], - z_lsq_coeff_4_e2c[:, 0], - ) - * z_quad_vector_sum0_4 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_5_e2c[:, 1], - z_lsq_coeff_5_e2c[:, 0], - ) - * z_quad_vector_sum0_5 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_6_e2c[:, 1], - z_lsq_coeff_6_e2c[:, 0], - ) - * z_quad_vector_sum0_6 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_7_e2c[:, 1], - z_lsq_coeff_7_e2c[:, 0], - ) - * z_quad_vector_sum0_7 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_8_e2c[:, 1], - z_lsq_coeff_8_e2c[:, 0], - ) - * z_quad_vector_sum0_8 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_9_e2c[:, 1], - z_lsq_coeff_9_e2c[:, 0], - ) - * z_quad_vector_sum0_9 - + np.where( - patch0_cell_rel_idx_dsl == int32(1), - z_lsq_coeff_10_e2c[:, 1], - z_lsq_coeff_10_e2c[:, 0], - ) - * z_quad_vector_sum0_10 - ) +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + constant_field, + random_field, + zero_field, +) - return p_out_e_hybrid_1a +class TestHfluxFfslHybridStencil01a(StencilTest): + PROGRAM = hflux_ffsl_hybrid_stencil_01a + OUTPUTS = ("p_out_e_hybrid_1a",) -def test_hflux_ffsl_hybrid_stencil_01a(backend): - grid = SimpleGrid() - z_lsq_coeff_1 = random_field(grid, CellDim, KDim) - z_lsq_coeff_2 = random_field(grid, CellDim, KDim) - z_lsq_coeff_3 = random_field(grid, CellDim, KDim) - z_lsq_coeff_4 = random_field(grid, CellDim, KDim) - z_lsq_coeff_5 = random_field(grid, CellDim, KDim) - z_lsq_coeff_6 = random_field(grid, CellDim, KDim) - z_lsq_coeff_7 = random_field(grid, CellDim, KDim) - z_lsq_coeff_8 = random_field(grid, CellDim, KDim) - z_lsq_coeff_9 = random_field(grid, CellDim, KDim) - z_lsq_coeff_10 = random_field(grid, CellDim, KDim) - z_quad_vector_sum0_1 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_2 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_3 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_4 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_5 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_6 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_7 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_8 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_9 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum0_10 = random_field(grid, EdgeDim, KDim) - patch0_cell_rel_idx_dsl = constant_field(grid, 1, EdgeDim, KDim, dtype=int32) - p_out_e_hybrid_1a = zero_field(grid, EdgeDim, KDim) + @staticmethod + def reference( + grid, + z_lsq_coeff_1: np.array, + z_lsq_coeff_2: np.array, + z_lsq_coeff_3: np.array, + z_lsq_coeff_4: np.array, + z_lsq_coeff_5: np.array, + z_lsq_coeff_6: np.array, + z_lsq_coeff_7: np.array, + z_lsq_coeff_8: np.array, + z_lsq_coeff_9: np.array, + z_lsq_coeff_10: np.array, + z_quad_vector_sum0_1: np.array, + z_quad_vector_sum0_2: np.array, + z_quad_vector_sum0_3: np.array, + z_quad_vector_sum0_4: np.array, + z_quad_vector_sum0_5: np.array, + z_quad_vector_sum0_6: np.array, + z_quad_vector_sum0_7: np.array, + z_quad_vector_sum0_8: np.array, + z_quad_vector_sum0_9: np.array, + z_quad_vector_sum0_10: np.array, + patch0_cell_rel_idx_dsl: np.array, + **kwargs, + ): + e2c = grid.connectivities[E2CDim] + z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] + z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] + z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] + z_lsq_coeff_4_e2c = z_lsq_coeff_4[e2c] + z_lsq_coeff_5_e2c = z_lsq_coeff_5[e2c] + z_lsq_coeff_6_e2c = z_lsq_coeff_6[e2c] + z_lsq_coeff_7_e2c = z_lsq_coeff_7[e2c] + z_lsq_coeff_8_e2c = z_lsq_coeff_8[e2c] + z_lsq_coeff_9_e2c = z_lsq_coeff_9[e2c] + z_lsq_coeff_10_e2c = z_lsq_coeff_10[e2c] - ref = hflux_ffsl_hybrid_stencil_01a_numpy( - grid.connectivities[E2CDim], - z_lsq_coeff_1.asnumpy(), - z_lsq_coeff_2.asnumpy(), - z_lsq_coeff_3.asnumpy(), - z_lsq_coeff_4.asnumpy(), - z_lsq_coeff_5.asnumpy(), - z_lsq_coeff_6.asnumpy(), - z_lsq_coeff_7.asnumpy(), - z_lsq_coeff_8.asnumpy(), - z_lsq_coeff_9.asnumpy(), - z_lsq_coeff_10.asnumpy(), - z_quad_vector_sum0_1.asnumpy(), - z_quad_vector_sum0_2.asnumpy(), - z_quad_vector_sum0_3.asnumpy(), - z_quad_vector_sum0_4.asnumpy(), - z_quad_vector_sum0_5.asnumpy(), - z_quad_vector_sum0_6.asnumpy(), - z_quad_vector_sum0_7.asnumpy(), - z_quad_vector_sum0_8.asnumpy(), - z_quad_vector_sum0_9.asnumpy(), - z_quad_vector_sum0_10.asnumpy(), - patch0_cell_rel_idx_dsl.asnumpy(), - ) + p_out_e_hybrid_1a = ( + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_e2c[:, 1], + z_lsq_coeff_1_e2c[:, 0], + ) + * z_quad_vector_sum0_1 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_e2c[:, 1], + z_lsq_coeff_2_e2c[:, 0], + ) + * z_quad_vector_sum0_2 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_e2c[:, 1], + z_lsq_coeff_3_e2c[:, 0], + ) + * z_quad_vector_sum0_3 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_4_e2c[:, 1], + z_lsq_coeff_4_e2c[:, 0], + ) + * z_quad_vector_sum0_4 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_5_e2c[:, 1], + z_lsq_coeff_5_e2c[:, 0], + ) + * z_quad_vector_sum0_5 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_6_e2c[:, 1], + z_lsq_coeff_6_e2c[:, 0], + ) + * z_quad_vector_sum0_6 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_7_e2c[:, 1], + z_lsq_coeff_7_e2c[:, 0], + ) + * z_quad_vector_sum0_7 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_8_e2c[:, 1], + z_lsq_coeff_8_e2c[:, 0], + ) + * z_quad_vector_sum0_8 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_9_e2c[:, 1], + z_lsq_coeff_9_e2c[:, 0], + ) + * z_quad_vector_sum0_9 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_10_e2c[:, 1], + z_lsq_coeff_10_e2c[:, 0], + ) + * z_quad_vector_sum0_10 + ) - hflux_ffsl_hybrid_stencil_01a.with_backend(backend)( - z_lsq_coeff_1, - z_lsq_coeff_2, - z_lsq_coeff_3, - z_lsq_coeff_4, - z_lsq_coeff_5, - z_lsq_coeff_6, - z_lsq_coeff_7, - z_lsq_coeff_8, - z_lsq_coeff_9, - z_lsq_coeff_10, - z_quad_vector_sum0_1, - z_quad_vector_sum0_2, - z_quad_vector_sum0_3, - z_quad_vector_sum0_4, - z_quad_vector_sum0_5, - z_quad_vector_sum0_6, - z_quad_vector_sum0_7, - z_quad_vector_sum0_8, - z_quad_vector_sum0_9, - z_quad_vector_sum0_10, - patch0_cell_rel_idx_dsl, - p_out_e_hybrid_1a, - offset_provider={ - "E2C": grid.get_offset_provider("E2C"), - }, - ) + return dict(p_out_e_hybrid_1a=p_out_e_hybrid_1a) - assert np.allclose(p_out_e_hybrid_1a.asnumpy(), ref) + @pytest.fixture + def input_data(self, grid): + z_lsq_coeff_1 = random_field(grid, CellDim, KDim) + z_lsq_coeff_2 = random_field(grid, CellDim, KDim) + z_lsq_coeff_3 = random_field(grid, CellDim, KDim) + z_lsq_coeff_4 = random_field(grid, CellDim, KDim) + z_lsq_coeff_5 = random_field(grid, CellDim, KDim) + z_lsq_coeff_6 = random_field(grid, CellDim, KDim) + z_lsq_coeff_7 = random_field(grid, CellDim, KDim) + z_lsq_coeff_8 = random_field(grid, CellDim, KDim) + z_lsq_coeff_9 = random_field(grid, CellDim, KDim) + z_lsq_coeff_10 = random_field(grid, CellDim, KDim) + z_quad_vector_sum0_1 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_2 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_3 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_4 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_5 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_6 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_7 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_8 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_9 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum0_10 = random_field(grid, EdgeDim, KDim) + patch0_cell_rel_idx_dsl = constant_field(grid, 1, EdgeDim, KDim, dtype=int32) + p_out_e_hybrid_1a = zero_field(grid, EdgeDim, KDim) + return dict( + z_lsq_coeff_1=z_lsq_coeff_1, + z_lsq_coeff_2=z_lsq_coeff_2, + z_lsq_coeff_3=z_lsq_coeff_3, + z_lsq_coeff_4=z_lsq_coeff_4, + z_lsq_coeff_5=z_lsq_coeff_5, + z_lsq_coeff_6=z_lsq_coeff_6, + z_lsq_coeff_7=z_lsq_coeff_7, + z_lsq_coeff_8=z_lsq_coeff_8, + z_lsq_coeff_9=z_lsq_coeff_9, + z_lsq_coeff_10=z_lsq_coeff_10, + z_quad_vector_sum0_1=z_quad_vector_sum0_1, + z_quad_vector_sum0_2=z_quad_vector_sum0_2, + z_quad_vector_sum0_3=z_quad_vector_sum0_3, + z_quad_vector_sum0_4=z_quad_vector_sum0_4, + z_quad_vector_sum0_5=z_quad_vector_sum0_5, + z_quad_vector_sum0_6=z_quad_vector_sum0_6, + z_quad_vector_sum0_7=z_quad_vector_sum0_7, + z_quad_vector_sum0_8=z_quad_vector_sum0_8, + z_quad_vector_sum0_9=z_quad_vector_sum0_9, + z_quad_vector_sum0_10=z_quad_vector_sum0_10, + patch0_cell_rel_idx_dsl=patch0_cell_rel_idx_dsl, + p_out_e_hybrid_1a=p_out_e_hybrid_1a, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py index 6cce67d1fd..5c57bf79b9 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflux_ffsl_hybrid_stencil_02.py @@ -12,42 +12,38 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.hflux_ffsl_hybrid_stencil_02 import ( hflux_ffsl_hybrid_stencil_02, ) from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field - - -def hflux_ffsl_hybrid_stencil_02_numpy( - p_out_e_hybrid_2: np.ndarray, - p_mass_flx_e: np.ndarray, - z_dreg_area: np.ndarray, -): - p_out_e_hybrid_2 = p_mass_flx_e * p_out_e_hybrid_2 / z_dreg_area - - return p_out_e_hybrid_2 - - -def test_hflux_ffsl_hybrid_stencil_02(backend): - grid = SimpleGrid() - p_out_e_hybrid_2 = random_field(grid, EdgeDim, KDim) - p_mass_flx_e = random_field(grid, EdgeDim, KDim) - z_dreg_area = random_field(grid, EdgeDim, KDim) - - ref = hflux_ffsl_hybrid_stencil_02_numpy( - p_out_e_hybrid_2.asnumpy(), - p_mass_flx_e.asnumpy(), - z_dreg_area.asnumpy(), - ) - - hflux_ffsl_hybrid_stencil_02.with_backend(backend)( - p_out_e_hybrid_2, - p_mass_flx_e, - z_dreg_area, - offset_provider={}, - ) - - assert np.allclose(p_out_e_hybrid_2.asnumpy(), ref) +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestHfluxFfslHybridStencil02(StencilTest): + PROGRAM = hflux_ffsl_hybrid_stencil_02 + OUTPUTS = ("p_out_e_hybrid_2",) + + @staticmethod + def reference( + grid, + p_out_e_hybrid_2: np.array, + p_mass_flx_e: np.array, + z_dreg_area: np.array, + **kwargs, + ): + p_out_e_hybrid_2 = p_mass_flx_e * p_out_e_hybrid_2 / z_dreg_area + + return dict(p_out_e_hybrid_2=p_out_e_hybrid_2) + + @pytest.fixture + def input_data(self, grid): + p_out_e_hybrid_2 = random_field(grid, EdgeDim, KDim) + p_mass_flx_e = random_field(grid, EdgeDim, KDim) + z_dreg_area = random_field(grid, EdgeDim, KDim) + return dict( + p_mass_flx_e=p_mass_flx_e, + z_dreg_area=z_dreg_area, + p_out_e_hybrid_2=p_out_e_hybrid_2, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py index 8d6e241496..ba9ba72d47 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py @@ -19,415 +19,199 @@ prep_gauss_quadrature_c_list_stencil, ) from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import constant_field, random_field, zero_field - - -def prep_gauss_quadrature_c_list_stencil_numpy( - famask_int: np.array, - p_coords_dreg_v_1_x: np.array, - p_coords_dreg_v_2_x: np.array, - p_coords_dreg_v_3_x: np.array, - p_coords_dreg_v_4_x: np.array, - p_coords_dreg_v_1_y: np.array, - p_coords_dreg_v_2_y: np.array, - p_coords_dreg_v_3_y: np.array, - p_coords_dreg_v_4_y: np.array, - shape_func_1_1: float, - shape_func_2_1: float, - shape_func_3_1: float, - shape_func_4_1: float, - shape_func_1_2: float, - shape_func_2_2: float, - shape_func_3_2: float, - shape_func_4_2: float, - shape_func_1_3: float, - shape_func_2_3: float, - shape_func_3_3: float, - shape_func_4_3: float, - shape_func_1_4: float, - shape_func_2_4: float, - shape_func_3_4: float, - shape_func_4_4: float, - zeta_1: float, - zeta_2: float, - zeta_3: float, - zeta_4: float, - eta_1: float, - eta_2: float, - eta_3: float, - eta_4: float, - wgt_zeta_1: float, - wgt_zeta_2: float, - wgt_eta_1: float, - wgt_eta_2: float, - dbl_eps: float, - eps: float, - p_dreg_area_in: np.array, -) -> tuple[np.ndarray]: - z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 - z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 - z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 - z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 - - z_eta_1_1 = 1.0 - eta_1 - z_eta_2_1 = 1.0 - eta_2 - z_eta_3_1 = 1.0 - eta_3 - z_eta_4_1 = 1.0 - eta_4 - z_eta_1_2 = 1.0 + eta_1 - z_eta_2_2 = 1.0 + eta_2 - z_eta_3_2 = 1.0 + eta_3 - z_eta_4_2 = 1.0 + eta_4 - z_eta_1_3 = 1.0 - zeta_1 - z_eta_2_3 = 1.0 - zeta_2 - z_eta_3_3 = 1.0 - zeta_3 - z_eta_4_3 = 1.0 - zeta_4 - z_eta_1_4 = 1.0 + zeta_1 - z_eta_2_4 = 1.0 + zeta_2 - z_eta_3_4 = 1.0 + zeta_3 - z_eta_4_4 = 1.0 + zeta_4 +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + constant_field, + random_field, + zero_field, +) - famask_bool = np.where(famask_int == int32(1), True, False) - p_coords_dreg_v_1_x = np.where(famask_bool, p_coords_dreg_v_1_x, 0.0) - p_coords_dreg_v_2_x = np.where(famask_bool, p_coords_dreg_v_2_x, 0.0) - p_coords_dreg_v_3_x = np.where(famask_bool, p_coords_dreg_v_3_x, 0.0) - p_coords_dreg_v_4_x = np.where(famask_bool, p_coords_dreg_v_4_x, 0.0) - p_coords_dreg_v_1_y = np.where(famask_bool, p_coords_dreg_v_1_y, 0.0) - p_coords_dreg_v_2_y = np.where(famask_bool, p_coords_dreg_v_2_y, 0.0) - p_coords_dreg_v_3_y = np.where(famask_bool, p_coords_dreg_v_3_y, 0.0) - p_coords_dreg_v_4_y = np.where(famask_bool, p_coords_dreg_v_4_y, 0.0) - wgt_t_detjac_1 = np.where( - famask_bool, - dbl_eps - + z_wgt_1 - * ( - ( - z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) - ), - 0.0, - ) - wgt_t_detjac_2 = np.where( - famask_bool, - dbl_eps - + z_wgt_2 - * ( - ( - z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) - ), - 0.0, - ) - wgt_t_detjac_3 = np.where( - famask_bool, - dbl_eps - + z_wgt_3 - * ( - ( - z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) - ), - 0.0, - ) - wgt_t_detjac_4 = np.where( - famask_bool, - dbl_eps - + z_wgt_4 - * ( - ( - z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) - ), - 0.0, +class TestPrepGaussQuadratureCListStencil(StencilTest): + PROGRAM = prep_gauss_quadrature_c_list_stencil + OUTPUTS = ( + "p_quad_vector_sum_1", + "p_quad_vector_sum_2", + "p_quad_vector_sum_3", + "p_quad_vector_sum_4", + "p_quad_vector_sum_5", + "p_quad_vector_sum_6", + "p_quad_vector_sum_7", + "p_quad_vector_sum_8", + "p_quad_vector_sum_9", + "p_quad_vector_sum_10", + "p_dreg_area", ) - z_gauss_pts_1_x = ( - shape_func_1_1 * p_coords_dreg_v_1_x - + shape_func_2_1 * p_coords_dreg_v_2_x - + shape_func_3_1 * p_coords_dreg_v_3_x - + shape_func_4_1 * p_coords_dreg_v_4_x - ) - z_gauss_pts_1_y = ( - shape_func_1_1 * p_coords_dreg_v_1_y - + shape_func_2_1 * p_coords_dreg_v_2_y - + shape_func_3_1 * p_coords_dreg_v_3_y - + shape_func_4_1 * p_coords_dreg_v_4_y - ) - z_gauss_pts_2_x = ( - shape_func_1_2 * p_coords_dreg_v_1_x - + shape_func_2_2 * p_coords_dreg_v_2_x - + shape_func_3_2 * p_coords_dreg_v_3_x - + shape_func_4_2 * p_coords_dreg_v_4_x - ) - z_gauss_pts_2_y = ( - shape_func_1_2 * p_coords_dreg_v_1_y - + shape_func_2_2 * p_coords_dreg_v_2_y - + shape_func_3_2 * p_coords_dreg_v_3_y - + shape_func_4_2 * p_coords_dreg_v_4_y - ) - z_gauss_pts_3_x = ( - shape_func_1_3 * p_coords_dreg_v_1_x - + shape_func_2_3 * p_coords_dreg_v_2_x - + shape_func_3_3 * p_coords_dreg_v_3_x - + shape_func_4_3 * p_coords_dreg_v_4_x - ) - z_gauss_pts_3_y = ( - shape_func_1_3 * p_coords_dreg_v_1_y - + shape_func_2_3 * p_coords_dreg_v_2_y - + shape_func_3_3 * p_coords_dreg_v_3_y - + shape_func_4_3 * p_coords_dreg_v_4_y - ) - z_gauss_pts_4_x = ( - shape_func_1_4 * p_coords_dreg_v_1_x - + shape_func_2_4 * p_coords_dreg_v_2_x - + shape_func_3_4 * p_coords_dreg_v_3_x - + shape_func_4_4 * p_coords_dreg_v_4_x - ) - z_gauss_pts_4_y = ( - shape_func_1_4 * p_coords_dreg_v_1_y - + shape_func_2_4 * p_coords_dreg_v_2_y - + shape_func_3_4 * p_coords_dreg_v_3_y - + shape_func_4_4 * p_coords_dreg_v_4_y - ) + @staticmethod + def _compute_wgt_t_detjac( + wgt_zeta_1, + wgt_eta_1, + wgt_zeta_2, + wgt_eta_2, + eta_1, + eta_2, + eta_3, + eta_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + famask_int, + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + dbl_eps, + ): + z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 + z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 + z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 + z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 - p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 - p_quad_vector_sum_2 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x - ) - p_quad_vector_sum_3 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y - ) - p_quad_vector_sum_4 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x - ) - p_quad_vector_sum_5 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y - ) - p_quad_vector_sum_6 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y - ) - p_quad_vector_sum_7 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x - ) - p_quad_vector_sum_8 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y - ) - p_quad_vector_sum_9 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y - ) - p_quad_vector_sum_10 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y - ) + z_eta_1_1, z_eta_2_1, z_eta_3_1, z_eta_4_1 = ( + 1.0 - eta_1, + 1.0 - eta_2, + 1.0 - eta_3, + 1.0 - eta_4, + ) + z_eta_1_2, z_eta_2_2, z_eta_3_2, z_eta_4_2 = ( + 1.0 + eta_1, + 1.0 + eta_2, + 1.0 + eta_3, + 1.0 + eta_4, + ) + z_eta_1_3, z_eta_2_3, z_eta_3_3, z_eta_4_3 = ( + 1.0 - zeta_1, + 1.0 - zeta_2, + 1.0 - zeta_3, + 1.0 - zeta_4, + ) + z_eta_1_4, z_eta_2_4, z_eta_3_4, z_eta_4_4 = ( + 1.0 + zeta_1, + 1.0 + zeta_2, + 1.0 + zeta_3, + 1.0 + zeta_4, + ) - p_dreg_area = p_dreg_area_in + p_quad_vector_sum_1 - return ( - p_quad_vector_sum_1, - p_quad_vector_sum_2, - p_quad_vector_sum_3, - p_quad_vector_sum_4, - p_quad_vector_sum_5, - p_quad_vector_sum_6, - p_quad_vector_sum_7, - p_quad_vector_sum_8, - p_quad_vector_sum_9, - p_quad_vector_sum_10, - p_dreg_area, - ) + famask_bool = np.where(famask_int == np.int32(1), True, False) + p_coords_dreg_v_1_x = np.where(famask_bool, p_coords_dreg_v_1_x, 0.0) + p_coords_dreg_v_2_x = np.where(famask_bool, p_coords_dreg_v_2_x, 0.0) + p_coords_dreg_v_3_x = np.where(famask_bool, p_coords_dreg_v_3_x, 0.0) + p_coords_dreg_v_4_x = np.where(famask_bool, p_coords_dreg_v_4_x, 0.0) + p_coords_dreg_v_1_y = np.where(famask_bool, p_coords_dreg_v_1_y, 0.0) + p_coords_dreg_v_2_y = np.where(famask_bool, p_coords_dreg_v_2_y, 0.0) + p_coords_dreg_v_3_y = np.where(famask_bool, p_coords_dreg_v_3_y, 0.0) + p_coords_dreg_v_4_y = np.where(famask_bool, p_coords_dreg_v_4_y, 0.0) -@pytest.mark.slow_tests -def test_prep_gauss_quadrature_c_list_stencil(backend): - grid = SimpleGrid() + wgt_t_detjac_1 = np.where( + famask_bool, + dbl_eps + + z_wgt_1 + * ( + ( + z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) - famask_int = constant_field(grid, 1, EdgeDim, KDim, dtype=int32) - p_coords_dreg_v_1_x = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_2_x = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_3_x = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_4_x = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_1_y = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_2_y = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_3_y = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_4_y = random_field(grid, EdgeDim, KDim) - shape_func_1_1 = 0.001 - shape_func_2_1 = 0.001 - shape_func_3_1 = 0.001 - shape_func_4_1 = 0.001 - shape_func_1_2 = 0.001 - shape_func_2_2 = 0.001 - shape_func_3_2 = 0.001 - shape_func_4_2 = 0.001 - shape_func_1_3 = 0.001 - shape_func_2_3 = 0.001 - shape_func_3_3 = 0.001 - shape_func_4_3 = 0.001 - shape_func_1_4 = 0.001 - shape_func_2_4 = 0.001 - shape_func_3_4 = 0.001 - shape_func_4_4 = 0.001 - zeta_1 = 0.002 - zeta_2 = 0.002 - zeta_3 = 0.002 - zeta_4 = 0.002 - eta_1 = 0.5 - eta_2 = 0.5 - eta_3 = 0.5 - eta_4 = 0.5 - wgt_zeta_1 = 0.003 - wgt_zeta_2 = 0.003 - wgt_eta_1 = 0.002 - wgt_eta_2 = 0.007 - dbl_eps = np.float64(0.1) - eps = 0.1 - p_dreg_area_in = random_field(grid, EdgeDim, KDim) - p_quad_vector_sum_1 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_2 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_3 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_4 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_5 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_6 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_7 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_8 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_9 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_10 = zero_field(grid, EdgeDim, KDim) - p_dreg_area = zero_field(grid, EdgeDim, KDim) + wgt_t_detjac_2 = np.where( + famask_bool, + dbl_eps + + z_wgt_2 + * ( + ( + z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) - ( - ref_1, - ref_2, - ref_3, - ref_4, - ref_5, - ref_6, - ref_7, - ref_8, - ref_9, - ref_10, - ref_11, - ) = prep_gauss_quadrature_c_list_stencil_numpy( - np.asarray(famask_int), - np.asarray(p_coords_dreg_v_1_x), - np.asarray(p_coords_dreg_v_2_x), - np.asarray(p_coords_dreg_v_3_x), - np.asarray(p_coords_dreg_v_4_x), - np.asarray(p_coords_dreg_v_1_y), - np.asarray(p_coords_dreg_v_2_y), - np.asarray(p_coords_dreg_v_3_y), - np.asarray(p_coords_dreg_v_4_y), - shape_func_1_1, - shape_func_2_1, - shape_func_3_1, - shape_func_4_1, - shape_func_1_2, - shape_func_2_2, - shape_func_3_2, - shape_func_4_2, - shape_func_1_3, - shape_func_2_3, - shape_func_3_3, - shape_func_4_3, - shape_func_1_4, - shape_func_2_4, - shape_func_3_4, - shape_func_4_4, - zeta_1, - zeta_2, - zeta_3, - zeta_4, - eta_1, - eta_2, - eta_3, - eta_4, - wgt_zeta_1, - wgt_zeta_2, - wgt_eta_1, - wgt_eta_2, - dbl_eps, - eps, - np.asarray(p_dreg_area_in), - ) + wgt_t_detjac_3 = np.where( + famask_bool, + dbl_eps + + z_wgt_3 + * ( + ( + z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + wgt_t_detjac_4 = np.where( + famask_bool, + dbl_eps + + z_wgt_4 + * ( + ( + z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) - prep_gauss_quadrature_c_list_stencil.with_backend(backend)( - famask_int, + return wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 + + @staticmethod + def _compute_z_gauss_points( p_coords_dreg_v_1_x, p_coords_dreg_v_2_x, p_coords_dreg_v_3_x, @@ -452,53 +236,401 @@ def test_prep_gauss_quadrature_c_list_stencil(backend): shape_func_2_4, shape_func_3_4, shape_func_4_4, - zeta_1, - zeta_2, - zeta_3, - zeta_4, - eta_1, - eta_2, - eta_3, - eta_4, - wgt_zeta_1, - wgt_zeta_2, - wgt_eta_1, - wgt_eta_2, - dbl_eps, - eps, - p_dreg_area_in, - p_quad_vector_sum_1, - p_quad_vector_sum_2, - p_quad_vector_sum_3, - p_quad_vector_sum_4, - p_quad_vector_sum_5, - p_quad_vector_sum_6, - p_quad_vector_sum_7, - p_quad_vector_sum_8, - p_quad_vector_sum_9, - p_quad_vector_sum_10, - p_dreg_area, - offset_provider={}, - ) - co1 = np.asarray(p_quad_vector_sum_1) - co2 = np.asarray(p_quad_vector_sum_2) - co3 = np.asarray(p_quad_vector_sum_3) - co4 = np.asarray(p_quad_vector_sum_4) - co5 = np.asarray(p_quad_vector_sum_5) - co6 = np.asarray(p_quad_vector_sum_6) - co7 = np.asarray(p_quad_vector_sum_7) - co8 = np.asarray(p_quad_vector_sum_8) - co9 = np.asarray(p_quad_vector_sum_9) - co10 = np.asarray(p_quad_vector_sum_10) - co11 = np.asarray(p_dreg_area) - assert np.allclose(ref_1, co1) - assert np.allclose(ref_2, co2) - assert np.allclose(ref_3, co3) - assert np.allclose(ref_4, co4) - assert np.allclose(ref_5, co5) - assert np.allclose(ref_6, co6) - assert np.allclose(ref_7, co7) - assert np.allclose(ref_8, co8) - assert np.allclose(ref_9, co9) - assert np.allclose(ref_10, co10) - assert np.allclose(ref_11, co11) + ): + z_gauss_pts_1_x = ( + shape_func_1_1 * p_coords_dreg_v_1_x + + shape_func_2_1 * p_coords_dreg_v_2_x + + shape_func_3_1 * p_coords_dreg_v_3_x + + shape_func_4_1 * p_coords_dreg_v_4_x + ) + z_gauss_pts_1_y = ( + shape_func_1_1 * p_coords_dreg_v_1_y + + shape_func_2_1 * p_coords_dreg_v_2_y + + shape_func_3_1 * p_coords_dreg_v_3_y + + shape_func_4_1 * p_coords_dreg_v_4_y + ) + z_gauss_pts_2_x = ( + shape_func_1_2 * p_coords_dreg_v_1_x + + shape_func_2_2 * p_coords_dreg_v_2_x + + shape_func_3_2 * p_coords_dreg_v_3_x + + shape_func_4_2 * p_coords_dreg_v_4_x + ) + z_gauss_pts_2_y = ( + shape_func_1_2 * p_coords_dreg_v_1_y + + shape_func_2_2 * p_coords_dreg_v_2_y + + shape_func_3_2 * p_coords_dreg_v_3_y + + shape_func_4_2 * p_coords_dreg_v_4_y + ) + z_gauss_pts_3_x = ( + shape_func_1_3 * p_coords_dreg_v_1_x + + shape_func_2_3 * p_coords_dreg_v_2_x + + shape_func_3_3 * p_coords_dreg_v_3_x + + shape_func_4_3 * p_coords_dreg_v_4_x + ) + z_gauss_pts_3_y = ( + shape_func_1_3 * p_coords_dreg_v_1_y + + shape_func_2_3 * p_coords_dreg_v_2_y + + shape_func_3_3 * p_coords_dreg_v_3_y + + shape_func_4_3 * p_coords_dreg_v_4_y + ) + z_gauss_pts_4_x = ( + shape_func_1_4 * p_coords_dreg_v_1_x + + shape_func_2_4 * p_coords_dreg_v_2_x + + shape_func_3_4 * p_coords_dreg_v_3_x + + shape_func_4_4 * p_coords_dreg_v_4_x + ) + z_gauss_pts_4_y = ( + shape_func_1_4 * p_coords_dreg_v_1_y + + shape_func_2_4 * p_coords_dreg_v_2_y + + shape_func_3_4 * p_coords_dreg_v_3_y + + shape_func_4_4 * p_coords_dreg_v_4_y + ) + return ( + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, + ) + + @staticmethod + def _compute_vector_sums( + wgt_t_detjac_1, + wgt_t_detjac_2, + wgt_t_detjac_3, + wgt_t_detjac_4, + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, + ): + p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 + p_quad_vector_sum_2 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x + ) + p_quad_vector_sum_3 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y + ) + p_quad_vector_sum_4 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_5 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_6 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_7 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_8 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_9 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_10 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + return ( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + ) + + @classmethod + def reference( + cls, + grid, + famask_int: np.array, + p_coords_dreg_v_1_x: np.array, + p_coords_dreg_v_2_x: np.array, + p_coords_dreg_v_3_x: np.array, + p_coords_dreg_v_4_x: np.array, + p_coords_dreg_v_1_y: np.array, + p_coords_dreg_v_2_y: np.array, + p_coords_dreg_v_3_y: np.array, + p_coords_dreg_v_4_y: np.array, + shape_func_1_1: float, + shape_func_2_1: float, + shape_func_3_1: float, + shape_func_4_1: float, + shape_func_1_2: float, + shape_func_2_2: float, + shape_func_3_2: float, + shape_func_4_2: float, + shape_func_1_3: float, + shape_func_2_3: float, + shape_func_3_3: float, + shape_func_4_3: float, + shape_func_1_4: float, + shape_func_2_4: float, + shape_func_3_4: float, + shape_func_4_4: float, + zeta_1: float, + zeta_2: float, + zeta_3: float, + zeta_4: float, + eta_1: float, + eta_2: float, + eta_3: float, + eta_4: float, + wgt_zeta_1: float, + wgt_zeta_2: float, + wgt_eta_1: float, + wgt_eta_2: float, + dbl_eps: float, + eps: float, + p_dreg_area_in: np.array, + **kwargs, + ): + wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 = cls._compute_wgt_t_detjac( + wgt_zeta_1, + wgt_eta_1, + wgt_zeta_2, + wgt_eta_2, + eta_1, + eta_2, + eta_3, + eta_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + famask_int, + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + dbl_eps, + ) + + ( + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, + ) = cls._compute_z_gauss_points( + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + ) + + ( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + ) = cls._compute_vector_sums( + wgt_t_detjac_1, + wgt_t_detjac_2, + wgt_t_detjac_3, + wgt_t_detjac_4, + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, + ) + + p_dreg_area = p_dreg_area_in + p_quad_vector_sum_1 + return dict( + p_quad_vector_sum_1=p_quad_vector_sum_1, + p_quad_vector_sum_2=p_quad_vector_sum_2, + p_quad_vector_sum_3=p_quad_vector_sum_3, + p_quad_vector_sum_4=p_quad_vector_sum_4, + p_quad_vector_sum_5=p_quad_vector_sum_5, + p_quad_vector_sum_6=p_quad_vector_sum_6, + p_quad_vector_sum_7=p_quad_vector_sum_7, + p_quad_vector_sum_8=p_quad_vector_sum_8, + p_quad_vector_sum_9=p_quad_vector_sum_9, + p_quad_vector_sum_10=p_quad_vector_sum_10, + p_dreg_area=p_dreg_area, + ) + + @pytest.mark.slow_tests + @pytest.fixture + def input_data(self, grid): + famask_int = constant_field(grid, 1, EdgeDim, KDim, dtype=int32) + p_coords_dreg_v_1_x = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_2_x = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_3_x = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_4_x = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_1_y = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_2_y = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_3_y = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_4_y = random_field(grid, EdgeDim, KDim) + shape_func_1_1 = 0.001 + shape_func_2_1 = 0.001 + shape_func_3_1 = 0.001 + shape_func_4_1 = 0.001 + shape_func_1_2 = 0.001 + shape_func_2_2 = 0.001 + shape_func_3_2 = 0.001 + shape_func_4_2 = 0.001 + shape_func_1_3 = 0.001 + shape_func_2_3 = 0.001 + shape_func_3_3 = 0.001 + shape_func_4_3 = 0.001 + shape_func_1_4 = 0.001 + shape_func_2_4 = 0.001 + shape_func_3_4 = 0.001 + shape_func_4_4 = 0.001 + zeta_1 = 0.002 + zeta_2 = 0.002 + zeta_3 = 0.002 + zeta_4 = 0.002 + eta_1 = 0.5 + eta_2 = 0.5 + eta_3 = 0.5 + eta_4 = 0.5 + wgt_zeta_1 = 0.003 + wgt_zeta_2 = 0.003 + wgt_eta_1 = 0.002 + wgt_eta_2 = 0.007 + dbl_eps = np.float64(0.1) + eps = 0.1 + p_dreg_area_in = random_field(grid, EdgeDim, KDim) + p_quad_vector_sum_1 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_2 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_3 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_4 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_5 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_6 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_7 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_8 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_9 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_10 = zero_field(grid, EdgeDim, KDim) + p_dreg_area = zero_field(grid, EdgeDim, KDim) + return dict( + famask_int=famask_int, + p_coords_dreg_v_1_x=p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x=p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x=p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x=p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y=p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y=p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y=p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y=p_coords_dreg_v_4_y, + shape_func_1_1=shape_func_1_1, + shape_func_2_1=shape_func_2_1, + shape_func_3_1=shape_func_3_1, + shape_func_4_1=shape_func_4_1, + shape_func_1_2=shape_func_1_2, + shape_func_2_2=shape_func_2_2, + shape_func_3_2=shape_func_3_2, + shape_func_4_2=shape_func_4_2, + shape_func_1_3=shape_func_1_3, + shape_func_2_3=shape_func_2_3, + shape_func_3_3=shape_func_3_3, + shape_func_4_3=shape_func_4_3, + shape_func_1_4=shape_func_1_4, + shape_func_2_4=shape_func_2_4, + shape_func_3_4=shape_func_3_4, + shape_func_4_4=shape_func_4_4, + zeta_1=zeta_1, + zeta_2=zeta_2, + zeta_3=zeta_3, + zeta_4=zeta_4, + eta_1=eta_1, + eta_2=eta_2, + eta_3=eta_3, + eta_4=eta_4, + wgt_zeta_1=wgt_zeta_1, + wgt_zeta_2=wgt_zeta_2, + wgt_eta_1=wgt_eta_1, + wgt_eta_2=wgt_eta_2, + dbl_eps=dbl_eps, + eps=eps, + p_dreg_area_in=p_dreg_area_in, + p_quad_vector_sum_1=p_quad_vector_sum_1, + p_quad_vector_sum_2=p_quad_vector_sum_2, + p_quad_vector_sum_3=p_quad_vector_sum_3, + p_quad_vector_sum_4=p_quad_vector_sum_4, + p_quad_vector_sum_5=p_quad_vector_sum_5, + p_quad_vector_sum_6=p_quad_vector_sum_6, + p_quad_vector_sum_7=p_quad_vector_sum_7, + p_quad_vector_sum_8=p_quad_vector_sum_8, + p_quad_vector_sum_9=p_quad_vector_sum_9, + p_quad_vector_sum_10=p_quad_vector_sum_10, + p_dreg_area=p_dreg_area, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_stencil.py index 9c521c56c3..07c2e16a7e 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_stencil.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_stencil.py @@ -18,362 +18,40 @@ prep_gauss_quadrature_c_stencil, ) from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field -def prep_gauss_quadrature_c_stencil_numpy( - p_coords_dreg_v_1_x: np.array, - p_coords_dreg_v_2_x: np.array, - p_coords_dreg_v_3_x: np.array, - p_coords_dreg_v_4_x: np.array, - p_coords_dreg_v_1_y: np.array, - p_coords_dreg_v_2_y: np.array, - p_coords_dreg_v_3_y: np.array, - p_coords_dreg_v_4_y: np.array, - shape_func_1_1: float, - shape_func_2_1: float, - shape_func_3_1: float, - shape_func_4_1: float, - shape_func_1_2: float, - shape_func_2_2: float, - shape_func_3_2: float, - shape_func_4_2: float, - shape_func_1_3: float, - shape_func_2_3: float, - shape_func_3_3: float, - shape_func_4_3: float, - shape_func_1_4: float, - shape_func_2_4: float, - shape_func_3_4: float, - shape_func_4_4: float, - zeta_1: float, - zeta_2: float, - zeta_3: float, - zeta_4: float, - eta_1: float, - eta_2: float, - eta_3: float, - eta_4: float, - wgt_zeta_1: float, - wgt_zeta_2: float, - wgt_eta_1: float, - wgt_eta_2: float, - dbl_eps: float, - eps: float, -) -> tuple[np.ndarray]: - z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 - z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 - z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 - z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 - - z_eta_1_1 = 1.0 - eta_1 - z_eta_2_1 = 1.0 - eta_2 - z_eta_3_1 = 1.0 - eta_3 - z_eta_4_1 = 1.0 - eta_4 - z_eta_1_2 = 1.0 + eta_1 - z_eta_2_2 = 1.0 + eta_2 - z_eta_3_2 = 1.0 + eta_3 - z_eta_4_2 = 1.0 + eta_4 - z_eta_1_3 = 1.0 - zeta_1 - z_eta_2_3 = 1.0 - zeta_2 - z_eta_3_3 = 1.0 - zeta_3 - z_eta_4_3 = 1.0 - zeta_4 - z_eta_1_4 = 1.0 + zeta_1 - z_eta_2_4 = 1.0 + zeta_2 - z_eta_3_4 = 1.0 + zeta_3 - z_eta_4_4 = 1.0 + zeta_4 - - wgt_t_detjac_1 = dbl_eps + z_wgt_1 * ( - ( - z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) - ) - wgt_t_detjac_2 = dbl_eps + z_wgt_2 * ( - ( - z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) - ) - wgt_t_detjac_3 = dbl_eps + z_wgt_3 * ( - ( - z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) - ) - wgt_t_detjac_4 = dbl_eps + z_wgt_4 * ( - ( - z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) - ) - - z_gauss_pts_1_x = ( - shape_func_1_1 * p_coords_dreg_v_1_x - + shape_func_2_1 * p_coords_dreg_v_2_x - + shape_func_3_1 * p_coords_dreg_v_3_x - + shape_func_4_1 * p_coords_dreg_v_4_x - ) - z_gauss_pts_1_y = ( - shape_func_1_1 * p_coords_dreg_v_1_y - + shape_func_2_1 * p_coords_dreg_v_2_y - + shape_func_3_1 * p_coords_dreg_v_3_y - + shape_func_4_1 * p_coords_dreg_v_4_y - ) - z_gauss_pts_2_x = ( - shape_func_1_2 * p_coords_dreg_v_1_x - + shape_func_2_2 * p_coords_dreg_v_2_x - + shape_func_3_2 * p_coords_dreg_v_3_x - + shape_func_4_2 * p_coords_dreg_v_4_x - ) - z_gauss_pts_2_y = ( - shape_func_1_2 * p_coords_dreg_v_1_y - + shape_func_2_2 * p_coords_dreg_v_2_y - + shape_func_3_2 * p_coords_dreg_v_3_y - + shape_func_4_2 * p_coords_dreg_v_4_y - ) - z_gauss_pts_3_x = ( - shape_func_1_3 * p_coords_dreg_v_1_x - + shape_func_2_3 * p_coords_dreg_v_2_x - + shape_func_3_3 * p_coords_dreg_v_3_x - + shape_func_4_3 * p_coords_dreg_v_4_x - ) - z_gauss_pts_3_y = ( - shape_func_1_3 * p_coords_dreg_v_1_y - + shape_func_2_3 * p_coords_dreg_v_2_y - + shape_func_3_3 * p_coords_dreg_v_3_y - + shape_func_4_3 * p_coords_dreg_v_4_y - ) - z_gauss_pts_4_x = ( - shape_func_1_4 * p_coords_dreg_v_1_x - + shape_func_2_4 * p_coords_dreg_v_2_x - + shape_func_3_4 * p_coords_dreg_v_3_x - + shape_func_4_4 * p_coords_dreg_v_4_x - ) - z_gauss_pts_4_y = ( - shape_func_1_4 * p_coords_dreg_v_1_y - + shape_func_2_4 * p_coords_dreg_v_2_y - + shape_func_3_4 * p_coords_dreg_v_3_y - + shape_func_4_4 * p_coords_dreg_v_4_y - ) - - p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 - p_quad_vector_sum_2 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x +class TestPrepGaussQuadratureCStencil(StencilTest): + PROGRAM = prep_gauss_quadrature_c_stencil + OUTPUTS = ( + "p_quad_vector_sum_1", + "p_quad_vector_sum_2", + "p_quad_vector_sum_3", + "p_quad_vector_sum_4", + "p_quad_vector_sum_5", + "p_quad_vector_sum_6", + "p_quad_vector_sum_7", + "p_quad_vector_sum_8", + "p_quad_vector_sum_9", + "p_quad_vector_sum_10", + "p_dreg_area_out", ) - p_quad_vector_sum_3 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y - ) - p_quad_vector_sum_4 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x - ) - p_quad_vector_sum_5 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y - ) - p_quad_vector_sum_6 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y - ) - p_quad_vector_sum_7 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x - ) - p_quad_vector_sum_8 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y - ) - p_quad_vector_sum_9 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y - ) - p_quad_vector_sum_10 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y - ) - - z_area = p_quad_vector_sum_1 - p_dreg_area_out = np.where( - z_area >= 0.0, - np.maximum(eps, np.absolute(z_area)), - -np.maximum(eps, np.absolute(z_area)), - ) - return ( - p_quad_vector_sum_1, - p_quad_vector_sum_2, - p_quad_vector_sum_3, - p_quad_vector_sum_4, - p_quad_vector_sum_5, - p_quad_vector_sum_6, - p_quad_vector_sum_7, - p_quad_vector_sum_8, - p_quad_vector_sum_9, - p_quad_vector_sum_10, - p_dreg_area_out, - ) - - -@pytest.mark.slow_tests -def test_prep_gauss_quadrature_c_stencil(backend): - grid = SimpleGrid() - - p_coords_dreg_v_1_x = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_2_x = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_3_x = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_4_x = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_1_y = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_2_y = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_3_y = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_4_y = random_field(grid, EdgeDim, KDim) - shape_func_1_1 = 0.001 - shape_func_2_1 = 0.001 - shape_func_3_1 = 0.001 - shape_func_4_1 = 0.001 - shape_func_1_2 = 0.001 - shape_func_2_2 = 0.001 - shape_func_3_2 = 0.001 - shape_func_4_2 = 0.001 - shape_func_1_3 = 0.001 - shape_func_2_3 = 0.001 - shape_func_3_3 = 0.001 - shape_func_4_3 = 0.001 - shape_func_1_4 = 0.001 - shape_func_2_4 = 0.001 - shape_func_3_4 = 0.001 - shape_func_4_4 = 0.001 - zeta_1 = 0.002 - zeta_2 = 0.002 - zeta_3 = 0.002 - zeta_4 = 0.002 - eta_1 = 0.5 - eta_2 = 0.5 - eta_3 = 0.5 - eta_4 = 0.5 - wgt_zeta_1 = 0.003 - wgt_zeta_2 = 0.003 - wgt_eta_1 = 0.002 - wgt_eta_2 = 0.007 - dbl_eps = np.float64(0.1) - eps = 0.1 - p_quad_vector_sum_1 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_2 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_3 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_4 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_5 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_6 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_7 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_8 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_9 = zero_field(grid, EdgeDim, KDim) - p_quad_vector_sum_10 = zero_field(grid, EdgeDim, KDim) - p_dreg_area_out = zero_field(grid, EdgeDim, KDim) - ( - ref_1, - ref_2, - ref_3, - ref_4, - ref_5, - ref_6, - ref_7, - ref_8, - ref_9, - ref_10, - ref_11, - ) = prep_gauss_quadrature_c_stencil_numpy( - np.asarray(p_coords_dreg_v_1_x), - np.asarray(p_coords_dreg_v_2_x), - np.asarray(p_coords_dreg_v_3_x), - np.asarray(p_coords_dreg_v_4_x), - np.asarray(p_coords_dreg_v_1_y), - np.asarray(p_coords_dreg_v_2_y), - np.asarray(p_coords_dreg_v_3_y), - np.asarray(p_coords_dreg_v_4_y), - shape_func_1_1, - shape_func_2_1, - shape_func_3_1, - shape_func_4_1, - shape_func_1_2, - shape_func_2_2, - shape_func_3_2, - shape_func_4_2, - shape_func_1_3, - shape_func_2_3, - shape_func_3_3, - shape_func_4_3, - shape_func_1_4, - shape_func_2_4, - shape_func_3_4, - shape_func_4_4, + @staticmethod + def _compute_wgt_t_detjac( + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, zeta_1, zeta_2, zeta_3, @@ -382,15 +60,115 @@ def test_prep_gauss_quadrature_c_stencil(backend): eta_2, eta_3, eta_4, - wgt_zeta_1, - wgt_zeta_2, - wgt_eta_1, - wgt_eta_2, - dbl_eps, - eps, - ) + ): + # Compute z_wgt + z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 + z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 + z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 + z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 + + # Compute z_eta + z_eta_1_1, z_eta_2_1, z_eta_3_1, z_eta_4_1 = ( + 1.0 - eta_1, + 1.0 - eta_2, + 1.0 - eta_3, + 1.0 - eta_4, + ) + z_eta_1_2, z_eta_2_2, z_eta_3_2, z_eta_4_2 = ( + 1.0 + eta_1, + 1.0 + eta_2, + 1.0 + eta_3, + 1.0 + eta_4, + ) + z_eta_1_3, z_eta_2_3, z_eta_3_3, z_eta_4_3 = ( + 1.0 - zeta_1, + 1.0 - zeta_2, + 1.0 - zeta_3, + 1.0 - zeta_4, + ) + z_eta_1_4, z_eta_2_4, z_eta_3_4, z_eta_4_4 = ( + 1.0 + zeta_1, + 1.0 + zeta_2, + 1.0 + zeta_3, + 1.0 + zeta_4, + ) - prep_gauss_quadrature_c_stencil.with_backend(backend)( + wgt_t_detjac_1 = dbl_eps + z_wgt_1 * ( + ( + z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + wgt_t_detjac_2 = dbl_eps + z_wgt_2 * ( + ( + z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + wgt_t_detjac_3 = dbl_eps + z_wgt_3 * ( + ( + z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + wgt_t_detjac_4 = dbl_eps + z_wgt_4 * ( + ( + z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + return wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 + + @staticmethod + def _compute_z_gauss_points( p_coords_dreg_v_1_x, p_coords_dreg_v_2_x, p_coords_dreg_v_3_x, @@ -415,52 +193,400 @@ def test_prep_gauss_quadrature_c_stencil(backend): shape_func_2_4, shape_func_3_4, shape_func_4_4, - zeta_1, - zeta_2, - zeta_3, - zeta_4, - eta_1, - eta_2, - eta_3, - eta_4, - wgt_zeta_1, - wgt_zeta_2, - wgt_eta_1, - wgt_eta_2, - dbl_eps, - eps, - p_quad_vector_sum_1, - p_quad_vector_sum_2, - p_quad_vector_sum_3, - p_quad_vector_sum_4, - p_quad_vector_sum_5, - p_quad_vector_sum_6, - p_quad_vector_sum_7, - p_quad_vector_sum_8, - p_quad_vector_sum_9, - p_quad_vector_sum_10, - p_dreg_area_out, - offset_provider={}, - ) - co1 = np.asarray(p_quad_vector_sum_1) - co2 = np.asarray(p_quad_vector_sum_2) - co3 = np.asarray(p_quad_vector_sum_3) - co4 = np.asarray(p_quad_vector_sum_4) - co5 = np.asarray(p_quad_vector_sum_5) - co6 = np.asarray(p_quad_vector_sum_6) - co7 = np.asarray(p_quad_vector_sum_7) - co8 = np.asarray(p_quad_vector_sum_8) - co9 = np.asarray(p_quad_vector_sum_9) - co10 = np.asarray(p_quad_vector_sum_10) - co11 = np.asarray(p_dreg_area_out) - assert np.allclose(ref_1, co1) - assert np.allclose(ref_2, co2) - assert np.allclose(ref_3, co3) - assert np.allclose(ref_4, co4) - assert np.allclose(ref_5, co5) - assert np.allclose(ref_6, co6) - assert np.allclose(ref_7, co7) - assert np.allclose(ref_8, co8) - assert np.allclose(ref_9, co9) - assert np.allclose(ref_10, co10) - assert np.allclose(ref_11, co11) + ): + z_gauss_pts_1_x = ( + shape_func_1_1 * p_coords_dreg_v_1_x + + shape_func_2_1 * p_coords_dreg_v_2_x + + shape_func_3_1 * p_coords_dreg_v_3_x + + shape_func_4_1 * p_coords_dreg_v_4_x + ) + z_gauss_pts_1_y = ( + shape_func_1_1 * p_coords_dreg_v_1_y + + shape_func_2_1 * p_coords_dreg_v_2_y + + shape_func_3_1 * p_coords_dreg_v_3_y + + shape_func_4_1 * p_coords_dreg_v_4_y + ) + z_gauss_pts_2_x = ( + shape_func_1_2 * p_coords_dreg_v_1_x + + shape_func_2_2 * p_coords_dreg_v_2_x + + shape_func_3_2 * p_coords_dreg_v_3_x + + shape_func_4_2 * p_coords_dreg_v_4_x + ) + z_gauss_pts_2_y = ( + shape_func_1_2 * p_coords_dreg_v_1_y + + shape_func_2_2 * p_coords_dreg_v_2_y + + shape_func_3_2 * p_coords_dreg_v_3_y + + shape_func_4_2 * p_coords_dreg_v_4_y + ) + z_gauss_pts_3_x = ( + shape_func_1_3 * p_coords_dreg_v_1_x + + shape_func_2_3 * p_coords_dreg_v_2_x + + shape_func_3_3 * p_coords_dreg_v_3_x + + shape_func_4_3 * p_coords_dreg_v_4_x + ) + z_gauss_pts_3_y = ( + shape_func_1_3 * p_coords_dreg_v_1_y + + shape_func_2_3 * p_coords_dreg_v_2_y + + shape_func_3_3 * p_coords_dreg_v_3_y + + shape_func_4_3 * p_coords_dreg_v_4_y + ) + z_gauss_pts_4_x = ( + shape_func_1_4 * p_coords_dreg_v_1_x + + shape_func_2_4 * p_coords_dreg_v_2_x + + shape_func_3_4 * p_coords_dreg_v_3_x + + shape_func_4_4 * p_coords_dreg_v_4_x + ) + z_gauss_pts_4_y = ( + shape_func_1_4 * p_coords_dreg_v_1_y + + shape_func_2_4 * p_coords_dreg_v_2_y + + shape_func_3_4 * p_coords_dreg_v_3_y + + shape_func_4_4 * p_coords_dreg_v_4_y + ) + return ( + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, + ) + + @staticmethod + def _compute_vector_sums( + wgt_t_detjac_1, + wgt_t_detjac_2, + wgt_t_detjac_3, + wgt_t_detjac_4, + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, + ): + p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 + p_quad_vector_sum_2 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x + ) + p_quad_vector_sum_3 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y + ) + p_quad_vector_sum_4 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_5 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_6 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_7 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_8 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_9 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_10 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + return ( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + ) + + @classmethod + def reference( + cls, + grid, + p_coords_dreg_v_1_x: np.array, + p_coords_dreg_v_2_x: np.array, + p_coords_dreg_v_3_x: np.array, + p_coords_dreg_v_4_x: np.array, + p_coords_dreg_v_1_y: np.array, + p_coords_dreg_v_2_y: np.array, + p_coords_dreg_v_3_y: np.array, + p_coords_dreg_v_4_y: np.array, + shape_func_1_1: float, + shape_func_2_1: float, + shape_func_3_1: float, + shape_func_4_1: float, + shape_func_1_2: float, + shape_func_2_2: float, + shape_func_3_2: float, + shape_func_4_2: float, + shape_func_1_3: float, + shape_func_2_3: float, + shape_func_3_3: float, + shape_func_4_3: float, + shape_func_1_4: float, + shape_func_2_4: float, + shape_func_3_4: float, + shape_func_4_4: float, + zeta_1: float, + zeta_2: float, + zeta_3: float, + zeta_4: float, + eta_1: float, + eta_2: float, + eta_3: float, + eta_4: float, + wgt_zeta_1: float, + wgt_zeta_2: float, + wgt_eta_1: float, + wgt_eta_2: float, + dbl_eps: float, + eps: float, + **kwargs, + ): + + wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 = cls._compute_wgt_t_detjac( + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, + ) + + ( + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, + ) = cls._compute_z_gauss_points( + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + ) + + ( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + ) = cls._compute_vector_sums( + wgt_t_detjac_1, + wgt_t_detjac_2, + wgt_t_detjac_3, + wgt_t_detjac_4, + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, + ) + + z_area = p_quad_vector_sum_1 + p_dreg_area_out = np.where( + z_area >= 0.0, + np.maximum(eps, np.absolute(z_area)), + -np.maximum(eps, np.absolute(z_area)), + ) + return dict( + p_quad_vector_sum_1=p_quad_vector_sum_1, + p_quad_vector_sum_2=p_quad_vector_sum_2, + p_quad_vector_sum_3=p_quad_vector_sum_3, + p_quad_vector_sum_4=p_quad_vector_sum_4, + p_quad_vector_sum_5=p_quad_vector_sum_5, + p_quad_vector_sum_6=p_quad_vector_sum_6, + p_quad_vector_sum_7=p_quad_vector_sum_7, + p_quad_vector_sum_8=p_quad_vector_sum_8, + p_quad_vector_sum_9=p_quad_vector_sum_9, + p_quad_vector_sum_10=p_quad_vector_sum_10, + p_dreg_area_out=p_dreg_area_out, + ) + + @pytest.mark.slow_tests + @pytest.fixture + def input_data(self, grid): + p_coords_dreg_v_1_x = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_2_x = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_3_x = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_4_x = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_1_y = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_2_y = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_3_y = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_4_y = random_field(grid, EdgeDim, KDim) + shape_func_1_1 = 0.001 + shape_func_2_1 = 0.001 + shape_func_3_1 = 0.001 + shape_func_4_1 = 0.001 + shape_func_1_2 = 0.001 + shape_func_2_2 = 0.001 + shape_func_3_2 = 0.001 + shape_func_4_2 = 0.001 + shape_func_1_3 = 0.001 + shape_func_2_3 = 0.001 + shape_func_3_3 = 0.001 + shape_func_4_3 = 0.001 + shape_func_1_4 = 0.001 + shape_func_2_4 = 0.001 + shape_func_3_4 = 0.001 + shape_func_4_4 = 0.001 + zeta_1 = 0.002 + zeta_2 = 0.002 + zeta_3 = 0.002 + zeta_4 = 0.002 + eta_1 = 0.5 + eta_2 = 0.5 + eta_3 = 0.5 + eta_4 = 0.5 + wgt_zeta_1 = 0.003 + wgt_zeta_2 = 0.003 + wgt_eta_1 = 0.002 + wgt_eta_2 = 0.007 + dbl_eps = np.float64(0.1) + eps = 0.1 + p_quad_vector_sum_1 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_2 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_3 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_4 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_5 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_6 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_7 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_8 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_9 = zero_field(grid, EdgeDim, KDim) + p_quad_vector_sum_10 = zero_field(grid, EdgeDim, KDim) + p_dreg_area_out = zero_field(grid, EdgeDim, KDim) + return dict( + p_coords_dreg_v_1_x=p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x=p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x=p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x=p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y=p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y=p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y=p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y=p_coords_dreg_v_4_y, + shape_func_1_1=shape_func_1_1, + shape_func_2_1=shape_func_2_1, + shape_func_3_1=shape_func_3_1, + shape_func_4_1=shape_func_4_1, + shape_func_1_2=shape_func_1_2, + shape_func_2_2=shape_func_2_2, + shape_func_3_2=shape_func_3_2, + shape_func_4_2=shape_func_4_2, + shape_func_1_3=shape_func_1_3, + shape_func_2_3=shape_func_2_3, + shape_func_3_3=shape_func_3_3, + shape_func_4_3=shape_func_4_3, + shape_func_1_4=shape_func_1_4, + shape_func_2_4=shape_func_2_4, + shape_func_3_4=shape_func_3_4, + shape_func_4_4=shape_func_4_4, + zeta_1=zeta_1, + zeta_2=zeta_2, + zeta_3=zeta_3, + zeta_4=zeta_4, + eta_1=eta_1, + eta_2=eta_2, + eta_3=eta_3, + eta_4=eta_4, + wgt_zeta_1=wgt_zeta_1, + wgt_zeta_2=wgt_zeta_2, + wgt_eta_1=wgt_eta_1, + wgt_eta_2=wgt_eta_2, + dbl_eps=dbl_eps, + eps=eps, + p_quad_vector_sum_1=p_quad_vector_sum_1, + p_quad_vector_sum_2=p_quad_vector_sum_2, + p_quad_vector_sum_3=p_quad_vector_sum_3, + p_quad_vector_sum_4=p_quad_vector_sum_4, + p_quad_vector_sum_5=p_quad_vector_sum_5, + p_quad_vector_sum_6=p_quad_vector_sum_6, + p_quad_vector_sum_7=p_quad_vector_sum_7, + p_quad_vector_sum_8=p_quad_vector_sum_8, + p_quad_vector_sum_9=p_quad_vector_sum_9, + p_quad_vector_sum_10=p_quad_vector_sum_10, + p_dreg_area_out=p_dreg_area_out, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_c_svd_stencil.py index 26df554d28..03a627840c 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_c_svd_stencil.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_c_svd_stencil.py @@ -18,321 +18,301 @@ recon_lsq_cell_c_svd_stencil, ) from icon4py.model.common.dimension import C2E2C2E2CDim, CECECDim, CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field - +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + reshape, + zero_field, +) -def recon_lsq_cell_c_svd_stencil_numpy( - c2e2c2e2c: np.ndarray, - p_cc: np.ndarray, - lsq_pseudoinv_1: np.ndarray, - lsq_pseudoinv_2: np.ndarray, - lsq_pseudoinv_3: np.ndarray, - lsq_pseudoinv_4: np.ndarray, - lsq_pseudoinv_5: np.ndarray, - lsq_pseudoinv_6: np.ndarray, - lsq_pseudoinv_7: np.ndarray, - lsq_pseudoinv_8: np.ndarray, - lsq_pseudoinv_9: np.ndarray, - lsq_moments_1: np.ndarray, - lsq_moments_2: np.ndarray, - lsq_moments_3: np.ndarray, - lsq_moments_4: np.ndarray, - lsq_moments_5: np.ndarray, - lsq_moments_6: np.ndarray, - lsq_moments_7: np.ndarray, - lsq_moments_8: np.ndarray, - lsq_moments_9: np.ndarray, -) -> tuple[np.ndarray]: - lsq_moments_1 = np.expand_dims(lsq_moments_1, axis=-1) - lsq_moments_2 = np.expand_dims(lsq_moments_2, axis=-1) - lsq_moments_3 = np.expand_dims(lsq_moments_3, axis=-1) - lsq_moments_4 = np.expand_dims(lsq_moments_4, axis=-1) - lsq_moments_5 = np.expand_dims(lsq_moments_5, axis=-1) - lsq_moments_6 = np.expand_dims(lsq_moments_6, axis=-1) - lsq_moments_7 = np.expand_dims(lsq_moments_7, axis=-1) - lsq_moments_8 = np.expand_dims(lsq_moments_8, axis=-1) - lsq_moments_9 = np.expand_dims(lsq_moments_9, axis=-1) - lsq_moments_1 = np.broadcast_to(lsq_moments_1, p_cc.shape) - lsq_moments_2 = np.broadcast_to(lsq_moments_2, p_cc.shape) - lsq_moments_3 = np.broadcast_to(lsq_moments_3, p_cc.shape) - lsq_moments_4 = np.broadcast_to(lsq_moments_4, p_cc.shape) - lsq_moments_5 = np.broadcast_to(lsq_moments_5, p_cc.shape) - lsq_moments_6 = np.broadcast_to(lsq_moments_6, p_cc.shape) - lsq_moments_7 = np.broadcast_to(lsq_moments_7, p_cc.shape) - lsq_moments_8 = np.broadcast_to(lsq_moments_8, p_cc.shape) - lsq_moments_9 = np.broadcast_to(lsq_moments_9, p_cc.shape) - lsq_pseudoinv_9 = np.expand_dims(lsq_pseudoinv_9, axis=-1) - lsq_pseudoinv_8 = np.expand_dims(lsq_pseudoinv_8, axis=-1) - lsq_pseudoinv_7 = np.expand_dims(lsq_pseudoinv_7, axis=-1) - lsq_pseudoinv_6 = np.expand_dims(lsq_pseudoinv_6, axis=-1) - lsq_pseudoinv_5 = np.expand_dims(lsq_pseudoinv_5, axis=-1) - lsq_pseudoinv_4 = np.expand_dims(lsq_pseudoinv_4, axis=-1) - lsq_pseudoinv_3 = np.expand_dims(lsq_pseudoinv_3, axis=-1) - lsq_pseudoinv_2 = np.expand_dims(lsq_pseudoinv_2, axis=-1) - lsq_pseudoinv_1 = np.expand_dims(lsq_pseudoinv_1, axis=-1) - p_coeff_10_dsl = ( - lsq_pseudoinv_9[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_9[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_9[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_9[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_9[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_9[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_9[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_9[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_9[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) +class TestReconLsqCellCSvdStencil(StencilTest): + PROGRAM = recon_lsq_cell_c_svd_stencil + OUTPUTS = ( + "p_coeff_1_dsl", + "p_coeff_2_dsl", + "p_coeff_3_dsl", + "p_coeff_4_dsl", + "p_coeff_5_dsl", + "p_coeff_6_dsl", + "p_coeff_7_dsl", + "p_coeff_8_dsl", + "p_coeff_9_dsl", + "p_coeff_10_dsl", ) - p_coeff_9_dsl = ( - lsq_pseudoinv_8[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_8[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_8[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_8[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_8[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_8[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_8[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_8[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_8[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) + @staticmethod + def reference( + grid, + p_cc: np.array, + lsq_pseudoinv_1: np.array, + lsq_pseudoinv_2: np.array, + lsq_pseudoinv_3: np.array, + lsq_pseudoinv_4: np.array, + lsq_pseudoinv_5: np.array, + lsq_pseudoinv_6: np.array, + lsq_pseudoinv_7: np.array, + lsq_pseudoinv_8: np.array, + lsq_pseudoinv_9: np.array, + lsq_moments_1: np.array, + lsq_moments_2: np.array, + lsq_moments_3: np.array, + lsq_moments_4: np.array, + lsq_moments_5: np.array, + lsq_moments_6: np.array, + lsq_moments_7: np.array, + lsq_moments_8: np.array, + lsq_moments_9: np.array, + **kwargs, + ): + c2e2c2e2c = grid.connectivities[C2E2C2E2CDim] + lsq_moments_1 = np.expand_dims(lsq_moments_1, axis=-1) + lsq_moments_2 = np.expand_dims(lsq_moments_2, axis=-1) + lsq_moments_3 = np.expand_dims(lsq_moments_3, axis=-1) + lsq_moments_4 = np.expand_dims(lsq_moments_4, axis=-1) + lsq_moments_5 = np.expand_dims(lsq_moments_5, axis=-1) + lsq_moments_6 = np.expand_dims(lsq_moments_6, axis=-1) + lsq_moments_7 = np.expand_dims(lsq_moments_7, axis=-1) + lsq_moments_8 = np.expand_dims(lsq_moments_8, axis=-1) + lsq_moments_9 = np.expand_dims(lsq_moments_9, axis=-1) + lsq_moments_1 = np.broadcast_to(lsq_moments_1, p_cc.shape) + lsq_moments_2 = np.broadcast_to(lsq_moments_2, p_cc.shape) + lsq_moments_3 = np.broadcast_to(lsq_moments_3, p_cc.shape) + lsq_moments_4 = np.broadcast_to(lsq_moments_4, p_cc.shape) + lsq_moments_5 = np.broadcast_to(lsq_moments_5, p_cc.shape) + lsq_moments_6 = np.broadcast_to(lsq_moments_6, p_cc.shape) + lsq_moments_7 = np.broadcast_to(lsq_moments_7, p_cc.shape) + lsq_moments_8 = np.broadcast_to(lsq_moments_8, p_cc.shape) + lsq_pseudoinv_9 = reshape(lsq_pseudoinv_9, c2e2c2e2c.shape) + lsq_pseudoinv_9 = np.expand_dims(lsq_pseudoinv_9, axis=-1) + lsq_pseudoinv_8 = reshape(lsq_pseudoinv_8, c2e2c2e2c.shape) + lsq_pseudoinv_8 = np.expand_dims(lsq_pseudoinv_8, axis=-1) + lsq_pseudoinv_7 = reshape(lsq_pseudoinv_7, c2e2c2e2c.shape) + lsq_pseudoinv_7 = np.expand_dims(lsq_pseudoinv_7, axis=-1) + lsq_pseudoinv_6 = reshape(lsq_pseudoinv_6, c2e2c2e2c.shape) + lsq_pseudoinv_6 = np.expand_dims(lsq_pseudoinv_6, axis=-1) + lsq_pseudoinv_5 = reshape(lsq_pseudoinv_5, c2e2c2e2c.shape) + lsq_pseudoinv_5 = np.expand_dims(lsq_pseudoinv_5, axis=-1) + lsq_pseudoinv_4 = reshape(lsq_pseudoinv_4, c2e2c2e2c.shape) + lsq_pseudoinv_4 = np.expand_dims(lsq_pseudoinv_4, axis=-1) + lsq_pseudoinv_3 = reshape(lsq_pseudoinv_3, c2e2c2e2c.shape) + lsq_pseudoinv_3 = np.expand_dims(lsq_pseudoinv_3, axis=-1) + lsq_pseudoinv_2 = reshape(lsq_pseudoinv_2, c2e2c2e2c.shape) + lsq_pseudoinv_2 = np.expand_dims(lsq_pseudoinv_2, axis=-1) + lsq_pseudoinv_1 = reshape(lsq_pseudoinv_1, c2e2c2e2c.shape) + lsq_pseudoinv_1 = np.expand_dims(lsq_pseudoinv_1, axis=-1) - p_coeff_8_dsl = ( - lsq_pseudoinv_7[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_7[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_7[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_7[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_7[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_7[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_7[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_7[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_7[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) + p_coeff_10_dsl = ( + lsq_pseudoinv_9[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_9[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_9[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_9[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_9[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_9[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_9[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_9[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_9[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) - p_coeff_7_dsl = ( - lsq_pseudoinv_6[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_6[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_6[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_6[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_6[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_6[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_6[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_6[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_6[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) + p_coeff_9_dsl = ( + lsq_pseudoinv_8[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_8[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_8[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_8[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_8[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_8[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_8[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_8[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_8[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) - p_coeff_6_dsl = ( - lsq_pseudoinv_5[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_5[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_5[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_5[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_5[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_5[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_5[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_5[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_5[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) + p_coeff_8_dsl = ( + lsq_pseudoinv_7[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_7[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_7[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_7[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_7[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_7[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_7[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_7[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_7[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) - p_coeff_5_dsl = ( - lsq_pseudoinv_4[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_4[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_4[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_4[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_4[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_4[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_4[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_4[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_4[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) + p_coeff_7_dsl = ( + lsq_pseudoinv_6[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_6[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_6[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_6[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_6[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_6[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_6[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_6[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_6[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) - p_coeff_4_dsl = ( - lsq_pseudoinv_3[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_3[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_3[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_3[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_3[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_3[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_3[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_3[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_3[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) + p_coeff_6_dsl = ( + lsq_pseudoinv_5[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_5[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_5[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_5[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_5[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_5[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_5[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_5[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_5[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) - p_coeff_3_dsl = ( - lsq_pseudoinv_2[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_2[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_2[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_2[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_2[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_2[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_2[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_2[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_2[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) + p_coeff_5_dsl = ( + lsq_pseudoinv_4[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_4[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_4[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_4[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_4[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_4[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_4[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_4[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_4[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) - p_coeff_2_dsl = ( - lsq_pseudoinv_1[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) - + lsq_pseudoinv_1[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) - + lsq_pseudoinv_1[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) - + lsq_pseudoinv_1[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) - + lsq_pseudoinv_1[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) - + lsq_pseudoinv_1[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) - + lsq_pseudoinv_1[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) - + lsq_pseudoinv_1[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) - + lsq_pseudoinv_1[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) - ) + p_coeff_4_dsl = ( + lsq_pseudoinv_3[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_3[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_3[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_3[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_3[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_3[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_3[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_3[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_3[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) - p_coeff_1_dsl = p_cc - ( - p_coeff_2_dsl * lsq_moments_1 - + p_coeff_3_dsl * lsq_moments_2 - + p_coeff_4_dsl * lsq_moments_3 - + p_coeff_5_dsl * lsq_moments_4 - + p_coeff_6_dsl * lsq_moments_5 - + p_coeff_7_dsl * lsq_moments_6 - + p_coeff_8_dsl * lsq_moments_7 - + p_coeff_9_dsl * lsq_moments_8 - + p_coeff_10_dsl * lsq_moments_9 - ) - return ( - p_coeff_1_dsl, - p_coeff_2_dsl, - p_coeff_3_dsl, - p_coeff_4_dsl, - p_coeff_5_dsl, - p_coeff_6_dsl, - p_coeff_7_dsl, - p_coeff_8_dsl, - p_coeff_9_dsl, - p_coeff_10_dsl, - ) + p_coeff_3_dsl = ( + lsq_pseudoinv_2[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_2[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_2[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_2[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_2[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_2[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_2[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_2[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_2[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + p_coeff_2_dsl = ( + lsq_pseudoinv_1[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_1[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_1[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_1[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_1[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_1[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_1[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_1[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_1[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) -@pytest.mark.slow_tests -def test_recon_lsq_cell_c_svd_stencil(backend): - grid = SimpleGrid() - p_cc = random_field(grid, CellDim, KDim) - lsq_pseudoinv_1 = random_field(grid, CellDim, C2E2C2E2CDim) - lsq_pseudoinv_2 = random_field(grid, CellDim, C2E2C2E2CDim) - lsq_pseudoinv_3 = random_field(grid, CellDim, C2E2C2E2CDim) - lsq_pseudoinv_4 = random_field(grid, CellDim, C2E2C2E2CDim) - lsq_pseudoinv_5 = random_field(grid, CellDim, C2E2C2E2CDim) - lsq_pseudoinv_6 = random_field(grid, CellDim, C2E2C2E2CDim) - lsq_pseudoinv_7 = random_field(grid, CellDim, C2E2C2E2CDim) - lsq_pseudoinv_8 = random_field(grid, CellDim, C2E2C2E2CDim) - lsq_pseudoinv_9 = random_field(grid, CellDim, C2E2C2E2CDim) - lsq_pseudoinv_1_field = as_1D_sparse_field(lsq_pseudoinv_1, CECECDim) - lsq_pseudoinv_2_field = as_1D_sparse_field(lsq_pseudoinv_2, CECECDim) - lsq_pseudoinv_3_field = as_1D_sparse_field(lsq_pseudoinv_3, CECECDim) - lsq_pseudoinv_4_field = as_1D_sparse_field(lsq_pseudoinv_4, CECECDim) - lsq_pseudoinv_5_field = as_1D_sparse_field(lsq_pseudoinv_5, CECECDim) - lsq_pseudoinv_6_field = as_1D_sparse_field(lsq_pseudoinv_6, CECECDim) - lsq_pseudoinv_7_field = as_1D_sparse_field(lsq_pseudoinv_7, CECECDim) - lsq_pseudoinv_8_field = as_1D_sparse_field(lsq_pseudoinv_8, CECECDim) - lsq_pseudoinv_9_field = as_1D_sparse_field(lsq_pseudoinv_9, CECECDim) - lsq_moments_1 = random_field(grid, CellDim) - lsq_moments_2 = random_field(grid, CellDim) - lsq_moments_3 = random_field(grid, CellDim) - lsq_moments_4 = random_field(grid, CellDim) - lsq_moments_5 = random_field(grid, CellDim) - lsq_moments_6 = random_field(grid, CellDim) - lsq_moments_7 = random_field(grid, CellDim) - lsq_moments_8 = random_field(grid, CellDim) - lsq_moments_9 = random_field(grid, CellDim) - p_coeff_1_dsl = zero_field(grid, CellDim, KDim) - p_coeff_2_dsl = zero_field(grid, CellDim, KDim) - p_coeff_3_dsl = zero_field(grid, CellDim, KDim) - p_coeff_4_dsl = zero_field(grid, CellDim, KDim) - p_coeff_5_dsl = zero_field(grid, CellDim, KDim) - p_coeff_6_dsl = zero_field(grid, CellDim, KDim) - p_coeff_7_dsl = zero_field(grid, CellDim, KDim) - p_coeff_8_dsl = zero_field(grid, CellDim, KDim) - p_coeff_9_dsl = zero_field(grid, CellDim, KDim) - p_coeff_10_dsl = zero_field(grid, CellDim, KDim) + p_coeff_1_dsl = p_cc - ( + p_coeff_2_dsl * lsq_moments_1 + + p_coeff_3_dsl * lsq_moments_2 + + p_coeff_4_dsl * lsq_moments_3 + + p_coeff_5_dsl * lsq_moments_4 + + p_coeff_6_dsl * lsq_moments_5 + + p_coeff_7_dsl * lsq_moments_6 + + p_coeff_8_dsl * lsq_moments_7 + + p_coeff_9_dsl * lsq_moments_8 + + p_coeff_10_dsl * lsq_moments_9 + ) + return dict( + p_coeff_1_dsl=p_coeff_1_dsl, + p_coeff_2_dsl=p_coeff_2_dsl, + p_coeff_3_dsl=p_coeff_3_dsl, + p_coeff_4_dsl=p_coeff_4_dsl, + p_coeff_5_dsl=p_coeff_5_dsl, + p_coeff_6_dsl=p_coeff_6_dsl, + p_coeff_7_dsl=p_coeff_7_dsl, + p_coeff_8_dsl=p_coeff_8_dsl, + p_coeff_9_dsl=p_coeff_9_dsl, + p_coeff_10_dsl=p_coeff_10_dsl, + ) - ( - ref_1, - ref_2, - ref_3, - ref_4, - ref_5, - ref_6, - ref_7, - ref_8, - ref_9, - ref_10, - ) = recon_lsq_cell_c_svd_stencil_numpy( - grid.connectivities[C2E2C2E2CDim], - p_cc.asnumpy(), - lsq_pseudoinv_1.asnumpy(), - lsq_pseudoinv_2.asnumpy(), - lsq_pseudoinv_3.asnumpy(), - lsq_pseudoinv_4.asnumpy(), - lsq_pseudoinv_5.asnumpy(), - lsq_pseudoinv_6.asnumpy(), - lsq_pseudoinv_7.asnumpy(), - lsq_pseudoinv_8.asnumpy(), - lsq_pseudoinv_9.asnumpy(), - lsq_moments_1.asnumpy(), - lsq_moments_2.asnumpy(), - lsq_moments_3.asnumpy(), - lsq_moments_4.asnumpy(), - lsq_moments_5.asnumpy(), - lsq_moments_6.asnumpy(), - lsq_moments_7.asnumpy(), - lsq_moments_8.asnumpy(), - lsq_moments_9.asnumpy(), - ) - - recon_lsq_cell_c_svd_stencil.with_backend(backend)( - p_cc, - lsq_pseudoinv_1_field, - lsq_pseudoinv_2_field, - lsq_pseudoinv_3_field, - lsq_pseudoinv_4_field, - lsq_pseudoinv_5_field, - lsq_pseudoinv_6_field, - lsq_pseudoinv_7_field, - lsq_pseudoinv_8_field, - lsq_pseudoinv_9_field, - lsq_moments_1, - lsq_moments_2, - lsq_moments_3, - lsq_moments_4, - lsq_moments_5, - lsq_moments_6, - lsq_moments_7, - lsq_moments_8, - lsq_moments_9, - p_coeff_1_dsl, - p_coeff_2_dsl, - p_coeff_3_dsl, - p_coeff_4_dsl, - p_coeff_5_dsl, - p_coeff_6_dsl, - p_coeff_7_dsl, - p_coeff_8_dsl, - p_coeff_9_dsl, - p_coeff_10_dsl, - offset_provider={ - "C2E2C2E2C": grid.get_offset_provider("C2E2C2E2C"), - "C2CECEC": grid.get_offset_provider("C2CECEC"), - }, - ) - co1 = p_coeff_1_dsl.asnumpy() - co2 = p_coeff_2_dsl.asnumpy() - co3 = p_coeff_3_dsl.asnumpy() - co4 = p_coeff_4_dsl.asnumpy() - co5 = p_coeff_5_dsl.asnumpy() - co6 = p_coeff_6_dsl.asnumpy() - co7 = p_coeff_7_dsl.asnumpy() - co8 = p_coeff_8_dsl.asnumpy() - co9 = p_coeff_9_dsl.asnumpy() - co10 = p_coeff_10_dsl.asnumpy() - assert np.allclose(ref_1, co1) - assert np.allclose(ref_2, co2) - assert np.allclose(ref_3, co3) - assert np.allclose(ref_4, co4) - assert np.allclose(ref_5, co5) - assert np.allclose(ref_6, co6) - assert np.allclose(ref_7, co7) - assert np.allclose(ref_8, co8) - assert np.allclose(ref_9, co9) - assert np.allclose(ref_10, co10) + @pytest.mark.slow_tests + @pytest.fixture + def input_data(self, grid): + p_cc = random_field(grid, CellDim, KDim) + lsq_pseudoinv_1_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_pseudoinv_2_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_pseudoinv_3_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_pseudoinv_4_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_pseudoinv_5_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_pseudoinv_6_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_pseudoinv_7_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_pseudoinv_8_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_pseudoinv_9_field = as_1D_sparse_field( + random_field(grid, CellDim, C2E2C2E2CDim), CECECDim + ) + lsq_moments_1 = random_field(grid, CellDim) + lsq_moments_2 = random_field(grid, CellDim) + lsq_moments_3 = random_field(grid, CellDim) + lsq_moments_4 = random_field(grid, CellDim) + lsq_moments_5 = random_field(grid, CellDim) + lsq_moments_6 = random_field(grid, CellDim) + lsq_moments_7 = random_field(grid, CellDim) + lsq_moments_8 = random_field(grid, CellDim) + lsq_moments_9 = random_field(grid, CellDim) + p_coeff_1_dsl = zero_field(grid, CellDim, KDim) + p_coeff_2_dsl = zero_field(grid, CellDim, KDim) + p_coeff_3_dsl = zero_field(grid, CellDim, KDim) + p_coeff_4_dsl = zero_field(grid, CellDim, KDim) + p_coeff_5_dsl = zero_field(grid, CellDim, KDim) + p_coeff_6_dsl = zero_field(grid, CellDim, KDim) + p_coeff_7_dsl = zero_field(grid, CellDim, KDim) + p_coeff_8_dsl = zero_field(grid, CellDim, KDim) + p_coeff_9_dsl = zero_field(grid, CellDim, KDim) + p_coeff_10_dsl = zero_field(grid, CellDim, KDim) + return dict( + p_cc=p_cc, + lsq_pseudoinv_1=lsq_pseudoinv_1_field, + lsq_pseudoinv_2=lsq_pseudoinv_2_field, + lsq_pseudoinv_3=lsq_pseudoinv_3_field, + lsq_pseudoinv_4=lsq_pseudoinv_4_field, + lsq_pseudoinv_5=lsq_pseudoinv_5_field, + lsq_pseudoinv_6=lsq_pseudoinv_6_field, + lsq_pseudoinv_7=lsq_pseudoinv_7_field, + lsq_pseudoinv_8=lsq_pseudoinv_8_field, + lsq_pseudoinv_9=lsq_pseudoinv_9_field, + lsq_moments_1=lsq_moments_1, + lsq_moments_2=lsq_moments_2, + lsq_moments_3=lsq_moments_3, + lsq_moments_4=lsq_moments_4, + lsq_moments_5=lsq_moments_5, + lsq_moments_6=lsq_moments_6, + lsq_moments_7=lsq_moments_7, + lsq_moments_8=lsq_moments_8, + lsq_moments_9=lsq_moments_9, + p_coeff_1_dsl=p_coeff_1_dsl, + p_coeff_2_dsl=p_coeff_2_dsl, + p_coeff_3_dsl=p_coeff_3_dsl, + p_coeff_4_dsl=p_coeff_4_dsl, + p_coeff_5_dsl=p_coeff_5_dsl, + p_coeff_6_dsl=p_coeff_6_dsl, + p_coeff_7_dsl=p_coeff_7_dsl, + p_coeff_8_dsl=p_coeff_8_dsl, + p_coeff_9_dsl=p_coeff_9_dsl, + p_coeff_10_dsl=p_coeff_10_dsl, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_l_svd_stencil.py b/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_l_svd_stencil.py index cc0a59d247..298b6e71f3 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_l_svd_stencil.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_l_svd_stencil.py @@ -12,65 +12,66 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.recon_lsq_cell_l_svd_stencil import ( recon_lsq_cell_l_svd_stencil, ) from icon4py.model.common.dimension import C2E2CDim, CECDim, CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field - - -def recon_lsq_cell_l_svd_stencil_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 = p_cc - p_coeff_2 = np.sum(lsq_pseudoinv_1 * n_diff, axis=1) - p_coeff_3 = np.sum(lsq_pseudoinv_2 * n_diff, axis=1) - return p_coeff_1, p_coeff_2, p_coeff_3 - - -def test_recon_lsq_cell_l_svd_stencil(backend): - grid = SimpleGrid() - p_cc = random_field(grid, CellDim, KDim) - lsq_pseudoinv_1 = random_field(grid, CellDim, C2E2CDim) - lsq_pseudoinv_1_field = as_1D_sparse_field(lsq_pseudoinv_1, CECDim) +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + reshape, + zero_field, +) - lsq_pseudoinv_2 = random_field(grid, CellDim, C2E2CDim) - lsq_pseudoinv_2_field = as_1D_sparse_field(lsq_pseudoinv_2, CECDim) - p_coeff_1 = zero_field(grid, CellDim, KDim) - p_coeff_2 = zero_field(grid, CellDim, KDim) - p_coeff_3 = zero_field(grid, CellDim, KDim) - ref_1, ref_2, ref_3 = recon_lsq_cell_l_svd_stencil_numpy( - grid.connectivities[C2E2CDim], - p_cc.asnumpy(), - lsq_pseudoinv_1.asnumpy(), - lsq_pseudoinv_2.asnumpy(), +class TestReconLsqCellLSvdStencil(StencilTest): + PROGRAM = recon_lsq_cell_l_svd_stencil + OUTPUTS = ( + "p_coeff_1_dsl", + "p_coeff_2_dsl", + "p_coeff_3_dsl", ) - recon_lsq_cell_l_svd_stencil.with_backend(backend)( - p_cc, - lsq_pseudoinv_1_field, - lsq_pseudoinv_2_field, - p_coeff_1, - p_coeff_2, - p_coeff_3, - offset_provider={ - "C2E2C": grid.get_offset_provider("C2E2C"), - "C2CEC": grid.get_offset_provider("C2CEC"), - }, - ) - co1 = p_coeff_1.asnumpy() - co2 = p_coeff_2.asnumpy() - co3 = p_coeff_3.asnumpy() - assert np.allclose(ref_1, co1) - assert np.allclose(ref_2, co2) - assert np.allclose(ref_3, co3) + @staticmethod + def reference( + grid, + p_cc: np.array, + lsq_pseudoinv_1: np.array, + lsq_pseudoinv_2: np.array, + **kwargs, + ): + c2e2c = grid.connectivities[C2E2CDim] + p_cc_e = np.expand_dims(p_cc, axis=1) + n_diff = p_cc[c2e2c] - p_cc_e + lsq_pseudoinv_1 = reshape(lsq_pseudoinv_1, c2e2c.shape) + lsq_pseudoinv_1 = np.expand_dims(lsq_pseudoinv_1, axis=-1) + lsq_pseudoinv_2 = reshape(lsq_pseudoinv_2, c2e2c.shape) + lsq_pseudoinv_2 = np.expand_dims(lsq_pseudoinv_2, axis=-1) + p_coeff_1_dsl = p_cc + p_coeff_2_dsl = np.sum(lsq_pseudoinv_1 * n_diff, axis=1) + p_coeff_3_dsl = np.sum(lsq_pseudoinv_2 * n_diff, axis=1) + return dict( + p_coeff_1_dsl=p_coeff_1_dsl, + p_coeff_2_dsl=p_coeff_2_dsl, + p_coeff_3_dsl=p_coeff_3_dsl, + ) + + @pytest.fixture + def input_data(self, grid): + p_cc = random_field(grid, CellDim, KDim) + lsq_pseudoinv_1_field = as_1D_sparse_field(random_field(grid, CellDim, C2E2CDim), CECDim) + lsq_pseudoinv_2_field = as_1D_sparse_field(random_field(grid, CellDim, C2E2CDim), CECDim) + p_coeff_1_dsl = zero_field(grid, CellDim, KDim) + p_coeff_2_dsl = zero_field(grid, CellDim, KDim) + p_coeff_3_dsl = zero_field(grid, CellDim, KDim) + return dict( + p_cc=p_cc, + lsq_pseudoinv_1=lsq_pseudoinv_1_field, + lsq_pseudoinv_2=lsq_pseudoinv_2_field, + p_coeff_1_dsl=p_coeff_1_dsl, + p_coeff_2_dsl=p_coeff_2_dsl, + p_coeff_3_dsl=p_coeff_3_dsl, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_01.py index 9e8cb8b2b1..f45130e824 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_01.py @@ -12,51 +12,53 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.step_advection_stencil_01 import step_advection_stencil_01 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.helpers import Output, StencilTest, 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: - tmp = pd_time * ( - p_mflx_contra_v[:, 1:] * deepatmo_divzl - p_mflx_contra_v[:, :-1] * deepatmo_divzu +class TestStepAdvectionStencil01(StencilTest): + PROGRAM = step_advection_stencil_01 + OUTPUTS = ( + Output( + "rhodz_ast2", + refslice=(slice(None), slice(None, -1)), + gtslice=(slice(None), slice(None, -1)), + ), ) - return rhodz_ast + tmp - -def test_step_advection_stencil_01(backend): - grid = SimpleGrid() - rhodz_ast = random_field(grid, CellDim, KDim) - p_mflx_contra = random_field(grid, CellDim, KDim, extend={KDim: 1}) - deepatmo_divzl = random_field(grid, KDim) - deepatmo_divzu = random_field(grid, KDim) - result = zero_field(grid, CellDim, KDim) - p_dtime = 0.1 - - ref = step_advection_stencil_01_numpy( - rhodz_ast.asnumpy(), - p_mflx_contra.asnumpy(), - deepatmo_divzl.asnumpy(), - deepatmo_divzu.asnumpy(), + @staticmethod + def reference( + grid, + rhodz_ast: np.array, + p_mflx_contra_v: np.array, + deepatmo_divzl: np.array, + deepatmo_divzu: np.array, p_dtime, - ) + **kwargs, + ): + tmp = p_dtime * ( + p_mflx_contra_v[:, 1:] * deepatmo_divzl - p_mflx_contra_v[:, :-1] * deepatmo_divzu + ) + rhodz_ast2 = rhodz_ast + tmp - step_advection_stencil_01.with_backend(backend)( - rhodz_ast, - p_mflx_contra, - deepatmo_divzl, - deepatmo_divzu, - p_dtime, - result, - offset_provider={"Koff": KDim}, - ) + return dict(rhodz_ast2=rhodz_ast2) - assert np.allclose(ref[:, :-1], result.asnumpy()[:, :-1]) + @pytest.fixture + def input_data(self, grid): + rhodz_ast = random_field(grid, CellDim, KDim) + p_mflx_contra_v = random_field(grid, CellDim, KDim, extend={KDim: 1}) + deepatmo_divzl = random_field(grid, KDim) + deepatmo_divzu = random_field(grid, KDim) + rhodz_ast2 = zero_field(grid, CellDim, KDim) + p_dtime = 0.1 + return dict( + rhodz_ast=rhodz_ast, + p_mflx_contra_v=p_mflx_contra_v, + deepatmo_divzl=deepatmo_divzl, + deepatmo_divzu=deepatmo_divzu, + p_dtime=p_dtime, + rhodz_ast2=rhodz_ast2, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_02.py index 88f32eea42..c34cd6f037 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_02.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_02.py @@ -12,49 +12,50 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.step_advection_stencil_02 import step_advection_stencil_02 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.helpers import Output, StencilTest, 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 = p_mflx_contra_v[:, 1:] * deepatmo_divzl - p_mflx_contra_v[:, :-1] * deepatmo_divzu - return np.maximum(0.1 * rhodz_new, rhodz_new) - pd_time * tmp - - -def test_step_advection_stencil_02(backend): - grid = SimpleGrid() - rhodz_ast = random_field(grid, CellDim, KDim) - p_mflx_contra = random_field(grid, CellDim, KDim, extend={KDim: 1}) - deepatmo_divzl = random_field(grid, KDim) - deepatmo_divzu = random_field(grid, KDim) - result = zero_field(grid, CellDim, KDim) - p_dtime = 0.1 - - ref = step_advection_stencil_02_numpy( - rhodz_ast.asnumpy(), - p_mflx_contra.asnumpy(), - deepatmo_divzl.asnumpy(), - deepatmo_divzu.asnumpy(), - p_dtime, - ) - - step_advection_stencil_02.with_backend(backend)( - rhodz_ast, - p_mflx_contra, - deepatmo_divzl, - deepatmo_divzu, - p_dtime, - result, - offset_provider={"Koff": KDim}, +class TestStepAdvectionStencil02(StencilTest): + PROGRAM = step_advection_stencil_02 + OUTPUTS = ( + Output( + "rhodz_ast2", + refslice=(slice(None), slice(None, -1)), + gtslice=(slice(None), slice(None, -1)), + ), ) - assert np.allclose(ref[:, :-1], result.asnumpy()[:, :-1]) + @staticmethod + def reference( + grid, + p_rhodz_new: np.array, + p_mflx_contra_v: np.array, + deepatmo_divzl: np.array, + deepatmo_divzu: np.array, + p_dtime: float, + **kwargs, + ): + tmp = p_mflx_contra_v[:, 1:] * deepatmo_divzl - p_mflx_contra_v[:, :-1] * deepatmo_divzu + rhodz_ast2 = np.maximum(0.1 * p_rhodz_new, p_rhodz_new) - p_dtime * tmp + return dict(rhodz_ast2=rhodz_ast2) + + @pytest.fixture + def input_data(self, grid): + p_rhodz_new = random_field(grid, CellDim, KDim) + p_mflx_contra_v = random_field(grid, CellDim, KDim, extend={KDim: 1}) + deepatmo_divzl = random_field(grid, KDim) + deepatmo_divzu = random_field(grid, KDim) + p_dtime = 0.1 + rhodz_ast2 = zero_field(grid, CellDim, KDim) + return dict( + p_rhodz_new=p_rhodz_new, + p_mflx_contra_v=p_mflx_contra_v, + deepatmo_divzl=deepatmo_divzl, + deepatmo_divzu=deepatmo_divzu, + p_dtime=p_dtime, + rhodz_ast2=rhodz_ast2, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_03.py b/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_03.py index cb001eda6f..2e1f83ab08 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_03.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_03.py @@ -12,43 +12,39 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.step_advection_stencil_03 import step_advection_stencil_03 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field -def step_advection_stencil_03_numpy( - p_tracer_now: np.array, - p_grf_tend_tracer: np.array, - p_dtime, -) -> np.array: - p_tracer_new = p_tracer_now + p_dtime * p_grf_tend_tracer - p_tracer_new = np.where(p_tracer_new < 0.0, 0.0, p_tracer_new) +class TestStepAdvectionStencil03(StencilTest): + PROGRAM = step_advection_stencil_03 + OUTPUTS = ("p_tracer_new",) - return p_tracer_new - - -def test_step_advection_stencil_03(backend): - grid = SimpleGrid() - - p_tracer_now = random_field(grid, CellDim, KDim) - p_grf_tend_tracer = random_field(grid, CellDim, KDim) - p_tracer_new = random_field(grid, CellDim, KDim) - - p_dtime = np.float64(5.0) - - ref = step_advection_stencil_03_numpy( - p_tracer_now.asnumpy(), - p_grf_tend_tracer.asnumpy(), - p_dtime, - ) - step_advection_stencil_03.with_backend(backend)( - p_tracer_now, - p_grf_tend_tracer, - p_tracer_new, + @staticmethod + def reference( + grid, + p_tracer_now: np.array, + p_grf_tend_tracer: np.array, p_dtime, - offset_provider={}, - ) - assert np.allclose(p_tracer_new.asnumpy(), ref) + **kwargs, + ): + p_tracer_new = p_tracer_now + p_dtime * p_grf_tend_tracer + p_tracer_new = np.where(p_tracer_new < 0.0, 0.0, p_tracer_new) + + return dict(p_tracer_new=p_tracer_new) + + @pytest.fixture + def input_data(self, grid): + p_tracer_now = random_field(grid, CellDim, KDim) + p_grf_tend_tracer = random_field(grid, CellDim, KDim) + p_tracer_new = random_field(grid, CellDim, KDim) + p_dtime = np.float64(5.0) + return dict( + p_tracer_now=p_tracer_now, + p_grf_tend_tracer=p_grf_tend_tracer, + p_dtime=p_dtime, + p_tracer_new=p_tracer_new, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_04.py b/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_04.py index 66c443bc2e..6541824d3a 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_04.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_step_advection_stencil_04.py @@ -12,40 +12,38 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.step_advection_stencil_04 import step_advection_stencil_04 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field -def step_advection_stencil_04_numpy( - p_tracer_now: np.array, - p_tracer_new: np.array, - p_dtime, -) -> np.array: - opt_ddt_tracer_adv = (p_tracer_new - p_tracer_now) / p_dtime - return opt_ddt_tracer_adv +class TestStepAdvectionStencil04(StencilTest): + PROGRAM = step_advection_stencil_04 + OUTPUTS = ("opt_ddt_tracer_adv",) - -def test_step_advection_stencil_04(backend): - grid = SimpleGrid() - - p_tracer_now = random_field(grid, CellDim, KDim) - p_tracer_new = random_field(grid, CellDim, KDim) - opt_ddt_tracer_adv = zero_field(grid, CellDim, KDim) - p_dtime = np.float64(5.0) - - ref = step_advection_stencil_04_numpy( - p_tracer_now.asnumpy(), - p_tracer_new.asnumpy(), - p_dtime, - ) - step_advection_stencil_04.with_backend(backend)( - p_tracer_now, - p_tracer_new, - opt_ddt_tracer_adv, + @staticmethod + def reference( + grid, + p_tracer_now: np.array, + p_tracer_new: np.array, p_dtime, - offset_provider={}, - ) - assert np.allclose(opt_ddt_tracer_adv.asnumpy(), ref) + **kwargs, + ): + opt_ddt_tracer_adv = (p_tracer_new - p_tracer_now) / p_dtime + + return dict(opt_ddt_tracer_adv=opt_ddt_tracer_adv) + + @pytest.fixture + def input_data(self, grid): + p_tracer_now = random_field(grid, CellDim, KDim) + p_tracer_new = random_field(grid, CellDim, KDim) + opt_ddt_tracer_adv = zero_field(grid, CellDim, KDim) + p_dtime = np.float64(5.0) + return dict( + p_tracer_now=p_tracer_now, + p_tracer_new=p_tracer_new, + p_dtime=p_dtime, + opt_ddt_tracer_adv=opt_ddt_tracer_adv, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura3_stencil_01.py index 57c6304ef9..a538479e1e 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura3_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura3_stencil_01.py @@ -12,205 +12,183 @@ # 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.advection.upwind_hflux_miura3_stencil_01 import ( upwind_hflux_miura3_stencil_01, ) from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, random_mask, zero_field +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + random_mask, + zero_field, +) -def upwind_hflux_miura3_stencil_01_numpy( - e2c: np.array, - z_lsq_coeff_1: np.array, - z_lsq_coeff_2: np.array, - z_lsq_coeff_3: np.array, - z_lsq_coeff_4: np.array, - z_lsq_coeff_5: np.array, - z_lsq_coeff_6: np.array, - z_lsq_coeff_7: np.array, - z_lsq_coeff_8: np.array, - z_lsq_coeff_9: np.array, - z_lsq_coeff_10: np.array, - z_quad_vector_sum_1: np.array, - z_quad_vector_sum_2: np.array, - z_quad_vector_sum_3: np.array, - z_quad_vector_sum_4: np.array, - z_quad_vector_sum_5: np.array, - z_quad_vector_sum_6: np.array, - z_quad_vector_sum_7: np.array, - z_quad_vector_sum_8: np.array, - z_quad_vector_sum_9: np.array, - z_quad_vector_sum_10: np.array, - z_dreg_area: np.array, - p_mass_flx_e: np.array, - cell_rel_idx_dsl: np.array, -) -> np.array: - z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] - z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] - z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] - z_lsq_coeff_4_e2c = z_lsq_coeff_4[e2c] - z_lsq_coeff_5_e2c = z_lsq_coeff_5[e2c] - z_lsq_coeff_6_e2c = z_lsq_coeff_6[e2c] - z_lsq_coeff_7_e2c = z_lsq_coeff_7[e2c] - z_lsq_coeff_8_e2c = z_lsq_coeff_8[e2c] - z_lsq_coeff_9_e2c = z_lsq_coeff_9[e2c] - z_lsq_coeff_10_e2c = z_lsq_coeff_10[e2c] +class TestUpwindHfluxMiura3Stencil01(StencilTest): + PROGRAM = upwind_hflux_miura3_stencil_01 + OUTPUTS = ("p_out_e_miura3",) - p_out_e_miura3 = ( - ( - np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_1_e2c[:, 1], - z_lsq_coeff_1_e2c[:, 0], - ) - * z_quad_vector_sum_1 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_2_e2c[:, 1], - z_lsq_coeff_2_e2c[:, 0], - ) - * z_quad_vector_sum_2 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_3_e2c[:, 1], - z_lsq_coeff_3_e2c[:, 0], - ) - * z_quad_vector_sum_3 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_4_e2c[:, 1], - z_lsq_coeff_4_e2c[:, 0], - ) - * z_quad_vector_sum_4 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_5_e2c[:, 1], - z_lsq_coeff_5_e2c[:, 0], - ) - * z_quad_vector_sum_5 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_6_e2c[:, 1], - z_lsq_coeff_6_e2c[:, 0], - ) - * z_quad_vector_sum_6 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_7_e2c[:, 1], - z_lsq_coeff_7_e2c[:, 0], - ) - * z_quad_vector_sum_7 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_8_e2c[:, 1], - z_lsq_coeff_8_e2c[:, 0], - ) - * z_quad_vector_sum_8 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_9_e2c[:, 1], - z_lsq_coeff_9_e2c[:, 0], - ) - * z_quad_vector_sum_9 - + np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_10_e2c[:, 1], - z_lsq_coeff_10_e2c[:, 0], + @staticmethod + def reference( + grid, + z_lsq_coeff_1: np.array, + z_lsq_coeff_2: np.array, + z_lsq_coeff_3: np.array, + z_lsq_coeff_4: np.array, + z_lsq_coeff_5: np.array, + z_lsq_coeff_6: np.array, + z_lsq_coeff_7: np.array, + z_lsq_coeff_8: np.array, + z_lsq_coeff_9: np.array, + z_lsq_coeff_10: np.array, + z_quad_vector_sum_1: np.array, + z_quad_vector_sum_2: np.array, + z_quad_vector_sum_3: np.array, + z_quad_vector_sum_4: np.array, + z_quad_vector_sum_5: np.array, + z_quad_vector_sum_6: np.array, + z_quad_vector_sum_7: np.array, + z_quad_vector_sum_8: np.array, + z_quad_vector_sum_9: np.array, + z_quad_vector_sum_10: np.array, + z_dreg_area: np.array, + p_mass_flx_e: np.array, + cell_rel_idx_dsl: np.array, + **kwargs, + ): + e2c = grid.connectivities[E2CDim] + z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] + z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] + z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] + z_lsq_coeff_4_e2c = z_lsq_coeff_4[e2c] + z_lsq_coeff_5_e2c = z_lsq_coeff_5[e2c] + z_lsq_coeff_6_e2c = z_lsq_coeff_6[e2c] + z_lsq_coeff_7_e2c = z_lsq_coeff_7[e2c] + z_lsq_coeff_8_e2c = z_lsq_coeff_8[e2c] + z_lsq_coeff_9_e2c = z_lsq_coeff_9[e2c] + z_lsq_coeff_10_e2c = z_lsq_coeff_10[e2c] + + p_out_e_miura3 = ( + ( + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_e2c[:, 1], + z_lsq_coeff_1_e2c[:, 0], + ) + * z_quad_vector_sum_1 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_e2c[:, 1], + z_lsq_coeff_2_e2c[:, 0], + ) + * z_quad_vector_sum_2 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_e2c[:, 1], + z_lsq_coeff_3_e2c[:, 0], + ) + * z_quad_vector_sum_3 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_4_e2c[:, 1], + z_lsq_coeff_4_e2c[:, 0], + ) + * z_quad_vector_sum_4 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_5_e2c[:, 1], + z_lsq_coeff_5_e2c[:, 0], + ) + * z_quad_vector_sum_5 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_6_e2c[:, 1], + z_lsq_coeff_6_e2c[:, 0], + ) + * z_quad_vector_sum_6 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_7_e2c[:, 1], + z_lsq_coeff_7_e2c[:, 0], + ) + * z_quad_vector_sum_7 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_8_e2c[:, 1], + z_lsq_coeff_8_e2c[:, 0], + ) + * z_quad_vector_sum_8 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_9_e2c[:, 1], + z_lsq_coeff_9_e2c[:, 0], + ) + * z_quad_vector_sum_9 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_10_e2c[:, 1], + z_lsq_coeff_10_e2c[:, 0], + ) + * z_quad_vector_sum_10 ) - * z_quad_vector_sum_10 + / z_dreg_area + * p_mass_flx_e ) - / z_dreg_area - * p_mass_flx_e - ) - - return p_out_e_miura3 + return dict(p_out_e_miura3=p_out_e_miura3) -def test_upwind_hflux_miura3_stencil_01(backend): - grid = SimpleGrid() - - z_lsq_coeff_1 = random_field(grid, CellDim, KDim) - z_lsq_coeff_2 = random_field(grid, CellDim, KDim) - z_lsq_coeff_3 = random_field(grid, CellDim, KDim) - z_lsq_coeff_4 = random_field(grid, CellDim, KDim) - z_lsq_coeff_5 = random_field(grid, CellDim, KDim) - z_lsq_coeff_6 = random_field(grid, CellDim, KDim) - z_lsq_coeff_7 = random_field(grid, CellDim, KDim) - z_lsq_coeff_8 = random_field(grid, CellDim, KDim) - z_lsq_coeff_9 = random_field(grid, CellDim, KDim) - z_lsq_coeff_10 = random_field(grid, CellDim, KDim) - z_quad_vector_sum_1 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_2 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_3 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_4 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_5 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_6 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_7 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_8 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_9 = random_field(grid, EdgeDim, KDim) - z_quad_vector_sum_10 = random_field(grid, EdgeDim, KDim) - p_mass_flx_e = random_field(grid, EdgeDim, KDim) - z_dreg_area = random_field(grid, EdgeDim, KDim) - cell_rel_idx_dsl = random_mask(grid, EdgeDim, KDim, dtype=int32) - p_out_e_miura3 = zero_field(grid, EdgeDim, KDim) - - ref = upwind_hflux_miura3_stencil_01_numpy( - grid.connectivities[E2CDim], - z_lsq_coeff_1.asnumpy(), - z_lsq_coeff_2.asnumpy(), - z_lsq_coeff_3.asnumpy(), - z_lsq_coeff_4.asnumpy(), - z_lsq_coeff_5.asnumpy(), - z_lsq_coeff_6.asnumpy(), - z_lsq_coeff_7.asnumpy(), - z_lsq_coeff_8.asnumpy(), - z_lsq_coeff_9.asnumpy(), - z_lsq_coeff_10.asnumpy(), - z_quad_vector_sum_1.asnumpy(), - z_quad_vector_sum_2.asnumpy(), - z_quad_vector_sum_3.asnumpy(), - z_quad_vector_sum_4.asnumpy(), - z_quad_vector_sum_5.asnumpy(), - z_quad_vector_sum_6.asnumpy(), - z_quad_vector_sum_7.asnumpy(), - z_quad_vector_sum_8.asnumpy(), - z_quad_vector_sum_9.asnumpy(), - z_quad_vector_sum_10.asnumpy(), - z_dreg_area.asnumpy(), - p_mass_flx_e.asnumpy(), - cell_rel_idx_dsl.asnumpy(), - ) - - upwind_hflux_miura3_stencil_01.with_backend(backend)( - z_lsq_coeff_1, - z_lsq_coeff_2, - z_lsq_coeff_3, - z_lsq_coeff_4, - z_lsq_coeff_5, - z_lsq_coeff_6, - z_lsq_coeff_7, - z_lsq_coeff_8, - z_lsq_coeff_9, - z_lsq_coeff_10, - z_quad_vector_sum_1, - z_quad_vector_sum_2, - z_quad_vector_sum_3, - z_quad_vector_sum_4, - z_quad_vector_sum_5, - z_quad_vector_sum_6, - z_quad_vector_sum_7, - z_quad_vector_sum_8, - z_quad_vector_sum_9, - z_quad_vector_sum_10, - z_dreg_area, - p_mass_flx_e, - cell_rel_idx_dsl, - p_out_e_miura3, - offset_provider={ - "E2C": grid.get_offset_provider("E2C"), - }, - ) - assert np.allclose(p_out_e_miura3.asnumpy(), ref) + @pytest.fixture + def input_data(self, grid): + z_lsq_coeff_1 = random_field(grid, CellDim, KDim) + z_lsq_coeff_2 = random_field(grid, CellDim, KDim) + z_lsq_coeff_3 = random_field(grid, CellDim, KDim) + z_lsq_coeff_4 = random_field(grid, CellDim, KDim) + z_lsq_coeff_5 = random_field(grid, CellDim, KDim) + z_lsq_coeff_6 = random_field(grid, CellDim, KDim) + z_lsq_coeff_7 = random_field(grid, CellDim, KDim) + z_lsq_coeff_8 = random_field(grid, CellDim, KDim) + z_lsq_coeff_9 = random_field(grid, CellDim, KDim) + z_lsq_coeff_10 = random_field(grid, CellDim, KDim) + z_quad_vector_sum_1 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_2 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_3 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_4 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_5 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_6 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_7 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_8 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_9 = random_field(grid, EdgeDim, KDim) + z_quad_vector_sum_10 = random_field(grid, EdgeDim, KDim) + p_mass_flx_e = random_field(grid, EdgeDim, KDim) + z_dreg_area = random_field(grid, EdgeDim, KDim) + cell_rel_idx_dsl = random_mask(grid, EdgeDim, KDim, dtype=int32) + p_out_e_miura3 = zero_field(grid, EdgeDim, KDim) + return dict( + z_lsq_coeff_1=z_lsq_coeff_1, + z_lsq_coeff_2=z_lsq_coeff_2, + z_lsq_coeff_3=z_lsq_coeff_3, + z_lsq_coeff_4=z_lsq_coeff_4, + z_lsq_coeff_5=z_lsq_coeff_5, + z_lsq_coeff_6=z_lsq_coeff_6, + z_lsq_coeff_7=z_lsq_coeff_7, + z_lsq_coeff_8=z_lsq_coeff_8, + z_lsq_coeff_9=z_lsq_coeff_9, + z_lsq_coeff_10=z_lsq_coeff_10, + z_quad_vector_sum_1=z_quad_vector_sum_1, + z_quad_vector_sum_2=z_quad_vector_sum_2, + z_quad_vector_sum_3=z_quad_vector_sum_3, + z_quad_vector_sum_4=z_quad_vector_sum_4, + z_quad_vector_sum_5=z_quad_vector_sum_5, + z_quad_vector_sum_6=z_quad_vector_sum_6, + z_quad_vector_sum_7=z_quad_vector_sum_7, + z_quad_vector_sum_8=z_quad_vector_sum_8, + z_quad_vector_sum_9=z_quad_vector_sum_9, + z_quad_vector_sum_10=z_quad_vector_sum_10, + z_dreg_area=z_dreg_area, + p_mass_flx_e=p_mass_flx_e, + cell_rel_idx_dsl=cell_rel_idx_dsl, + p_out_e_miura3=p_out_e_miura3, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py index d8ca5dd519..1e33e6a7a7 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_01.py @@ -12,87 +12,81 @@ # 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.advection.upwind_hflux_miura_cycl_stencil_01 import ( upwind_hflux_miura_cycl_stencil_01, ) from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, random_mask, zero_field - - -def upwind_hflux_miura_cycl_stencil_01_numpy( - e2c: np.array, - z_lsq_coeff_1_dsl: np.array, - z_lsq_coeff_2_dsl: np.array, - z_lsq_coeff_3_dsl: np.array, - distv_bary_1: np.array, - distv_bary_2: np.array, - p_mass_flx_e: np.array, - cell_rel_idx_dsl: int32, -): - z_lsq_coeff_1_dsl_e2c = z_lsq_coeff_1_dsl[e2c] - z_lsq_coeff_2_dsl_e2c = z_lsq_coeff_2_dsl[e2c] - z_lsq_coeff_3_dsl_e2c = z_lsq_coeff_3_dsl[e2c] - - z_tracer_mflx_dsl = ( - np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_1_dsl_e2c[:, 1], - z_lsq_coeff_1_dsl_e2c[:, 0], - ) - + distv_bary_1 - * np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_2_dsl_e2c[:, 1], - z_lsq_coeff_2_dsl_e2c[:, 0], - ) - + distv_bary_2 - * np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_3_dsl_e2c[:, 1], - z_lsq_coeff_3_dsl_e2c[:, 0], - ) - ) * p_mass_flx_e +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + random_mask, + zero_field, +) - return z_tracer_mflx_dsl +class TestUpwindHfluxMiuraCyclStencil01(StencilTest): + PROGRAM = upwind_hflux_miura_cycl_stencil_01 + OUTPUTS = ("z_tracer_mflx_dsl",) -def test_upwind_hflux_miura_cycl_stencil_01(backend): - grid = SimpleGrid() + @staticmethod + def reference( + grid, + z_lsq_coeff_1_dsl: np.array, + z_lsq_coeff_2_dsl: np.array, + z_lsq_coeff_3_dsl: np.array, + distv_bary_1: np.array, + distv_bary_2: np.array, + p_mass_flx_e: np.array, + cell_rel_idx_dsl: np.array, + **kwargs, + ): + e2c = grid.connectivities[E2CDim] + z_lsq_coeff_1_dsl_e2c = z_lsq_coeff_1_dsl[e2c] + z_lsq_coeff_2_dsl_e2c = z_lsq_coeff_2_dsl[e2c] + z_lsq_coeff_3_dsl_e2c = z_lsq_coeff_3_dsl[e2c] - z_lsq_coeff_1_dsl = random_field(grid, CellDim, KDim) - z_lsq_coeff_2_dsl = random_field(grid, CellDim, KDim) - z_lsq_coeff_3_dsl = random_field(grid, CellDim, KDim) - distv_bary_1 = random_field(grid, EdgeDim, KDim) - distv_bary_2 = random_field(grid, EdgeDim, KDim) - p_mass_flx_e = random_field(grid, EdgeDim, KDim) - cell_rel_idx_dsl = random_mask(grid, EdgeDim, KDim, dtype=int32) - z_tracer_mflx_dsl = zero_field(grid, EdgeDim, KDim) + z_tracer_mflx_dsl = ( + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_dsl_e2c[:, 1], + z_lsq_coeff_1_dsl_e2c[:, 0], + ) + + distv_bary_1 + * np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_dsl_e2c[:, 1], + z_lsq_coeff_2_dsl_e2c[:, 0], + ) + + distv_bary_2 + * np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_dsl_e2c[:, 1], + z_lsq_coeff_3_dsl_e2c[:, 0], + ) + ) * p_mass_flx_e - ref = upwind_hflux_miura_cycl_stencil_01_numpy( - grid.connectivities[E2CDim], - z_lsq_coeff_1_dsl.asnumpy(), - z_lsq_coeff_2_dsl.asnumpy(), - z_lsq_coeff_3_dsl.asnumpy(), - distv_bary_1.asnumpy(), - distv_bary_2.asnumpy(), - p_mass_flx_e.asnumpy(), - cell_rel_idx_dsl.asnumpy(), - ) + return dict(z_tracer_mflx_dsl=z_tracer_mflx_dsl) - upwind_hflux_miura_cycl_stencil_01.with_backend(backend)( - z_lsq_coeff_1_dsl, - z_lsq_coeff_2_dsl, - z_lsq_coeff_3_dsl, - distv_bary_1, - distv_bary_2, - p_mass_flx_e, - cell_rel_idx_dsl, - z_tracer_mflx_dsl, - offset_provider={ - "E2C": grid.get_offset_provider("E2C"), - }, - ) - assert np.allclose(ref, z_tracer_mflx_dsl.asnumpy()) + @pytest.fixture + def input_data(self, grid): + z_lsq_coeff_1_dsl = random_field(grid, CellDim, KDim) + z_lsq_coeff_2_dsl = random_field(grid, CellDim, KDim) + z_lsq_coeff_3_dsl = random_field(grid, CellDim, KDim) + distv_bary_1 = random_field(grid, EdgeDim, KDim) + distv_bary_2 = random_field(grid, EdgeDim, KDim) + p_mass_flx_e = random_field(grid, EdgeDim, KDim) + cell_rel_idx_dsl = random_mask(grid, EdgeDim, KDim, dtype=int32) + z_tracer_mflx_dsl = zero_field(grid, EdgeDim, KDim) + return dict( + z_lsq_coeff_1_dsl=z_lsq_coeff_1_dsl, + z_lsq_coeff_2_dsl=z_lsq_coeff_2_dsl, + z_lsq_coeff_3_dsl=z_lsq_coeff_3_dsl, + distv_bary_1=distv_bary_1, + distv_bary_2=distv_bary_2, + p_mass_flx_e=p_mass_flx_e, + cell_rel_idx_dsl=cell_rel_idx_dsl, + z_tracer_mflx_dsl=z_tracer_mflx_dsl, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py index 6edc412c24..dbb64d6807 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_02.py @@ -12,89 +12,82 @@ # 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.advection.upwind_hflux_miura_cycl_stencil_02 import ( upwind_hflux_miura_cycl_stencil_02, ) from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field -def upwind_hflux_miura_cycl_stencil_02_numpy( - c2e: np.array, - nsub: int32, - p_mass_flx_e: np.array, - geofac_div: np.array, - z_rhofluxdiv_c: np.array, - z_tracer_mflx: np.array, - z_rho_now: np.array, - z_tracer_now: np.array, - z_dtsub: float, -): - p_mass_flx_e_c2e = p_mass_flx_e[c2e] - geofac_div = np.expand_dims(geofac_div, axis=-1) - z_tracer_mflx_c2e = z_tracer_mflx[c2e] - - z_rhofluxdiv_c_out = ( - np.sum(p_mass_flx_e_c2e * geofac_div, axis=1) if nsub == int32(1) else z_rhofluxdiv_c +class TestUpwindHfluxMiuraCyclStencil02(StencilTest): + PROGRAM = upwind_hflux_miura_cycl_stencil_02 + OUTPUTS = ( + "z_rhofluxdiv_c_out", + "z_fluxdiv_c_dsl", + "z_rho_new_dsl", + "z_tracer_new_dsl", ) - z_fluxdiv_c_dsl = np.sum(z_tracer_mflx_c2e * geofac_div, axis=1) - - z_rho_new_dsl = z_rho_now - z_dtsub * z_rhofluxdiv_c_out - - z_tracer_new_dsl = (z_tracer_now * z_rho_now - z_dtsub * z_fluxdiv_c_dsl) / z_rho_new_dsl - return (z_rhofluxdiv_c_out, z_fluxdiv_c_dsl, z_rho_new_dsl, z_tracer_new_dsl) + @staticmethod + def reference( + grid, + nsub: int32, + p_mass_flx_e: np.array, + geofac_div: np.array, + z_rhofluxdiv_c: np.array, + z_tracer_mflx: np.array, + z_rho_now: np.array, + z_tracer_now: np.array, + z_dtsub: float, + **kwargs, + ): + c2e = grid.connectivities[C2EDim] + p_mass_flx_e_c2e = p_mass_flx_e[c2e] + geofac_div = np.expand_dims(geofac_div, axis=-1) + z_tracer_mflx_c2e = z_tracer_mflx[c2e] + z_rhofluxdiv_c_out = ( + np.sum(p_mass_flx_e_c2e * geofac_div, axis=1) if nsub == int32(1) else z_rhofluxdiv_c + ) + z_fluxdiv_c_dsl = np.sum(z_tracer_mflx_c2e * geofac_div, axis=1) + z_rho_new_dsl = z_rho_now - z_dtsub * z_rhofluxdiv_c_out + z_tracer_new_dsl = (z_tracer_now * z_rho_now - z_dtsub * z_fluxdiv_c_dsl) / z_rho_new_dsl -def test_upwind_hflux_miura_cycl_stencil_02(backend): - grid = SimpleGrid() - nsub = int32(1) - p_mass_flx_e = random_field(grid, EdgeDim, KDim) - geofac_div = random_field(grid, CellDim, C2EDim) - z_rhofluxdiv_c = random_field(grid, CellDim, KDim) - z_tracer_mflx = random_field(grid, EdgeDim, KDim) - z_rho_now = random_field(grid, CellDim, KDim) - z_tracer_now = random_field(grid, CellDim, KDim) - z_dtsub = 0.5 - z_rhofluxdiv_c_out = random_field(grid, CellDim, KDim) - z_fluxdiv_c_dsl = random_field(grid, CellDim, KDim) - z_rho_new_dsl = random_field(grid, CellDim, KDim) - z_tracer_new_dsl = random_field(grid, CellDim, KDim) + return dict( + z_rhofluxdiv_c_out=z_rhofluxdiv_c_out, + z_fluxdiv_c_dsl=z_fluxdiv_c_dsl, + z_rho_new_dsl=z_rho_new_dsl, + z_tracer_new_dsl=z_tracer_new_dsl, + ) - ref_1, ref_2, ref_3, ref_4 = upwind_hflux_miura_cycl_stencil_02_numpy( - grid.connectivities[C2EDim], - nsub, - p_mass_flx_e.asnumpy(), - geofac_div.asnumpy(), - z_rhofluxdiv_c.asnumpy(), - z_tracer_mflx.asnumpy(), - z_rho_now.asnumpy(), - z_tracer_now.asnumpy(), - z_dtsub, - ) - - upwind_hflux_miura_cycl_stencil_02.with_backend(backend)( - nsub, - p_mass_flx_e, - geofac_div, - z_rhofluxdiv_c, - z_tracer_mflx, - z_rho_now, - z_tracer_now, - z_dtsub, - z_rhofluxdiv_c_out, - z_fluxdiv_c_dsl, - z_rho_new_dsl, - z_tracer_new_dsl, - offset_provider={ - "C2CE": grid.get_offset_provider("C2CE"), - "C2E": grid.get_offset_provider("C2E"), - }, - ) - assert np.allclose(ref_1, z_rhofluxdiv_c_out.asnumpy()) - assert np.allclose(ref_2, z_fluxdiv_c_dsl.asnumpy()) - assert np.allclose(ref_3, z_rho_new_dsl.asnumpy()) - assert np.allclose(ref_4, z_tracer_new_dsl.asnumpy()) + @pytest.fixture + def input_data(self, grid): + nsub = int32(1) + p_mass_flx_e = random_field(grid, EdgeDim, KDim) + geofac_div = random_field(grid, CellDim, C2EDim) + z_rhofluxdiv_c = random_field(grid, CellDim, KDim) + z_tracer_mflx = random_field(grid, EdgeDim, KDim) + z_rho_now = random_field(grid, CellDim, KDim) + z_tracer_now = random_field(grid, CellDim, KDim) + z_dtsub = 0.5 + z_rhofluxdiv_c_out = zero_field(grid, CellDim, KDim) + z_fluxdiv_c_dsl = zero_field(grid, CellDim, KDim) + z_rho_new_dsl = zero_field(grid, CellDim, KDim) + z_tracer_new_dsl = zero_field(grid, CellDim, KDim) + return dict( + nsub=nsub, + p_mass_flx_e=p_mass_flx_e, + geofac_div=geofac_div, + z_rhofluxdiv_c=z_rhofluxdiv_c, + z_tracer_mflx=z_tracer_mflx, + z_rho_now=z_rho_now, + z_tracer_now=z_tracer_now, + z_dtsub=z_dtsub, + z_rhofluxdiv_c_out=z_rhofluxdiv_c_out, + z_fluxdiv_c_dsl=z_fluxdiv_c_dsl, + z_rho_new_dsl=z_rho_new_dsl, + z_tracer_new_dsl=z_tracer_new_dsl, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py index 0e19320edd..8c1a99e509 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03a.py @@ -12,38 +12,36 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_03a import ( upwind_hflux_miura_cycl_stencil_03a, ) from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, zero_field - - -def upwind_hflux_miura_cycl_stencil_03a_numpy( - z_tracer_mflx_1_dsl: np.array, - z_tracer_mflx_2_dsl: np.array, -): - p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl) / float(2) - return p_out_e - - -def test_upwind_hflux_miura_cycl_stencil_03a(backend): - grid = SimpleGrid() - z_tracer_mflx_1_dsl = random_field(grid, EdgeDim, KDim) - z_tracer_mflx_2_dsl = random_field(grid, EdgeDim, KDim) - p_out_e = zero_field(grid, EdgeDim, KDim) - - ref = upwind_hflux_miura_cycl_stencil_03a_numpy( - z_tracer_mflx_1_dsl.asnumpy(), - z_tracer_mflx_2_dsl.asnumpy(), - ) - - upwind_hflux_miura_cycl_stencil_03a.with_backend(backend)( - z_tracer_mflx_1_dsl, - z_tracer_mflx_2_dsl, - p_out_e, - offset_provider={}, - ) - assert np.allclose(ref, p_out_e.asnumpy()) +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestUpwindHfluxMiuraCyclStencil03a(StencilTest): + PROGRAM = upwind_hflux_miura_cycl_stencil_03a + OUTPUTS = ("p_out_e",) + + @staticmethod + def reference( + grid, + z_tracer_mflx_1_dsl: np.array, + z_tracer_mflx_2_dsl: np.array, + **kwargs, + ): + p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl) / float(2) + return dict(p_out_e=p_out_e) + + @pytest.fixture + def input_data(self, grid): + z_tracer_mflx_1_dsl = random_field(grid, EdgeDim, KDim) + z_tracer_mflx_2_dsl = random_field(grid, EdgeDim, KDim) + p_out_e = zero_field(grid, EdgeDim, KDim) + return dict( + z_tracer_mflx_1_dsl=z_tracer_mflx_1_dsl, + z_tracer_mflx_2_dsl=z_tracer_mflx_2_dsl, + p_out_e=p_out_e, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py index 9ac1ee6017..481fd93c84 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_cycl_stencil_03b.py @@ -12,42 +12,39 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_03b import ( upwind_hflux_miura_cycl_stencil_03b, ) from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, zero_field - - -def upwind_hflux_miura_cycl_stencil_03b_numpy( - z_tracer_mflx_1_dsl: np.array, - z_tracer_mflx_2_dsl: np.array, - z_tracer_mflx_3_dsl: np.array, -): - p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl + z_tracer_mflx_3_dsl) / float(3) - return p_out_e - - -def test_upwind_hflux_miura_cycl_stencil_03b(backend): - grid = SimpleGrid() - z_tracer_mflx_1_dsl = random_field(grid, EdgeDim, KDim) - z_tracer_mflx_2_dsl = random_field(grid, EdgeDim, KDim) - z_tracer_mflx_3_dsl = random_field(grid, EdgeDim, KDim) - p_out_e = zero_field(grid, EdgeDim, KDim) - - ref = upwind_hflux_miura_cycl_stencil_03b_numpy( - z_tracer_mflx_1_dsl.asnumpy(), - z_tracer_mflx_2_dsl.asnumpy(), - z_tracer_mflx_3_dsl.asnumpy(), - ) - - upwind_hflux_miura_cycl_stencil_03b.with_backend(backend)( - z_tracer_mflx_1_dsl, - z_tracer_mflx_2_dsl, - z_tracer_mflx_3_dsl, - p_out_e, - offset_provider={}, - ) - assert np.allclose(ref, p_out_e.asnumpy()) +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestUpwindHfluxMiuraCyclStencil03b(StencilTest): + PROGRAM = upwind_hflux_miura_cycl_stencil_03b + OUTPUTS = ("p_out_e",) + + @staticmethod + def reference( + grid, + z_tracer_mflx_1_dsl: np.array, + z_tracer_mflx_2_dsl: np.array, + z_tracer_mflx_3_dsl: np.array, + **kwargs, + ): + p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl + z_tracer_mflx_3_dsl) / float(3) + return dict(p_out_e=p_out_e) + + @pytest.fixture + def input_data(self, grid): + z_tracer_mflx_1_dsl = random_field(grid, EdgeDim, KDim) + z_tracer_mflx_2_dsl = random_field(grid, EdgeDim, KDim) + z_tracer_mflx_3_dsl = random_field(grid, EdgeDim, KDim) + p_out_e = zero_field(grid, EdgeDim, KDim) + return dict( + z_tracer_mflx_1_dsl=z_tracer_mflx_1_dsl, + z_tracer_mflx_2_dsl=z_tracer_mflx_2_dsl, + z_tracer_mflx_3_dsl=z_tracer_mflx_3_dsl, + p_out_e=p_out_e, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_stencil_01.py index 0b9d47d818..59e3c026fe 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura_stencil_01.py @@ -12,87 +12,81 @@ # 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.advection.upwind_hflux_miura_stencil_01 import ( upwind_hflux_miura_stencil_01, ) from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import constant_field, random_field, zero_field - - -def upwind_hflux_miura_stencil_01_numpy( - e2c: np.array, - z_lsq_coeff_1: np.array, - z_lsq_coeff_2: np.array, - z_lsq_coeff_3: np.array, - distv_bary_1: np.array, - distv_bary_2: np.array, - p_mass_flx_e: np.array, - cell_rel_idx_dsl: np.array, -) -> np.array: - z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] - z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] - z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] - - p_out_e = ( - np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_1_e2c[:, 1], - z_lsq_coeff_1_e2c[:, 0], - ) - + distv_bary_1 - * np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_2_e2c[:, 1], - z_lsq_coeff_2_e2c[:, 0], - ) - + distv_bary_2 - * np.where( - cell_rel_idx_dsl == int32(1), - z_lsq_coeff_3_e2c[:, 1], - z_lsq_coeff_3_e2c[:, 0], - ) - ) * p_mass_flx_e +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + constant_field, + random_field, + zero_field, +) - return p_out_e +class TestUpwindHfluxMiuraStencil01(StencilTest): + PROGRAM = upwind_hflux_miura_stencil_01 + OUTPUTS = ("p_out_e",) -def test_upwind_hflux_miura_stencil_01(backend): - grid = SimpleGrid() + @staticmethod + def reference( + grid, + z_lsq_coeff_1: np.array, + z_lsq_coeff_2: np.array, + z_lsq_coeff_3: np.array, + distv_bary_1: np.array, + distv_bary_2: np.array, + p_mass_flx_e: np.array, + cell_rel_idx_dsl: np.array, + **kwargs, + ): + e2c = grid.connectivities[E2CDim] + z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] + z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] + z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] - z_lsq_coeff_1 = random_field(grid, CellDim, KDim) - z_lsq_coeff_2 = random_field(grid, CellDim, KDim) - z_lsq_coeff_3 = random_field(grid, CellDim, KDim) - distv_bary_1 = random_field(grid, EdgeDim, KDim) - distv_bary_2 = random_field(grid, EdgeDim, KDim) - p_mass_flx_e = random_field(grid, EdgeDim, KDim) - cell_rel_idx_dsl = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) - p_out_e = zero_field(grid, EdgeDim, KDim) + p_out_e = ( + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_e2c[:, 1], + z_lsq_coeff_1_e2c[:, 0], + ) + + distv_bary_1 + * np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_e2c[:, 1], + z_lsq_coeff_2_e2c[:, 0], + ) + + distv_bary_2 + * np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_e2c[:, 1], + z_lsq_coeff_3_e2c[:, 0], + ) + ) * p_mass_flx_e - ref = upwind_hflux_miura_stencil_01_numpy( - grid.connectivities[E2CDim], - z_lsq_coeff_1.asnumpy(), - z_lsq_coeff_2.asnumpy(), - z_lsq_coeff_3.asnumpy(), - distv_bary_1.asnumpy(), - distv_bary_2.asnumpy(), - p_mass_flx_e.asnumpy(), - cell_rel_idx_dsl.asnumpy(), - ) + return dict(p_out_e=p_out_e) - upwind_hflux_miura_stencil_01.with_backend(backend)( - z_lsq_coeff_1, - z_lsq_coeff_2, - z_lsq_coeff_3, - distv_bary_1, - distv_bary_2, - p_mass_flx_e, - cell_rel_idx_dsl, - p_out_e, - offset_provider={ - "E2C": grid.get_offset_provider("E2C"), - }, - ) - assert np.allclose(p_out_e.asnumpy(), ref) + @pytest.fixture + def input_data(self, grid): + z_lsq_coeff_1 = random_field(grid, CellDim, KDim) + z_lsq_coeff_2 = random_field(grid, CellDim, KDim) + z_lsq_coeff_3 = random_field(grid, CellDim, KDim) + distv_bary_1 = random_field(grid, EdgeDim, KDim) + distv_bary_2 = random_field(grid, EdgeDim, KDim) + p_mass_flx_e = random_field(grid, EdgeDim, KDim) + cell_rel_idx_dsl = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) + p_out_e = zero_field(grid, EdgeDim, KDim) + return dict( + z_lsq_coeff_1=z_lsq_coeff_1, + z_lsq_coeff_2=z_lsq_coeff_2, + z_lsq_coeff_3=z_lsq_coeff_3, + distv_bary_1=distv_bary_1, + distv_bary_2=distv_bary_2, + p_mass_flx_e=p_mass_flx_e, + cell_rel_idx_dsl=cell_rel_idx_dsl, + p_out_e=p_out_e, + ) diff --git a/model/common/src/icon4py/model/common/grid/icon.py b/model/common/src/icon4py/model/common/grid/icon.py index c99ad6c93a..09f564fa6b 100644 --- a/model/common/src/icon4py/model/common/grid/icon.py +++ b/model/common/src/icon4py/model/common/grid/icon.py @@ -16,11 +16,13 @@ from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.common.dimension import ( + C2E2C2E2CDim, C2E2CDim, C2E2CODim, C2EDim, C2VDim, CECDim, + CECECDim, CEDim, CellDim, E2C2EDim, @@ -63,6 +65,12 @@ def __init__(self): "E2C2E": (self._get_offset_provider, E2C2EDim, EdgeDim, EdgeDim), "E2C2EO": (self._get_offset_provider, E2C2EODim, EdgeDim, EdgeDim), "Koff": (lambda: KDim,), # Koff is a special case + "C2CECEC ": ( + self._get_offset_provider_for_sparse_fields, + C2E2C2E2CDim, + CellDim, + CECECDim, + ), } @builder diff --git a/model/common/src/icon4py/model/common/grid/simple.py b/model/common/src/icon4py/model/common/grid/simple.py index 5e5736cc73..ab611d67c7 100644 --- a/model/common/src/icon4py/model/common/grid/simple.py +++ b/model/common/src/icon4py/model/common/grid/simple.py @@ -22,6 +22,7 @@ C2EDim, C2VDim, CECDim, + CECECDim, CEDim, CellDim, E2C2EDim, @@ -425,6 +426,12 @@ def __init__(self): "E2ECV": (self._get_offset_provider_for_sparse_fields, E2C2VDim, EdgeDim, ECVDim), "E2EC": (self._get_offset_provider_for_sparse_fields, E2CDim, EdgeDim, ECDim), "C2CEC": (self._get_offset_provider_for_sparse_fields, C2E2CDim, CellDim, CECDim), + "C2CECEC": ( + self._get_offset_provider_for_sparse_fields, + C2E2C2E2CDim, + CellDim, + CECECDim, + ), } @property