From 8c6e2178dadc3182be84d7e8a50c03e829de0ff3 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Mon, 15 Jan 2024 15:27:16 +0100 Subject: [PATCH 01/14] Refactor test_hflux_ffsl_hybrid_stencil_02.py --- .../advection/hflux_ffsl_hybrid_stencil_02.py | 3 +- .../test_hflux_ffsl_hybrid_stencil_02.py | 63 +++++++++---------- 2 files changed, 31 insertions(+), 35 deletions(-) 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/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..4ecc80a792 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,37 @@ # 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, + ) From d3df7c10763c50d8d5ecc8f1bfad18452c53a4dc Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Tue, 16 Jan 2024 17:16:26 +0100 Subject: [PATCH 02/14] Refactor hflux_ffsl_hybrid_stencil_01a.py --- .../hflux_ffsl_hybrid_stencil_01a.py | 6 +- .../test_hflux_ffsl_hybrid_stencil_01a.py | 331 ++++++++---------- 2 files changed, 156 insertions(+), 181 deletions(-) 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/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..53bee2da7d 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,168 @@ # 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 +from icon4py.model.common.test_utils.helpers import StencilTest, 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] +class TestHfluxFfslHybridStencil01a(StencilTest): + PROGRAM = hflux_ffsl_hybrid_stencil_01a + OUTPUTS = ("p_out_e_hybrid_1a",) - 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 - ) - - return p_out_e_hybrid_1a + @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] + 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 + ) -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) - - 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(), - ) - - 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, + ) \ No newline at end of file From 3f6bc86924192c723cf3eef51bf21108d045af10 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Tue, 16 Jan 2024 19:43:52 +0100 Subject: [PATCH 03/14] Refactor six advection tests --- ...st_prep_gauss_quadrature_c_list_stencil.py | 875 ++++++++---------- .../test_prep_gauss_quadrature_c_stencil.py | 786 ++++++++-------- .../test_step_advection_stencil_01.py | 80 +- .../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 | 350 ++++--- ...test_upwind_hflux_miura_cycl_stencil_01.py | 134 ++- ...test_upwind_hflux_miura_cycl_stencil_02.py | 139 ++- ...est_upwind_hflux_miura_cycl_stencil_03a.py | 57 +- ...est_upwind_hflux_miura_cycl_stencil_03b.py | 64 +- .../test_upwind_hflux_miura_stencil_01.py | 134 ++- 12 files changed, 1309 insertions(+), 1509 deletions(-) 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..601537e91a 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,486 +19,421 @@ 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 +from icon4py.model.common.test_utils.helpers import StencilTest, constant_field, random_field, zero_field +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",) -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 + @staticmethod + def reference( + 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, + ): + 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 + 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 - 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) + 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) - ) + wgt_t_detjac_1 = np.where( + famask_bool, + dbl_eps + + z_wgt_1 * ( - 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_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_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_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_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_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_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, - ) - - 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 - ) - 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 - ) - - 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, - ) - + ( + 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, + ) -@pytest.mark.slow_tests -def test_prep_gauss_quadrature_c_list_stencil(backend): - grid = SimpleGrid() + 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 + ) - 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) + 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 + ) - ( - 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), - ) + 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, + ) - prep_gauss_quadrature_c_list_stencil.with_backend(backend)( - 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, - 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, - 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) + @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_6, + p_quad_vector_sum_8=p_quad_vector_sum_7, + p_quad_vector_sum_9=p_quad_vector_sum_8, + p_quad_vector_sum_10=p_quad_vector_sum_10, + p_dreg_area=p_dreg_area, + ) \ No newline at end of file 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..739aec885e 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,449 +18,387 @@ 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 +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", + ) -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 + @staticmethod + def reference( + 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, + ): + 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 + 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) + 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) + ) ) - * ( - 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) + 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) + ) ) - - ( - 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) + 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) + ) ) - * ( - 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_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) + ) ) - ) - 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_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_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_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_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_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_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) + 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 ) - ) - 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_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_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_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_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_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_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) + 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 ) - ) - 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) + + 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 ) - * ( - 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) + 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 ) - - ( - 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) + 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 ) - * ( - 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) + 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_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 - ) - 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, - 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, - ) + 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, + ) - prep_gauss_quadrature_c_stencil.with_backend(backend)( - 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, - 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) + @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, + ) \ No newline at end of file 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..0dda6ce204 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,51 @@ # 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, - ) - - step_advection_stencil_01.with_backend(backend)( - rhodz_ast, - p_mflx_contra, - deepatmo_divzl, - deepatmo_divzu, - p_dtime, - result, - offset_provider={"Koff": KDim}, - ) - - assert np.allclose(ref[:, :-1], result.asnumpy()[:, :-1]) + **kwargs, + ): + tmp = p_dtime * ( + p_mflx_contra_v[:, 1:] * deepatmo_divzl - p_mflx_contra_v[:, :-1] * deepatmo_divzu + ) + rhodz_ast2 = rhodz_ast + tmp + + return dict(rhodz_ast2=rhodz_ast2) + + @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, + ) \ No newline at end of file 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..356aa84c9d 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,48 @@ # 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 StencilTest, Output, 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, +class TestStepAdvectionStencil02(StencilTest): + PROGRAM = step_advection_stencil_02 + OUTPUTS = ( + Output( + "rhodz_ast2", refslice=(slice(None), slice(None, -1)), gtslice=(slice(None), slice(None, -1)) + ), ) - - step_advection_stencil_02.with_backend(backend)( - rhodz_ast, - p_mflx_contra, - deepatmo_divzl, - deepatmo_divzu, - p_dtime, - result, - offset_provider={"Koff": KDim}, - ) - - 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, + ) \ No newline at end of file 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..96c3d56d5a 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, + ) \ No newline at end of file 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..ffa389a5b8 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..727208e5e9 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,177 @@ # 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 +class TestUpwindHfluxMiura3Stencil01(StencilTest): + PROGRAM = upwind_hflux_miura3_stencil_01 + OUTPUTS = ("p_out_e_miura3",) -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] + @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], + 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_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_dreg_area + * p_mass_flx_e ) - / z_dreg_area - * p_mass_flx_e - ) - - return 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(), - ) + return dict(p_out_e_miura3=p_out_e_miura3) - 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_1, + z_quad_vector_sum_3=z_quad_vector_sum_1, + z_quad_vector_sum_4=z_quad_vector_sum_1, + z_quad_vector_sum_5=z_quad_vector_sum_1, + z_quad_vector_sum_6=z_quad_vector_sum_1, + z_quad_vector_sum_7=z_quad_vector_sum_1, + z_quad_vector_sum_8=z_quad_vector_sum_1, + z_quad_vector_sum_9=z_quad_vector_sum_1, + z_quad_vector_sum_10=z_quad_vector_sum_1, + 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, + ) \ No newline at end of file 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..94d304cc26 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,75 @@ # 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 +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, random_mask, zero_field +class TestUpwindHfluxMiuraCyclStencil01(StencilTest): + PROGRAM = upwind_hflux_miura_cycl_stencil_01 + OUTPUTS = ("z_tracer_mflx_dsl",) -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] + @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_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 + 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 - return z_tracer_mflx_dsl + return dict(z_tracer_mflx_dsl=z_tracer_mflx_dsl) - -def test_upwind_hflux_miura_cycl_stencil_01(backend): - grid = SimpleGrid() - - 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) - - 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(), - ) - - 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, + ) \ No newline at end of file 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..5f348452cd 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, zero_field, random_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, + ) \ No newline at end of file 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..4822a031e7 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,35 @@ # 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, + ) \ No newline at end of file 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..cd290693a5 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,38 @@ # 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..7eddfc68b6 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,75 @@ # 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 +from icon4py.model.common.test_utils.helpers import StencilTest, constant_field, random_field, zero_field +class TestUpwindHfluxMiuraStencil01(StencilTest): + PROGRAM = upwind_hflux_miura_stencil_01 + OUTPUTS = ("p_out_e",) -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] + @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] - 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 - - return p_out_e + 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 + return dict(p_out_e=p_out_e) -def test_upwind_hflux_miura_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) - 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) - - 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(), - ) - - 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_1, + z_lsq_coeff_3=z_lsq_coeff_1, + distv_bary_1=distv_bary_1, + distv_bary_2=distv_bary_1, + p_mass_flx_e=p_mass_flx_e, + cell_rel_idx_dsl=cell_rel_idx_dsl, + p_out_e=p_out_e, + ) From 77f29455da157dca6c49c3ac0656b5d1a835de2f Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Tue, 16 Jan 2024 19:45:51 +0100 Subject: [PATCH 04/14] Add grid_type=GridType.UNSTRUCTURED in advection stencils --- .../atmosphere/advection/divide_flux_area_list_stencil_01.py | 3 ++- .../atmosphere/advection/divide_flux_area_list_stencil_02.py | 3 ++- .../model/atmosphere/advection/hflx_limiter_pd_stencil_02.py | 5 +++-- .../advection/prep_gauss_quadrature_c_list_stencil.py | 5 +++-- .../atmosphere/advection/prep_gauss_quadrature_c_stencil.py | 3 ++- .../atmosphere/advection/recon_lsq_cell_c_svd_stencil.py | 3 ++- .../atmosphere/advection/recon_lsq_cell_l_svd_stencil.py | 3 ++- .../model/atmosphere/advection/step_advection_stencil_01.py | 5 +++-- .../model/atmosphere/advection/step_advection_stencil_02.py | 5 +++-- .../model/atmosphere/advection/step_advection_stencil_03.py | 3 ++- .../model/atmosphere/advection/step_advection_stencil_04.py | 3 ++- .../atmosphere/advection/upwind_hflux_miura3_stencil_01.py | 3 ++- .../advection/upwind_hflux_miura_cycl_stencil_01.py | 3 ++- .../advection/upwind_hflux_miura_cycl_stencil_02.py | 3 ++- .../advection/upwind_hflux_miura_cycl_stencil_03a.py | 3 ++- .../advection/upwind_hflux_miura_cycl_stencil_03b.py | 3 ++- .../atmosphere/advection/upwind_hflux_miura_stencil_01.py | 3 ++- 17 files changed, 38 insertions(+), 21 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/hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py index 9d1be46771..dda82c1d9e 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py @@ -36,17 +36,18 @@ def _hflx_limiter_pd_stencil_02( return p_mflx_tracer_h_out -@program +@program(grid_type="unstructured") def hflx_limiter_pd_stencil_02( refin_ctrl: Field[[EdgeDim], int32], r_m: Field[[CellDim, KDim], float], + p_mflx_tracer_h_in: Field[[EdgeDim, KDim], float], p_mflx_tracer_h: Field[[EdgeDim, KDim], float], bound: int32, ): _hflx_limiter_pd_stencil_02( refin_ctrl, r_m, - p_mflx_tracer_h, + p_mflx_tracer_h_in, bound, out=p_mflx_tracer_h, ) 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..5cc21f6cf4 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], @@ -431,4 +432,4 @@ def prep_gauss_quadrature_c_list_stencil( p_quad_vector_sum_10, p_dreg_area, ), - ) + ) \ No newline at end of file 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..c2c54b7069 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,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program @@ -176,7 +177,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..af82ad223e 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,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program @@ -41,7 +42,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..45cec475e7 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,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program @@ -31,7 +32,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], @@ -47,4 +48,4 @@ def step_advection_stencil_01( deepatmo_divzu, p_dtime, out=rhodz_ast2, - ) + ) \ No newline at end of file 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..4423dc4192 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,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import maximum @@ -31,7 +32,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], @@ -47,4 +48,4 @@ def step_advection_stencil_02( deepatmo_divzu, p_dtime, out=rhodz_ast2, - ) + ) \ No newline at end of file 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..afcd59126d 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,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, where @@ -114,7 +115,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..592d0955ad 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,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, where @@ -51,7 +52,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..c29b987e08 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,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import float64 @@ -27,7 +28,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..72c07a1427 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,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program @@ -27,7 +28,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..b8ce740080 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,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import GridType from gt4py.next.common import Field from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, where @@ -39,7 +40,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], From d6dd89a6ece012560ddf81e7ae44945bedfed1be Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Tue, 16 Jan 2024 19:57:46 +0100 Subject: [PATCH 05/14] Run pre-commit --- .../prep_gauss_quadrature_c_list_stencil.py | 2 +- .../advection/recon_lsq_cell_c_svd_stencil.py | 3 +-- .../advection/recon_lsq_cell_l_svd_stencil.py | 3 +-- .../advection/step_advection_stencil_01.py | 5 ++--- .../advection/step_advection_stencil_02.py | 5 ++--- .../advection/upwind_hflux_miura3_stencil_01.py | 3 +-- .../advection/upwind_hflux_miura_cycl_stencil_01.py | 3 +-- .../upwind_hflux_miura_cycl_stencil_03a.py | 3 +-- .../upwind_hflux_miura_cycl_stencil_03b.py | 3 +-- .../advection/upwind_hflux_miura_stencil_01.py | 3 +-- .../test_hflux_ffsl_hybrid_stencil_01a.py | 9 +++++++-- .../test_hflux_ffsl_hybrid_stencil_02.py | 3 ++- .../test_prep_gauss_quadrature_c_list_stencil.py | 13 ++++++++++--- .../test_prep_gauss_quadrature_c_stencil.py | 3 ++- .../stencil_tests/test_step_advection_stencil_01.py | 8 +++++--- .../stencil_tests/test_step_advection_stencil_02.py | 10 ++++++---- .../stencil_tests/test_step_advection_stencil_03.py | 4 ++-- .../stencil_tests/test_step_advection_stencil_04.py | 6 +++--- .../test_upwind_hflux_miura3_stencil_01.py | 10 ++++++++-- .../test_upwind_hflux_miura_cycl_stencil_01.py | 12 +++++++++--- .../test_upwind_hflux_miura_cycl_stencil_02.py | 10 +++++----- .../test_upwind_hflux_miura_cycl_stencil_03a.py | 3 ++- .../test_upwind_hflux_miura_cycl_stencil_03b.py | 3 ++- .../test_upwind_hflux_miura_stencil_01.py | 8 +++++++- 24 files changed, 82 insertions(+), 53 deletions(-) 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 5cc21f6cf4..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 @@ -432,4 +432,4 @@ def prep_gauss_quadrature_c_list_stencil( p_quad_vector_sum_10, p_dreg_area, ), - ) \ No newline at end of file + ) 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 c2c54b7069..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,8 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import GridType -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 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 af82ad223e..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,8 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import GridType -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 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 45cec475e7..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,8 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import GridType -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 @@ -48,4 +47,4 @@ def step_advection_stencil_01( deepatmo_divzu, p_dtime, out=rhodz_ast2, - ) \ No newline at end of file + ) 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 4423dc4192..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,8 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import GridType -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 @@ -48,4 +47,4 @@ def step_advection_stencil_02( deepatmo_divzu, p_dtime, out=rhodz_ast2, - ) \ No newline at end of file + ) 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 afcd59126d..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,8 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import GridType -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 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 592d0955ad..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,8 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import GridType -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 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 c29b987e08..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,8 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import GridType -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 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 72c07a1427..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,8 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import GridType -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 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 b8ce740080..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,8 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import GridType -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 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 53bee2da7d..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 @@ -19,7 +19,12 @@ hflux_ffsl_hybrid_stencil_01a, ) from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, constant_field, random_field, zero_field +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + constant_field, + random_field, + zero_field, +) class TestHfluxFfslHybridStencil01a(StencilTest): @@ -176,4 +181,4 @@ def input_data(self, grid): 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, - ) \ No newline at end of file + ) 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 4ecc80a792..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 @@ -20,6 +20,7 @@ from icon4py.model.common.dimension import EdgeDim, KDim 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",) @@ -45,4 +46,4 @@ def input_data(self, grid): 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 601537e91a..4c8024e210 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,7 +19,13 @@ prep_gauss_quadrature_c_list_stencil, ) from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, constant_field, random_field, zero_field +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + constant_field, + random_field, + zero_field, +) + class TestPrepGaussQuadratureCListStencil(StencilTest): PROGRAM = prep_gauss_quadrature_c_list_stencil @@ -34,7 +40,8 @@ class TestPrepGaussQuadratureCListStencil(StencilTest): "p_quad_vector_sum_8", "p_quad_vector_sum_9", "p_quad_vector_sum_10", - "p_dreg_area",) + "p_dreg_area", + ) @staticmethod def reference( @@ -436,4 +443,4 @@ def input_data(self, grid): p_quad_vector_sum_9=p_quad_vector_sum_8, p_quad_vector_sum_10=p_quad_vector_sum_10, p_dreg_area=p_dreg_area, - ) \ No newline at end of file + ) 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 739aec885e..f98bd8560f 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 @@ -20,6 +20,7 @@ from icon4py.model.common.dimension import EdgeDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + class TestPrepGaussQuadratureCStencil(StencilTest): PROGRAM = prep_gauss_quadrature_c_stencil OUTPUTS = ( @@ -401,4 +402,4 @@ def input_data(self, grid): 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, - ) \ No newline at end of file + ) 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 0dda6ce204..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 @@ -23,7 +23,9 @@ class TestStepAdvectionStencil01(StencilTest): PROGRAM = step_advection_stencil_01 OUTPUTS = ( Output( - "rhodz_ast2", refslice=(slice(None), slice(None, -1)), gtslice=(slice(None), slice(None, -1)) + "rhodz_ast2", + refslice=(slice(None), slice(None, -1)), + gtslice=(slice(None), slice(None, -1)), ), ) @@ -38,7 +40,7 @@ def reference( **kwargs, ): tmp = p_dtime * ( - p_mflx_contra_v[:, 1:] * deepatmo_divzl - p_mflx_contra_v[:, :-1] * deepatmo_divzu + p_mflx_contra_v[:, 1:] * deepatmo_divzl - p_mflx_contra_v[:, :-1] * deepatmo_divzu ) rhodz_ast2 = rhodz_ast + tmp @@ -59,4 +61,4 @@ def input_data(self, grid): deepatmo_divzu=deepatmo_divzu, p_dtime=p_dtime, rhodz_ast2=rhodz_ast2, - ) \ No newline at end of file + ) 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 356aa84c9d..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 @@ -16,17 +16,19 @@ 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.test_utils.helpers import StencilTest, Output, random_field, zero_field +from icon4py.model.common.test_utils.helpers import Output, StencilTest, random_field, zero_field class TestStepAdvectionStencil02(StencilTest): PROGRAM = step_advection_stencil_02 OUTPUTS = ( Output( - "rhodz_ast2", refslice=(slice(None), slice(None, -1)), gtslice=(slice(None), slice(None, -1)) + "rhodz_ast2", + refslice=(slice(None), slice(None, -1)), + gtslice=(slice(None), slice(None, -1)), ), ) - + @staticmethod def reference( grid, @@ -56,4 +58,4 @@ def input_data(self, grid): deepatmo_divzu=deepatmo_divzu, p_dtime=p_dtime, rhodz_ast2=rhodz_ast2, - ) \ No newline at end of file + ) 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 96c3d56d5a..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 @@ -21,7 +21,7 @@ class TestStepAdvectionStencil03(StencilTest): PROGRAM = step_advection_stencil_03 - OUTPUTS = ("p_tracer_new", ) + OUTPUTS = ("p_tracer_new",) @staticmethod def reference( @@ -47,4 +47,4 @@ def input_data(self, grid): p_grf_tend_tracer=p_grf_tend_tracer, p_dtime=p_dtime, p_tracer_new=p_tracer_new, - ) \ No newline at end of file + ) 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 ffa389a5b8..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 @@ -23,7 +23,7 @@ class TestStepAdvectionStencil04(StencilTest): PROGRAM = step_advection_stencil_04 OUTPUTS = ("opt_ddt_tracer_adv",) - @staticmethod + @staticmethod def reference( grid, p_tracer_now: np.array, @@ -34,7 +34,7 @@ def reference( 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) @@ -45,5 +45,5 @@ def input_data(self, grid): p_tracer_now=p_tracer_now, p_tracer_new=p_tracer_new, p_dtime=p_dtime, - opt_ddt_tracer_adv=opt_ddt_tracer_adv + 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 727208e5e9..be9f837410 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 @@ -19,7 +19,13 @@ upwind_hflux_miura3_stencil_01, ) from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, random_mask, zero_field +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + random_mask, + zero_field, +) + class TestUpwindHfluxMiura3Stencil01(StencilTest): PROGRAM = upwind_hflux_miura3_stencil_01 @@ -185,4 +191,4 @@ def input_data(self, grid): p_mass_flx_e=p_mass_flx_e, cell_rel_idx_dsl=cell_rel_idx_dsl, p_out_e_miura3=p_out_e_miura3, - ) \ No newline at end of file + ) 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 94d304cc26..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 @@ -19,13 +19,19 @@ upwind_hflux_miura_cycl_stencil_01, ) from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, random_mask, zero_field +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + random_mask, + zero_field, +) + class TestUpwindHfluxMiuraCyclStencil01(StencilTest): PROGRAM = upwind_hflux_miura_cycl_stencil_01 OUTPUTS = ("z_tracer_mflx_dsl",) - @staticmethod + @staticmethod def reference( grid, z_lsq_coeff_1_dsl: np.array, @@ -83,4 +89,4 @@ def input_data(self, grid): p_mass_flx_e=p_mass_flx_e, cell_rel_idx_dsl=cell_rel_idx_dsl, z_tracer_mflx_dsl=z_tracer_mflx_dsl, - ) \ No newline at end of file + ) 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 5f348452cd..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 @@ -19,15 +19,15 @@ upwind_hflux_miura_cycl_stencil_02, ) from icon4py.model.common.dimension import C2EDim, CellDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, zero_field, random_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field class TestUpwindHfluxMiuraCyclStencil02(StencilTest): PROGRAM = upwind_hflux_miura_cycl_stencil_02 OUTPUTS = ( "z_rhofluxdiv_c_out", - "z_fluxdiv_c_dsl", - "z_rho_new_dsl", + "z_fluxdiv_c_dsl", + "z_rho_new_dsl", "z_tracer_new_dsl", ) @@ -60,7 +60,7 @@ def reference( 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 + z_tracer_new_dsl=z_tracer_new_dsl, ) @pytest.fixture @@ -90,4 +90,4 @@ def input_data(self, grid): z_fluxdiv_c_dsl=z_fluxdiv_c_dsl, z_rho_new_dsl=z_rho_new_dsl, z_tracer_new_dsl=z_tracer_new_dsl, - ) \ No newline at end of file + ) 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 4822a031e7..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 @@ -20,6 +20,7 @@ from icon4py.model.common.dimension import EdgeDim, KDim 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",) @@ -43,4 +44,4 @@ def input_data(self, grid): 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, - ) \ No newline at end of file + ) 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 cd290693a5..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 @@ -20,6 +20,7 @@ from icon4py.model.common.dimension import EdgeDim, KDim 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",) @@ -45,5 +46,5 @@ def input_data(self, grid): 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 + 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 7eddfc68b6..21d3e229d6 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 @@ -19,7 +19,13 @@ upwind_hflux_miura_stencil_01, ) from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, constant_field, random_field, zero_field +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + constant_field, + random_field, + zero_field, +) + class TestUpwindHfluxMiuraStencil01(StencilTest): PROGRAM = upwind_hflux_miura_stencil_01 From 1062b5630b1ca0f8f2d47aca7c210191ef91923e Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Tue, 16 Jan 2024 20:14:38 +0100 Subject: [PATCH 06/14] Pre-commit fixes --- ...est_prep_gauss_quadrature_c_list_stencil.py | 6 +++--- .../test_upwind_hflux_miura3_stencil_01.py | 18 +++++++++--------- .../test_upwind_hflux_miura_stencil_01.py | 4 ++-- 3 files changed, 14 insertions(+), 14 deletions(-) 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 4c8024e210..60f4419dbb 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 @@ -438,9 +438,9 @@ def input_data(self, grid): 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_6, - p_quad_vector_sum_8=p_quad_vector_sum_7, - p_quad_vector_sum_9=p_quad_vector_sum_8, + 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_upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_upwind_hflux_miura3_stencil_01.py index be9f837410..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 @@ -178,15 +178,15 @@ def input_data(self, grid): 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_1, - z_quad_vector_sum_3=z_quad_vector_sum_1, - z_quad_vector_sum_4=z_quad_vector_sum_1, - z_quad_vector_sum_5=z_quad_vector_sum_1, - z_quad_vector_sum_6=z_quad_vector_sum_1, - z_quad_vector_sum_7=z_quad_vector_sum_1, - z_quad_vector_sum_8=z_quad_vector_sum_1, - z_quad_vector_sum_9=z_quad_vector_sum_1, - z_quad_vector_sum_10=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, 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 21d3e229d6..71ecfb861a 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 @@ -82,8 +82,8 @@ def input_data(self, grid): 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_1, - z_lsq_coeff_3=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_1, p_mass_flx_e=p_mass_flx_e, From 5f8f2ec58052b825626d6935e65fc5ec99d39111 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Tue, 16 Jan 2024 20:18:08 +0100 Subject: [PATCH 07/14] Pre-commit fixes - missing variable --- .../tests/stencil_tests/test_upwind_hflux_miura_stencil_01.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 71ecfb861a..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 @@ -85,7 +85,7 @@ def input_data(self, grid): 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_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, From 6e105e50a91f5882b481f24b956adb13dcc8a52a Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 18 Jan 2024 09:54:27 +0100 Subject: [PATCH 08/14] Add C2CECEC offset in Simple- and IconGrid --- model/common/src/icon4py/model/common/grid/icon.py | 3 +++ model/common/src/icon4py/model/common/grid/simple.py | 2 ++ 2 files changed, 5 insertions(+) diff --git a/model/common/src/icon4py/model/common/grid/icon.py b/model/common/src/icon4py/model/common/grid/icon.py index 5672ea9886..6044ec0680 100644 --- a/model/common/src/icon4py/model/common/grid/icon.py +++ b/model/common/src/icon4py/model/common/grid/icon.py @@ -16,10 +16,12 @@ from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.common.dimension import ( + C2E2C2E2CDim, C2E2CDim, C2E2CODim, C2EDim, CECDim, + CECECDim, CEDim, CellDim, E2C2EDim, @@ -61,6 +63,7 @@ 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..05981e6edd 100644 --- a/model/common/src/icon4py/model/common/grid/simple.py +++ b/model/common/src/icon4py/model/common/grid/simple.py @@ -21,6 +21,7 @@ C2E2CODim, C2EDim, C2VDim, + CECECDim, CECDim, CEDim, CellDim, @@ -425,6 +426,7 @@ 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 From 1bc4be07d6432ca1d87ae0549e62b42032c86e65 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 18 Jan 2024 10:08:27 +0100 Subject: [PATCH 09/14] Refactor slow_tests advection --- .../test_divide_flux_area_list_stencil_02.py | 563 ++++++++--------- .../test_recon_lsq_cell_c_svd_stencil.py | 585 +++++++++--------- .../test_recon_lsq_cell_l_svd_stencil.py | 111 ++-- 3 files changed, 602 insertions(+), 657 deletions(-) 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..d6827483dd 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,280 @@ # 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, + numpy_to_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_recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/tests/stencil_tests/test_recon_lsq_cell_c_svd_stencil.py index 26df554d28..1d69f9292a 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,302 @@ 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 ( + reshape, + StencilTest, + as_1D_sparse_field, + random_field, + 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_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_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_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_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_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_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_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_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_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_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_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_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_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_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_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) + ) -@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_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) + ) - ( - 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(), - ) + 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, + ) - 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..d0dbd26364 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,70 @@ # 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 ( + reshape, + StencilTest, + as_1D_sparse_field, + random_field, + 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, + ) From b1a91a312a735190a6ac7d21f77d3c82d4de02cf Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 18 Jan 2024 10:20:08 +0100 Subject: [PATCH 10/14] Run pre-commit --- .../stencil_tests/test_recon_lsq_cell_c_svd_stencil.py | 5 ++--- .../stencil_tests/test_recon_lsq_cell_l_svd_stencil.py | 10 +++------- model/common/src/icon4py/model/common/grid/icon.py | 7 ++++++- model/common/src/icon4py/model/common/grid/simple.py | 9 +++++++-- 4 files changed, 18 insertions(+), 13 deletions(-) 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 1d69f9292a..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 @@ -19,10 +19,10 @@ ) from icon4py.model.common.dimension import C2E2C2E2CDim, CECECDim, CellDim, KDim from icon4py.model.common.test_utils.helpers import ( - reshape, StencilTest, as_1D_sparse_field, random_field, + reshape, zero_field, ) @@ -103,7 +103,6 @@ def reference( lsq_pseudoinv_1 = reshape(lsq_pseudoinv_1, c2e2c2e2c.shape) 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) @@ -114,7 +113,7 @@ def reference( + 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_9_dsl = ( lsq_pseudoinv_8[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) 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 d0dbd26364..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 @@ -19,10 +19,10 @@ ) from icon4py.model.common.dimension import C2E2CDim, CECDim, CellDim, KDim from icon4py.model.common.test_utils.helpers import ( - reshape, StencilTest, as_1D_sparse_field, random_field, + reshape, zero_field, ) @@ -62,12 +62,8 @@ def reference( @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 - ) + 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) diff --git a/model/common/src/icon4py/model/common/grid/icon.py b/model/common/src/icon4py/model/common/grid/icon.py index 6044ec0680..3acf4dfc41 100644 --- a/model/common/src/icon4py/model/common/grid/icon.py +++ b/model/common/src/icon4py/model/common/grid/icon.py @@ -63,7 +63,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), + "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 05981e6edd..ab611d67c7 100644 --- a/model/common/src/icon4py/model/common/grid/simple.py +++ b/model/common/src/icon4py/model/common/grid/simple.py @@ -21,8 +21,8 @@ C2E2CODim, C2EDim, C2VDim, - CECECDim, CECDim, + CECECDim, CEDim, CellDim, E2C2EDim, @@ -426,7 +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), + "C2CECEC": ( + self._get_offset_provider_for_sparse_fields, + C2E2C2E2CDim, + CellDim, + CECECDim, + ), } @property From d771bf54e24f1830e75820200237865f9d369d0e Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 18 Jan 2024 13:47:48 +0100 Subject: [PATCH 11/14] Fix pre-commit - imported helpers not used --- .../tests/stencil_tests/test_divide_flux_area_list_stencil_02.py | 1 - 1 file changed, 1 deletion(-) 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 d6827483dd..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 @@ -22,7 +22,6 @@ from icon4py.model.common.test_utils.helpers import ( StencilTest, as_1D_sparse_field, - numpy_to_1D_sparse_field, random_field, random_mask, reshape, From 369abf554b8cc5954652a052b3fff7591ed7eb6a Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 18 Jan 2024 15:08:40 +0100 Subject: [PATCH 12/14] Add last slow test divide_flux_area_list_stencil_01 --- .../test_divide_flux_area_list_stencil_01.py | 1473 ++++++++--------- 1 file changed, 735 insertions(+), 738 deletions(-) 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..645f53deca 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,726 +19,144 @@ 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: - dx1 = p1_lon - p0_lon - dy1 = p1_lat - p0_lat - - dx2 = p2_lon - p0_lon - dy2 = p2_lat - p0_lat - - dx1dy2 = dx1 * dy2 - dy1dx2 = dy1 * dx2 - - lccw = np.where(dx1dy2 > dy1dx2, True, False) - ccw_out = np.where(lccw, int32(1), int32(-1)) # 1: clockwise, -1: counterclockwise - 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( +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", + ) + + # Check whether lines inters. + @staticmethod + 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 + + dx2 = p2_lon - p0_lon + dy2 = p2_lat - p0_lat + + dx1dy2 = dx1 * dy2 + dy1dx2 = dy1 * dx2 + + lccw = np.where(dx1dy2 > dy1dx2, True, False) + ccw_out = np.where(lccw, int32(1), int32(-1)) # 1: clockwise, -1: counterclockwise + return ccw_out + + # Check whether two lines intersect + @staticmethod + def _lintersect_numpy( line1_p1_lon, line1_p1_lat, line1_p2_lon, line1_p2_lat, line2_p1_lon, line2_p1_lat, - ) * ccw( - line1_p1_lon, - line1_p1_lat, - line1_p2_lon, - line1_p2_lat, - line2_p2_lon, - line2_p2_lat, - ) - intersect2 = ccw( - line2_p1_lon, - line2_p1_lat, line2_p2_lon, line2_p2_lat, + ): + intersect1 = TestDivideFluxAreaListStencil01._ccw_numpy( + line1_p1_lon, + line1_p1_lat, + line1_p2_lon, + line1_p2_lat, + line2_p1_lon, + line2_p1_lat, + ) * TestDivideFluxAreaListStencil01._ccw_numpy( + line1_p1_lon, + line1_p1_lat, + line1_p2_lon, + line1_p2_lat, + line2_p2_lon, + line2_p2_lat, + ) + intersect2 = TestDivideFluxAreaListStencil01._ccw_numpy( + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, + line1_p1_lon, + line1_p1_lat, + ) * TestDivideFluxAreaListStencil01._ccw_numpy( + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, + line1_p2_lon, + line1_p2_lat, + ) + lintersect_out = np.where((intersect1 + intersect2) == -2, True, False) + + return lintersect_out + + # Compute intersection point of two lines in 2D + @staticmethod + def _line_intersect_numpy( line1_p1_lon, line1_p1_lat, - ) * ccw( + line1_p2_lon, + line1_p2_lat, line2_p1_lon, line2_p1_lat, line2_p2_lon, line2_p2_lat, - line1_p2_lon, - line1_p2_lat, - ) - lintersect_out = np.where((intersect1 + intersect2) == -2, True, False) - - return lintersect_out - - -# 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]: - 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) - - intersect_1 = (line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon) / (m1 - m2) - intersect_2 = line1_p1_lat + m1 * (intersect_1 - line1_p1_lon) - - 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), - ) - - # 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), - ) - - # 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_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, - ) - - # 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( - 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( - 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( - 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 - ) - - # ------------------------------------------------- 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( - 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, - ) - # 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( - 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, - ) - # 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, - ) + ): + 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) - # 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 + intersect_1 = (line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon) / ( + m1 - m2 + ) + intersect_2 = line1_p1_lat + m1 * (intersect_1 - line1_p1_lon) - # ----------------------------------------------- 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) + return intersect_1, intersect_2 - # ------------------------------------------------ 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, - ) - - # --------------------------------------------- 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 ( - 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, - ) - - -@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), - ) - - divide_flux_area_list_stencil_01.with_backend(backend)( + @classmethod + def reference( + TestDivideFluxAreaListStencil01, + grid, famask_int, p_vn, - ptr_v3_lon_field, - ptr_v3_lat_field, + ptr_v3_lon, + ptr_v3_lat, tangent_orientation_dsl, dreg_patch0_1_lon_dsl, dreg_patch0_1_lat_dsl, @@ -748,48 +166,627 @@ 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_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, - 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) + **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) + + 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), + ) + + # 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), + ) + + # Create first mask does departure-line segment intersects with A1V3 + lintersect_line1 = TestDivideFluxAreaListStencil01._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 = TestDivideFluxAreaListStencil01._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) + # ------------------------------------------------- Case 1 + mask_case1 = np.logical_and.reduce([lintersect_line1, lintersect_line2, famask_bool]) + ps1_x, ps1_y = TestDivideFluxAreaListStencil01._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 = TestDivideFluxAreaListStencil01._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 - 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( + 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( + 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( + 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 + ) + + # ------------------------------------------------- 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( + 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, + ) + # 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( + 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, + ) + # 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, + ) + + # 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 = TestDivideFluxAreaListStencil01._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 = TestDivideFluxAreaListStencil01._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 + 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 = TestDivideFluxAreaListStencil01._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 = TestDivideFluxAreaListStencil01._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 = 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, + ) + + # --------------------------------------------- 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.fixture + @pytest.mark.slow_tests + 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, + ) From f96a3330784de5d44cff972d38175f024010ea71 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Thu, 18 Jan 2024 16:27:06 +0100 Subject: [PATCH 13/14] restore hflx_limiter_pd_stencil_02.py (was on Sam's list) --- .../model/atmosphere/advection/hflx_limiter_pd_stencil_02.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py index 5f38d0b648..2380296e9e 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py @@ -36,18 +36,17 @@ def _hflx_limiter_pd_stencil_02( return p_mflx_tracer_h_out -@program(grid_type="unstructured") +@program def hflx_limiter_pd_stencil_02( refin_ctrl: Field[[EdgeDim], int32], r_m: Field[[CellDim, KDim], float], - p_mflx_tracer_h_in: Field[[EdgeDim, KDim], float], p_mflx_tracer_h: Field[[EdgeDim, KDim], float], bound: int32, ): _hflx_limiter_pd_stencil_02( refin_ctrl, r_m, - p_mflx_tracer_h_in, + p_mflx_tracer_h, bound, out=p_mflx_tracer_h, ) From f07e537367dee882c60fe59a6976f2bb5bceb163 Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Wed, 24 Jan 2024 14:30:06 +0100 Subject: [PATCH 14/14] Refactor long advection tests with functions --- .../test_divide_flux_area_list_stencil_01.py | 1281 +++++++++++++---- ...st_prep_gauss_quadrature_c_list_stencil.py | 310 +++- .../test_prep_gauss_quadrature_c_stencil.py | 301 +++- 3 files changed, 1511 insertions(+), 381 deletions(-) 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 645f53deca..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 @@ -29,6 +29,95 @@ ) +# 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 + + dx2 = p2_lon - p0_lon + dy2 = p2_lat - p0_lat + + dx1dy2 = dx1 * dy2 + dy1dx2 = dy1 * dx2 + + lccw = np.where(dx1dy2 > dy1dx2, True, False) + ccw_out = np.where(lccw, int32(1), int32(-1)) # 1: clockwise, -1: counterclockwise + return ccw_out + + +# 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_numpy( + line1_p1_lon, + line1_p1_lat, + line1_p2_lon, + line1_p2_lat, + line2_p2_lon, + line2_p2_lat, + ) + intersect2 = _ccw_numpy( + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, + line1_p1_lon, + line1_p1_lat, + ) * _ccw_numpy( + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, + line1_p2_lon, + line1_p2_lat, + ) + lintersect_out = np.where((intersect1 + intersect2) == -2, True, False) + + return lintersect_out + + +# Compute intersection point of two lines in 2D +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) + + intersect_1 = (line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon) / (m1 - m2) + intersect_2 = line1_p1_lat + m1 * (intersect_1 - line1_p1_lon) + + return intersect_1, intersect_2 + + class TestDivideFluxAreaListStencil01(StencilTest): PROGRAM = divide_flux_area_list_stencil_01 OUTPUTS = ( @@ -58,106 +147,8 @@ class TestDivideFluxAreaListStencil01(StencilTest): "dreg_patch2_4_lat_vmask", ) - # Check whether lines inters. - @staticmethod - 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 - - dx2 = p2_lon - p0_lon - dy2 = p2_lat - p0_lat - - dx1dy2 = dx1 * dy2 - dy1dx2 = dy1 * dx2 - - lccw = np.where(dx1dy2 > dy1dx2, True, False) - ccw_out = np.where(lccw, int32(1), int32(-1)) # 1: clockwise, -1: counterclockwise - return ccw_out - - # Check whether two lines intersect - @staticmethod - 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 = TestDivideFluxAreaListStencil01._ccw_numpy( - line1_p1_lon, - line1_p1_lat, - line1_p2_lon, - line1_p2_lat, - line2_p1_lon, - line2_p1_lat, - ) * TestDivideFluxAreaListStencil01._ccw_numpy( - line1_p1_lon, - line1_p1_lat, - line1_p2_lon, - line1_p2_lat, - line2_p2_lon, - line2_p2_lat, - ) - intersect2 = TestDivideFluxAreaListStencil01._ccw_numpy( - line2_p1_lon, - line2_p1_lat, - line2_p2_lon, - line2_p2_lat, - line1_p1_lon, - line1_p1_lat, - ) * TestDivideFluxAreaListStencil01._ccw_numpy( - line2_p1_lon, - line2_p1_lat, - line2_p2_lon, - line2_p2_lat, - line1_p2_lon, - line1_p2_lat, - ) - lintersect_out = np.where((intersect1 + intersect2) == -2, True, False) - - return lintersect_out - - # Compute intersection point of two lines in 2D @staticmethod - 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) - - intersect_1 = (line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon) / ( - m1 - m2 - ) - intersect_2 = line1_p1_lat + m1 * (intersect_1 - line1_p1_lon) - - return intersect_1, intersect_2 - - @classmethod - def reference( - TestDivideFluxAreaListStencil01, - grid, - famask_int, - p_vn, - ptr_v3_lon, - ptr_v3_lat, - tangent_orientation_dsl, + def _generate_flux_area_geometry( dreg_patch0_1_lon_dsl, dreg_patch0_1_lat_dsl, dreg_patch0_2_lon_dsl, @@ -166,33 +157,26 @@ def reference( dreg_patch0_3_lat_dsl, dreg_patch0_4_lon_dsl, dreg_patch0_4_lat_dsl, - **kwargs, + p_vn, + ptr_v3_lon_e, + ptr_v3_lat_e, ): - 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) - 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_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 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( @@ -206,7 +190,6 @@ def reference( np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), ) - # 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( @@ -220,38 +203,7 @@ def reference( np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), ) - # Create first mask does departure-line segment intersects with A1V3 - lintersect_line1 = TestDivideFluxAreaListStencil01._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 = TestDivideFluxAreaListStencil01._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) - # ------------------------------------------------- Case 1 - mask_case1 = np.logical_and.reduce([lintersect_line1, lintersect_line2, famask_bool]) - ps1_x, ps1_y = TestDivideFluxAreaListStencil01._line_intersect_numpy( + return ( fl_line_p1_lon, fl_line_p1_lat, fl_line_p2_lon, @@ -260,44 +212,84 @@ def reference( tri_line1_p1_lat, tri_line1_p2_lon, tri_line1_p2_lat, - ) - ps2_x, ps2_y = TestDivideFluxAreaListStencil01._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, + 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, ) - # 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) + @staticmethod + def _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, + ): + 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), - dreg_patch0_2_lon_dsl, + 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), - dreg_patch0_2_lat_dsl, + arrival_pts_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_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), - dreg_patch0_4_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, ) - # Case 1 - patch 1 + + @staticmethod + def _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, + ): 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) @@ -314,6 +306,29 @@ def reference( 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, + 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) @@ -332,11 +347,38 @@ def reference( 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] + 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, ) - # Case 2a - patch 0 + + @staticmethod + def _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 = 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( @@ -361,7 +403,37 @@ def reference( np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), dreg_patch0_4_lat_dsl, ) - # Case 2a - patch 1 + + 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, + 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 = np.where( mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask ) @@ -394,21 +466,39 @@ def reference( 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] + 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 2b - patch 0 + + @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, + 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 = 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( @@ -433,16 +523,72 @@ def reference( 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 + + 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, + 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, + ): + 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_1_lon_vmask = np.where( mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask ) @@ -476,41 +622,38 @@ def reference( dreg_patch2_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 + 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, + ) - # ----------------------------------------------- Case 3a - # Check whether flux area edge 2 intersects with triangle edge 1 - lintersect_e2_line1 = TestDivideFluxAreaListStencil01._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 = TestDivideFluxAreaListStencil01._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 + @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, + 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 = 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( @@ -535,7 +678,39 @@ def reference( np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), dreg_patch0_4_lat_dsl, ) - # Case 3a - patch 1 + + 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_1_lon_vmask = np.where( mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask ) @@ -568,40 +743,37 @@ def reference( 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 = TestDivideFluxAreaListStencil01._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 = TestDivideFluxAreaListStencil01._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, + 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 3b - patch 0 + + @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) @@ -626,16 +798,39 @@ def reference( 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 + + 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, + 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_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask ) @@ -669,6 +864,564 @@ def reference( 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. @@ -725,8 +1478,8 @@ def reference( dreg_patch2_4_lat_vmask=dreg_patch2_4_lat_vmask, ) - @pytest.fixture @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) 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 60f4419dbb..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 @@ -44,73 +44,62 @@ class TestPrepGaussQuadratureCListStencil(StencilTest): ) @staticmethod - def reference( - 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, + 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 - 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 + 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, + ) + + famask_bool = np.where(famask_int == np.int32(1), True, False) - 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) @@ -144,6 +133,7 @@ def reference( ), 0.0, ) + wgt_t_detjac_2 = np.where( famask_bool, dbl_eps @@ -168,6 +158,7 @@ def reference( ), 0.0, ) + wgt_t_detjac_3 = np.where( famask_bool, dbl_eps @@ -217,6 +208,35 @@ def reference( 0.0, ) + 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, + 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, + ): 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 @@ -265,7 +285,32 @@ def reference( + 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 @@ -321,6 +366,151 @@ def reference( + 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( 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 f98bd8560f..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 @@ -38,69 +38,60 @@ class TestPrepGaussQuadratureCStencil(StencilTest): ) @staticmethod - def reference( - 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, + 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, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, ): + # 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 - 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 + # 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, + ) wgt_t_detjac_1 = dbl_eps + z_wgt_1 * ( ( @@ -174,7 +165,35 @@ def reference( - 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, + 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, + ): 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 @@ -223,7 +242,32 @@ def reference( + 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 @@ -279,6 +323,149 @@ def reference( + 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(