diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_apply_horizontal_density_increment.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_apply_horizontal_density_increment.py index 2b82a73a20..52820b534f 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_apply_horizontal_density_increment.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_apply_horizontal_density_increment.py @@ -15,7 +15,7 @@ from icon4py.model.atmosphere.advection.stencils import apply_horizontal_density_increment from icon4py.model.common import dimension as dims -from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.common.grid import base from icon4py.model.testing import stencil_tests @@ -29,9 +29,9 @@ class TestApplyHorizontalDensityIncrement(stencil_tests.StencilTest): ), ) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_rhodz_new: np.ndarray, p_mflx_contra_v: np.ndarray, deepatmo_divzl: np.ndarray, @@ -43,15 +43,15 @@ def reference( 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) -> dict: - p_rhodz_new = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_mflx_contra_v = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid) -> dict: + p_rhodz_new = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_mflx_contra_v = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - deepatmo_divzl = data_alloc.random_field(grid, dims.KDim) - deepatmo_divzu = data_alloc.random_field(grid, dims.KDim) - rhodz_ast2 = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + deepatmo_divzl = self.data_alloc.random_field(dims.KDim) + deepatmo_divzu = self.data_alloc.random_field(dims.KDim) + rhodz_ast2 = self.data_alloc.zero_field(dims.CellDim, dims.KDim) p_dtime = 0.1 return dict( p_rhodz_new=p_rhodz_new, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_apply_interpolated_tracer_time_tendency.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_apply_interpolated_tracer_time_tendency.py index 4484a4159a..9a506670e2 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_apply_interpolated_tracer_time_tendency.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_apply_interpolated_tracer_time_tendency.py @@ -16,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -24,9 +23,9 @@ class TestApplyInterpolatedTracerTimeTendency(stencil_tests.StencilTest): PROGRAM = apply_interpolated_tracer_time_tendency OUTPUTS = ("p_tracer_new",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_tracer_now: np.ndarray, p_grf_tend_tracer: np.ndarray, p_dtime: float, @@ -37,11 +36,11 @@ def reference( return dict(p_tracer_new=p_tracer_new) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - p_tracer_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_grf_tend_tracer = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_tracer_new = data_alloc.random_field(grid, dims.CellDim, dims.KDim) + p_tracer_now = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_grf_tend_tracer = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_tracer_new = self.data_alloc.random_field(dims.CellDim, dims.KDim) p_dtime = np.float64(5.0) return dict( p_tracer_now=p_tracer_now, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_apply_vertical_density_increment.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_apply_vertical_density_increment.py index 9c428c9c60..cfe248ff6a 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_apply_vertical_density_increment.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_apply_vertical_density_increment.py @@ -16,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -30,9 +29,9 @@ class TestApplyVerticalDensityIncrement(stencil_tests.StencilTest): ), ) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, rhodz_ast: np.ndarray, p_mflx_contra_v: np.ndarray, deepatmo_divzl: np.ndarray, @@ -47,15 +46,15 @@ def reference( return dict(rhodz_ast2=rhodz_ast2) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - rhodz_ast = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_mflx_contra_v = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + rhodz_ast = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_mflx_contra_v = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - deepatmo_divzl = data_alloc.random_field(grid, dims.KDim) - deepatmo_divzu = data_alloc.random_field(grid, dims.KDim) - rhodz_ast2 = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + deepatmo_divzl = self.data_alloc.random_field(dims.KDim) + deepatmo_divzu = self.data_alloc.random_field(dims.KDim) + rhodz_ast2 = self.data_alloc.zero_field(dims.CellDim, dims.KDim) p_dtime = 0.1 return dict( rhodz_ast=rhodz_ast, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_average_horizontal_flux_subcycling_2.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_average_horizontal_flux_subcycling_2.py index d73e3354ac..73dabf5ad2 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_average_horizontal_flux_subcycling_2.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_average_horizontal_flux_subcycling_2.py @@ -16,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -24,9 +23,9 @@ class TestAverageHorizontalFluxSubcycling2(stencil_tests.StencilTest): PROGRAM = average_horizontal_flux_subcycling_2 OUTPUTS = ("p_out_e",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_tracer_mflx_1_dsl: np.ndarray, z_tracer_mflx_2_dsl: np.ndarray, **kwargs: Any, @@ -34,11 +33,11 @@ def reference( 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 + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - z_tracer_mflx_1_dsl = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - z_tracer_mflx_2_dsl = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_out_e = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) + z_tracer_mflx_1_dsl = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + z_tracer_mflx_2_dsl = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_out_e = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) return dict( z_tracer_mflx_1_dsl=z_tracer_mflx_1_dsl, z_tracer_mflx_2_dsl=z_tracer_mflx_2_dsl, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_average_horizontal_flux_subcycling_3.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_average_horizontal_flux_subcycling_3.py index 9883b6fc5a..60620fc273 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_average_horizontal_flux_subcycling_3.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_average_horizontal_flux_subcycling_3.py @@ -16,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -24,9 +23,9 @@ class TestAverageHorizontalFluxSubcycling3(stencil_tests.StencilTest): PROGRAM = average_horizontal_flux_subcycling_3 OUTPUTS = ("p_out_e",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_tracer_mflx_1_dsl: np.ndarray, z_tracer_mflx_2_dsl: np.ndarray, z_tracer_mflx_3_dsl: np.ndarray, @@ -35,12 +34,12 @@ def reference( 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 + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - z_tracer_mflx_1_dsl = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - z_tracer_mflx_2_dsl = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - z_tracer_mflx_3_dsl = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_out_e = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) + z_tracer_mflx_1_dsl = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + z_tracer_mflx_2_dsl = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + z_tracer_mflx_3_dsl = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_out_e = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) return dict( z_tracer_mflx_1_dsl=z_tracer_mflx_1_dsl, z_tracer_mflx_2_dsl=z_tracer_mflx_2_dsl, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_antidiffusive_cell_fluxes_and_min_max.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_antidiffusive_cell_fluxes_and_min_max.py index 63dc972850..703ae2a652 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_antidiffusive_cell_fluxes_and_min_max.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_antidiffusive_cell_fluxes_and_min_max.py @@ -6,15 +6,18 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from collections.abc import Mapping +from typing import cast + import gt4py.next as gtx import numpy as np import pytest -import icon4py.model.common.utils.data_allocation as data_alloc from icon4py.model.atmosphere.advection.stencils.compute_antidiffusive_cell_fluxes_and_min_max import ( compute_antidiffusive_cell_fluxes_and_min_max, ) from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.testing import stencil_tests @@ -28,9 +31,9 @@ class TestComputeAntidiffusiveCellFluxesAndMinMax(stencil_tests.StencilTest): "z_tracer_min", ) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, geofac_div: np.ndarray, p_rhodz_now: np.ndarray, p_rhodz_new: np.ndarray, @@ -40,7 +43,8 @@ def reference( p_dtime: float, **kwargs, ) -> dict: - c2e = connectivities[dims.C2EDim] + connectivities = stencil_tests.connectivities_asnumpy(grid) + c2e = connectivities[dims.C2E] z_anti_c2e = z_anti[c2e] geofac_div = np.expand_dims(geofac_div, axis=-1) @@ -77,21 +81,21 @@ def reference( z_tracer_min=z_tracer_min, ) - @pytest.fixture - def input_data(self, grid) -> dict: - geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) - p_rhodz_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_rhodz_new = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - z_mflx_low = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - z_anti = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_cc = data_alloc.random_field(grid, dims.CellDim, dims.KDim) + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid) -> dict: + geofac_div = self.data_alloc.random_field(dims.CellDim, dims.C2EDim) + p_rhodz_now = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_rhodz_new = self.data_alloc.random_field(dims.CellDim, dims.KDim) + z_mflx_low = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + z_anti = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_cc = self.data_alloc.random_field(dims.CellDim, dims.KDim) p_dtime = 5.0 - z_mflx_anti_in = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - z_mflx_anti_out = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - z_tracer_new_low = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - z_tracer_max = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - z_tracer_min = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + z_mflx_anti_in = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + z_mflx_anti_out = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + z_tracer_new_low = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + z_tracer_max = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + z_tracer_min = self.data_alloc.zero_field(dims.CellDim, dims.KDim) return dict( geofac_div=geofac_div, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_barycentric_backtrajectory.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_barycentric_backtrajectory.py index 63054b2a19..cf17f35abc 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_barycentric_backtrajectory.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_barycentric_backtrajectory.py @@ -10,11 +10,11 @@ import numpy as np import pytest -import icon4py.model.common.utils.data_allocation as data_alloc from icon4py.model.atmosphere.advection.stencils.compute_barycentric_backtrajectory import ( compute_barycentric_backtrajectory, ) from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.testing import stencil_tests @@ -22,9 +22,9 @@ class TestComputeBarycentricBacktrajectory(stencil_tests.StencilTest): PROGRAM = compute_barycentric_backtrajectory OUTPUTS = ("p_cell_idx", "p_distv_bary_1", "p_distv_bary_2") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_vn: np.ndarray, p_vt: np.ndarray, cell_idx: np.ndarray, @@ -79,24 +79,24 @@ def reference( p_distv_bary_2=p_distv_bary_2, ) - @pytest.fixture - def input_data(self, grid) -> dict: - p_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_vt = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid) -> dict: + p_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_vt = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) cell_idx = grid.get_connectivity("E2C") - pos_on_tplane_e_1 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - pos_on_tplane_e_2 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - primal_normal_cell_1 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) + pos_on_tplane_e_1 = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + pos_on_tplane_e_2 = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + primal_normal_cell_1 = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) - dual_normal_cell_1 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - primal_normal_cell_2 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) + dual_normal_cell_1 = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + primal_normal_cell_2 = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) - dual_normal_cell_2 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) + dual_normal_cell_2 = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) - p_cell_idx = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=gtx.int32) - p_cell_rel_idx_dsl = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=gtx.int32) - p_distv_bary_1 = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_distv_bary_2 = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) + p_cell_idx = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=gtx.int32) + p_cell_rel_idx_dsl = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=gtx.int32) + p_distv_bary_1 = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_distv_bary_2 = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) p_dthalf = 2.0 return dict( diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_barycentric_backtrajectory_alt.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_barycentric_backtrajectory_alt.py index 1adb8af3aa..a7996f909c 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_barycentric_backtrajectory_alt.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_barycentric_backtrajectory_alt.py @@ -10,11 +10,11 @@ import numpy as np import pytest -import icon4py.model.common.utils.data_allocation as data_alloc from icon4py.model.atmosphere.advection.stencils.compute_barycentric_backtrajectory_alt import ( compute_barycentric_backtrajectory_alt, ) from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.testing import stencil_tests @@ -22,9 +22,9 @@ class TestComputeBarycentricBacktrajectoryAlt(stencil_tests.StencilTest): PROGRAM = compute_barycentric_backtrajectory_alt OUTPUTS = ("p_distv_bary_1", "p_distv_bary_2") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_vn: np.ndarray, p_vt: np.ndarray, pos_on_tplane_e_1: np.ndarray, @@ -72,18 +72,18 @@ def reference( p_distv_bary_2=p_distv_bary_2, ) - @pytest.fixture - def input_data(self, grid) -> dict: - p_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_vt = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - pos_on_tplane_e_1 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - pos_on_tplane_e_2 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - primal_normal_cell_1 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - dual_normal_cell_1 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - primal_normal_cell_2 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - dual_normal_cell_2 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - p_distv_bary_1 = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_distv_bary_2 = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid) -> dict: + p_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_vt = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + pos_on_tplane_e_1 = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + pos_on_tplane_e_2 = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + primal_normal_cell_1 = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + dual_normal_cell_1 = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + primal_normal_cell_2 = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + dual_normal_cell_2 = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + p_distv_bary_1 = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_distv_bary_2 = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) p_dthalf = 2.0 return dict( diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ffsl_backtrajectory.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ffsl_backtrajectory.py index 60bb15e9cf..8878f8e487 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ffsl_backtrajectory.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ffsl_backtrajectory.py @@ -10,11 +10,11 @@ import numpy as np import pytest -import icon4py.model.common.utils.data_allocation as data_alloc from icon4py.model.atmosphere.advection.stencils.compute_ffsl_backtrajectory import ( compute_ffsl_backtrajectory, ) from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.testing import stencil_tests @@ -34,9 +34,9 @@ class TestComputeFfslBacktrajectory(stencil_tests.StencilTest): "p_coords_dreg_v_4_lat_dsl", ) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_vn: np.ndarray, p_vt: np.ndarray, cell_idx: np.ndarray, @@ -145,40 +145,40 @@ def reference( "p_coords_dreg_v_4_lat_dsl": p_coords_dreg_v_4_lat_dsl, } - @pytest.fixture - def input_data(self, grid) -> dict: - p_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_vt = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid) -> dict: + p_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_vt = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) cell_idx = grid.get_connectivity("E2C") - cell_blk = data_alloc.constant_field(grid, 1, dims.EdgeDim, dims.E2CDim, dtype=gtx.int32) + cell_blk = self.data_alloc.constant_field(1, dims.EdgeDim, dims.E2CDim, dtype=gtx.int32) - edge_verts_1_x = data_alloc.random_field(grid, dims.EdgeDim) - edge_verts_2_x = data_alloc.random_field(grid, dims.EdgeDim) - edge_verts_1_y = data_alloc.random_field(grid, dims.EdgeDim) - edge_verts_2_y = data_alloc.random_field(grid, dims.EdgeDim) - pos_on_tplane_e_1_x = data_alloc.random_field(grid, dims.EdgeDim) - pos_on_tplane_e_2_x = data_alloc.random_field(grid, dims.EdgeDim) - pos_on_tplane_e_1_y = data_alloc.random_field(grid, dims.EdgeDim) - pos_on_tplane_e_2_y = data_alloc.random_field(grid, dims.EdgeDim) - primal_normal_cell_x = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - dual_normal_cell_x = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - primal_normal_cell_y = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - dual_normal_cell_y = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - lvn_sys_pos = data_alloc.constant_field(grid, True, dims.EdgeDim, dims.KDim, dtype=bool) + edge_verts_1_x = self.data_alloc.random_field(dims.EdgeDim) + edge_verts_2_x = self.data_alloc.random_field(dims.EdgeDim) + edge_verts_1_y = self.data_alloc.random_field(dims.EdgeDim) + edge_verts_2_y = self.data_alloc.random_field(dims.EdgeDim) + pos_on_tplane_e_1_x = self.data_alloc.random_field(dims.EdgeDim) + pos_on_tplane_e_2_x = self.data_alloc.random_field(dims.EdgeDim) + pos_on_tplane_e_1_y = self.data_alloc.random_field(dims.EdgeDim) + pos_on_tplane_e_2_y = self.data_alloc.random_field(dims.EdgeDim) + primal_normal_cell_x = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + dual_normal_cell_x = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + primal_normal_cell_y = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + dual_normal_cell_y = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + lvn_sys_pos = self.data_alloc.constant_field(True, dims.EdgeDim, dims.KDim, dtype=bool) p_dt = 2.0 - p_cell_idx = data_alloc.constant_field(grid, 0, dims.EdgeDim, dims.KDim, dtype=gtx.int32) - p_cell_rel_idx_dsl = data_alloc.constant_field( - grid, 0, dims.EdgeDim, dims.KDim, dtype=gtx.int32 + p_cell_idx = self.data_alloc.constant_field(0, dims.EdgeDim, dims.KDim, dtype=gtx.int32) + p_cell_rel_idx_dsl = self.data_alloc.constant_field( + 0, dims.EdgeDim, dims.KDim, dtype=gtx.int32 ) - p_cell_blk = data_alloc.constant_field(grid, 0, dims.EdgeDim, dims.KDim, dtype=gtx.int32) - p_coords_dreg_v_1_lon_dsl = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_2_lon_dsl = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_3_lon_dsl = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_4_lon_dsl = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_1_lat_dsl = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_2_lat_dsl = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_3_lat_dsl = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_4_lat_dsl = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) + p_cell_blk = self.data_alloc.constant_field(0, dims.EdgeDim, dims.KDim, dtype=gtx.int32) + p_coords_dreg_v_1_lon_dsl = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_2_lon_dsl = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_3_lon_dsl = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_4_lon_dsl = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_1_lat_dsl = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_2_lat_dsl = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_3_lat_dsl = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_4_lat_dsl = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) return dict( p_vn=p_vn, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ffsl_backtrajectory_counterclockwise_indicator.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ffsl_backtrajectory_counterclockwise_indicator.py index 064c915c6c..c51d32744e 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ffsl_backtrajectory_counterclockwise_indicator.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ffsl_backtrajectory_counterclockwise_indicator.py @@ -16,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -24,9 +23,9 @@ class TestComputeFfslBacktrajectoryCounterclockwiseIndicator(stencil_tests.Stenc PROGRAM = compute_ffsl_backtrajectory_counterclockwise_indicator OUTPUTS = ("lvn_sys_pos",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_vn: np.ndarray, tangent_orientation: np.ndarray, lcounterclock: bool, @@ -44,11 +43,11 @@ def reference( return dict(lvn_sys_pos=lvn_sys_pos) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - p_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - tangent_orientation = data_alloc.random_field(grid, dims.EdgeDim) - lvn_sys_pos = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=bool) + p_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + tangent_orientation = self.data_alloc.random_field(dims.EdgeDim) + lvn_sys_pos = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=bool) lcounterclock = True return dict( p_vn=p_vn, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ffsl_backtrajectory_length_indicator.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ffsl_backtrajectory_length_indicator.py index 326d040463..ab5d951b51 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ffsl_backtrajectory_length_indicator.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ffsl_backtrajectory_length_indicator.py @@ -14,7 +14,7 @@ compute_ffsl_backtrajectory_length_indicator, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.common.grid import base from icon4py.model.testing import stencil_tests @@ -22,9 +22,9 @@ class TestComputeFfslBacktrajectoryLengthIndicator(stencil_tests.StencilTest): PROGRAM = compute_ffsl_backtrajectory_length_indicator OUTPUTS = ("opt_famask_dsl",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_vn: np.ndarray, p_vt: np.ndarray, edge_cell_length: np.ndarray, @@ -46,12 +46,12 @@ def reference( return dict(opt_famask_dsl=opt_famask_dsl) - @pytest.fixture - def input_data(self, grid) -> dict: - p_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_vt = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - edge_cell_length = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - opt_famask_dsl = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=gtx.int32) + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid) -> dict: + p_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_vt = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + edge_cell_length = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + opt_famask_dsl = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=gtx.int32) p_dt = 1.0 return dict( diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_horizontal_tracer_flux_from_cubic_coefficients.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_horizontal_tracer_flux_from_cubic_coefficients.py index cdf27e4979..a5adb5f6e1 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_horizontal_tracer_flux_from_cubic_coefficients.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_horizontal_tracer_flux_from_cubic_coefficients.py @@ -14,7 +14,7 @@ compute_horizontal_tracer_flux_from_cubic_coefficients, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.common.grid import base from icon4py.model.testing import stencil_tests @@ -22,9 +22,9 @@ class TestComputeHorizontalTracerFluxFromCubicCoefficients(stencil_tests.Stencil PROGRAM = compute_horizontal_tracer_flux_from_cubic_coefficients OUTPUTS = ("p_out_e_hybrid_2",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_out_e_hybrid_2: np.ndarray, p_mass_flx_e: np.ndarray, z_dreg_area: np.ndarray, @@ -34,11 +34,11 @@ def reference( return dict(p_out_e_hybrid_2=p_out_e_hybrid_2) - @pytest.fixture - def input_data(self, grid) -> dict: - p_out_e_hybrid_2 = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_mass_flx_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - z_dreg_area = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid) -> dict: + p_out_e_hybrid_2 = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_mass_flx_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + z_dreg_area = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) return dict( p_mass_flx_e=p_mass_flx_e, z_dreg_area=z_dreg_area, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients.py index 51e4cc49d2..804b9136c9 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -16,7 +17,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -24,9 +24,9 @@ class TestComputeHorizontalTracerFluxFromLinearCoefficients(stencil_tests.Stenci PROGRAM = compute_horizontal_tracer_flux_from_linear_coefficients OUTPUTS = ("p_out_e",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_lsq_coeff_1: np.ndarray, z_lsq_coeff_2: np.ndarray, z_lsq_coeff_3: np.ndarray, @@ -37,8 +37,9 @@ def reference( p_out_e: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) p_out_e_cp = p_out_e.copy() - e2c = connectivities[dims.E2CDim] + e2c = connectivities[dims.E2C] 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] @@ -57,18 +58,18 @@ def reference( return dict(p_out_e=p_out_e) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - z_lsq_coeff_1 = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - z_lsq_coeff_2 = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - z_lsq_coeff_3 = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - distv_bary_1 = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - distv_bary_2 = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_mass_flx_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - cell_rel_idx_dsl = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim, low=0.0, high=2.0, dtype=gtx.int32 + z_lsq_coeff_1 = self.data_alloc.random_field(dims.CellDim, dims.KDim) + z_lsq_coeff_2 = self.data_alloc.random_field(dims.CellDim, dims.KDim) + z_lsq_coeff_3 = self.data_alloc.random_field(dims.CellDim, dims.KDim) + distv_bary_1 = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + distv_bary_2 = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_mass_flx_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + cell_rel_idx_dsl = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim, low=0.0, high=2.0, dtype=gtx.int32 ) - p_out_e = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) + p_out_e = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) edge_domain = h_grid.domain(dims.EdgeDim) horizontal_start = grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5)) diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients_alt.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients_alt.py index f490a7bf82..70fed8741d 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients_alt.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients_alt.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -16,7 +17,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -24,9 +24,9 @@ class TestComputeHorizontalTracerFluxFromLinearCoefficientsAlt(stencil_tests.Ste PROGRAM = compute_horizontal_tracer_flux_from_linear_coefficients_alt OUTPUTS = ("p_out_e",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_lsq_coeff_1: np.ndarray, z_lsq_coeff_2: np.ndarray, z_lsq_coeff_3: np.ndarray, @@ -37,8 +37,9 @@ def reference( p_out_e: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) p_out_e_cp = p_out_e.copy() - e2c = connectivities[dims.E2CDim] + e2c = connectivities[dims.E2C] 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] @@ -57,16 +58,16 @@ def reference( return dict(p_out_e=p_out_e) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - z_lsq_coeff_1 = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - z_lsq_coeff_2 = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - z_lsq_coeff_3 = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - distv_bary_1 = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - distv_bary_2 = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_mass_flx_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_out_e = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) + z_lsq_coeff_1 = self.data_alloc.random_field(dims.CellDim, dims.KDim) + z_lsq_coeff_2 = self.data_alloc.random_field(dims.CellDim, dims.KDim) + z_lsq_coeff_3 = self.data_alloc.random_field(dims.CellDim, dims.KDim) + distv_bary_1 = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + distv_bary_2 = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_mass_flx_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_out_e = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) edge_domain = h_grid.domain(dims.EdgeDim) horizontal_start = grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5)) diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_positive_definite_horizontal_multiplicative_flux_factor.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_positive_definite_horizontal_multiplicative_flux_factor.py index c6ee1ad78e..62e5f190c0 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_positive_definite_horizontal_multiplicative_flux_factor.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_positive_definite_horizontal_multiplicative_flux_factor.py @@ -6,15 +6,18 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from collections.abc import Mapping +from typing import cast + import gt4py.next as gtx import numpy as np import pytest -import icon4py.model.common.utils.data_allocation as data_alloc from icon4py.model.atmosphere.advection.stencils.compute_positive_definite_horizontal_multiplicative_flux_factor import ( compute_positive_definite_horizontal_multiplicative_flux_factor, ) from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.testing import stencil_tests @@ -22,9 +25,9 @@ class TestComputePositiveDefiniteHorizontalMultiplicativeFluxFactor(stencil_test PROGRAM = compute_positive_definite_horizontal_multiplicative_flux_factor OUTPUTS = ("r_m",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, geofac_div: np.ndarray, p_cc: np.ndarray, p_rhodz_now: np.ndarray, @@ -33,7 +36,8 @@ def reference( dbl_eps, **kwargs, ) -> dict: - c2e = connectivities[dims.C2EDim] + connectivities = stencil_tests.connectivities_asnumpy(grid) + c2e = connectivities[dims.C2E] geofac_div = np.expand_dims(geofac_div, axis=-1) p_m_0 = np.maximum( 0.0, @@ -53,13 +57,13 @@ def reference( return dict(r_m=r_m) - @pytest.fixture - def input_data(self, grid) -> dict: - geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) - p_cc = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_rhodz_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_mflx_tracer_h = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - r_m = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid) -> dict: + geofac_div = self.data_alloc.random_field(dims.CellDim, dims.C2EDim) + p_cc = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_rhodz_now = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_mflx_tracer_h = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + r_m = self.data_alloc.zero_field(dims.CellDim, dims.KDim) p_dtime = np.float64(5) dbl_eps = np.float64(1e-9) return dict( diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm4gpu_parabola_coefficients.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm4gpu_parabola_coefficients.py index ffaeefc5c6..b4a4b23567 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm4gpu_parabola_coefficients.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm4gpu_parabola_coefficients.py @@ -16,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -24,9 +23,9 @@ class TestComputePpm4gpuParabolaCoefficients(stencil_tests.StencilTest): PROGRAM = compute_ppm4gpu_parabola_coefficients OUTPUTS = ("z_delta_q", "z_a1") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_face_up: np.ndarray, z_face_low: np.ndarray, p_cc: np.ndarray, @@ -36,13 +35,13 @@ def reference( z_a1 = p_cc - 0.5 * (z_face_up + z_face_low) return dict(z_delta_q=z_delta_q, z_a1=z_a1) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - z_face_up = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - z_face_low = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_cc = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - z_delta_q = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - z_a1 = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + z_face_up = self.data_alloc.random_field(dims.CellDim, dims.KDim) + z_face_low = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_cc = self.data_alloc.random_field(dims.CellDim, dims.KDim) + z_delta_q = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + z_a1 = self.data_alloc.zero_field(dims.CellDim, dims.KDim) return dict( z_face_up=z_face_up, z_face_low=z_face_low, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm_all_face_values.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm_all_face_values.py index 054cb659fe..28964860e7 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm_all_face_values.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm_all_face_values.py @@ -16,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -25,9 +24,9 @@ class TestComputePpmAllFaceValues(stencil_tests.StencilTest): PROGRAM = compute_ppm_all_face_values OUTPUTS = ("p_face",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_cc: np.ndarray, p_cellhgt_mc_now: np.ndarray, p_face_in: np.ndarray, @@ -49,12 +48,12 @@ def reference( p_face[:, 1:] = np.where((k[1:] == elevp1), p_cc[:, :-1], p_face[:, 1:]) return dict(p_face=p_face) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - p_cc = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_cellhgt_mc_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_face_in = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_face = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + p_cc = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_cellhgt_mc_now = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_face_in = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_face = self.data_alloc.zero_field(dims.CellDim, dims.KDim) slev = gtx.int32(1) slevp1 = gtx.int32(2) elev = grid.num_levels - 2 diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm_quadratic_face_values.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm_quadratic_face_values.py index 98095dc0f1..de5df28763 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm_quadratic_face_values.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm_quadratic_face_values.py @@ -16,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -27,9 +26,9 @@ class TestComputePpmQuadraticFaceValues(stencil_tests.StencilTest): PROGRAM = compute_ppm_quadratic_face_values OUTPUTS = (stencil_tests.Output("p_face", refslice=outslice, gtslice=outslice),) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_cc: np.ndarray, p_cellhgt_mc_now: np.ndarray, **kwargs: Any, @@ -42,11 +41,11 @@ def reference( ) return dict(p_face=p_face) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - p_face = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_cc = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_cellhgt_mc_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim) + p_face = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_cc = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_cellhgt_mc_now = self.data_alloc.random_field(dims.CellDim, dims.KDim) return dict( p_cc=p_cc, p_cellhgt_mc_now=p_cellhgt_mc_now, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm_quartic_face_values.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm_quartic_face_values.py index b6192bf216..9edc1fafc0 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm_quartic_face_values.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm_quartic_face_values.py @@ -16,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -24,9 +23,9 @@ class TestComputePpmQuarticFaceValues(stencil_tests.StencilTest): PROGRAM = compute_ppm_quartic_face_values OUTPUTS = (stencil_tests.Output("p_face", gtslice=(slice(None), slice(2, None))),) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_cc: np.ndarray, p_cellhgt_mc_now: np.ndarray, z_slope: np.ndarray, @@ -68,14 +67,14 @@ def reference( ) return dict(p_face=p_face) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - p_cc = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_cellhgt_mc_now = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + p_cc = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_cellhgt_mc_now = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - z_slope = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_face = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + z_slope = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_face = self.data_alloc.zero_field(dims.CellDim, dims.KDim) return dict( p_cc=p_cc, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm_slope.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm_slope.py index d875fd8f2c..2a79682255 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm_slope.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_ppm_slope.py @@ -14,7 +14,6 @@ from icon4py.model.atmosphere.advection.stencils.compute_ppm_slope import compute_ppm_slope from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -27,9 +26,9 @@ class TestComputePpmSlope(stencil_tests.StencilTest): ), ) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_cc: np.ndarray, p_cellhgt_mc_now: np.ndarray, elev: gtx.int32, @@ -64,12 +63,12 @@ def reference( z_slope = np.where(k[1:-1] < elev, z_slope_a, z_slope_b) return dict(z_slope=z_slope) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - z_slope = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - p_cc = data_alloc.random_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) - p_cellhgt_mc_now = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + z_slope = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + p_cc = self.data_alloc.random_field(dims.CellDim, dims.KDim, extend={dims.KDim: 1}) + p_cellhgt_mc_now = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) elev = grid.num_levels - 2 diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_tendency.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_tendency.py index 88810d6269..960b191707 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_tendency.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_tendency.py @@ -14,7 +14,6 @@ from icon4py.model.atmosphere.advection.stencils.compute_tendency import compute_tendency from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -22,9 +21,9 @@ class TestComputeTendency(stencil_tests.StencilTest): PROGRAM = compute_tendency OUTPUTS = ("opt_ddt_tracer_adv",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_tracer_now: np.ndarray, p_tracer_new: np.ndarray, p_dtime, @@ -34,11 +33,11 @@ def reference( return dict(opt_ddt_tracer_adv=opt_ddt_tracer_adv) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - p_tracer_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_tracer_new = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - opt_ddt_tracer_adv = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + p_tracer_now = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_tracer_new = self.data_alloc.random_field(dims.CellDim, dims.KDim) + opt_ddt_tracer_adv = self.data_alloc.zero_field(dims.CellDim, dims.KDim) p_dtime = np.float64(5.0) return dict( p_tracer_now=p_tracer_now, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_vertical_parabola_limiter_condition.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_vertical_parabola_limiter_condition.py index e881c1f702..250cca3c37 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_vertical_parabola_limiter_condition.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_vertical_parabola_limiter_condition.py @@ -16,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -24,9 +23,9 @@ class TestComputeVerticalParabolaLimiterCondition(stencil_tests.StencilTest): PROGRAM = compute_vertical_parabola_limiter_condition OUTPUTS = ("l_limit",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_face: np.ndarray, p_cc: np.ndarray, **kwargs: Any, @@ -36,11 +35,11 @@ def reference( l_limit = np.where(np.abs(z_delta) < -1 * z_a6i, 1, 0) return dict(l_limit=l_limit) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - p_cc = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_face = data_alloc.random_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) - l_limit = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=gtx.int32) + p_cc = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_face = self.data_alloc.random_field(dims.CellDim, dims.KDim, extend={dims.KDim: 1}) + l_limit = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=gtx.int32) return dict( p_face=p_face, p_cc=p_cc, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_vertical_tracer_flux_upwind.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_vertical_tracer_flux_upwind.py index 03d40aaa52..497d1c5373 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_vertical_tracer_flux_upwind.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_vertical_tracer_flux_upwind.py @@ -16,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -27,9 +26,9 @@ class TestComputeVerticalTracerFluxUpwind(stencil_tests.StencilTest): PROGRAM = compute_vertical_tracer_flux_upwind OUTPUTS = (stencil_tests.Output("p_upflux", refslice=outslice, gtslice=outslice),) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_cc: np.ndarray, p_mflx_contra_v: np.ndarray, **kwargs: Any, @@ -41,14 +40,14 @@ def reference( ) return dict(p_upflux=p_upflux) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - p_cc = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_mflx_contra_v = data_alloc.random_field( - grid, dims.CellDim, dims.KDim + p_cc = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_mflx_contra_v = self.data_alloc.random_field( + dims.CellDim, dims.KDim ) # TODO(dastrm): should be KHalfDim - p_upflux = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim + p_upflux = self.data_alloc.zero_field( + dims.CellDim, dims.KDim ) # TODO(dastrm): should be KHalfDim return dict( p_cc=p_cc, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_integrate_tracer_density_horizontally.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_integrate_tracer_density_horizontally.py index 6bb07ba7cb..6ddb507eae 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_integrate_tracer_density_horizontally.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_integrate_tracer_density_horizontally.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -16,7 +17,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -29,9 +29,9 @@ class TestIntegrateTracerDensityHorizontally(stencil_tests.StencilTest): "z_tracer_new_dsl", ) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_mass_flx_e: np.ndarray, geofac_div: np.ndarray, z_rhofluxdiv_c: np.ndarray, @@ -42,7 +42,8 @@ def reference( nsub: gtx.int32, **kwargs: Any, ) -> dict: - c2e = connectivities[dims.C2EDim] + connectivities = stencil_tests.connectivities_asnumpy(grid) + c2e = connectivities[dims.C2E] 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] @@ -61,18 +62,18 @@ def reference( z_tracer_new_dsl=z_tracer_new_dsl, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - p_mass_flx_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) - z_rhofluxdiv_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - z_tracer_mflx = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - z_rho_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - z_tracer_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - z_rhofluxdiv_c_out = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - z_fluxdiv_c_dsl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - z_rho_new_dsl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - z_tracer_new_dsl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + p_mass_flx_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + geofac_div = self.data_alloc.random_field(dims.CellDim, dims.C2EDim) + z_rhofluxdiv_c = self.data_alloc.random_field(dims.CellDim, dims.KDim) + z_tracer_mflx = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + z_rho_now = self.data_alloc.random_field(dims.CellDim, dims.KDim) + z_tracer_now = self.data_alloc.random_field(dims.CellDim, dims.KDim) + z_rhofluxdiv_c_out = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + z_fluxdiv_c_dsl = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + z_rho_new_dsl = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + z_tracer_new_dsl = self.data_alloc.zero_field(dims.CellDim, dims.KDim) z_dtsub = 0.5 nsub = 1 return dict( diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_integrate_tracer_horizontally.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_integrate_tracer_horizontally.py index f1f1134f66..eb432aad10 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_integrate_tracer_horizontally.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_integrate_tracer_horizontally.py @@ -5,13 +5,13 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np import pytest -import icon4py.model.common.utils.data_allocation as data_alloc from icon4py.model.atmosphere.advection.stencils.integrate_tracer_horizontally import ( integrate_tracer_horizontally, ) @@ -24,9 +24,9 @@ class TestIntegrateTracerHorizontally(stencil_tests.StencilTest): PROGRAM = integrate_tracer_horizontally OUTPUTS = ("tracer_new_hor",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_mflx_tracer_h: np.ndarray, deepatmo_divh: np.ndarray, tracer_now: np.ndarray, @@ -36,25 +36,26 @@ def reference( p_dtime: float, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) geofac_div = np.expand_dims(geofac_div, axis=-1) tracer_new_hor = ( tracer_now * rhodz_now - p_dtime * deepatmo_divh - * np.sum(p_mflx_tracer_h[connectivities[dims.C2EDim]] * geofac_div, axis=1) + * np.sum(p_mflx_tracer_h[connectivities[dims.C2E]] * geofac_div, axis=1) ) / rhodz_new return dict(tracer_new_hor=tracer_new_hor) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - p_mflx_tracer_h = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - deepatmo_divh = data_alloc.random_field(grid, dims.KDim) - tracer_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - rhodz_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - rhodz_new = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) + p_mflx_tracer_h = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + deepatmo_divh = self.data_alloc.random_field(dims.KDim) + tracer_now = self.data_alloc.random_field(dims.CellDim, dims.KDim) + rhodz_now = self.data_alloc.random_field(dims.CellDim, dims.KDim) + rhodz_new = self.data_alloc.random_field(dims.CellDim, dims.KDim) + geofac_div = self.data_alloc.random_field(dims.CellDim, dims.C2EDim) p_dtime = np.float64(5.0) - tracer_new_hor = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + tracer_new_hor = self.data_alloc.zero_field(dims.CellDim, dims.KDim) return dict( p_mflx_tracer_h=p_mflx_tracer_h, deepatmo_divh=deepatmo_divh, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_integrate_tracer_vertically.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_integrate_tracer_vertically.py index 6ab7a36d13..e5d856d8d2 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_integrate_tracer_vertically.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_integrate_tracer_vertically.py @@ -16,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -24,9 +23,9 @@ class TestIntegrateTracerVertically(stencil_tests.StencilTest): PROGRAM = integrate_tracer_vertically OUTPUTS = ("tracer_new",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, tracer_now: np.ndarray, rhodz_now: np.ndarray, p_mflx_tracer_v: np.ndarray, @@ -58,18 +57,18 @@ def reference( return dict(tracer_new=tracer_new) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - tracer_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - rhodz_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_mflx_tracer_v = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + tracer_now = self.data_alloc.random_field(dims.CellDim, dims.KDim) + rhodz_now = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_mflx_tracer_v = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - deepatmo_divzl = data_alloc.random_field(grid, dims.KDim) - deepatmo_divzu = data_alloc.random_field(grid, dims.KDim) - rhodz_new = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - tracer_new = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - k = data_alloc.index_field(grid, dims.KDim) + deepatmo_divzl = self.data_alloc.random_field(dims.KDim) + deepatmo_divzu = self.data_alloc.random_field(dims.KDim) + rhodz_new = self.data_alloc.random_field(dims.CellDim, dims.KDim) + tracer_new = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + k = self.data_alloc.index_field(dims.KDim) p_dtime = np.float64(5.0) ivadv_tracer = 1 iadv_slev_jt = 4 diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_limit_vertical_parabola_semi_monotonically.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_limit_vertical_parabola_semi_monotonically.py index b7079d0aca..b021cac853 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_limit_vertical_parabola_semi_monotonically.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_limit_vertical_parabola_semi_monotonically.py @@ -16,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -24,9 +23,9 @@ class TestLimitVerticalParabolaSemiMonotonically(stencil_tests.StencilTest): PROGRAM = limit_vertical_parabola_semi_monotonically OUTPUTS = ("p_face_up", "p_face_low") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, l_limit: np.ndarray, p_face: np.ndarray, p_cc: np.ndarray, @@ -47,13 +46,13 @@ def reference( ) return dict(p_face_up=q_face_up, p_face_low=q_face_low) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - l_limit = data_alloc.random_mask(grid, dims.CellDim, dims.KDim, dtype=gtx.int32) - p_cc = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - p_face = data_alloc.random_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) - p_face_up = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - p_face_low = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + l_limit = self.data_alloc.random_mask(dims.CellDim, dims.KDim, dtype=gtx.int32) + p_cc = self.data_alloc.random_field(dims.CellDim, dims.KDim) + p_face = self.data_alloc.random_field(dims.CellDim, dims.KDim, extend={dims.KDim: 1}) + p_face_up = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + p_face_low = self.data_alloc.zero_field(dims.CellDim, dims.KDim) return dict( l_limit=l_limit, p_cc=p_cc, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_limit_vertical_slope_semi_monotonically.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_limit_vertical_slope_semi_monotonically.py index fb32ae1f24..d4f66a0057 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_limit_vertical_slope_semi_monotonically.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_limit_vertical_slope_semi_monotonically.py @@ -16,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -24,9 +23,9 @@ class TestLimitVerticalSlopeSemiMonotonically(stencil_tests.StencilTest): PROGRAM = limit_vertical_slope_semi_monotonically OUTPUTS = (stencil_tests.Output("z_slope", gtslice=(slice(None), slice(1, -1))),) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_cc: np.ndarray, z_slope: np.ndarray, k: np.ndarray, @@ -39,11 +38,11 @@ def reference( slope = np.where(z_slope[:, 1:-1] >= 0.0, slope_l, -slope_l) return dict(z_slope=slope) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - p_cc = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - z_slope = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - k = data_alloc.index_field(grid, dims.KDim) + p_cc = self.data_alloc.random_field(dims.CellDim, dims.KDim) + z_slope = self.data_alloc.random_field(dims.CellDim, dims.KDim) + k = self.data_alloc.index_field(dims.KDim) elev = k[-2].as_scalar() return dict( diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_postprocess_antidiffusive_cell_fluxes_and_min_max.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_postprocess_antidiffusive_cell_fluxes_and_min_max.py index 9fe005d986..27180930e5 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_postprocess_antidiffusive_cell_fluxes_and_min_max.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_postprocess_antidiffusive_cell_fluxes_and_min_max.py @@ -16,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -24,9 +23,9 @@ class TestPostprocessAntidiffusiveCellFluxesAndMinMax(stencil_tests.StencilTest) PROGRAM = postprocess_antidiffusive_cell_fluxes_and_min_max OUTPUTS = ("z_tracer_new_low", "z_tracer_max", "z_tracer_min") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, refin_ctrl: np.ndarray, p_cc: np.ndarray, z_tracer_new_low: np.ndarray, @@ -54,18 +53,18 @@ def reference( z_tracer_min=z_tracer_min_out, ) - @pytest.fixture() + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: hi_bound, lo_bound = 3, 1 - refin_ctrl = data_alloc.constant_field(grid, 2, dims.CellDim, dtype=gtx.int32) - p_cc = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - z_tracer_new_low_in = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - z_tracer_max_in = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - z_tracer_min_in = data_alloc.random_field(grid, dims.CellDim, dims.KDim) + refin_ctrl = self.data_alloc.constant_field(2, dims.CellDim, dtype=gtx.int32) + p_cc = self.data_alloc.random_field(dims.CellDim, dims.KDim) + z_tracer_new_low_in = self.data_alloc.random_field(dims.CellDim, dims.KDim) + z_tracer_max_in = self.data_alloc.random_field(dims.CellDim, dims.KDim) + z_tracer_min_in = self.data_alloc.random_field(dims.CellDim, dims.KDim) - z_tracer_new_low_out = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - z_tracer_max_out = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - z_tracer_min_out = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + z_tracer_new_low_out = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + z_tracer_max_out = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + z_tracer_min_out = self.data_alloc.zero_field(dims.CellDim, dims.KDim) return dict( refin_ctrl=refin_ctrl, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_ffsl_flux_area_patches_list.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_ffsl_flux_area_patches_list.py index bc2d3c6933..0af461942b 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_ffsl_flux_area_patches_list.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_ffsl_flux_area_patches_list.py @@ -10,11 +10,11 @@ import numpy as np import pytest -import icon4py.model.common.utils.data_allocation as data_alloc from icon4py.model.atmosphere.advection.stencils.prepare_ffsl_flux_area_patches_list import ( prepare_ffsl_flux_area_patches_list, ) from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.testing import stencil_tests @@ -113,259 +113,264 @@ def _line_intersect_numpy( return intersect_1, intersect_2 -@pytest.mark.gtfn_too_slow -class TestPrepareFfslFluxAreaPatchesList(stencil_tests.StencilTest): - PROGRAM = prepare_ffsl_flux_area_patches_list - 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", +# -- Reference helpers -- +def _generate_flux_area_geometry( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + p_vn, + ptr_v3_lon_e, + ptr_v3_lat_e, +): + arrival_pts_1_lon_dsl = dreg_patch0_1_lon_dsl + arrival_pts_1_lat_dsl = dreg_patch0_1_lat_dsl + arrival_pts_2_lon_dsl = dreg_patch0_2_lon_dsl + arrival_pts_2_lat_dsl = dreg_patch0_2_lat_dsl + depart_pts_1_lon_dsl = dreg_patch0_4_lon_dsl + depart_pts_1_lat_dsl = dreg_patch0_4_lat_dsl + depart_pts_2_lon_dsl = dreg_patch0_3_lon_dsl + depart_pts_2_lat_dsl = dreg_patch0_3_lat_dsl + + lvn_pos = p_vn >= 0.0 + + 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 + + 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), ) - @staticmethod - def _generate_flux_area_geometry( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - p_vn, - ptr_v3_lon_e, - ptr_v3_lat_e, - ): - arrival_pts_1_lon_dsl = dreg_patch0_1_lon_dsl - arrival_pts_1_lat_dsl = dreg_patch0_1_lat_dsl - arrival_pts_2_lon_dsl = dreg_patch0_2_lon_dsl - arrival_pts_2_lat_dsl = dreg_patch0_2_lat_dsl - depart_pts_1_lon_dsl = dreg_patch0_4_lon_dsl - depart_pts_1_lat_dsl = dreg_patch0_4_lat_dsl - depart_pts_2_lon_dsl = dreg_patch0_3_lon_dsl - depart_pts_2_lat_dsl = dreg_patch0_3_lat_dsl - - lvn_pos = p_vn >= 0.0 - - 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 - - tri_line1_p1_lon = arrival_pts_1_lon_dsl - tri_line1_p1_lat = arrival_pts_1_lat_dsl - tri_line1_p2_lon = np.where( - lvn_pos, - np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), - np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), - ) - tri_line1_p2_lat = np.where( - lvn_pos, - np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), - np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), - ) - - tri_line2_p1_lon = arrival_pts_2_lon_dsl - tri_line2_p1_lat = arrival_pts_2_lat_dsl - tri_line2_p2_lon = np.where( - lvn_pos, - np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), - np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), - ) - tri_line2_p2_lat = np.where( - lvn_pos, - np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), - np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), - ) - - return ( - fl_line_p1_lon, - fl_line_p1_lat, - fl_line_p2_lon, - fl_line_p2_lat, - tri_line1_p1_lon, - tri_line1_p1_lat, - tri_line1_p2_lon, - tri_line1_p2_lat, - tri_line2_p1_lon, - tri_line2_p1_lat, - tri_line2_p2_lon, - tri_line2_p2_lat, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, - arrival_pts_2_lon_dsl, - arrival_pts_2_lat_dsl, - depart_pts_1_lon_dsl, - depart_pts_1_lat_dsl, - depart_pts_2_lon_dsl, - depart_pts_2_lat_dsl, - ) + tri_line2_p1_lon = arrival_pts_2_lon_dsl + tri_line2_p1_lat = arrival_pts_2_lat_dsl + tri_line2_p2_lon = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), + ) + tri_line2_p2_lat = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), + ) - @staticmethod - def _apply_case1_patch0( - mask_case1, - lvn_sys_pos, + return ( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, arrival_pts_1_lon_dsl, arrival_pts_1_lat_dsl, arrival_pts_2_lon_dsl, arrival_pts_2_lat_dsl, - ps1_x, - ps1_y, - ps2_x, - ps2_y, depart_pts_1_lon_dsl, depart_pts_1_lat_dsl, depart_pts_2_lon_dsl, depart_pts_2_lat_dsl, - ): - dreg_patch0_1_lon_dsl = arrival_pts_1_lon_dsl - dreg_patch0_1_lat_dsl = arrival_pts_1_lat_dsl - dreg_patch0_2_lon_dsl = np.where( - mask_case1, - np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), - arrival_pts_2_lon_dsl, - ) - dreg_patch0_2_lat_dsl = np.where( - mask_case1, - np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), - arrival_pts_2_lat_dsl, - ) - dreg_patch0_3_lon_dsl = np.where(mask_case1, ps2_x, depart_pts_2_lon_dsl) - dreg_patch0_3_lat_dsl = np.where(mask_case1, ps2_y, depart_pts_2_lat_dsl) - dreg_patch0_4_lon_dsl = np.where( - mask_case1, - np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), - depart_pts_1_lon_dsl, - ) - dreg_patch0_4_lat_dsl = np.where( - mask_case1, - np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), - depart_pts_1_lat_dsl, - ) + ) - return ( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ) - @staticmethod - def _apply_case1_patch1( +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, - lvn_sys_pos, - arrival_pts_1_lon_dsl, - arrival_pts_1_lat_dsl, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + arrival_pts_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case1, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + arrival_pts_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case1, ps2_x, depart_pts_2_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case1, ps2_y, depart_pts_2_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case1, + np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), depart_pts_1_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case1, + np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), depart_pts_1_lat_dsl, - ps1_x, - ps1_y, - ): - dreg_patch1_1_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) - dreg_patch1_1_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) - dreg_patch1_4_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) - dreg_patch1_4_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) - dreg_patch1_2_lon_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0 - ) - dreg_patch1_2_lat_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0 - ) - dreg_patch1_3_lon_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0 - ) - dreg_patch1_3_lat_vmask = np.where( - mask_case1, np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0 - ) + ) - return ( - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - ) + return ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + ) - @staticmethod - def _apply_case1_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) - 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 - ) - 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, - ) +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) + dreg_patch1_4_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) + dreg_patch1_2_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0 + ) + dreg_patch1_2_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0 + ) + dreg_patch1_3_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0 + ) + dreg_patch1_3_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0 + ) + + return ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + ) + + +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) + 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 + ) + + return ( + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + ) - @staticmethod - def _apply_case2a_patch0( + +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( 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, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case2a, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case2a, + np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case2a, + np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + + return ( dreg_patch0_1_lon_dsl, dreg_patch0_1_lat_dsl, dreg_patch0_2_lon_dsl, @@ -374,53 +379,53 @@ def _apply_case2a_patch0( 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( - mask_case2a, - np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), - dreg_patch0_2_lon_dsl, - ) - dreg_patch0_2_lat_dsl = np.where( - mask_case2a, - np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), - dreg_patch0_2_lat_dsl, - ) - dreg_patch0_3_lon_dsl = np.where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = np.where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) - dreg_patch0_4_lon_dsl = np.where( - mask_case2a, - np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), - dreg_patch0_4_lon_dsl, - ) - dreg_patch0_4_lat_dsl = np.where( - mask_case2a, - np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), - dreg_patch0_4_lat_dsl, - ) + ) - return ( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ) - @staticmethod - def _apply_case2a_patch1( +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) + 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, - 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, + 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, + ) + + return ( dreg_patch1_1_lon_vmask, dreg_patch1_1_lat_vmask, dreg_patch1_4_lon_vmask, @@ -429,63 +434,55 @@ def _apply_case2a_patch1( 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 - ) - 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, - ) + ) - 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_case2b_patch0( +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( 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, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case2b, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case2b, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + + return ( dreg_patch0_1_lon_dsl, dreg_patch0_1_lat_dsl, dreg_patch0_2_lon_dsl, @@ -494,46 +491,32 @@ def _apply_case2b_patch0( dreg_patch0_3_lat_dsl, dreg_patch0_4_lon_dsl, dreg_patch0_4_lat_dsl, - ): - dreg_patch0_1_lon_dsl = np.where(mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = np.where(mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_2_lon_dsl = np.where( - mask_case2b, - np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), - dreg_patch0_2_lon_dsl, - ) - dreg_patch0_2_lat_dsl = np.where( - mask_case2b, - np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), - dreg_patch0_2_lat_dsl, - ) - dreg_patch0_3_lon_dsl = np.where(mask_case2b, ps2_x, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = np.where(mask_case2b, ps2_y, dreg_patch0_3_lat_dsl) - dreg_patch0_4_lon_dsl = np.where( - mask_case2b, - np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), - dreg_patch0_4_lon_dsl, - ) - dreg_patch0_4_lat_dsl = np.where( - mask_case2b, - np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), - dreg_patch0_4_lat_dsl, - ) + ) - return ( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ) - @staticmethod - def _apply_case2b_patch1( - mask_case2b, +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, @@ -542,39 +525,53 @@ def _apply_case2b_patch1( 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( + +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) + 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, - 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, + 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, + ) + + return ( dreg_patch2_1_lon_vmask, dreg_patch2_1_lat_vmask, dreg_patch2_4_lon_vmask, @@ -583,63 +580,55 @@ def _apply_case2b_patch2( 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 - ) - 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, - ) + ) - return ( - dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask, - dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask, - dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask, - ) - @staticmethod - def _apply_case3a_patch0( +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( 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, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case3a, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case3a, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + + return ( dreg_patch0_1_lon_dsl, dreg_patch0_1_lat_dsl, dreg_patch0_2_lon_dsl, @@ -648,55 +637,55 @@ def _apply_case3a_patch0( dreg_patch0_3_lat_dsl, dreg_patch0_4_lon_dsl, dreg_patch0_4_lat_dsl, - ): - dreg_patch0_1_lon_dsl = np.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = np.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_2_lon_dsl = np.where( - mask_case3a, - np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), - dreg_patch0_2_lon_dsl, - ) - dreg_patch0_2_lat_dsl = np.where( - mask_case3a, - np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), - dreg_patch0_2_lat_dsl, - ) - dreg_patch0_3_lon_dsl = np.where(mask_case3a, ps2_x, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = np.where(mask_case3a, ps2_y, dreg_patch0_3_lat_dsl) - dreg_patch0_4_lon_dsl = np.where( - mask_case3a, - np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), - dreg_patch0_4_lon_dsl, - ) - dreg_patch0_4_lat_dsl = np.where( - mask_case3a, - np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), - dreg_patch0_4_lat_dsl, - ) + ) - return ( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - ) - @staticmethod - def _apply_case3a_patch1( +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) + 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, - 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, + np.where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), + dreg_patch1_2_lon_vmask, + ) + dreg_patch1_2_lat_vmask = np.where( + mask_case3a, + np.where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), + dreg_patch1_2_lat_vmask, + ) + dreg_patch1_3_lon_vmask = np.where(mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), + dreg_patch1_4_lon_vmask, + ) + dreg_patch1_4_lat_vmask = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), + dreg_patch1_4_lat_vmask, + ) + + return ( dreg_patch1_1_lon_vmask, dreg_patch1_1_lat_vmask, dreg_patch1_4_lon_vmask, @@ -705,61 +694,53 @@ def _apply_case3a_patch1( 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 - ) - dreg_patch1_1_lat_vmask = np.where( - mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask - ) - dreg_patch1_2_lon_vmask = np.where( - mask_case3a, - np.where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), - dreg_patch1_2_lon_vmask, - ) - dreg_patch1_2_lat_vmask = np.where( - mask_case3a, - np.where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), - dreg_patch1_2_lat_vmask, - ) - dreg_patch1_3_lon_vmask = np.where( - mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask - ) - dreg_patch1_3_lat_vmask = np.where( - mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask - ) - dreg_patch1_4_lon_vmask = np.where( - mask_case3a, - np.where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), - dreg_patch1_4_lon_vmask, - ) - dreg_patch1_4_lat_vmask = np.where( - mask_case3a, - np.where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), - dreg_patch1_4_lat_vmask, - ) + ) - return ( - dreg_patch1_1_lon_vmask, - dreg_patch1_1_lat_vmask, - dreg_patch1_4_lon_vmask, - dreg_patch1_4_lat_vmask, - dreg_patch1_2_lon_vmask, - dreg_patch1_2_lat_vmask, - dreg_patch1_3_lon_vmask, - dreg_patch1_3_lat_vmask, - ) - @staticmethod - def _apply_case3b_patch0( +def _apply_case3b_patch0( + mask_case3b, + arrival_pts_1_lon_dsl, + arrival_pts_1_lat_dsl, + arrival_pts_2_lon_dsl, + arrival_pts_2_lat_dsl, + pi2_x, + pi2_y, + lvn_sys_pos, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, +): + dreg_patch0_1_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_4_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) + dreg_patch0_4_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( mask_case3b, - 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, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case3b, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where( + mask_case3b, + np.where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), + dreg_patch0_3_lon_dsl, + ) + dreg_patch0_3_lat_dsl = np.where( + mask_case3b, + np.where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), + dreg_patch0_3_lat_dsl, + ) + + return ( dreg_patch0_1_lon_dsl, dreg_patch0_1_lat_dsl, dreg_patch0_4_lon_dsl, @@ -768,55 +749,55 @@ def _apply_case3b_patch0( dreg_patch0_2_lat_dsl, dreg_patch0_3_lon_dsl, dreg_patch0_3_lat_dsl, - ): - dreg_patch0_1_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_4_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) - dreg_patch0_4_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) - dreg_patch0_2_lon_dsl = np.where( - mask_case3b, - np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), - dreg_patch0_2_lon_dsl, - ) - dreg_patch0_2_lat_dsl = np.where( - mask_case3b, - np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), - dreg_patch0_2_lat_dsl, - ) - dreg_patch0_3_lon_dsl = np.where( - mask_case3b, - np.where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), - dreg_patch0_3_lon_dsl, - ) - dreg_patch0_3_lat_dsl = np.where( - mask_case3b, - np.where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), - dreg_patch0_3_lat_dsl, - ) + ) - return ( - dreg_patch0_1_lon_dsl, - dreg_patch0_1_lat_dsl, - dreg_patch0_4_lon_dsl, - dreg_patch0_4_lat_dsl, - dreg_patch0_2_lon_dsl, - dreg_patch0_2_lat_dsl, - dreg_patch0_3_lon_dsl, - dreg_patch0_3_lat_dsl, - ) - @staticmethod - def _apply_case3b_patch2( +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) + 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, - 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, + np.where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), + dreg_patch2_2_lon_vmask, + ) + dreg_patch2_2_lat_vmask = np.where( + mask_case3b, + np.where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), + dreg_patch2_2_lat_vmask, + ) + dreg_patch2_3_lon_vmask = np.where(mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where( + mask_case3b, + np.where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), + dreg_patch2_4_lon_vmask, + ) + dreg_patch2_4_lat_vmask = np.where( + mask_case3b, + np.where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), + dreg_patch2_4_lat_vmask, + ) + + return ( dreg_patch2_1_lon_vmask, dreg_patch2_1_lat_vmask, dreg_patch2_2_lon_vmask, @@ -825,55 +806,42 @@ def _apply_case3b_patch2( 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 - ) - dreg_patch2_1_lat_vmask = np.where( - mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask - ) - dreg_patch2_2_lon_vmask = np.where( - mask_case3b, - np.where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), - dreg_patch2_2_lon_vmask, - ) - dreg_patch2_2_lat_vmask = np.where( - mask_case3b, - np.where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), - dreg_patch2_2_lat_vmask, - ) - dreg_patch2_3_lon_vmask = np.where( - mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask - ) - dreg_patch2_3_lat_vmask = np.where( - mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask - ) - dreg_patch2_4_lon_vmask = np.where( - mask_case3b, - np.where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), - dreg_patch2_4_lon_vmask, - ) - dreg_patch2_4_lat_vmask = np.where( - mask_case3b, - np.where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), - dreg_patch2_4_lat_vmask, - ) + ) - return ( - dreg_patch2_1_lon_vmask, - dreg_patch2_1_lat_vmask, - dreg_patch2_2_lon_vmask, - dreg_patch2_2_lat_vmask, - dreg_patch2_3_lon_vmask, - dreg_patch2_3_lat_vmask, - dreg_patch2_4_lon_vmask, - dreg_patch2_4_lat_vmask, - ) - @classmethod +@pytest.mark.gtfn_too_slow +class TestPrepareFfslFluxAreaPatchesList(stencil_tests.StencilTest): + PROGRAM = prepare_ffsl_flux_area_patches_list + 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", + ) + + @stencil_tests.static_reference def reference( - cls, - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, famask_int, p_vn, ptr_v3_lon, @@ -893,7 +861,7 @@ def reference( 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( + result_tuple = _generate_flux_area_geometry( dreg_patch0_1_lon_dsl, dreg_patch0_1_lat_dsl, dreg_patch0_2_lon_dsl, @@ -992,7 +960,7 @@ def reference( dreg_patch0_3_lat_dsl, dreg_patch0_4_lon_dsl, dreg_patch0_4_lat_dsl, - ) = cls._apply_case1_patch0( + ) = _apply_case1_patch0( mask_case1, lvn_sys_pos, arrival_pts_1_lon_dsl, @@ -1018,7 +986,7 @@ def reference( dreg_patch1_2_lat_vmask, dreg_patch1_3_lon_vmask, dreg_patch1_3_lat_vmask, - ) = cls._apply_case1_patch1( + ) = _apply_case1_patch1( mask_case1, lvn_sys_pos, arrival_pts_1_lon_dsl, @@ -1038,7 +1006,7 @@ def reference( dreg_patch2_2_lat_vmask, dreg_patch2_3_lon_vmask, dreg_patch2_3_lat_vmask, - ) = cls._apply_case1_patch2( + ) = _apply_case1_patch2( mask_case1, lvn_sys_pos, arrival_pts_2_lon_dsl, @@ -1051,10 +1019,14 @@ def reference( # ------------------------------------------------- Case 2a mask_case2a = np.logical_and.reduce( - [lintersect_line1, np.logical_not(lintersect_line2), famask_bool] + [ + lintersect_line1, + np.logical_not(lintersect_line2), + famask_bool, + ] ) # Case 2a - patch 0 - result_tuple_patch0 = cls._apply_case2a_patch0( + result_tuple_patch0 = _apply_case2a_patch0( mask_case2a, lvn_sys_pos, arrival_pts_1_lon_dsl, @@ -1086,7 +1058,7 @@ def reference( dreg_patch0_4_lat_dsl, ) = result_tuple_patch0 # Case 2a - patch 1 - result_tuple_patch1 = cls._apply_case2a_patch1( + result_tuple_patch1 = _apply_case2a_patch1( mask_case2a, lvn_sys_pos, arrival_pts_1_lon_dsl, @@ -1127,10 +1099,14 @@ def reference( # -------------------------------------------------- Case 2b mask_case2b = np.logical_and.reduce( - [lintersect_line2, np.logical_not(lintersect_line1), famask_bool] + [ + lintersect_line2, + np.logical_not(lintersect_line1), + famask_bool, + ] ) # Case 2b - patch 0 - result_tuple_patch0_case2b = cls._apply_case2b_patch0( + result_tuple_patch0_case2b = _apply_case2b_patch0( mask_case2b, lvn_sys_pos, arrival_pts_1_lon_dsl, @@ -1163,7 +1139,7 @@ def reference( ) = result_tuple_patch0_case2b # Case 2b - patch 1 - result_tuple_patch1_case2b = cls._apply_case2b_patch1( + result_tuple_patch1_case2b = _apply_case2b_patch1( mask_case2b, dreg_patch1_1_lon_vmask, dreg_patch1_1_lat_vmask, @@ -1187,7 +1163,7 @@ def reference( ) = result_tuple_patch1_case2b # Case 2b - patch 2 - result_tuple_patch2_case2b = cls._apply_case2b_patch2( + result_tuple_patch2_case2b = _apply_case2b_patch2( mask_case2b, lvn_sys_pos, arrival_pts_2_lon_dsl, @@ -1251,7 +1227,7 @@ def reference( tri_line1_p2_lat, ) # Case 3a - patch 0 - result = cls._apply_case3a_patch0( + result = _apply_case3a_patch0( mask_case3a, arrival_pts_1_lon_dsl, arrival_pts_1_lat_dsl, @@ -1293,7 +1269,7 @@ def reference( dreg_patch1_2_lat_vmask, dreg_patch1_3_lon_vmask, dreg_patch1_3_lat_vmask, - ) = cls._apply_case3a_patch1( + ) = _apply_case3a_patch1( mask_case3a, lvn_sys_pos, arrival_pts_1_lon_dsl, @@ -1356,7 +1332,7 @@ def reference( dreg_patch0_2_lat_dsl, dreg_patch0_3_lon_dsl, dreg_patch0_3_lat_dsl, - ) = cls._apply_case3b_patch0( + ) = _apply_case3b_patch0( mask_case3b, arrival_pts_1_lon_dsl, arrival_pts_1_lat_dsl, @@ -1394,7 +1370,7 @@ def reference( dreg_patch2_3_lat_vmask, dreg_patch2_4_lon_vmask, dreg_patch2_4_lat_vmask, - ) = cls._apply_case3b_patch2( + ) = _apply_case3b_patch2( mask_case3b, arrival_pts_2_lon_dsl, arrival_pts_2_lat_dsl, @@ -1418,11 +1394,20 @@ def reference( # 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_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)] + [ + famask_bool, + np.logical_not(indices_previously_matched), + ] ) # Case 4 - patch 0 - no change # Case 4 - patch 1 @@ -1470,41 +1455,41 @@ def reference( dreg_patch2_4_lat_vmask=dreg_patch2_4_lat_vmask, ) - @pytest.fixture - def input_data(self, grid) -> dict: - famask_int = data_alloc.random_mask(grid, dims.EdgeDim, dims.KDim, dtype=gtx.int32) - p_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - ptr_v3_lon_field = data_alloc.random_field( - grid, dims.EdgeDim, dims.E2CDim, low=0.1, high=1.0 - ) - ptr_v3_lat_field = data_alloc.random_field( - grid, dims.EdgeDim, dims.E2CDim, low=0.1, high=1.0 - ) - tangent_orientation_dsl = data_alloc.random_field(grid, dims.EdgeDim, low=0.1, high=1.0) - dreg_patch0_1_lon_dsl = data_alloc.constant_field(grid, 1.0, dims.EdgeDim, dims.KDim) - dreg_patch0_1_lat_dsl = data_alloc.constant_field(grid, 1.0, dims.EdgeDim, dims.KDim) - dreg_patch0_2_lon_dsl = data_alloc.constant_field(grid, 2.0, dims.EdgeDim, dims.KDim) - dreg_patch0_2_lat_dsl = data_alloc.constant_field(grid, 2.0, dims.EdgeDim, dims.KDim) - dreg_patch0_3_lon_dsl = data_alloc.constant_field(grid, 3.0, dims.EdgeDim, dims.KDim) - dreg_patch0_3_lat_dsl = data_alloc.constant_field(grid, 3.0, dims.EdgeDim, dims.KDim) - dreg_patch0_4_lon_dsl = data_alloc.constant_field(grid, 4.0, dims.EdgeDim, dims.KDim) - dreg_patch0_4_lat_dsl = data_alloc.constant_field(grid, 4.0, dims.EdgeDim, dims.KDim) - dreg_patch1_1_lon_vmask = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - dreg_patch1_1_lat_vmask = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - dreg_patch1_2_lon_vmask = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - dreg_patch1_2_lat_vmask = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - dreg_patch1_3_lon_vmask = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - dreg_patch1_3_lat_vmask = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - dreg_patch1_4_lon_vmask = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - dreg_patch1_4_lat_vmask = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - dreg_patch2_1_lon_vmask = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - dreg_patch2_1_lat_vmask = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - dreg_patch2_2_lon_vmask = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - dreg_patch2_2_lat_vmask = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - dreg_patch2_3_lon_vmask = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - dreg_patch2_3_lat_vmask = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - dreg_patch2_4_lon_vmask = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - dreg_patch2_4_lat_vmask = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid) -> dict: + famask_int = self.data_alloc.random_mask(dims.EdgeDim, dims.KDim, dtype=gtx.int32) + p_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + ptr_v3_lon_field = self.data_alloc.random_field( + dims.EdgeDim, dims.E2CDim, low=0.1, high=1.0 + ) + ptr_v3_lat_field = self.data_alloc.random_field( + dims.EdgeDim, dims.E2CDim, low=0.1, high=1.0 + ) + tangent_orientation_dsl = self.data_alloc.random_field(dims.EdgeDim, low=0.1, high=1.0) + dreg_patch0_1_lon_dsl = self.data_alloc.constant_field(1.0, dims.EdgeDim, dims.KDim) + dreg_patch0_1_lat_dsl = self.data_alloc.constant_field(1.0, dims.EdgeDim, dims.KDim) + dreg_patch0_2_lon_dsl = self.data_alloc.constant_field(2.0, dims.EdgeDim, dims.KDim) + dreg_patch0_2_lat_dsl = self.data_alloc.constant_field(2.0, dims.EdgeDim, dims.KDim) + dreg_patch0_3_lon_dsl = self.data_alloc.constant_field(3.0, dims.EdgeDim, dims.KDim) + dreg_patch0_3_lat_dsl = self.data_alloc.constant_field(3.0, dims.EdgeDim, dims.KDim) + dreg_patch0_4_lon_dsl = self.data_alloc.constant_field(4.0, dims.EdgeDim, dims.KDim) + dreg_patch0_4_lat_dsl = self.data_alloc.constant_field(4.0, dims.EdgeDim, dims.KDim) + dreg_patch1_1_lon_vmask = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + dreg_patch1_1_lat_vmask = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + dreg_patch1_2_lon_vmask = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + dreg_patch1_2_lat_vmask = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + dreg_patch1_3_lon_vmask = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + dreg_patch1_3_lat_vmask = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + dreg_patch1_4_lon_vmask = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + dreg_patch1_4_lat_vmask = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + dreg_patch2_1_lon_vmask = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + dreg_patch2_1_lat_vmask = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + dreg_patch2_2_lon_vmask = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + dreg_patch2_2_lat_vmask = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + dreg_patch2_3_lon_vmask = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + dreg_patch2_3_lat_vmask = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + dreg_patch2_4_lon_vmask = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + dreg_patch2_4_lat_vmask = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) return dict( famask_int=famask_int, p_vn=p_vn, diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_numerical_quadrature_for_cubic_reconstruction.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_numerical_quadrature_for_cubic_reconstruction.py index efe12e6b23..19e6315bf3 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_numerical_quadrature_for_cubic_reconstruction.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_numerical_quadrature_for_cubic_reconstruction.py @@ -16,248 +16,217 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests -class TestPrepareNumericalQuadratureForCubicReconstruction(stencil_tests.StencilTest): - PROGRAM = prepare_numerical_quadrature_for_cubic_reconstruction - 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", - ) +# -- Reference helpers -- - @staticmethod - def _compute_wgt_t_detjac( - wgt_zeta_1, - wgt_zeta_2, - wgt_eta_1, - wgt_eta_2, - dbl_eps, - p_coords_dreg_v_1_x, - p_coords_dreg_v_2_x, - p_coords_dreg_v_3_x, - p_coords_dreg_v_4_x, - p_coords_dreg_v_1_y, - p_coords_dreg_v_2_y, - p_coords_dreg_v_3_y, - p_coords_dreg_v_4_y, - zeta_1, - zeta_2, - zeta_3, - 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 - # 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, +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 + + # 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 * ( + ( + 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, 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_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_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, + - ( + 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_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_x - p_coords_dreg_v_1_x) + - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) ) - wgt_t_detjac_2 = dbl_eps + z_wgt_2 * ( - ( - z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) + ) + wgt_t_detjac_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) ) - 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_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) ) - wgt_t_detjac_4 = dbl_eps + z_wgt_4 * ( - ( - z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) - + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) - ) - * ( - z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) - - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) - ) - - ( - z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) - + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) - ) - * ( - z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) - - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) - ) + - ( + z_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) ) - 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 - + 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_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_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 + ) + 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_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_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_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_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_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_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_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 + ) + 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_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_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_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 + - ( + 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) ) - 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, + * ( + z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) ) + ) + return wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 + - @staticmethod - def _compute_vector_sums( - wgt_t_detjac_1, - wgt_t_detjac_2, - wgt_t_detjac_3, - wgt_t_detjac_4, +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 + + shape_func_3_1 * p_coords_dreg_v_3_x + + shape_func_4_1 * p_coords_dreg_v_4_x + ) + z_gauss_pts_1_y = ( + shape_func_1_1 * p_coords_dreg_v_1_y + + shape_func_2_1 * p_coords_dreg_v_2_y + + shape_func_3_1 * p_coords_dreg_v_3_y + + shape_func_4_1 * p_coords_dreg_v_4_y + ) + z_gauss_pts_2_x = ( + shape_func_1_2 * p_coords_dreg_v_1_x + + shape_func_2_2 * p_coords_dreg_v_2_x + + shape_func_3_2 * p_coords_dreg_v_3_x + + shape_func_4_2 * p_coords_dreg_v_4_x + ) + z_gauss_pts_2_y = ( + shape_func_1_2 * p_coords_dreg_v_1_y + + shape_func_2_2 * p_coords_dreg_v_2_y + + shape_func_3_2 * p_coords_dreg_v_3_y + + shape_func_4_2 * p_coords_dreg_v_4_y + ) + z_gauss_pts_3_x = ( + shape_func_1_3 * p_coords_dreg_v_1_x + + shape_func_2_3 * p_coords_dreg_v_2_x + + shape_func_3_3 * p_coords_dreg_v_3_x + + shape_func_4_3 * p_coords_dreg_v_4_x + ) + z_gauss_pts_3_y = ( + shape_func_1_3 * p_coords_dreg_v_1_y + + shape_func_2_3 * p_coords_dreg_v_2_y + + shape_func_3_3 * p_coords_dreg_v_3_y + + shape_func_4_3 * p_coords_dreg_v_4_y + ) + z_gauss_pts_4_x = ( + shape_func_1_4 * p_coords_dreg_v_1_x + + shape_func_2_4 * p_coords_dreg_v_2_x + + shape_func_3_4 * p_coords_dreg_v_3_x + + shape_func_4_4 * p_coords_dreg_v_4_x + ) + z_gauss_pts_4_y = ( + shape_func_1_4 * p_coords_dreg_v_1_y + + shape_func_2_4 * p_coords_dreg_v_2_y + + shape_func_3_4 * p_coords_dreg_v_3_y + + shape_func_4_4 * p_coords_dreg_v_4_y + ) + return ( z_gauss_pts_1_x, z_gauss_pts_1_y, z_gauss_pts_2_x, @@ -266,79 +235,111 @@ def _compute_vector_sums( z_gauss_pts_3_y, z_gauss_pts_4_x, z_gauss_pts_4_y, - ): - p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 - p_quad_vector_sum_2 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x - ) - p_quad_vector_sum_3 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y - ) - p_quad_vector_sum_4 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x - ) - p_quad_vector_sum_5 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y - ) - p_quad_vector_sum_6 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y - ) - p_quad_vector_sum_7 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x - ) - p_quad_vector_sum_8 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y - ) - p_quad_vector_sum_9 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y - ) - p_quad_vector_sum_10 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y - ) - return ( - p_quad_vector_sum_1, - p_quad_vector_sum_2, - p_quad_vector_sum_3, - p_quad_vector_sum_4, - p_quad_vector_sum_5, - p_quad_vector_sum_6, - p_quad_vector_sum_7, - p_quad_vector_sum_8, - p_quad_vector_sum_9, - p_quad_vector_sum_10, - ) + ) + + +def _compute_vector_sums( + wgt_t_detjac_1, + wgt_t_detjac_2, + wgt_t_detjac_3, + wgt_t_detjac_4, + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, +): + p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 + p_quad_vector_sum_2 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x + ) + p_quad_vector_sum_3 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y + ) + p_quad_vector_sum_4 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_5 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_6 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_7 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_8 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_9 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_10 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + return ( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + ) + + +class TestPrepareNumericalQuadratureForCubicReconstruction(stencil_tests.StencilTest): + PROGRAM = prepare_numerical_quadrature_for_cubic_reconstruction + 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", + ) - @classmethod + @stencil_tests.static_reference def reference( - cls, - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_coords_dreg_v_1_x: np.ndarray, p_coords_dreg_v_2_x: np.ndarray, p_coords_dreg_v_3_x: np.ndarray, @@ -379,7 +380,7 @@ def reference( eps: float, **kwargs: Any, ) -> dict: - wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 = cls._compute_wgt_t_detjac( + wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 = _compute_wgt_t_detjac( wgt_zeta_1, wgt_zeta_2, wgt_eta_1, @@ -412,7 +413,7 @@ def reference( z_gauss_pts_3_y, z_gauss_pts_4_x, z_gauss_pts_4_y, - ) = cls._compute_z_gauss_points( + ) = _compute_z_gauss_points( p_coords_dreg_v_1_x, p_coords_dreg_v_2_x, p_coords_dreg_v_3_x, @@ -450,7 +451,7 @@ def reference( p_quad_vector_sum_8, p_quad_vector_sum_9, p_quad_vector_sum_10, - ) = cls._compute_vector_sums( + ) = _compute_vector_sums( wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, @@ -485,27 +486,27 @@ def reference( p_dreg_area_out=p_dreg_area_out, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - p_coords_dreg_v_1_x = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_2_x = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_3_x = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_4_x = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_1_y = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_2_y = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_3_y = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_4_y = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_1 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_2 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_3 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_4 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_5 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_6 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_7 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_8 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_9 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_10 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_dreg_area_out = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) + p_coords_dreg_v_1_x = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_2_x = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_3_x = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_4_x = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_1_y = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_2_y = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_3_y = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_4_y = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_1 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_2 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_3 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_4 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_5 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_6 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_7 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_8 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_9 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_10 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_dreg_area_out = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) shape_func_1_1 = 0.001 shape_func_2_1 = 0.001 shape_func_3_1 = 0.001 diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_numerical_quadrature_list_for_cubic_reconstruction.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_numerical_quadrature_list_for_cubic_reconstruction.py index 66b376c37a..5110158cf5 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_numerical_quadrature_list_for_cubic_reconstruction.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_numerical_quadrature_list_for_cubic_reconstruction.py @@ -16,285 +16,254 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests -class TestPrepareNumericalQuadratureListForCubicReconstruction(stencil_tests.StencilTest): - PROGRAM = prepare_numerical_quadrature_list_for_cubic_reconstruction - 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", - ) +# -- Reference helpers -- - @staticmethod - def _compute_wgt_t_detjac( - 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, - 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, - 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, 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, - ) +def _compute_wgt_t_detjac( + 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, + 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, + 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, 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 == 1, True, False) + famask_bool = np.where(famask_int == 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) + 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 + wgt_t_detjac_1 = np.where( + famask_bool, + dbl_eps + + z_wgt_1 + * ( + ( + z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) * ( - ( - z_eta_1_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, - ) + 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 + 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, - ) + 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 + 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_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_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_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_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, + ) - return wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 + 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 - + shape_func_3_1 * p_coords_dreg_v_3_x - + shape_func_4_1 * p_coords_dreg_v_4_x - ) - z_gauss_pts_1_y = ( - shape_func_1_1 * p_coords_dreg_v_1_y - + shape_func_2_1 * p_coords_dreg_v_2_y - + shape_func_3_1 * p_coords_dreg_v_3_y - + shape_func_4_1 * p_coords_dreg_v_4_y - ) - z_gauss_pts_2_x = ( - shape_func_1_2 * p_coords_dreg_v_1_x - + shape_func_2_2 * p_coords_dreg_v_2_x - + shape_func_3_2 * p_coords_dreg_v_3_x - + shape_func_4_2 * p_coords_dreg_v_4_x - ) - z_gauss_pts_2_y = ( - shape_func_1_2 * p_coords_dreg_v_1_y - + shape_func_2_2 * p_coords_dreg_v_2_y - + shape_func_3_2 * p_coords_dreg_v_3_y - + shape_func_4_2 * p_coords_dreg_v_4_y - ) - z_gauss_pts_3_x = ( - shape_func_1_3 * p_coords_dreg_v_1_x - + shape_func_2_3 * p_coords_dreg_v_2_x - + shape_func_3_3 * p_coords_dreg_v_3_x - + shape_func_4_3 * p_coords_dreg_v_4_x - ) - z_gauss_pts_3_y = ( - shape_func_1_3 * p_coords_dreg_v_1_y - + shape_func_2_3 * p_coords_dreg_v_2_y - + shape_func_3_3 * p_coords_dreg_v_3_y - + shape_func_4_3 * p_coords_dreg_v_4_y - ) - z_gauss_pts_4_x = ( - shape_func_1_4 * p_coords_dreg_v_1_x - + shape_func_2_4 * p_coords_dreg_v_2_x - + shape_func_3_4 * p_coords_dreg_v_3_x - + shape_func_4_4 * p_coords_dreg_v_4_x - ) - z_gauss_pts_4_y = ( - shape_func_1_4 * p_coords_dreg_v_1_y - + shape_func_2_4 * p_coords_dreg_v_2_y - + shape_func_3_4 * p_coords_dreg_v_3_y - + shape_func_4_4 * p_coords_dreg_v_4_y - ) - return ( - z_gauss_pts_1_x, - z_gauss_pts_1_y, - z_gauss_pts_2_x, - z_gauss_pts_2_y, - z_gauss_pts_3_x, - z_gauss_pts_3_y, - z_gauss_pts_4_x, - z_gauss_pts_4_y, - ) - @staticmethod - def _compute_vector_sums( - wgt_t_detjac_1, - wgt_t_detjac_2, - wgt_t_detjac_3, - wgt_t_detjac_4, +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 + + shape_func_3_1 * p_coords_dreg_v_3_x + + shape_func_4_1 * p_coords_dreg_v_4_x + ) + z_gauss_pts_1_y = ( + shape_func_1_1 * p_coords_dreg_v_1_y + + shape_func_2_1 * p_coords_dreg_v_2_y + + shape_func_3_1 * p_coords_dreg_v_3_y + + shape_func_4_1 * p_coords_dreg_v_4_y + ) + z_gauss_pts_2_x = ( + shape_func_1_2 * p_coords_dreg_v_1_x + + shape_func_2_2 * p_coords_dreg_v_2_x + + shape_func_3_2 * p_coords_dreg_v_3_x + + shape_func_4_2 * p_coords_dreg_v_4_x + ) + z_gauss_pts_2_y = ( + shape_func_1_2 * p_coords_dreg_v_1_y + + shape_func_2_2 * p_coords_dreg_v_2_y + + shape_func_3_2 * p_coords_dreg_v_3_y + + shape_func_4_2 * p_coords_dreg_v_4_y + ) + z_gauss_pts_3_x = ( + shape_func_1_3 * p_coords_dreg_v_1_x + + shape_func_2_3 * p_coords_dreg_v_2_x + + shape_func_3_3 * p_coords_dreg_v_3_x + + shape_func_4_3 * p_coords_dreg_v_4_x + ) + z_gauss_pts_3_y = ( + shape_func_1_3 * p_coords_dreg_v_1_y + + shape_func_2_3 * p_coords_dreg_v_2_y + + shape_func_3_3 * p_coords_dreg_v_3_y + + shape_func_4_3 * p_coords_dreg_v_4_y + ) + z_gauss_pts_4_x = ( + shape_func_1_4 * p_coords_dreg_v_1_x + + shape_func_2_4 * p_coords_dreg_v_2_x + + shape_func_3_4 * p_coords_dreg_v_3_x + + shape_func_4_4 * p_coords_dreg_v_4_x + ) + z_gauss_pts_4_y = ( + shape_func_1_4 * p_coords_dreg_v_1_y + + shape_func_2_4 * p_coords_dreg_v_2_y + + shape_func_3_4 * p_coords_dreg_v_3_y + + shape_func_4_4 * p_coords_dreg_v_4_y + ) + return ( z_gauss_pts_1_x, z_gauss_pts_1_y, z_gauss_pts_2_x, @@ -303,79 +272,111 @@ def _compute_vector_sums( z_gauss_pts_3_y, z_gauss_pts_4_x, z_gauss_pts_4_y, - ): - p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 - p_quad_vector_sum_2 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x - ) - p_quad_vector_sum_3 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y - ) - p_quad_vector_sum_4 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x - ) - p_quad_vector_sum_5 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y - ) - p_quad_vector_sum_6 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y - ) - p_quad_vector_sum_7 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x - ) - p_quad_vector_sum_8 = ( - wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y - ) - p_quad_vector_sum_9 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y - ) - p_quad_vector_sum_10 = ( - wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y - + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y - + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y - + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y - ) - return ( - p_quad_vector_sum_1, - p_quad_vector_sum_2, - p_quad_vector_sum_3, - p_quad_vector_sum_4, - p_quad_vector_sum_5, - p_quad_vector_sum_6, - p_quad_vector_sum_7, - p_quad_vector_sum_8, - p_quad_vector_sum_9, - p_quad_vector_sum_10, - ) + ) + + +def _compute_vector_sums( + wgt_t_detjac_1, + wgt_t_detjac_2, + wgt_t_detjac_3, + wgt_t_detjac_4, + z_gauss_pts_1_x, + z_gauss_pts_1_y, + z_gauss_pts_2_x, + z_gauss_pts_2_y, + z_gauss_pts_3_x, + z_gauss_pts_3_y, + z_gauss_pts_4_x, + z_gauss_pts_4_y, +): + p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 + p_quad_vector_sum_2 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x + ) + p_quad_vector_sum_3 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y + ) + p_quad_vector_sum_4 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_5 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_6 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_7 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_8 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_9 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_10 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + return ( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + ) + + +class TestPrepareNumericalQuadratureListForCubicReconstruction(stencil_tests.StencilTest): + PROGRAM = prepare_numerical_quadrature_list_for_cubic_reconstruction + 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", + ) - @classmethod + @stencil_tests.static_reference def reference( - cls, - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, famask_int: np.ndarray, p_coords_dreg_v_1_x: np.ndarray, p_coords_dreg_v_2_x: np.ndarray, @@ -417,7 +418,7 @@ def reference( dbl_eps: float, **kwargs: Any, ) -> dict: - wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 = cls._compute_wgt_t_detjac( + wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 = _compute_wgt_t_detjac( p_coords_dreg_v_1_x, p_coords_dreg_v_2_x, p_coords_dreg_v_3_x, @@ -451,7 +452,7 @@ def reference( z_gauss_pts_3_y, z_gauss_pts_4_x, z_gauss_pts_4_y, - ) = cls._compute_z_gauss_points( + ) = _compute_z_gauss_points( p_coords_dreg_v_1_x, p_coords_dreg_v_2_x, p_coords_dreg_v_3_x, @@ -489,7 +490,7 @@ def reference( p_quad_vector_sum_8, p_quad_vector_sum_9, p_quad_vector_sum_10, - ) = cls._compute_vector_sums( + ) = _compute_vector_sums( wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, @@ -519,29 +520,29 @@ def reference( p_dreg_area=p_dreg_area, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - famask_int = data_alloc.constant_field(grid, 1, dims.EdgeDim, dims.KDim, dtype=gtx.int32) - p_coords_dreg_v_1_x = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_2_x = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_3_x = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_4_x = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_1_y = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_2_y = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_3_y = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_coords_dreg_v_4_y = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_dreg_area_in = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_1 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_2 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_3 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_4 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_5 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_6 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_7 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_8 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_9 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_quad_vector_sum_10 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - p_dreg_area = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) + famask_int = self.data_alloc.constant_field(1, dims.EdgeDim, dims.KDim, dtype=gtx.int32) + p_coords_dreg_v_1_x = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_2_x = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_3_x = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_4_x = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_1_y = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_2_y = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_3_y = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_coords_dreg_v_4_y = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_dreg_area_in = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_1 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_2 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_3 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_4 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_5 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_6 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_7 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_8 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_9 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_quad_vector_sum_10 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + p_dreg_area = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) shape_func_1_1 = 0.001 shape_func_2_1 = 0.001 shape_func_3_1 = 0.001 diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_reconstruct_cubic_coefficients_svd.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_reconstruct_cubic_coefficients_svd.py index ed85115d46..bd2d24b91c 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_reconstruct_cubic_coefficients_svd.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_reconstruct_cubic_coefficients_svd.py @@ -5,13 +5,13 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np import pytest -import icon4py.model.common.utils.data_allocation as data_alloc from icon4py.model.atmosphere.advection.stencils.reconstruct_cubic_coefficients_svd import ( reconstruct_cubic_coefficients_svd, ) @@ -36,9 +36,9 @@ class TestReconstructCubicCoefficientsSvd(stencil_tests.StencilTest): "p_coeff_10_dsl", ) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_cc: np.ndarray, lsq_pseudoinv_1: np.ndarray, lsq_pseudoinv_2: np.ndarray, @@ -70,6 +70,7 @@ def reference( p_coeff_10_dsl: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) p_coeff_1_dsl_cp = p_coeff_1_dsl.copy() p_coeff_2_dsl_cp = p_coeff_2_dsl.copy() p_coeff_3_dsl_cp = p_coeff_3_dsl.copy() @@ -81,7 +82,7 @@ def reference( p_coeff_9_dsl_cp = p_coeff_9_dsl.copy() p_coeff_10_dsl_cp = p_coeff_10_dsl.copy() - c2e2c2e2c = connectivities[dims.C2E2C2E2CDim] + c2e2c2e2c = connectivities[dims.C2E2C2E2C] 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) @@ -304,38 +305,38 @@ def reference( p_coeff_10_dsl=p_coeff_10_dsl, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - p_cc = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - lsq_pseudoinv_1_field = data_alloc.random_field(grid, dims.CellDim, dims.C2E2C2E2CDim) - lsq_pseudoinv_2_field = data_alloc.random_field(grid, dims.CellDim, dims.C2E2C2E2CDim) - lsq_pseudoinv_3_field = data_alloc.random_field(grid, dims.CellDim, dims.C2E2C2E2CDim) - lsq_pseudoinv_4_field = data_alloc.random_field(grid, dims.CellDim, dims.C2E2C2E2CDim) - lsq_pseudoinv_5_field = data_alloc.random_field(grid, dims.CellDim, dims.C2E2C2E2CDim) - lsq_pseudoinv_6_field = data_alloc.random_field(grid, dims.CellDim, dims.C2E2C2E2CDim) - lsq_pseudoinv_7_field = data_alloc.random_field(grid, dims.CellDim, dims.C2E2C2E2CDim) - lsq_pseudoinv_8_field = data_alloc.random_field(grid, dims.CellDim, dims.C2E2C2E2CDim) - lsq_pseudoinv_9_field = data_alloc.random_field(grid, dims.CellDim, dims.C2E2C2E2CDim) + p_cc = self.data_alloc.random_field(dims.CellDim, dims.KDim) + lsq_pseudoinv_1_field = self.data_alloc.random_field(dims.CellDim, dims.C2E2C2E2CDim) + lsq_pseudoinv_2_field = self.data_alloc.random_field(dims.CellDim, dims.C2E2C2E2CDim) + lsq_pseudoinv_3_field = self.data_alloc.random_field(dims.CellDim, dims.C2E2C2E2CDim) + lsq_pseudoinv_4_field = self.data_alloc.random_field(dims.CellDim, dims.C2E2C2E2CDim) + lsq_pseudoinv_5_field = self.data_alloc.random_field(dims.CellDim, dims.C2E2C2E2CDim) + lsq_pseudoinv_6_field = self.data_alloc.random_field(dims.CellDim, dims.C2E2C2E2CDim) + lsq_pseudoinv_7_field = self.data_alloc.random_field(dims.CellDim, dims.C2E2C2E2CDim) + lsq_pseudoinv_8_field = self.data_alloc.random_field(dims.CellDim, dims.C2E2C2E2CDim) + lsq_pseudoinv_9_field = self.data_alloc.random_field(dims.CellDim, dims.C2E2C2E2CDim) - lsq_moments_1 = data_alloc.random_field(grid, dims.CellDim) - lsq_moments_2 = data_alloc.random_field(grid, dims.CellDim) - lsq_moments_3 = data_alloc.random_field(grid, dims.CellDim) - lsq_moments_4 = data_alloc.random_field(grid, dims.CellDim) - lsq_moments_5 = data_alloc.random_field(grid, dims.CellDim) - lsq_moments_6 = data_alloc.random_field(grid, dims.CellDim) - lsq_moments_7 = data_alloc.random_field(grid, dims.CellDim) - lsq_moments_8 = data_alloc.random_field(grid, dims.CellDim) - lsq_moments_9 = data_alloc.random_field(grid, dims.CellDim) - p_coeff_1_dsl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - p_coeff_2_dsl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - p_coeff_3_dsl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - p_coeff_4_dsl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - p_coeff_5_dsl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - p_coeff_6_dsl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - p_coeff_7_dsl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - p_coeff_8_dsl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - p_coeff_9_dsl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - p_coeff_10_dsl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + lsq_moments_1 = self.data_alloc.random_field(dims.CellDim) + lsq_moments_2 = self.data_alloc.random_field(dims.CellDim) + lsq_moments_3 = self.data_alloc.random_field(dims.CellDim) + lsq_moments_4 = self.data_alloc.random_field(dims.CellDim) + lsq_moments_5 = self.data_alloc.random_field(dims.CellDim) + lsq_moments_6 = self.data_alloc.random_field(dims.CellDim) + lsq_moments_7 = self.data_alloc.random_field(dims.CellDim) + lsq_moments_8 = self.data_alloc.random_field(dims.CellDim) + lsq_moments_9 = self.data_alloc.random_field(dims.CellDim) + p_coeff_1_dsl = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + p_coeff_2_dsl = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + p_coeff_3_dsl = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + p_coeff_4_dsl = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + p_coeff_5_dsl = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + p_coeff_6_dsl = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + p_coeff_7_dsl = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + p_coeff_8_dsl = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + p_coeff_9_dsl = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + p_coeff_10_dsl = self.data_alloc.zero_field(dims.CellDim, dims.KDim) cell_domain = h_grid.domain(dims.CellDim) horizontal_start = grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_theta_and_exner.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_theta_and_exner.py index 883fabb34f..a408ffe026 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_theta_and_exner.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_theta_and_exner.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -17,8 +18,7 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.type_alias import vpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests from .test_calculate_nabla2_for_z import calculate_nabla2_for_z_numpy from .test_calculate_nabla2_of_theta import calculate_nabla2_of_theta_numpy @@ -31,13 +31,13 @@ @pytest.mark.skip_value_error @pytest.mark.uses_as_offset @pytest.mark.embedded_remap_error -class TestApplyDiffusionToThetaAndExner(StencilTest): +class TestApplyDiffusionToThetaAndExner(stencil_tests.StencilTest): PROGRAM = apply_diffusion_to_theta_and_exner OUTPUTS = ("theta_v", "exner") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, kh_smag_e: np.ndarray, inv_dual_edge_length: np.ndarray, theta_v_in: np.ndarray, @@ -53,6 +53,7 @@ def reference( apply_zdiffusion_t: bool, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) kwargs_2 = {k: v for k, v in kwargs.items() if k != "theta_v"} # remove unused kwargs z_nabla2_e = np.zeros_like(kh_smag_e) @@ -81,14 +82,14 @@ def reference( return dict(theta_v=theta_v, exner=exner) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid): - kh_smag_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - inv_dual_edge_length = data_alloc.random_field(grid, dims.EdgeDim) - theta_v_in = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) - zd_vertoffset = data_alloc.zero_field( - grid, dims.CellDim, dims.C2E2CDim, dims.KDim, dtype=gtx.int32 + kh_smag_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + inv_dual_edge_length = self.data_alloc.random_field(dims.EdgeDim) + theta_v_in = self.data_alloc.random_field(dims.CellDim, dims.KDim) + geofac_div = self.data_alloc.random_field(dims.CellDim, dims.C2EDim) + zd_vertoffset = self.data_alloc.zero_field( + dims.CellDim, dims.C2E2CDim, dims.KDim, dtype=gtx.int32 ) rng = np.random.default_rng() for k in range(grid.num_levels): @@ -98,13 +99,13 @@ def input_data(self, grid: base.Grid): high=grid.num_levels - k - 1, size=(zd_vertoffset.shape[0], zd_vertoffset.shape[1]), ) - zd_diffcoef = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - geofac_n2s_c = data_alloc.random_field(grid, dims.CellDim) - geofac_n2s_nbh = data_alloc.random_field(grid, dims.CellDim, dims.C2E2CDim) - vcoef = data_alloc.random_field(grid, dims.CellDim, dims.C2E2CDim, dims.KDim) - area = data_alloc.random_field(grid, dims.CellDim) - theta_v = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - exner = data_alloc.random_field(grid, dims.CellDim, dims.KDim) + zd_diffcoef = self.data_alloc.random_field(dims.CellDim, dims.KDim) + geofac_n2s_c = self.data_alloc.random_field(dims.CellDim) + geofac_n2s_nbh = self.data_alloc.random_field(dims.CellDim, dims.C2E2CDim) + vcoef = self.data_alloc.random_field(dims.CellDim, dims.C2E2CDim, dims.KDim) + area = self.data_alloc.random_field(dims.CellDim) + theta_v = self.data_alloc.random_field(dims.CellDim, dims.KDim) + exner = self.data_alloc.random_field(dims.CellDim, dims.KDim) rd_o_cvd = vpfloat("5.0") apply_zdiffusion_t = True diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py index debdae974f..acf187fd9e 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -14,8 +15,7 @@ from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_vn import apply_diffusion_to_vn from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest +from icon4py.model.testing import stencil_tests from .test_apply_nabla2_and_nabla4_global_to_vn import apply_nabla2_and_nabla4_global_to_vn_numpy from .test_apply_nabla2_and_nabla4_to_vn import apply_nabla2_and_nabla4_to_vn_numpy @@ -27,12 +27,12 @@ @pytest.mark.uses_concat_where @pytest.mark.continuous_benchmarking -class TestApplyDiffusionToVn(StencilTest): +class TestApplyDiffusionToVn(stencil_tests.StencilTest): PROGRAM = apply_diffusion_to_vn OUTPUTS = ("vn",) STATIC_PARAMS = { - StandardStaticVariants.NONE: (), - StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", "start_2nd_nudge_line_idx_e", @@ -40,16 +40,16 @@ class TestApplyDiffusionToVn(StencilTest): "vertical_end", "limited_area", ), - StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( "vertical_start", "vertical_end", "limited_area", ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, u_vert: np.ndarray, v_vert: np.ndarray, primal_normal_vert_v1: np.ndarray, @@ -68,6 +68,7 @@ def reference( limited_area: bool, **kwargs: Any, ): + connectivities = stencil_tests.connectivities_asnumpy(grid) edge = np.arange(area_edge.shape[0]) vn_cp = vn.copy() z_nabla4_e2 = calculate_nabla4_numpy( @@ -115,23 +116,23 @@ def reference( return dict(vn=vn) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - u_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim) - v_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim) + u_vert = self.data_alloc.random_field(dims.VertexDim, dims.KDim) + v_vert = self.data_alloc.random_field(dims.VertexDim, dims.KDim) - primal_normal_vert_v1 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2C2VDim) - primal_normal_vert_v2 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2C2VDim) + primal_normal_vert_v1 = self.data_alloc.random_field(dims.EdgeDim, dims.E2C2VDim) + primal_normal_vert_v2 = self.data_alloc.random_field(dims.EdgeDim, dims.E2C2VDim) - inv_vert_vert_length = data_alloc.random_field(grid, dims.EdgeDim) - inv_primal_edge_length = data_alloc.random_field(grid, dims.EdgeDim) + inv_vert_vert_length = self.data_alloc.random_field(dims.EdgeDim) + inv_primal_edge_length = self.data_alloc.random_field(dims.EdgeDim) - area_edge = data_alloc.random_field(grid, dims.EdgeDim) - kh_smag_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - z_nabla2_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - diff_multfac_vn = data_alloc.random_field(grid, dims.KDim) - vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - nudgecoeff_e = data_alloc.random_field(grid, dims.EdgeDim) + area_edge = self.data_alloc.random_field(dims.EdgeDim) + kh_smag_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + z_nabla2_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + diff_multfac_vn = self.data_alloc.random_field(dims.KDim) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + nudgecoeff_e = self.data_alloc.random_field(dims.EdgeDim) limited_area = grid.limited_area if hasattr(grid, "limited_area") else True fac_bdydiff_v = 5.0 diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py index 4e9652ac76..d61e67efff 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py @@ -5,6 +5,9 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from collections.abc import Mapping +from typing import cast + import gt4py.next as gtx import numpy as np import pytest @@ -15,9 +18,7 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing import definitions -from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest +from icon4py.model.testing import definitions, stencil_tests from .test_apply_nabla2_to_w import apply_nabla2_to_w_numpy from .test_apply_nabla2_to_w_in_upper_damping_layer import ( @@ -31,12 +32,12 @@ @pytest.mark.embedded_remap_error @pytest.mark.continuous_benchmarking -class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence(StencilTest): +class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence(stencil_tests.StencilTest): PROGRAM = apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence OUTPUTS = ("w", "dwdx", "dwdy") STATIC_PARAMS = { - StandardStaticVariants.NONE: (), - StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", "halo_idx", @@ -46,7 +47,7 @@ class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence(StencilTes "nrdmax", "type_shear", ), - StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( "vertical_start", "vertical_end", "nrdmax", @@ -54,9 +55,9 @@ class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence(StencilTes ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, area, geofac_n2s, geofac_grg_x, @@ -76,6 +77,7 @@ def reference( vertical_end, **kwargs, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) k = np.arange(w_old.shape[1]) cell = np.arange(w_old.shape[0]) reshaped_k = k[np.newaxis, :] @@ -118,7 +120,7 @@ def reference( out_dwdy[subset] = dwdy[subset] return dict(w=out_w, dwdx=out_dwdx, dwdy=out_dwdy) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: nrdmax = 13 cell_domain = h_grid.domain(dims.CellDim) @@ -138,17 +140,17 @@ def _get_start_index_for_w_diffusion() -> int32: horizontal_start = _get_start_index_for_w_diffusion() horizontal_end = grid.end_index(cell_domain(h_grid.Zone.HALO)) - geofac_grg_x = random_field(grid, dims.CellDim, dims.C2E2CODim) - geofac_grg_y = random_field(grid, dims.CellDim, dims.C2E2CODim) - diff_multfac_n2w = random_field(grid, dims.KDim) - area = random_field(grid, dims.CellDim) - geofac_n2s = random_field(grid, dims.CellDim, dims.C2E2CODim) - w_old = random_field(grid, dims.CellDim, dims.KDim) + geofac_grg_x = self.data_alloc.random_field(dims.CellDim, dims.C2E2CODim) + geofac_grg_y = self.data_alloc.random_field(dims.CellDim, dims.C2E2CODim) + diff_multfac_n2w = self.data_alloc.random_field(dims.KDim) + area = self.data_alloc.random_field(dims.CellDim) + geofac_n2s = self.data_alloc.random_field(dims.CellDim, dims.C2E2CODim) + w_old = self.data_alloc.random_field(dims.CellDim, dims.KDim) diff_multfac_w = 5.0 - w = zero_field(grid, dims.CellDim, dims.KDim) - dwdx = random_field(grid, dims.CellDim, dims.KDim) - dwdy = random_field(grid, dims.CellDim, dims.KDim) + w = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + dwdx = self.data_alloc.random_field(dims.CellDim, dims.KDim) + dwdy = self.data_alloc.random_field(dims.CellDim, dims.KDim) return dict( area=area, diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_and_nabla4_global_to_vn.py index c2ef62b392..e4631d36bd 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_and_nabla4_global_to_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_and_nabla4_global_to_vn.py @@ -17,8 +17,7 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def apply_nabla2_and_nabla4_global_to_vn_numpy( @@ -30,18 +29,18 @@ def apply_nabla2_and_nabla4_global_to_vn_numpy( return vn -class TestApplyNabla2AndNabla4GlobalToVn(StencilTest): +class TestApplyNabla2AndNabla4GlobalToVn(stencil_tests.StencilTest): PROGRAM = apply_nabla2_and_nabla4_global_to_vn OUTPUTS = ("vn",) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid): - area_edge = random_field(grid, dims.EdgeDim, dtype=wpfloat) - kh_smag_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - z_nabla2_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - z_nabla4_e2 = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - diff_multfac_vn = random_field(grid, dims.KDim, dtype=wpfloat) - vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + area_edge = self.data_alloc.random_field(dims.EdgeDim, dtype=wpfloat) + kh_smag_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + z_nabla2_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + z_nabla4_e2 = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + diff_multfac_vn = self.data_alloc.random_field(dims.KDim, dtype=wpfloat) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) return dict( area_edge=area_edge, @@ -56,9 +55,9 @@ def input_data(self, grid: base.Grid): vertical_end=gtx.int32(grid.num_levels), ) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, area_edge: np.ndarray, kh_smag_e: np.ndarray, z_nabla2_e: np.ndarray, diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_and_nabla4_to_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_and_nabla4_to_vn.py index 3b60be5291..41fabbd4f8 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_and_nabla4_to_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_and_nabla4_to_vn.py @@ -13,9 +13,9 @@ apply_nabla2_and_nabla4_to_vn, ) from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def apply_nabla2_and_nabla4_to_vn_numpy( @@ -38,19 +38,19 @@ def apply_nabla2_and_nabla4_to_vn_numpy( return vn -class TestApplyNabla2AndNabla4ToVn(StencilTest): +class TestApplyNabla2AndNabla4ToVn(stencil_tests.StencilTest): PROGRAM = apply_nabla2_and_nabla4_to_vn OUTPUTS = ("vn",) - @pytest.fixture - def input_data(self, grid): - area_edge = random_field(grid, dims.EdgeDim, dtype=wpfloat) - kh_smag_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - z_nabla2_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - z_nabla4_e2 = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - diff_multfac_vn = random_field(grid, dims.KDim, dtype=wpfloat) - nudgecoeff_e = random_field(grid, dims.EdgeDim, dtype=wpfloat) - vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): + area_edge = self.data_alloc.random_field(dims.EdgeDim, dtype=wpfloat) + kh_smag_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + z_nabla2_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + z_nabla4_e2 = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + diff_multfac_vn = self.data_alloc.random_field(dims.KDim, dtype=wpfloat) + nudgecoeff_e = self.data_alloc.random_field(dims.EdgeDim, dtype=wpfloat) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) nudgezone_diff = vpfloat("9.0") return dict( @@ -68,9 +68,9 @@ def input_data(self, grid): vertical_end=gtx.int32(grid.num_levels), ) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, area_edge: np.ndarray, kh_smag_e: np.ndarray, z_nabla2_e: np.ndarray, diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py index 0de1666efd..3cae8c7198 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py @@ -17,8 +17,7 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def apply_nabla2_to_vn_in_lateral_boundary_numpy( @@ -29,16 +28,16 @@ def apply_nabla2_to_vn_in_lateral_boundary_numpy( return vn -class TestApplyNabla2ToVnInLateralBoundary(StencilTest): +class TestApplyNabla2ToVnInLateralBoundary(stencil_tests.StencilTest): PROGRAM = apply_nabla2_to_vn_in_lateral_boundary OUTPUTS = ("vn",) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid): fac_bdydiff_v = wpfloat("5.0") - z_nabla2_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - area_edge = random_field(grid, dims.EdgeDim, dtype=wpfloat) - vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + z_nabla2_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + area_edge = self.data_alloc.random_field(dims.EdgeDim, dtype=wpfloat) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) return dict( fac_bdydiff_v=fac_bdydiff_v, z_nabla2_e=z_nabla2_e, @@ -50,9 +49,9 @@ def input_data(self, grid: base.Grid): vertical_end=gtx.int32(grid.num_levels), ) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_nabla2_e: np.ndarray, area_edge: np.ndarray, vn: np.ndarray, diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_to_w.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_to_w.py index cf505ee6ca..3512a03ae8 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_to_w.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_to_w.py @@ -5,26 +5,29 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from collections.abc import Mapping +from typing import cast + import gt4py.next as gtx import numpy as np import pytest from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w import apply_nabla2_to_w from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def apply_nabla2_to_w_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], area: np.ndarray, z_nabla2_c: np.ndarray, geofac_n2s: np.ndarray, w: np.ndarray, diff_multfac_w: float, ) -> np.ndarray: - c2e2cO = connectivities[dims.C2E2CODim] + c2e2cO = connectivities[dims.C2E2CO] geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) area = np.expand_dims(area, axis=-1) w = w - diff_multfac_w * area * area * np.sum( @@ -34,13 +37,13 @@ def apply_nabla2_to_w_numpy( @pytest.mark.embedded_remap_error -class TestMoApplyNabla2ToW(StencilTest): +class TestMoApplyNabla2ToW(stencil_tests.StencilTest): PROGRAM = apply_nabla2_to_w OUTPUTS = ("w",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, area: np.ndarray, z_nabla2_c: np.ndarray, geofac_n2s: np.ndarray, @@ -48,15 +51,16 @@ def reference( diff_multfac_w: float, **kwargs, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) w = apply_nabla2_to_w_numpy(connectivities, area, z_nabla2_c, geofac_n2s, w, diff_multfac_w) return dict(w=w) - @pytest.fixture - def input_data(self, grid): - area = random_field(grid, dims.CellDim, dtype=wpfloat) - z_nabla2_c = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - geofac_n2s = random_field(grid, dims.CellDim, dims.C2E2CODim, dtype=wpfloat) - w = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): + area = self.data_alloc.random_field(dims.CellDim, dtype=wpfloat) + z_nabla2_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + geofac_n2s = self.data_alloc.random_field(dims.CellDim, dims.C2E2CODim, dtype=wpfloat) + w = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) return dict( area=area, z_nabla2_c=z_nabla2_c, diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py index 74d532867b..22940587fa 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py @@ -15,8 +15,7 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def apply_nabla2_to_w_in_upper_damping_layer_numpy( @@ -30,16 +29,16 @@ def apply_nabla2_to_w_in_upper_damping_layer_numpy( return w -class TestApplyNabla2ToWInUpperDampingLayer(StencilTest): +class TestApplyNabla2ToWInUpperDampingLayer(stencil_tests.StencilTest): PROGRAM = apply_nabla2_to_w_in_upper_damping_layer OUTPUTS = ("w",) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid): - w = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - diff_multfac_n2w = random_field(grid, dims.KDim, dtype=wpfloat) - cell_area = random_field(grid, dims.CellDim, dtype=wpfloat) - z_nabla2_c = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + w = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + diff_multfac_n2w = self.data_alloc.random_field(dims.KDim, dtype=wpfloat) + cell_area = self.data_alloc.random_field(dims.CellDim, dtype=wpfloat) + z_nabla2_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) return dict( w=w, @@ -52,9 +51,9 @@ def input_data(self, grid: base.Grid): vertical_end=gtx.int32(grid.num_levels), ) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, w: np.ndarray, diff_multfac_n2w: np.ndarray, cell_area: np.ndarray, diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_diagnostics_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_diagnostics_for_turbulence.py index 142ee4d36a..18d53a4882 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_diagnostics_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_diagnostics_for_turbulence.py @@ -14,9 +14,9 @@ calculate_diagnostics_for_turbulence, ) from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import vpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def calculate_diagnostics_for_turbulence_numpy( @@ -29,13 +29,13 @@ def calculate_diagnostics_for_turbulence_numpy( return div_ic, hdef_ic -class TestCalculateDiagnosticsForTurbulence(StencilTest): +class TestCalculateDiagnosticsForTurbulence(stencil_tests.StencilTest): PROGRAM = calculate_diagnostics_for_turbulence OUTPUTS = ("div_ic", "hdef_ic") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, wgtfac_c: np.ndarray, div: np.ndarray, kh_c: np.ndarray, @@ -47,11 +47,11 @@ def reference( ) return dict(div_ic=div_ic, hdef_ic=hdef_ic) - @pytest.fixture - def input_data(self, grid): - wgtfac_c = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - div = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - kh_c = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - div_ic = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - hdef_ic = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): + wgtfac_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + div = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + kh_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + div_ic = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) + hdef_ic = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) return dict(wgtfac_c=wgtfac_c, div=div, kh_c=kh_c, div_ic=div_ic, hdef_ic=hdef_ic) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py index fa6659390b..6c1481276b 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py @@ -5,6 +5,9 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from collections.abc import Mapping +from typing import cast + import gt4py.next as gtx import numpy as np import pytest @@ -13,18 +16,18 @@ calculate_horizontal_gradients_for_turbulence, ) from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def calculate_horizontal_gradients_for_turbulence_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], w: np.ndarray, geofac_grg_x: np.ndarray, geofac_grg_y: np.ndarray, ) -> tuple[np.ndarray, np.ndarray]: - c2e2cO = connectivities[dims.C2E2CODim] + c2e2cO = connectivities[dims.C2E2CO] geofac_grg_x = np.expand_dims(geofac_grg_x, axis=-1) dwdx = np.sum(np.where((c2e2cO != -1)[:, :, np.newaxis], geofac_grg_x * w[c2e2cO], 0.0), axis=1) @@ -34,30 +37,31 @@ def calculate_horizontal_gradients_for_turbulence_numpy( @pytest.mark.embedded_remap_error -class TestCalculateHorizontalGradientsForTurbulence(StencilTest): +class TestCalculateHorizontalGradientsForTurbulence(stencil_tests.StencilTest): PROGRAM = calculate_horizontal_gradients_for_turbulence OUTPUTS = ("dwdx", "dwdy") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, w: np.ndarray, geofac_grg_x: np.ndarray, geofac_grg_y: np.ndarray, **kwargs, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) dwdx, dwdy = calculate_horizontal_gradients_for_turbulence_numpy( connectivities, w, geofac_grg_x, geofac_grg_y ) return dict(dwdx=dwdx, dwdy=dwdy) - @pytest.fixture - def input_data(self, grid): - w = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - geofac_grg_x = data_alloc.random_field(grid, dims.CellDim, dims.C2E2CODim, dtype=wpfloat) - geofac_grg_y = data_alloc.random_field(grid, dims.CellDim, dims.C2E2CODim, dtype=wpfloat) - dwdx = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - dwdy = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): + w = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + geofac_grg_x = self.data_alloc.random_field(dims.CellDim, dims.C2E2CODim, dtype=wpfloat) + geofac_grg_y = self.data_alloc.random_field(dims.CellDim, dims.C2E2CODim, dtype=wpfloat) + dwdx = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) + dwdy = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) return dict( w=w, diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index e9cd84bde4..e86eda2d40 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -6,6 +6,9 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from collections.abc import Mapping +from typing import cast + import gt4py.next as gtx import numpy as np import pytest @@ -15,8 +18,7 @@ ) from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base, horizontal as h_grid -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import definitions, stencil_tests +from icon4py.model.testing import stencil_tests @pytest.mark.continuous_benchmarking @@ -37,9 +39,9 @@ class TestCalculateNabla2AndSmagCoefficientsForVn(stencil_tests.StencilTest): ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, diff_multfac_smag: np.ndarray, tangent_orientation: np.ndarray, inv_primal_edge_length: np.ndarray, @@ -59,7 +61,8 @@ def reference( vertical_end: int, **kwargs, ) -> dict: - e2c2v = connectivities[dims.E2C2VDim] + connectivities = stencil_tests.connectivities_asnumpy(grid) + e2c2v = connectivities[dims.E2C2V] u_vert_e2c2v = u_vert[e2c2v] v_vert_e2c2v = v_vert[e2c2v] @@ -173,34 +176,34 @@ def reference( return dict(kh_smag_e=kh_smag_e_out, kh_smag_ec=kh_smag_ec_out, z_nabla2_e=z_nabla2_e_out) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - u_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim, dtype=ta.vpfloat) - v_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim, dtype=ta.vpfloat) + u_vert = self.data_alloc.random_field(dims.VertexDim, dims.KDim, dtype=ta.vpfloat) + v_vert = self.data_alloc.random_field(dims.VertexDim, dims.KDim, dtype=ta.vpfloat) smag_offset = ta.vpfloat("9.0") - diff_multfac_smag = data_alloc.random_field(grid, dims.KDim, dtype=ta.vpfloat) - tangent_orientation = data_alloc.random_sign(grid, dims.EdgeDim, dtype=ta.wpfloat) - vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) - smag_limit = data_alloc.random_field(grid, dims.KDim, dtype=ta.vpfloat) - inv_vert_vert_length = data_alloc.random_field(grid, dims.EdgeDim, dtype=ta.wpfloat) - inv_primal_edge_length = data_alloc.random_field(grid, dims.EdgeDim, dtype=ta.wpfloat) - - primal_normal_vert_x = data_alloc.random_field( - grid, dims.EdgeDim, dims.E2C2VDim, dtype=ta.wpfloat + diff_multfac_smag = self.data_alloc.random_field(dims.KDim, dtype=ta.vpfloat) + tangent_orientation = self.data_alloc.random_sign(dims.EdgeDim, dtype=ta.wpfloat) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + smag_limit = self.data_alloc.random_field(dims.KDim, dtype=ta.vpfloat) + inv_vert_vert_length = self.data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat) + inv_primal_edge_length = self.data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat) + + primal_normal_vert_x = self.data_alloc.random_field( + dims.EdgeDim, dims.E2C2VDim, dtype=ta.wpfloat ) - primal_normal_vert_y = data_alloc.random_field( - grid, dims.EdgeDim, dims.E2C2VDim, dtype=ta.wpfloat + primal_normal_vert_y = self.data_alloc.random_field( + dims.EdgeDim, dims.E2C2VDim, dtype=ta.wpfloat ) - dual_normal_vert_x = data_alloc.random_field( - grid, dims.EdgeDim, dims.E2C2VDim, dtype=ta.wpfloat + dual_normal_vert_x = self.data_alloc.random_field( + dims.EdgeDim, dims.E2C2VDim, dtype=ta.wpfloat ) - dual_normal_vert_y = data_alloc.random_field( - grid, dims.EdgeDim, dims.E2C2VDim, dtype=ta.wpfloat + dual_normal_vert_y = self.data_alloc.random_field( + dims.EdgeDim, dims.E2C2VDim, dtype=ta.wpfloat ) - z_nabla2_e = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) - kh_smag_e = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) - kh_smag_ec = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + z_nabla2_e = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + kh_smag_e = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + kh_smag_ec = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) edge_domain = h_grid.domain(dims.EdgeDim) horizontal_start = grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5)) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_for_w.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_for_w.py index eb405afab2..799318bc76 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_for_w.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_for_w.py @@ -5,6 +5,9 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from collections.abc import Mapping +from typing import cast + import gt4py.next as gtx import numpy as np import pytest @@ -14,14 +17,13 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils.data_allocation import constant_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def calculate_nabla2_for_w_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], w: np.ndarray, geofac_n2s: np.ndarray + connectivities: Mapping[gtx.FieldOffset, np.ndarray], w: np.ndarray, geofac_n2s: np.ndarray ) -> np.ndarray: - c2e2cO = connectivities[dims.C2E2CODim] + c2e2cO = connectivities[dims.C2E2CO] geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) z_nabla2_c = np.sum( np.where((c2e2cO != -1)[:, :, np.newaxis], w[c2e2cO] * geofac_n2s, 0), axis=1 @@ -30,25 +32,26 @@ def calculate_nabla2_for_w_numpy( @pytest.mark.embedded_remap_error -class TestCalculateNabla2ForW(StencilTest): +class TestCalculateNabla2ForW(stencil_tests.StencilTest): PROGRAM = calculate_nabla2_for_w OUTPUTS = ("z_nabla2_c",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, w: np.ndarray, geofac_n2s: np.ndarray, **kwargs, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) z_nabla2_c = calculate_nabla2_for_w_numpy(connectivities, w, geofac_n2s) return dict(z_nabla2_c=z_nabla2_c) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - w = constant_field(grid, 1.0, dims.CellDim, dims.KDim) - geofac_n2s = constant_field(grid, 2.0, dims.CellDim, dims.C2E2CODim) - z_nabla2_c = zero_field(grid, dims.CellDim, dims.KDim) + w = self.data_alloc.constant_field(1.0, dims.CellDim, dims.KDim) + geofac_n2s = self.data_alloc.constant_field(2.0, dims.CellDim, dims.C2E2CODim) + z_nabla2_c = self.data_alloc.zero_field(dims.CellDim, dims.KDim) return dict( w=w, diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_for_z.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_for_z.py index 8a4d9fb58e..0900810a7f 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_for_z.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_for_z.py @@ -5,6 +5,9 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from collections.abc import Mapping +from typing import cast + import gt4py.next as gtx import numpy as np import pytest @@ -13,14 +16,13 @@ calculate_nabla2_for_z, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import horizontal as h_grid +from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def calculate_nabla2_for_z_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], kh_smag_e: np.ndarray, inv_dual_edge_length: np.ndarray, theta_v: np.ndarray, @@ -29,7 +31,7 @@ def calculate_nabla2_for_z_numpy( ) -> np.ndarray: inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) - theta_v_e2c = theta_v[connectivities[dims.E2CDim]] + theta_v_e2c = theta_v[connectivities[dims.E2C]] theta_v_weighted = theta_v_e2c[:, 1] - theta_v_e2c[:, 0] z_nabla2_e = kh_smag_e * inv_dual_edge_length * theta_v_weighted @@ -37,30 +39,31 @@ def calculate_nabla2_for_z_numpy( return z_nabla2_e -class TestCalculateNabla2ForZ(StencilTest): +class TestCalculateNabla2ForZ(stencil_tests.StencilTest): PROGRAM = calculate_nabla2_for_z OUTPUTS = ("z_nabla2_e",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, kh_smag_e: np.ndarray, inv_dual_edge_length: np.ndarray, theta_v: np.ndarray, z_nabla2_e: np.ndarray, **kwargs, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) z_nabla2_e = calculate_nabla2_for_z_numpy( connectivities, kh_smag_e, inv_dual_edge_length, theta_v, z_nabla2_e, **kwargs ) return dict(z_nabla2_e=z_nabla2_e) - @pytest.fixture - def input_data(self, grid): - kh_smag_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - inv_dual_edge_length = random_field(grid, dims.EdgeDim, dtype=wpfloat) - theta_v = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - z_nabla2_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): + kh_smag_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + inv_dual_edge_length = self.data_alloc.random_field(dims.EdgeDim, dtype=wpfloat) + theta_v = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + z_nabla2_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) return dict( kh_smag_e=kh_smag_e, diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_of_theta.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_of_theta.py index 5355ee6371..0242c0d515 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_of_theta.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_of_theta.py @@ -5,6 +5,9 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from collections.abc import Mapping +from typing import cast + import gt4py.next as gtx import numpy as np import pytest @@ -14,14 +17,15 @@ ) from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests def calculate_nabla2_of_theta_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], z_nabla2_e: np.ndarray, geofac_div: np.ndarray + connectivities: Mapping[gtx.FieldOffset, np.ndarray], + z_nabla2_e: np.ndarray, + geofac_div: np.ndarray, ) -> np.ndarray: - c2e = connectivities[dims.C2EDim] + c2e = connectivities[dims.C2E] geofac_div = np.expand_dims(geofac_div, axis=-1) z_temp = np.sum(z_nabla2_e[c2e] * geofac_div, axis=1) # sum along edge dimension return z_temp @@ -31,22 +35,23 @@ class TestCalculateNabla2OfTheta(stencil_tests.StencilTest): PROGRAM = calculate_nabla2_of_theta OUTPUTS = ("z_temp",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_nabla2_e: np.ndarray, geofac_div: np.ndarray, **kwargs, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) z_temp = calculate_nabla2_of_theta_numpy(connectivities, z_nabla2_e, geofac_div) return dict(z_temp=z_temp) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - z_nabla2_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) - geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim, dtype=ta.wpfloat) + z_nabla2_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + geofac_div = self.data_alloc.random_field(dims.CellDim, dims.C2EDim, dtype=ta.wpfloat) - z_temp = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_temp = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) return dict( z_nabla2_e=z_nabla2_e, diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py index dc8f9ac4ba..806729f957 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py @@ -5,18 +5,21 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from collections.abc import Mapping +from typing import cast + import gt4py.next as gtx import numpy as np import pytest -import icon4py.model.common.utils.data_allocation as data_alloc from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import calculate_nabla4 from icon4py.model.common import dimension as dims, type_alias as ta -from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest +from icon4py.model.common.grid import base +from icon4py.model.testing import stencil_tests def calculate_nabla4_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], u_vert: np.ndarray, v_vert: np.ndarray, primal_normal_vert_v1: np.ndarray, @@ -25,7 +28,7 @@ def calculate_nabla4_numpy( inv_vert_vert_length: np.ndarray, inv_primal_edge_length: np.ndarray, ) -> np.ndarray: - e2c2v = connectivities[dims.E2C2VDim] + e2c2v = connectivities[dims.E2C2V] u_vert_e2c2v = u_vert[e2c2v] v_vert_e2c2v = v_vert[e2c2v] @@ -56,26 +59,26 @@ def calculate_nabla4_numpy( @pytest.mark.continuous_benchmarking -class TestCalculateNabla4(StencilTest): +class TestCalculateNabla4(stencil_tests.StencilTest): PROGRAM = calculate_nabla4 OUTPUTS = ("z_nabla4_e2",) STATIC_PARAMS = { - StandardStaticVariants.NONE: (), - StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", "vertical_start", "vertical_end", ), - StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( "vertical_start", "vertical_end", ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, u_vert: np.ndarray, v_vert: np.ndarray, primal_normal_vert_v1: np.ndarray, @@ -85,6 +88,7 @@ def reference( inv_primal_edge_length: np.ndarray, **kwargs, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) z_nabla4_e2 = calculate_nabla4_numpy( connectivities, u_vert, @@ -97,23 +101,23 @@ def reference( ) return dict(z_nabla4_e2=z_nabla4_e2) - @pytest.fixture - def input_data(self, grid) -> dict: - u_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim, dtype=ta.vpfloat) - v_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim, dtype=ta.vpfloat) + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid) -> dict: + u_vert = self.data_alloc.random_field(dims.VertexDim, dims.KDim, dtype=ta.vpfloat) + v_vert = self.data_alloc.random_field(dims.VertexDim, dims.KDim, dtype=ta.vpfloat) - primal_normal_vert_v1 = data_alloc.random_field( - grid, dims.EdgeDim, dims.E2C2VDim, dtype=ta.wpfloat + primal_normal_vert_v1 = self.data_alloc.random_field( + dims.EdgeDim, dims.E2C2VDim, dtype=ta.wpfloat ) - primal_normal_vert_v2 = data_alloc.random_field( - grid, dims.EdgeDim, dims.E2C2VDim, dtype=ta.wpfloat + primal_normal_vert_v2 = self.data_alloc.random_field( + dims.EdgeDim, dims.E2C2VDim, dtype=ta.wpfloat ) - z_nabla2_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) - inv_vert_vert_length = data_alloc.random_field(grid, dims.EdgeDim, dtype=ta.wpfloat) - inv_primal_edge_length = data_alloc.random_field(grid, dims.EdgeDim, dtype=ta.wpfloat) + z_nabla2_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + inv_vert_vert_length = self.data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat) + inv_primal_edge_length = self.data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat) - z_nabla4_e2 = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + z_nabla4_e2 = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) return dict( u_vert=u_vert, diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py index 50524d92d4..6413a25b7e 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py @@ -6,6 +6,8 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause import math +from collections.abc import Mapping +from typing import cast import gt4py.next as gtx import numpy as np @@ -17,22 +19,22 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base from icon4py.model.common.type_alias import vpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestEnhanceDiffusionCoefficientForGridPointColdPools(StencilTest): +class TestEnhanceDiffusionCoefficientForGridPointColdPools(stencil_tests.StencilTest): PROGRAM = enhance_diffusion_coefficient_for_grid_point_cold_pools OUTPUTS = ("kh_smag_e",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, kh_smag_e: np.ndarray, enh_diffu_3d: np.ndarray, **kwargs, ) -> dict: - e2c = connectivities[dims.E2CDim] + connectivities = stencil_tests.connectivities_asnumpy(grid) + e2c = connectivities[dims.E2C] kh_smag_e = np.maximum( kh_smag_e, np.max( @@ -42,10 +44,10 @@ def reference( ) return dict(kh_smag_e=kh_smag_e) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - kh_smag_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - enh_diffu_3d = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + kh_smag_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + enh_diffu_3d = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) return dict( kh_smag_e=kh_smag_e, diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py index fd0b4563a9..2fea8d1e5b 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py @@ -5,6 +5,9 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from collections.abc import Mapping +from typing import cast + import gt4py.next as gtx import numpy as np import pytest @@ -15,25 +18,25 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests @pytest.mark.embedded_remap_error -class TestTemporaryFieldForGridPointColdPoolsEnhancement(StencilTest): +class TestTemporaryFieldForGridPointColdPoolsEnhancement(stencil_tests.StencilTest): PROGRAM = temporary_field_for_grid_point_cold_pools_enhancement OUTPUTS = ("enh_diffu_3d",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, theta_v: np.ndarray, theta_ref_mc: np.ndarray, thresh_tdiff, smallest_vpfloat, **kwargs, ) -> dict: - c2e2c = connectivities[dims.C2E2CDim] + connectivities = stencil_tests.connectivities_asnumpy(grid) + c2e2c = connectivities[dims.C2E2C] tdiff = ( theta_v - np.sum(np.where((c2e2c != -1)[:, :, np.newaxis], theta_v[c2e2c], 0), axis=1) / 3 @@ -52,11 +55,11 @@ def reference( return dict(enh_diffu_3d=enh_diffu_3d) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - theta_v = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - theta_ref_mc = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - enh_diffu_3d = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + theta_v = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + theta_ref_mc = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + enh_diffu_3d = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) thresh_tdiff = wpfloat("5.0") smallest_vpfloat = -np.finfo(vpfloat).max diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_temporary_fields_for_turbulence_diagnostics.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_temporary_fields_for_turbulence_diagnostics.py index 0352652ad2..5e471447e3 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_temporary_fields_for_turbulence_diagnostics.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_temporary_fields_for_turbulence_diagnostics.py @@ -5,6 +5,9 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from collections.abc import Mapping +from typing import cast + import gt4py.next as gtx import numpy as np import pytest @@ -14,7 +17,6 @@ ) from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -22,9 +24,9 @@ class TestTemporaryFieldsForTurbulenceDiagnostics(stencil_tests.StencilTest): PROGRAM = temporary_fields_for_turbulence_diagnostics OUTPUTS = ("div", "kh_c") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, kh_smag_ec: np.ndarray, vn: np.ndarray, e_bln_c_s: np.ndarray, @@ -32,7 +34,8 @@ def reference( diff_multfac_smag: np.ndarray, **kwargs, ) -> dict: - c2e = connectivities[dims.C2EDim] + connectivities = stencil_tests.connectivities_asnumpy(grid) + c2e = connectivities[dims.C2E] geofac_div = np.expand_dims(geofac_div, axis=-1) e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) @@ -46,18 +49,18 @@ def reference( return dict(div=div, kh_c=kh_c) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid): - vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) - geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim, dtype=ta.wpfloat) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + geofac_div = self.data_alloc.random_field(dims.CellDim, dims.C2EDim, dtype=ta.wpfloat) - kh_smag_ec = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) - e_bln_c_s = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim, dtype=ta.wpfloat) + kh_smag_ec = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + e_bln_c_s = self.data_alloc.random_field(dims.CellDim, dims.C2EDim, dtype=ta.wpfloat) - diff_multfac_smag = data_alloc.random_field(grid, dims.KDim, dtype=ta.vpfloat) + diff_multfac_smag = self.data_alloc.random_field(dims.KDim, dtype=ta.vpfloat) - kh_c = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - div = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) + kh_c = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + div = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) return dict( kh_smag_ec=kh_smag_ec, diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index 37797d3b6a..9dc2ee3e7b 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -5,6 +5,9 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from collections.abc import Mapping +from typing import cast + import gt4py.next as gtx import numpy as np import pytest @@ -13,13 +16,13 @@ truly_horizontal_diffusion_nabla_of_theta_over_steep_points, ) from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def truly_horizontal_diffusion_nabla_of_theta_over_steep_points_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], zd_vertoffset: np.ndarray, zd_diffcoef: np.ndarray, geofac_n2s_c: np.ndarray, @@ -29,7 +32,7 @@ def truly_horizontal_diffusion_nabla_of_theta_over_steep_points_numpy( z_temp: np.ndarray, **kwargs, ) -> np.ndarray: - c2e2c = connectivities[dims.C2E2CDim] + c2e2c = connectivities[dims.C2E2C] full_shape = vcoef.shape geofac_n2s_nbh = np.expand_dims(geofac_n2s_nbh, axis=2) @@ -57,13 +60,13 @@ def truly_horizontal_diffusion_nabla_of_theta_over_steep_points_numpy( @pytest.mark.uses_as_offset -class TestTrulyHorizontalDiffusionNablaOfThetaOverSteepPoints(StencilTest): +class TestTrulyHorizontalDiffusionNablaOfThetaOverSteepPoints(stencil_tests.StencilTest): PROGRAM = truly_horizontal_diffusion_nabla_of_theta_over_steep_points OUTPUTS = ("z_temp",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, zd_vertoffset: np.ndarray, zd_diffcoef: np.ndarray, geofac_n2s_c: np.ndarray, @@ -73,6 +76,7 @@ def reference( z_temp: np.ndarray, **kwargs, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) z_temp = truly_horizontal_diffusion_nabla_of_theta_over_steep_points_numpy( connectivities, zd_vertoffset, @@ -85,9 +89,11 @@ def reference( ) return dict(z_temp=z_temp) - @pytest.fixture - def input_data(self, grid): - zd_vertoffset = zero_field(grid, dims.CellDim, dims.C2E2CDim, dims.KDim, dtype=gtx.int32) + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): + zd_vertoffset = self.data_alloc.zero_field( + dims.CellDim, dims.C2E2CDim, dims.KDim, dtype=gtx.int32 + ) rng = np.random.default_rng() for k in range(grid.num_levels): # construct offsets that reach all k-levels except the last (because we are using the entries of this field with `+1`) @@ -97,12 +103,12 @@ def input_data(self, grid): size=(zd_vertoffset.ndarray.shape[0], zd_vertoffset.ndarray.shape[1]), ) - zd_diffcoef = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - geofac_n2s_c = random_field(grid, dims.CellDim, dtype=wpfloat) - geofac_n2s_nbh = random_field(grid, dims.CellDim, dims.C2E2CDim, dtype=wpfloat) - vcoef = random_field(grid, dims.CellDim, dims.C2E2CDim, dims.KDim, dtype=wpfloat) - theta_v = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - z_temp = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + zd_diffcoef = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + geofac_n2s_c = self.data_alloc.random_field(dims.CellDim, dtype=wpfloat) + geofac_n2s_nbh = self.data_alloc.random_field(dims.CellDim, dims.C2E2CDim, dtype=wpfloat) + vcoef = self.data_alloc.random_field(dims.CellDim, dims.C2E2CDim, dims.KDim, dtype=wpfloat) + theta_v = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + z_temp = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) return dict( zd_vertoffset=zd_vertoffset, diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_update_theta_and_exner.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_update_theta_and_exner.py index d7bce46e8a..a8d7902c64 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_update_theta_and_exner.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_update_theta_and_exner.py @@ -15,8 +15,7 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def update_theta_and_exner_numpy( @@ -33,13 +32,13 @@ def update_theta_and_exner_numpy( return theta_v, exner -class TestUpdateThetaAndExner(StencilTest): +class TestUpdateThetaAndExner(stencil_tests.StencilTest): PROGRAM = update_theta_and_exner OUTPUTS = ("theta_v", "exner") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_temp: np.ndarray, area: np.ndarray, theta_v: np.ndarray, @@ -50,12 +49,12 @@ def reference( theta_v, exner = update_theta_and_exner_numpy(z_temp, area, theta_v, exner, rd_o_cvd) return dict(theta_v=theta_v, exner=exner) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - z_temp = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - area = random_field(grid, dims.CellDim, dtype=wpfloat) - theta_v = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - exner = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) + z_temp = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + area = self.data_alloc.random_field(dims.CellDim, dtype=wpfloat) + theta_v = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + exner = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) rd_o_cvd = vpfloat("5.0") return dict( diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_accumulate_prep_adv_fields.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_accumulate_prep_adv_fields.py index c58454e15d..39e7007219 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_accumulate_prep_adv_fields.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_accumulate_prep_adv_fields.py @@ -19,8 +19,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def accumulate_prep_adv_fields_numpy( @@ -36,13 +35,13 @@ def accumulate_prep_adv_fields_numpy( return vn_traj, mass_flx_me -class TestAccumulatePrepAdvFields(StencilTest): +class TestAccumulatePrepAdvFields(stencil_tests.StencilTest): PROGRAM = accumulate_prep_adv_fields OUTPUTS = ("vn_traj", "mass_flx_me") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_vn_avg: np.ndarray, mass_fl_e: np.ndarray, vn_traj: np.ndarray, @@ -60,12 +59,12 @@ def reference( return dict(vn_traj=vn_traj, mass_flx_me=mass_flx_me) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - mass_fl_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - mass_flx_me = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - z_vn_avg = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - vn_traj = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + mass_fl_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + mass_flx_me = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + z_vn_avg = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + vn_traj = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) r_nsubsteps = wpfloat("9.0") return dict( diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_analysis_increments_from_data_assimilation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_analysis_increments_from_data_assimilation.py index 5bf6bf9525..590fa80aeb 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_analysis_increments_from_data_assimilation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_analysis_increments_from_data_assimilation.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -17,12 +18,11 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def add_analysis_increments_from_data_assimilation_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], z_rho_expl: np.ndarray, rho_incr: np.ndarray, z_exner_expl: np.ndarray, @@ -34,13 +34,13 @@ def add_analysis_increments_from_data_assimilation_numpy( return (z_rho_expl, z_exner_expl) -class TestAddAnalysisIncrementsFromDataAssimilation(StencilTest): +class TestAddAnalysisIncrementsFromDataAssimilation(stencil_tests.StencilTest): PROGRAM = add_analysis_increments_from_data_assimilation OUTPUTS = ("z_rho_expl", "z_exner_expl") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_rho_expl: np.ndarray, rho_incr: np.ndarray, z_exner_expl: np.ndarray, @@ -48,6 +48,7 @@ def reference( iau_wgt_dyn: float, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) z_rho_expl, z_exner_expl = add_analysis_increments_from_data_assimilation_numpy( connectivities, z_rho_expl=z_rho_expl, @@ -58,12 +59,12 @@ def reference( ) return dict(z_rho_expl=z_rho_expl, z_exner_expl=z_exner_expl) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - z_exner_expl = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - exner_incr = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - z_rho_expl = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - rho_incr = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_exner_expl = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + exner_incr = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_rho_expl = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + rho_incr = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) iau_wgt_dyn = ta.wpfloat("8.0") return dict( diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_analysis_increments_to_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_analysis_increments_to_vn.py index 8f67332647..7dfa1d26f8 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_analysis_increments_to_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_analysis_increments_to_vn.py @@ -19,8 +19,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def add_analysis_increments_to_vn_numpy( @@ -30,13 +29,13 @@ def add_analysis_increments_to_vn_numpy( return vn -class TestAddAnalysisIncrementsToVn(StencilTest): +class TestAddAnalysisIncrementsToVn(stencil_tests.StencilTest): PROGRAM = add_analysis_increments_to_vn OUTPUTS = ("vn",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, vn_incr: np.ndarray, vn: np.ndarray, iau_wgt_dyn: ta.wpfloat, @@ -45,10 +44,10 @@ def reference( vn = add_analysis_increments_to_vn_numpy(vn_incr, vn, iau_wgt_dyn) return dict(vn=vn) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - vn_incr = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + vn_incr = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) iau_wgt_dyn = wpfloat("5.0") return dict( diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_extra_diffusion_for_normal_wind_tendency_approaching_cfl.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_extra_diffusion_for_normal_wind_tendency_approaching_cfl.py index 6c69b6b94e..90b7c19bfb 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_extra_diffusion_for_normal_wind_tendency_approaching_cfl.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_extra_diffusion_for_normal_wind_tendency_approaching_cfl.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -17,12 +18,11 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], levelmask: np.ndarray, c_lin_e: np.ndarray, z_w_con_c_full: np.ndarray, @@ -50,7 +50,7 @@ def add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_numpy( tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) inv_primal_edge_length = np.expand_dims(inv_primal_edge_length, axis=-1) - e2c = connectivities[dims.E2CDim] + e2c = connectivities[dims.E2C] w_con_e = np.where( (levelmask_offset_0) | (levelmask_offset_1), np.sum( @@ -73,8 +73,8 @@ def add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_numpy( ), difcoef, ) - e2v = connectivities[dims.E2VDim] - e2c2eo = connectivities[dims.E2C2EODim] + e2v = connectivities[dims.E2V] + e2c2eo = connectivities[dims.E2C2EO] ddt_vn_apc = np.where( ((levelmask_offset_0) | (levelmask_offset_1)) & (np.abs(w_con_e) > cfl_w_limit * ddqz_z_full_e), @@ -98,25 +98,25 @@ def add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_numpy( @pytest.mark.embedded_remap_error -class TestAddExtraDiffusionForNormalWindTendencyApproachingCfl(StencilTest): +class TestAddExtraDiffusionForNormalWindTendencyApproachingCfl(stencil_tests.StencilTest): PROGRAM = add_extra_diffusion_for_normal_wind_tendency_approaching_cfl OUTPUTS = ("ddt_vn_apc",) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - levelmask = data_alloc.random_mask(grid, dims.KDim, extend={dims.KDim: 1}) - c_lin_e = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat) - z_w_con_c_full = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - ddqz_z_full_e = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat, low=0.0 + levelmask = self.data_alloc.random_mask(dims.KDim, extend={dims.KDim: 1}) + c_lin_e = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat) + z_w_con_c_full = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + ddqz_z_full_e = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim, dtype=ta.vpfloat, low=0.0 ) - area_edge = data_alloc.random_field(grid, dims.EdgeDim) - tangent_orientation = data_alloc.random_field(grid, dims.EdgeDim) - inv_primal_edge_length = data_alloc.random_field(grid, dims.EdgeDim) - zeta = data_alloc.random_field(grid, dims.VertexDim, dims.KDim, dtype=ta.vpfloat) - geofac_grdiv = data_alloc.random_field(grid, dims.EdgeDim, dims.E2C2EODim) - vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - ddt_vn_apc = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + area_edge = self.data_alloc.random_field(dims.EdgeDim) + tangent_orientation = self.data_alloc.random_field(dims.EdgeDim) + inv_primal_edge_length = self.data_alloc.random_field(dims.EdgeDim) + zeta = self.data_alloc.random_field(dims.VertexDim, dims.KDim, dtype=ta.vpfloat) + geofac_grdiv = self.data_alloc.random_field(dims.EdgeDim, dims.E2C2EODim) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + ddt_vn_apc = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) cfl_w_limit = ta.vpfloat("4.0") scalfac_exdiff = 6.0 dtime = 2.0 @@ -141,9 +141,9 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala vertical_end=gtx.int32(grid.num_levels), ) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, levelmask: np.ndarray, c_lin_e: np.ndarray, z_w_con_c_full: np.ndarray, @@ -160,6 +160,7 @@ def reference( dtime: ta.wpfloat, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) ddt_vn_apc = add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_numpy( connectivities, levelmask, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_extra_diffusion_for_w_con_approaching_cfl.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_extra_diffusion_for_w_con_approaching_cfl.py index 1e855f11ac..d6d4070d81 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_extra_diffusion_for_w_con_approaching_cfl.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_extra_diffusion_for_w_con_approaching_cfl.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -19,12 +20,11 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, random_mask -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def add_extra_diffusion_for_w_con_approaching_cfl_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], cfl_clipping: np.ndarray, owner_mask: np.ndarray, z_w_con_c: np.ndarray, @@ -51,7 +51,7 @@ def add_extra_diffusion_for_w_con_approaching_cfl_numpy( 0, ) - c2e2cO = connectivities[dims.C2E2CODim] + c2e2cO = connectivities[dims.C2E2CO] ddt_w_adv = np.where( (cfl_clipping == 1) & (owner_mask == 1), ddt_w_adv @@ -71,13 +71,13 @@ def add_extra_diffusion_for_w_con_approaching_cfl_numpy( @pytest.mark.embedded_remap_error -class TestAddExtraDiffusionForWConApproachingCfl(StencilTest): +class TestAddExtraDiffusionForWConApproachingCfl(stencil_tests.StencilTest): PROGRAM = add_extra_diffusion_for_w_con_approaching_cfl OUTPUTS = ("ddt_w_adv",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, cfl_clipping: np.ndarray, owner_mask: np.ndarray, z_w_con_c: np.ndarray, @@ -91,6 +91,7 @@ def reference( dtime: ta.wpfloat, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) ddt_w_adv = add_extra_diffusion_for_w_con_approaching_cfl_numpy( connectivities, cfl_clipping, @@ -107,16 +108,16 @@ def reference( ) return dict(ddt_w_adv=ddt_w_adv) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - cfl_clipping = random_mask(grid, dims.CellDim, dims.KDim) - owner_mask = random_mask(grid, dims.CellDim) - z_w_con_c = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - ddqz_z_half = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - area = random_field(grid, dims.CellDim, dtype=wpfloat) - geofac_n2s = random_field(grid, dims.CellDim, dims.C2E2CODim, dtype=wpfloat) - w = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - ddt_w_adv = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + cfl_clipping = self.data_alloc.random_mask(dims.CellDim, dims.KDim) + owner_mask = self.data_alloc.random_mask(dims.CellDim) + z_w_con_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + ddqz_z_half = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + area = self.data_alloc.random_field(dims.CellDim, dtype=wpfloat) + geofac_n2s = self.data_alloc.random_field(dims.CellDim, dims.C2E2CODim, dtype=wpfloat) + w = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + ddt_w_adv = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) scalfac_exdiff = wpfloat("10.0") cfl_w_limit = vpfloat("3.0") dtime = wpfloat("2.0") diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_interpolated_horizontal_advection_of_w.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_interpolated_horizontal_advection_of_w.py index e27f6d9ee7..180101a8d8 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_interpolated_horizontal_advection_of_w.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_interpolated_horizontal_advection_of_w.py @@ -5,13 +5,13 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np import pytest -import icon4py.model.common.utils.data_allocation as data_alloc from icon4py.model.atmosphere.dycore.stencils.add_interpolated_horizontal_advection_of_w import ( add_interpolated_horizontal_advection_of_w, ) @@ -22,14 +22,14 @@ def add_interpolated_horizontal_advection_of_w_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], e_bln_c_s: np.ndarray, z_v_grad_w: np.ndarray, ddt_w_adv: np.ndarray, **kwargs: Any, ) -> np.ndarray: e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) - c2e = connectivities[dims.C2EDim] + c2e = connectivities[dims.C2E] ddt_w_adv = ddt_w_adv + np.sum( z_v_grad_w[c2e] * e_bln_c_s, @@ -42,24 +42,25 @@ class TestAddInterpolatedHorizontalAdvectionOfW(stencil_tests.StencilTest): PROGRAM = add_interpolated_horizontal_advection_of_w OUTPUTS = ("ddt_w_adv",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, e_bln_c_s: np.ndarray, z_v_grad_w: np.ndarray, ddt_w_adv: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) ddt_w_adv = add_interpolated_horizontal_advection_of_w_numpy( connectivities, e_bln_c_s, z_v_grad_w, ddt_w_adv ) return dict(ddt_w_adv=ddt_w_adv) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - z_v_grad_w = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) - e_bln_c_s = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim, dtype=ta.wpfloat) - ddt_w_adv = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_v_grad_w = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + e_bln_c_s = self.data_alloc.random_field(dims.CellDim, dims.C2EDim, dtype=ta.wpfloat) + ddt_w_adv = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) return dict( e_bln_c_s=e_bln_c_s, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_temporal_tendencies_to_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_temporal_tendencies_to_vn.py index 861d66d4ba..034f540381 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_temporal_tendencies_to_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_temporal_tendencies_to_vn.py @@ -18,8 +18,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests dycore_consts: Final = constants.PhysicsConstants() @@ -39,13 +38,13 @@ def add_temporal_tendencies_to_vn_numpy( return vn_nnew -class TestAddTemporalTendenciesToVn(StencilTest): +class TestAddTemporalTendenciesToVn(stencil_tests.StencilTest): PROGRAM = add_temporal_tendencies_to_vn OUTPUTS = ("vn_nnew",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, vn_nnow: np.ndarray, ddt_vn_apc_ntl1: np.ndarray, ddt_vn_phy: np.ndarray, @@ -59,15 +58,15 @@ def reference( ) return dict(vn_nnew=vn_nnew) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: dtime = wpfloat("10.0") - vn_nnow = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - ddt_vn_apc_ntl1 = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - ddt_vn_phy = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - z_theta_v_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - z_gradh_exner = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - vn_nnew = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + vn_nnow = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + ddt_vn_apc_ntl1 = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + ddt_vn_phy = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + z_theta_v_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + z_gradh_exner = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + vn_nnew = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) return dict( vn_nnow=vn_nnow, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_temporal_tendencies_to_vn_by_interpolating_between_time_levels.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_temporal_tendencies_to_vn_by_interpolating_between_time_levels.py index b84b1bbc74..98260056e0 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_temporal_tendencies_to_vn_by_interpolating_between_time_levels.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_temporal_tendencies_to_vn_by_interpolating_between_time_levels.py @@ -18,8 +18,7 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def add_temporal_tendencies_to_vn_by_interpolating_between_time_levels_numpy( @@ -43,13 +42,13 @@ def add_temporal_tendencies_to_vn_by_interpolating_between_time_levels_numpy( return vn_nnew -class TestAddTemporalTendenciesToVnByInterpolatingBetweenTimeLevels(StencilTest): +class TestAddTemporalTendenciesToVnByInterpolatingBetweenTimeLevels(stencil_tests.StencilTest): PROGRAM = add_temporal_tendencies_to_vn_by_interpolating_between_time_levels OUTPUTS = ("vn_nnew",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, vn_nnow: np.ndarray, ddt_vn_apc_ntl1: np.ndarray, ddt_vn_apc_ntl2: np.ndarray, @@ -76,15 +75,15 @@ def reference( ) return dict(vn_nnew=vn_nnew) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - vn_nnow = random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) - ddt_vn_apc_ntl1 = random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) - ddt_vn_apc_ntl2 = random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) - ddt_vn_phy = random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) - z_theta_v_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) - z_gradh_exner = random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) - vn_nnew = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + vn_nnow = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + ddt_vn_apc_ntl1 = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + ddt_vn_apc_ntl2 = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + ddt_vn_phy = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + z_theta_v_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + z_gradh_exner = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + vn_nnew = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) dtime = ta.wpfloat("5.0") wgt_nnow_vel = ta.wpfloat("8.0") wgt_nnew_vel = ta.wpfloat("7.0") diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_vertical_wind_derivative_to_divergence_damping.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_vertical_wind_derivative_to_divergence_damping.py index 723e7952e6..6d56c8946b 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_vertical_wind_derivative_to_divergence_damping.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_vertical_wind_derivative_to_divergence_damping.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -17,12 +18,11 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests def add_vertical_wind_derivative_to_divergence_damping_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], hmask_dd3d: np.ndarray, scalfac_dd3d: np.ndarray, inv_dual_edge_length: np.ndarray, @@ -33,7 +33,7 @@ def add_vertical_wind_derivative_to_divergence_damping_numpy( hmask_dd3d = np.expand_dims(hmask_dd3d, axis=-1) inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) - e2c = connectivities[dims.E2CDim] + e2c = connectivities[dims.E2C] z_dwdz_dd_e2c = z_dwdz_dd[e2c] z_dwdz_dd_weighted = z_dwdz_dd_e2c[:, 1] - z_dwdz_dd_e2c[:, 0] @@ -48,9 +48,9 @@ class TestAddVerticalWindDerivativeToDivergenceDamping(stencil_tests.StencilTest PROGRAM = add_vertical_wind_derivative_to_divergence_damping OUTPUTS = ("z_graddiv_vn",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, hmask_dd3d: np.ndarray, scalfac_dd3d: np.ndarray, inv_dual_edge_length: np.ndarray, @@ -58,6 +58,7 @@ def reference( z_graddiv_vn: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) z_graddiv_vn = add_vertical_wind_derivative_to_divergence_damping_numpy( connectivities, hmask_dd3d, @@ -68,13 +69,13 @@ def reference( ) return dict(z_graddiv_vn=z_graddiv_vn) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - hmask_dd3d = data_alloc.random_field(grid, dims.EdgeDim, dtype=ta.wpfloat) - scalfac_dd3d = data_alloc.random_field(grid, dims.KDim, dtype=ta.wpfloat) - inv_dual_edge_length = data_alloc.random_field(grid, dims.EdgeDim, dtype=ta.wpfloat) - z_dwdz_dd = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - z_graddiv_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + hmask_dd3d = self.data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat) + scalfac_dd3d = self.data_alloc.random_field(dims.KDim, dtype=ta.wpfloat) + inv_dual_edge_length = self.data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat) + z_dwdz_dd = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_graddiv_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) return dict( hmask_dd3d=hmask_dd3d, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_2nd_order_divergence_damping.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_2nd_order_divergence_damping.py index 08bd2c8d73..ff1d074473 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_2nd_order_divergence_damping.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_2nd_order_divergence_damping.py @@ -19,8 +19,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def apply_2nd_order_divergence_damping_numpy( @@ -30,13 +29,13 @@ def apply_2nd_order_divergence_damping_numpy( return vn -class TestApply2ndOrderDivergenceDamping(StencilTest): +class TestApply2ndOrderDivergenceDamping(stencil_tests.StencilTest): PROGRAM = apply_2nd_order_divergence_damping OUTPUTS = ("vn",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_graddiv_vn: np.ndarray, vn: np.ndarray, scal_divdamp_o2: ta.wpfloat, @@ -45,10 +44,10 @@ def reference( vn = apply_2nd_order_divergence_damping_numpy(z_graddiv_vn, vn, scal_divdamp_o2) return dict(vn=vn) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - z_graddiv_vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + z_graddiv_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) scal_divdamp_o2 = wpfloat("5.0") return dict( diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_4th_order_divergence_damping.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_4th_order_divergence_damping.py index b4ca1f1d56..edeaadfe76 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_4th_order_divergence_damping.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_4th_order_divergence_damping.py @@ -18,8 +18,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests from . import test_dycore_utils @@ -34,13 +33,13 @@ def apply_4th_order_divergence_damping_numpy( return vn -class TestApply4thOrderDivergenceDamping(StencilTest): +class TestApply4thOrderDivergenceDamping(stencil_tests.StencilTest): PROGRAM = apply_4th_order_divergence_damping OUTPUTS = ("vn",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, interpolated_fourth_order_divdamp_factor: np.ndarray, z_graddiv2_vn: np.ndarray, vn: np.ndarray, @@ -58,11 +57,11 @@ def reference( vn = apply_4th_order_divergence_damping_numpy(scal_divdamp, z_graddiv2_vn, vn) return dict(vn=vn) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - interpolated_fourth_order_divdamp_factor = data_alloc.random_field(grid, dims.KDim) - z_graddiv2_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + interpolated_fourth_order_divdamp_factor = self.data_alloc.random_field(dims.KDim) + z_graddiv2_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) divdamp_order = 24 mean_cell_area = 1000.0 diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py index 0ddbe7cf4e..81e4cb9831 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py @@ -6,19 +6,21 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from collections.abc import Mapping +from typing import cast + import gt4py.next as gtx import numpy as np import pytest import icon4py.model.common.type_alias as ta -import icon4py.model.testing.stencil_tests as test_helpers from icon4py.model.atmosphere.dycore.dycore_states import DivergenceDampingOrder from icon4py.model.atmosphere.dycore.stencils.compute_edge_diagnostics_for_dycore_and_update_vn import ( apply_divergence_damping_and_update_vn, ) from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid -from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.testing import stencil_tests from . import test_dycore_utils @@ -28,12 +30,12 @@ @pytest.mark.embedded_remap_error @pytest.mark.continuous_benchmarking -class TestApplyDivergenceDampingAndUpdateVn(test_helpers.StencilTest): +class TestApplyDivergenceDampingAndUpdateVn(stencil_tests.StencilTest): PROGRAM = apply_divergence_damping_and_update_vn OUTPUTS = ("next_vn",) STATIC_PARAMS = { - test_helpers.StandardStaticVariants.NONE: (), - test_helpers.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", "vertical_start", @@ -41,7 +43,7 @@ class TestApplyDivergenceDampingAndUpdateVn(test_helpers.StencilTest): "is_iau_active", "limited_area", ), - test_helpers.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( "vertical_start", "vertical_end", "is_iau_active", @@ -49,9 +51,9 @@ class TestApplyDivergenceDampingAndUpdateVn(test_helpers.StencilTest): ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, horizontal_gradient_of_normal_wind_divergence: np.ndarray, next_vn: np.ndarray, current_vn: np.ndarray, @@ -87,13 +89,14 @@ def reference( vertical_start: gtx.int32, vertical_end: gtx.int32, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) horz_idx = np.arange(horizontal_end)[:, np.newaxis] scaling_factor_for_3d_divdamp = np.expand_dims(scaling_factor_for_3d_divdamp, axis=0) horizontal_mask_for_3d_divdamp = np.expand_dims(horizontal_mask_for_3d_divdamp, axis=-1) inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) - e2c = connectivities[dims.E2CDim] + e2c = connectivities[dims.E2C] dwdz_at_edges_on_model_levels = dwdz_at_cells_on_model_levels[e2c] weighted_dwdz_at_edges_on_model_levels = ( dwdz_at_edges_on_model_levels[:, 1] - dwdz_at_edges_on_model_levels[:, 0] @@ -128,7 +131,7 @@ def reference( ) if apply_4th_order_divergence_damping: - e2c2eO = connectivities[dims.E2C2EODim] + e2c2eO = connectivities[dims.E2C2EO] # verified for e-10 squared_horizontal_gradient_of_total_divergence = np.where( (horizontal_start <= horz_idx) & (horz_idx < horizontal_end), @@ -186,7 +189,7 @@ def reference( return dict(next_vn=next_vn) - @pytest.fixture( + @stencil_tests.input_data_fixture( params=[ {"divdamp_order": do, "is_iau_active": ia, "second_order_divdamp_factor": sodf} for do, ia, sodf in [ @@ -206,31 +209,31 @@ def reference( f"divdamp_order[{param['divdamp_order']}]__is_iau_active[{param['is_iau_active']}]__second_order_divdamp_factor[{param['second_order_divdamp_factor']}]" ), ) - def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: - current_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - horizontal_mask_for_3d_divdamp = data_alloc.random_field(grid, dims.EdgeDim) - scaling_factor_for_3d_divdamp = data_alloc.random_field(grid, dims.KDim) - dwdz_at_cells_on_model_levels = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - inv_dual_edge_length = data_alloc.random_field(grid, dims.EdgeDim) - corrector_normal_wind_advective_tendency = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim + def input_data(self, grid: base.Grid, request: pytest.FixtureRequest) -> dict: + current_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + horizontal_mask_for_3d_divdamp = self.data_alloc.random_field(dims.EdgeDim) + scaling_factor_for_3d_divdamp = self.data_alloc.random_field(dims.KDim) + dwdz_at_cells_on_model_levels = self.data_alloc.random_field(dims.CellDim, dims.KDim) + inv_dual_edge_length = self.data_alloc.random_field(dims.EdgeDim) + corrector_normal_wind_advective_tendency = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim ) - predictor_normal_wind_advective_tendency = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim + predictor_normal_wind_advective_tendency = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim ) - normal_wind_tendency_due_to_slow_physics_process = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim + normal_wind_tendency_due_to_slow_physics_process = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim ) - horizontal_gradient_of_normal_wind_divergence = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim + horizontal_gradient_of_normal_wind_divergence = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim ) - normal_wind_iau_increment = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - next_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - theta_v_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - horizontal_pressure_gradient = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - geofac_grdiv = data_alloc.random_field(grid, dims.EdgeDim, dims.E2C2EODim) - interpolated_fourth_order_divdamp_factor = data_alloc.random_field(grid, dims.KDim) - nudgecoeff_e = data_alloc.random_field(grid, dims.EdgeDim) + normal_wind_iau_increment = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + next_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + theta_v_at_edges_on_model_levels = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + horizontal_pressure_gradient = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + geofac_grdiv = self.data_alloc.random_field(dims.EdgeDim, dims.E2C2EODim) + interpolated_fourth_order_divdamp_factor = self.data_alloc.random_field(dims.KDim) + nudgecoeff_e = self.data_alloc.random_field(dims.EdgeDim) mean_cell_area = 1000.0 max_nudging_coefficient = 0.3 diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_rayleigh_damping_mechanism.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_rayleigh_damping_mechanism.py index 4d2018f463..2a23de43e9 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_rayleigh_damping_mechanism.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_rayleigh_damping_mechanism.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -18,12 +19,11 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def apply_rayleigh_damping_mechanism_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], z_raylfac: np.ndarray, w: np.ndarray, ) -> np.ndarray: @@ -32,24 +32,25 @@ def apply_rayleigh_damping_mechanism_numpy( return w -class TestApplyRayleighDampingMechanism(StencilTest): +class TestApplyRayleighDampingMechanism(stencil_tests.StencilTest): PROGRAM = apply_rayleigh_damping_mechanism OUTPUTS = ("w",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_raylfac: np.ndarray, w: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) w = apply_rayleigh_damping_mechanism_numpy(connectivities, z_raylfac, w) return dict(w=w) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - z_raylfac = random_field(grid, dims.KDim, dtype=wpfloat) - w = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) + z_raylfac = self.data_alloc.random_field(dims.KDim, dtype=wpfloat) + w = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) return dict( z_raylfac=z_raylfac, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_weighted_2nd_and_4th_order_divergence_damping.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_weighted_2nd_and_4th_order_divergence_damping.py index 33505b461b..5f61f47207 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_weighted_2nd_and_4th_order_divergence_damping.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_weighted_2nd_and_4th_order_divergence_damping.py @@ -18,8 +18,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests from . import test_dycore_utils @@ -36,13 +35,13 @@ def apply_weighted_2nd_and_4th_order_divergence_damping_numpy( return vn -class TestApplyWeighted2ndAnd4thOrderDivergenceDamping(StencilTest): +class TestApplyWeighted2ndAnd4thOrderDivergenceDamping(stencil_tests.StencilTest): PROGRAM = apply_weighted_2nd_and_4th_order_divergence_damping OUTPUTS = ("vn",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, interpolated_fourth_order_divdamp_factor: np.ndarray, nudgecoeff_e: np.ndarray, z_graddiv2_vn: np.ndarray, @@ -74,12 +73,12 @@ def reference( ) return dict(vn=vn) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - interpolated_fourth_order_divdamp_factor = data_alloc.random_field(grid, dims.KDim) - nudgecoeff_e = data_alloc.random_field(grid, dims.EdgeDim, dtype=wpfloat) - z_graddiv2_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + interpolated_fourth_order_divdamp_factor = self.data_alloc.random_field(dims.KDim) + nudgecoeff_e = self.data_alloc.random_field(dims.EdgeDim, dtype=wpfloat) + z_graddiv2_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) divdamp_order = 24 mean_cell_area = 1000.0 diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py index 8ee6995109..a39825fd78 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py @@ -5,13 +5,13 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np import pytest -import icon4py.model.common.utils.data_allocation as data_alloc from icon4py.model.atmosphere.dycore.stencils.compute_advection_in_horizontal_momentum_equation import ( compute_advection_in_horizontal_momentum, ) @@ -25,7 +25,7 @@ def _compute_advective_normal_wind_tendency_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], horizontal_kinetic_energy_at_edges_on_model_levels: np.ndarray, coeff_gradekin: np.ndarray, horizontal_kinetic_energy_at_cells_on_model_levels: np.ndarray, @@ -37,7 +37,7 @@ def _compute_advective_normal_wind_tendency_numpy( vn_on_half_levels: np.ndarray, ddqz_z_full_e: np.ndarray, ) -> np.ndarray: - e2c = connectivities[dims.E2CDim] + e2c = connectivities[dims.E2C] horizontal_kinetic_energy_at_cells_on_model_levels_e2c = ( horizontal_kinetic_energy_at_cells_on_model_levels[e2c] ) @@ -55,7 +55,7 @@ def _compute_advective_normal_wind_tendency_numpy( + tangential_wind * ( coriolis_frequency - + 0.5 * np.sum(upward_vorticity_at_vertices[connectivities[dims.E2VDim]], axis=1) + + 0.5 * np.sum(upward_vorticity_at_vertices[connectivities[dims.E2V]], axis=1) ) + np.sum(contravariant_corrected_w_at_cells_on_model_levels[e2c] * c_lin_e, axis=1) * (vn_on_half_levels[:, :-1] - vn_on_half_levels[:, 1:]) @@ -65,7 +65,7 @@ def _compute_advective_normal_wind_tendency_numpy( def _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_without_levelmask_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], c_lin_e: np.ndarray, contravariant_corrected_w_at_cells_on_model_levels: np.ndarray, ddqz_z_full_e: np.ndarray, @@ -89,7 +89,7 @@ def _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_without_levelm tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) inv_primal_edge_length = np.expand_dims(inv_primal_edge_length, axis=-1) - e2c = connectivities[dims.E2CDim] + e2c = connectivities[dims.E2C] contravariant_corrected_w_at_edges_on_model_levels = np.sum( np.where( (e2c != -1)[:, :, np.newaxis], @@ -109,8 +109,8 @@ def _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_without_levelm ), difcoef, ) - e2v = connectivities[dims.E2VDim] - e2c2eo = connectivities[dims.E2C2EODim] + e2v = connectivities[dims.E2V] + e2c2eo = connectivities[dims.E2C2EO] normal_wind_advective_tendency = np.where( (np.abs(contravariant_corrected_w_at_edges_on_model_levels) > cfl_w_limit * ddqz_z_full_e), normal_wind_advective_tendency @@ -157,9 +157,9 @@ class TestFusedVelocityAdvectionStencilsHMomentum(stencil_tests.StencilTest): ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, normal_wind_advective_tendency: np.ndarray, vn: np.ndarray, horizontal_kinetic_energy_at_edges_on_model_levels: np.ndarray, @@ -183,6 +183,7 @@ def reference( end_index_of_damping_layer: int, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) normal_wind_advective_tendency_cp = normal_wind_advective_tendency.copy() nlev = kwargs["vertical_end"] k = np.arange(nlev) @@ -245,40 +246,40 @@ def reference( return dict(normal_wind_advective_tendency=normal_wind_advective_tendency) - @pytest.fixture( + @stencil_tests.input_data_fixture( params=[ {"apply_extra_diffusion_on_vn": value} for value in [True, False] ], # True for testing, False for benchmarking ids=lambda param: f"apply_extra_diffusion_on_vn[{param['apply_extra_diffusion_on_vn']}]", ) def input_data( - self, request: pytest.FixtureRequest, grid: base.Grid + self, grid: base.Grid, request: pytest.FixtureRequest ) -> dict[str, gtx.Field | state_utils.ScalarType]: - normal_wind_advective_tendency = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - horizontal_kinetic_energy_at_edges_on_model_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim + normal_wind_advective_tendency = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + horizontal_kinetic_energy_at_edges_on_model_levels = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim ) - tangential_wind = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - coriolis_frequency = data_alloc.random_field(grid, dims.EdgeDim) - contravariant_corrected_w_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim + tangential_wind = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + coriolis_frequency = self.data_alloc.random_field(dims.EdgeDim) + contravariant_corrected_w_at_cells_on_model_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim ) - vn_on_half_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim, extend={dims.KDim: 1} + vn_on_half_levels = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim, extend={dims.KDim: 1} ) - coeff_gradekin = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - e_bln_c_s = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) - c_lin_e = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - ddqz_z_full_e = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim, low=0.0 + coeff_gradekin = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + e_bln_c_s = self.data_alloc.random_field(dims.CellDim, dims.C2EDim) + c_lin_e = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + ddqz_z_full_e = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim, low=0.0 ) # this makes sure that the simplified stencil produces the same result as the numpy version - area_edge = data_alloc.random_field(grid, dims.EdgeDim) - tangent_orientation = data_alloc.random_field(grid, dims.EdgeDim) - inv_primal_edge_length = data_alloc.random_field(grid, dims.EdgeDim) - geofac_grdiv = data_alloc.random_field(grid, dims.EdgeDim, dims.E2C2EODim) + area_edge = self.data_alloc.random_field(dims.EdgeDim) + tangent_orientation = self.data_alloc.random_field(dims.EdgeDim) + inv_primal_edge_length = self.data_alloc.random_field(dims.EdgeDim) + geofac_grdiv = self.data_alloc.random_field(dims.EdgeDim, dims.E2C2EODim) - geofac_rot = data_alloc.random_field(grid, dims.VertexDim, dims.V2EDim) + geofac_rot = self.data_alloc.random_field(dims.VertexDim, dims.V2EDim) scalfac_exdiff = 0.6 dtime = 2.0 cfl_w_limit = 0.65 / dtime diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py index 3c32009b11..40886f1eb5 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -18,7 +19,6 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests from .test_add_interpolated_horizontal_advection_of_w import ( @@ -40,7 +40,7 @@ def interpolate_contravariant_correction_to_cells_on_half_levels_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], contravariant_correction_at_cells_on_half_levels: np.ndarray, contravariant_correction_at_edges_on_model_levels: np.ndarray, e_bln_c_s: np.ndarray, @@ -125,7 +125,7 @@ def compute_maximum_cfl_and_clip_contravariant_vertical_velocity_numpy( def compute_horizontal_advection_of_w( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], w: np.ndarray, tangential_wind_on_half_levels: np.ndarray, vn_on_half_levels: np.ndarray, @@ -155,7 +155,7 @@ def compute_horizontal_advection_of_w( def add_extra_diffusion_for_w_approaching_cfl_wihtout_levmask_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], cfl_clipping: np.ndarray, owner_mask: np.ndarray, contravariant_corrected_w_at_cells_on_half_levels: np.ndarray, @@ -183,7 +183,7 @@ def add_extra_diffusion_for_w_approaching_cfl_wihtout_levmask_numpy( 0, ) - c2e2cO = connectivities[dims.C2E2CODim] + c2e2cO = connectivities[dims.C2E2CO] vertical_wind_advective_tendency = np.where( (cfl_clipping == 1) & (owner_mask == 1), vertical_wind_advective_tendency @@ -203,7 +203,7 @@ def add_extra_diffusion_for_w_approaching_cfl_wihtout_levmask_numpy( def compute_advective_vertical_wind_tendency_and_apply_diffusion_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], vertical_wind_advective_tendency: np.ndarray, w: np.ndarray, horizontal_advection_of_w_at_edges_on_half_levels: np.ndarray, @@ -293,9 +293,9 @@ class TestFusedVelocityAdvectionStencilVMomentum(stencil_tests.StencilTest): ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, vertical_wind_advective_tendency: np.ndarray, contravariant_corrected_w_at_cells_on_model_levels: np.ndarray, vertical_cfl: np.ndarray, @@ -320,6 +320,7 @@ def reference( end_index_of_damping_layer: int, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) nlev = kwargs["vertical_end"] horizontal_advection_of_w_at_edges_on_half_levels = compute_horizontal_advection_of_w( @@ -409,37 +410,37 @@ def reference( vertical_cfl=vertical_cfl_ret, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - contravariant_corrected_w_at_cells_on_model_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim + contravariant_corrected_w_at_cells_on_model_levels = self.data_alloc.zero_field( + dims.CellDim, dims.KDim ) - vertical_wind_advective_tendency = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - w = data_alloc.random_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) - tangential_wind_on_half_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim, extend={dims.KDim: 1} + vertical_wind_advective_tendency = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + w = self.data_alloc.random_field(dims.CellDim, dims.KDim, extend={dims.KDim: 1}) + tangential_wind_on_half_levels = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim, extend={dims.KDim: 1} ) - vn_on_half_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim, extend={dims.KDim: 1} + vn_on_half_levels = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim, extend={dims.KDim: 1} ) - contravariant_correction_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + contravariant_correction_at_cells_on_half_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - coeff1_dwdz = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - coeff2_dwdz = data_alloc.random_field(grid, dims.CellDim, dims.KDim) + coeff1_dwdz = self.data_alloc.random_field(dims.CellDim, dims.KDim) + coeff2_dwdz = self.data_alloc.random_field(dims.CellDim, dims.KDim) - c_intp = data_alloc.random_field(grid, dims.VertexDim, dims.V2CDim) - inv_dual_edge_length = data_alloc.random_field(grid, dims.EdgeDim, low=1.0e-5) - inv_primal_edge_length = data_alloc.random_field(grid, dims.EdgeDim, low=1.0e-5) - tangent_orientation = data_alloc.random_field(grid, dims.EdgeDim, low=1.0e-5) - e_bln_c_s = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) + c_intp = self.data_alloc.random_field(dims.VertexDim, dims.V2CDim) + inv_dual_edge_length = self.data_alloc.random_field(dims.EdgeDim, low=1.0e-5) + inv_primal_edge_length = self.data_alloc.random_field(dims.EdgeDim, low=1.0e-5) + tangent_orientation = self.data_alloc.random_field(dims.EdgeDim, low=1.0e-5) + e_bln_c_s = self.data_alloc.random_field(dims.CellDim, dims.C2EDim) - vertical_cfl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - owner_mask = data_alloc.random_mask(grid, dims.CellDim) - ddqz_z_half = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - area = data_alloc.random_field(grid, dims.CellDim) - geofac_n2s = data_alloc.random_field(grid, dims.CellDim, dims.C2E2CODim) + vertical_cfl = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + owner_mask = self.data_alloc.random_mask(dims.CellDim) + ddqz_z_half = self.data_alloc.random_field(dims.CellDim, dims.KDim) + area = self.data_alloc.random_field(dims.CellDim) + geofac_n2s = self.data_alloc.random_field(dims.CellDim, dims.C2E2CODim) scalfac_exdiff = 10.0 dtime = 2.0 @@ -497,9 +498,9 @@ class TestFusedVelocityAdvectionStencilVMomentumAndContravariant(stencil_tests.S stencil_tests.StandardStaticVariants.NONE: (), # For now compile time variants triger error in gt4py } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, contravariant_correction_at_cells_on_half_levels: np.ndarray, vertical_wind_advective_tendency: np.ndarray, contravariant_corrected_w_at_cells_on_model_levels: np.ndarray, @@ -523,6 +524,7 @@ def reference( skip_compute_predictor_vertical_advection: bool, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) nlev = kwargs["vertical_end"] # We need to store the initial return field, because we only compute on a subdomain. @@ -622,7 +624,7 @@ def reference( vertical_cfl=vertical_cfl_ret, ) - @pytest.fixture( + @stencil_tests.input_data_fixture( params=[ {"skip_compute_predictor_vertical_advection": value} for value in [True, False] ], # True for benchmarking, False for testing @@ -633,32 +635,32 @@ def reference( def input_data( self, grid: base.Grid, request: pytest.FixtureRequest ) -> dict[str, gtx.Field | state_utils.ScalarType]: - contravariant_corrected_w_at_cells_on_model_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim + contravariant_corrected_w_at_cells_on_model_levels = self.data_alloc.zero_field( + dims.CellDim, dims.KDim ) - vertical_wind_advective_tendency = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - w = data_alloc.random_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) - horizontal_advection_of_w_at_edges_on_half_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim, extend={dims.KDim: 1} + vertical_wind_advective_tendency = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + w = self.data_alloc.random_field(dims.CellDim, dims.KDim, extend={dims.KDim: 1}) + horizontal_advection_of_w_at_edges_on_half_levels = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim, extend={dims.KDim: 1} ) - contravariant_correction_at_edges_on_model_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim + contravariant_correction_at_edges_on_model_levels = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim ) - contravariant_correction_at_cells_on_half_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + contravariant_correction_at_cells_on_half_levels = self.data_alloc.zero_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - coeff1_dwdz = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - coeff2_dwdz = data_alloc.random_field(grid, dims.CellDim, dims.KDim) + coeff1_dwdz = self.data_alloc.random_field(dims.CellDim, dims.KDim) + coeff2_dwdz = self.data_alloc.random_field(dims.CellDim, dims.KDim) - e_bln_c_s = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) - wgtfac_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim) + e_bln_c_s = self.data_alloc.random_field(dims.CellDim, dims.C2EDim) + wgtfac_c = self.data_alloc.random_field(dims.CellDim, dims.KDim) - vertical_cfl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - owner_mask = data_alloc.random_mask(grid, dims.CellDim) - ddqz_z_half = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - area = data_alloc.random_field(grid, dims.CellDim) - geofac_n2s = data_alloc.random_field(grid, dims.CellDim, dims.C2E2CODim) + vertical_cfl = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + owner_mask = self.data_alloc.random_mask(dims.CellDim) + ddqz_z_half = self.data_alloc.random_field(dims.CellDim, dims.KDim) + area = self.data_alloc.random_field(dims.CellDim) + geofac_n2s = self.data_alloc.random_field(dims.CellDim, dims.C2E2CODim) scalfac_exdiff = 10.0 dtime = 2.0 diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advective_vertical_wind_tendency.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advective_vertical_wind_tendency.py index 21fdb8c7ef..79978c83f5 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advective_vertical_wind_tendency.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advective_vertical_wind_tendency.py @@ -18,8 +18,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_advective_vertical_wind_tendency_numpy( @@ -38,13 +37,13 @@ def compute_advective_vertical_wind_tendency_numpy( return ddt_w_adv -class TestComputeAdvectiveVerticalWindTendency(StencilTest): +class TestComputeAdvectiveVerticalWindTendency(stencil_tests.StencilTest): PROGRAM = compute_advective_vertical_wind_tendency OUTPUTS = ("ddt_w_adv",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_w_con_c: np.ndarray, w: np.ndarray, coeff1_dwdz: np.ndarray, @@ -56,13 +55,15 @@ def reference( ) return dict(ddt_w_adv=ddt_w_adv) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - z_w_con_c = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - w = random_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=wpfloat) - coeff1_dwdz = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - coeff2_dwdz = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - ddt_w_adv = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + z_w_con_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + w = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=wpfloat + ) + coeff1_dwdz = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + coeff2_dwdz = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + ddt_w_adv = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) return dict( z_w_con_c=z_w_con_c, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_airmass.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_airmass.py index 61445dcf1e..c1ef715d54 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_airmass.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_airmass.py @@ -16,17 +16,16 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestComputeAirmass(StencilTest): +class TestComputeAirmass(stencil_tests.StencilTest): PROGRAM = compute_airmass OUTPUTS = ("airmass_out",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, rho_in: np.ndarray, ddqz_z_full_in: np.ndarray, deepatmo_t1mc_in: np.ndarray, @@ -35,12 +34,12 @@ def reference( airmass_out = rho_in * ddqz_z_full_in * deepatmo_t1mc_in return dict(airmass_out=airmass_out) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - rho_in = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - ddqz_z_full_in = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - deepatmo_t1mc_in = random_field(grid, dims.KDim, dtype=wpfloat) - airmass_out = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) + rho_in = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + ddqz_z_full_in = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + deepatmo_t1mc_in = self.data_alloc.random_field(dims.KDim, dtype=wpfloat) + airmass_out = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) return dict( rho_in=rho_in, ddqz_z_full_in=ddqz_z_full_in, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_approx_of_2nd_vertical_derivative_of_exner.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_approx_of_2nd_vertical_derivative_of_exner.py index d842690914..e767ebb97d 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_approx_of_2nd_vertical_derivative_of_exner.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_approx_of_2nd_vertical_derivative_of_exner.py @@ -18,8 +18,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_approx_of_2nd_vertical_derivative_of_exner_numpy( @@ -36,13 +35,13 @@ def compute_approx_of_2nd_vertical_derivative_of_exner_numpy( return z_dexner_dz_c_2 -class TestComputeApproxOf2ndVerticalDerivativeOfExner(StencilTest): +class TestComputeApproxOf2ndVerticalDerivativeOfExner(stencil_tests.StencilTest): PROGRAM = compute_approx_of_2nd_vertical_derivative_of_exner OUTPUTS = ("z_dexner_dz_c_2",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_theta_v_pr_ic: np.ndarray, d2dexdz2_fac1_mc: np.ndarray, d2dexdz2_fac2_mc: np.ndarray, @@ -57,16 +56,16 @@ def reference( ) return dict(z_dexner_dz_c_2=z_dexner_dz_c_2) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - z_theta_v_pr_ic = random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=vpfloat + z_theta_v_pr_ic = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=vpfloat ) - d2dexdz2_fac1_mc = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - z_rth_pr_2 = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - d2dexdz2_fac2_mc = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + d2dexdz2_fac1_mc = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + z_rth_pr_2 = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + d2dexdz2_fac2_mc = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) - z_dexner_dz_c_2 = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + z_dexner_dz_c_2 = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) return dict( z_theta_v_pr_ic=z_theta_v_pr_ic, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_averaged_vn_and_fluxes_and_prepare_tracer_advection.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_averaged_vn_and_fluxes_and_prepare_tracer_advection.py index cf503c38ca..1278b6e554 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_averaged_vn_and_fluxes_and_prepare_tracer_advection.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_averaged_vn_and_fluxes_and_prepare_tracer_advection.py @@ -5,13 +5,13 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np import pytest -import icon4py.model.common.utils.data_allocation as data_alloc from icon4py.model.atmosphere.dycore.stencils.compute_horizontal_velocity_quantities import ( compute_averaged_vn_and_fluxes, ) @@ -56,9 +56,9 @@ class TestComputeAveragedVnAndFluxesAndPrepareTracerAdvection(stencil_tests.Sten ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, spatially_averaged_vn: np.ndarray, mass_flux_at_edges_on_model_levels: np.ndarray, theta_v_flux_at_edges_on_model_levels: np.ndarray, @@ -76,6 +76,7 @@ def reference( horizontal_end: int, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) initial_spatially_averaged_vn = spatially_averaged_vn.copy() initial_mass_flux_at_edges_on_model_levels = mass_flux_at_edges_on_model_levels.copy() initial_theta_v_flux_at_edges_on_model_levels = theta_v_flux_at_edges_on_model_levels.copy() @@ -154,7 +155,7 @@ def reference( substep_averaged_mass_flux=substep_averaged_mass_flux, ) - @pytest.fixture( + @stencil_tests.input_data_fixture( params=[ {"prepare_advection": pa, "at_first_substep": afs} for pa, afs in [ @@ -167,19 +168,19 @@ def reference( ), ) def input_data( - self, request: pytest.FixtureRequest, grid: base.Grid + self, grid: base.Grid, request: pytest.FixtureRequest ) -> dict[str, gtx.Field | state_utils.ScalarType]: - spatially_averaged_vn = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - mass_fl_e = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - z_theta_v_fl_e = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - - substep_and_spatially_averaged_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - substep_averaged_mass_flux = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - e_flx_avg = data_alloc.random_field(grid, dims.EdgeDim, dims.E2C2EODim) - vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - z_rho_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - ddqz_z_full_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - z_theta_v_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) + spatially_averaged_vn = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + mass_fl_e = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + z_theta_v_fl_e = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + + substep_and_spatially_averaged_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + substep_averaged_mass_flux = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + e_flx_avg = self.data_alloc.random_field(dims.EdgeDim, dims.E2C2EODim) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + z_rho_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + ddqz_z_full_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + z_theta_v_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) prepare_advection = request.param["prepare_advection"] at_first_substep = request.param["at_first_substep"] r_nsubsteps = 0.5 diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_avg_vn_and_graddiv_vn_and_vt.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_avg_vn_and_graddiv_vn_and_vt.py index 5cf3efa33b..0510bbec14 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_avg_vn_and_graddiv_vn_and_vt.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_avg_vn_and_graddiv_vn_and_vt.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -18,19 +19,18 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_avg_vn_and_graddiv_vn_and_vt_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], e_flx_avg: np.ndarray, vn: np.ndarray, geofac_grdiv: np.ndarray, rbf_vec_coeff_e: np.ndarray, ) -> tuple[np.ndarray, ...]: - e2c2eO = connectivities[dims.E2C2EODim] - e2c2e = connectivities[dims.E2C2EDim] + e2c2eO = connectivities[dims.E2C2EO] + e2c2e = connectivities[dims.E2C2E] e_flx_avg = np.expand_dims(e_flx_avg, axis=-1) z_vn_avg = np.sum(vn[e2c2eO] * e_flx_avg, axis=1) geofac_grdiv = np.expand_dims(geofac_grdiv, axis=-1) @@ -43,19 +43,20 @@ def compute_avg_vn_and_graddiv_vn_and_vt_numpy( @pytest.mark.embedded_remap_error -class TestComputeAvgVnAndGraddivVnAndVt(StencilTest): +class TestComputeAvgVnAndGraddivVnAndVt(stencil_tests.StencilTest): PROGRAM = compute_avg_vn_and_graddiv_vn_and_vt OUTPUTS = ("z_vn_avg", "z_graddiv_vn", "vt") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, e_flx_avg: np.ndarray, vn: np.ndarray, geofac_grdiv: np.ndarray, rbf_vec_coeff_e: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) z_vn_avg, z_graddiv_vn, vt = compute_avg_vn_and_graddiv_vn_and_vt_numpy( connectivities, e_flx_avg, @@ -65,15 +66,15 @@ def reference( ) return dict(z_vn_avg=z_vn_avg, z_graddiv_vn=z_graddiv_vn, vt=vt) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - e_flx_avg = random_field(grid, dims.EdgeDim, dims.E2C2EODim, dtype=wpfloat) - geofac_grdiv = random_field(grid, dims.EdgeDim, dims.E2C2EODim, dtype=wpfloat) - rbf_vec_coeff_e = random_field(grid, dims.EdgeDim, dims.E2C2EDim, dtype=wpfloat) - vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - z_vn_avg = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - z_graddiv_vn = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - vt = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) + e_flx_avg = self.data_alloc.random_field(dims.EdgeDim, dims.E2C2EODim, dtype=wpfloat) + geofac_grdiv = self.data_alloc.random_field(dims.EdgeDim, dims.E2C2EODim, dtype=wpfloat) + rbf_vec_coeff_e = self.data_alloc.random_field(dims.EdgeDim, dims.E2C2EDim, dtype=wpfloat) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + z_vn_avg = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + z_graddiv_vn = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + vt = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) return dict( e_flx_avg=e_flx_avg, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_contravariant_correction.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_contravariant_correction.py index e557ec551d..a5a9d04b59 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_contravariant_correction.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_contravariant_correction.py @@ -18,8 +18,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_contravariant_correction_numpy( @@ -29,13 +28,13 @@ def compute_contravariant_correction_numpy( return z_w_concorr_me -class TestComputeContravariantCorrection(StencilTest): +class TestComputeContravariantCorrection(stencil_tests.StencilTest): PROGRAM = compute_contravariant_correction OUTPUTS = ("z_w_concorr_me",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, vn: np.ndarray, ddxn_z_full: np.ndarray, ddxt_z_full: np.ndarray, @@ -45,13 +44,13 @@ def reference( z_w_concorr_me = compute_contravariant_correction_numpy(vn, ddxn_z_full, ddxt_z_full, vt) return dict(z_w_concorr_me=z_w_concorr_me) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - ddxn_z_full = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - ddxt_z_full = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat, low=0.1) - vt = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - z_w_concorr_me = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + ddxn_z_full = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + ddxt_z_full = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat, low=0.1) + vt = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + z_w_concorr_me = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) return dict( vn=vn, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_contravariant_correction_of_w.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_contravariant_correction_of_w.py index 6a7d916250..93cd00731f 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_contravariant_correction_of_w.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_contravariant_correction_of_w.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -18,17 +19,16 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_contravariant_correction_of_w_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], e_bln_c_s: np.ndarray, z_w_concorr_me: np.ndarray, wgtfac_c: np.ndarray, ) -> np.ndarray: - c2e = connectivities[dims.C2EDim] + c2e = connectivities[dims.C2E] e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) z_w_concorr_me_offset_1 = np.roll(z_w_concorr_me, shift=1, axis=1) @@ -39,29 +39,30 @@ def compute_contravariant_correction_of_w_numpy( return w_concorr_c -class TestComputeContravariantCorrectionOfW(StencilTest): +class TestComputeContravariantCorrectionOfW(stencil_tests.StencilTest): PROGRAM = compute_contravariant_correction_of_w OUTPUTS = ("w_concorr_c",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, e_bln_c_s: np.ndarray, z_w_concorr_me: np.ndarray, wgtfac_c: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) w_concorr_c = compute_contravariant_correction_of_w_numpy( connectivities, e_bln_c_s, z_w_concorr_me, wgtfac_c ) return dict(w_concorr_c=w_concorr_c) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - e_bln_c_s = random_field(grid, dims.CellDim, dims.C2EDim, dtype=wpfloat) - z_w_concorr_me = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - wgtfac_c = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - w_concorr_c = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + e_bln_c_s = self.data_alloc.random_field(dims.CellDim, dims.C2EDim, dtype=wpfloat) + z_w_concorr_me = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + wgtfac_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + w_concorr_c = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) return dict( e_bln_c_s=e_bln_c_s, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_contravariant_correction_of_w_for_lower_boundary.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_contravariant_correction_of_w_for_lower_boundary.py index 1723d907a5..8f72e61d3e 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_contravariant_correction_of_w_for_lower_boundary.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_contravariant_correction_of_w_for_lower_boundary.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -18,17 +19,16 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_contravariant_correction_of_w_for_lower_boundary_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], e_bln_c_s: np.ndarray, z_w_concorr_me: np.ndarray, wgtfacq_c: np.ndarray, ) -> np.ndarray: - c2e = connectivities[dims.C2EDim] + c2e = connectivities[dims.C2E] e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) z_w_concorr_me_offset_1 = np.roll(z_w_concorr_me, shift=1, axis=1) @@ -48,30 +48,31 @@ def compute_contravariant_correction_of_w_for_lower_boundary_numpy( return w_concorr_c -class TestComputeContravariantCorrectionOfWForLowerBoundary(StencilTest): +class TestComputeContravariantCorrectionOfWForLowerBoundary(stencil_tests.StencilTest): PROGRAM = compute_contravariant_correction_of_w_for_lower_boundary OUTPUTS = ("w_concorr_c",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, e_bln_c_s: np.ndarray, z_w_concorr_me: np.ndarray, wgtfacq_c: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) w_concorr_c = compute_contravariant_correction_of_w_for_lower_boundary_numpy( connectivities, e_bln_c_s, z_w_concorr_me, wgtfacq_c ) return dict(w_concorr_c=w_concorr_c) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - e_bln_c_s = random_field(grid, dims.CellDim, dims.C2EDim, dtype=wpfloat) - z_w_concorr_me = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - wgtfacq_c = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - w_concorr_c = zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=vpfloat + e_bln_c_s = self.data_alloc.random_field(dims.CellDim, dims.C2EDim, dtype=wpfloat) + z_w_concorr_me = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + wgtfacq_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + w_concorr_c = self.data_alloc.zero_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=vpfloat ) return dict( diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_diagnostics_from_normal_wind.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_diagnostics_from_normal_wind.py index c47e4e5332..8bef1e13b0 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_diagnostics_from_normal_wind.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_diagnostics_from_normal_wind.py @@ -5,6 +5,9 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from collections.abc import Mapping +from typing import cast + import gt4py.next as gtx import numpy as np import pytest @@ -15,7 +18,6 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests from .test_compute_contravariant_correction import compute_contravariant_correction_numpy @@ -37,7 +39,7 @@ def compute_diagnostics_from_normal_wind_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], tangential_wind_on_half_levels: np.ndarray, tangential_wind: np.ndarray, vn_on_half_levels: np.ndarray, @@ -178,10 +180,9 @@ class TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContrava ), } - @classmethod + @stencil_tests.static_reference def reference( - cls, - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, tangential_wind: np.ndarray, tangential_wind_on_half_levels: np.ndarray, vn_on_half_levels: np.ndarray, @@ -206,6 +207,7 @@ def reference( vertical_start: int, vertical_end: int, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) initial_tangential_wind = tangential_wind.copy() initial_tangential_wind_on_half_levels = tangential_wind_on_half_levels.copy() initial_horizontal_kinetic_energy_at_edges_on_model_levels = ( @@ -290,7 +292,7 @@ def reference( horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, ) - @pytest.fixture( + @stencil_tests.input_data_fixture( params=[ {"skip_compute_predictor_vertical_advection": value} for value in [True, False] ], # True for benchmarking, False for testing @@ -301,31 +303,31 @@ def reference( def input_data( self, grid: base.Grid, request: pytest.FixtureRequest ) -> dict[str, gtx.Field | state_utils.ScalarType]: - horizontal_advection_of_w_at_edges_on_half_levels = data_alloc.zero_field( - grid, dims.EdgeDim, dims.KDim + horizontal_advection_of_w_at_edges_on_half_levels = self.data_alloc.zero_field( + dims.EdgeDim, dims.KDim ) - tangential_wind = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - tangential_wind_on_half_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - vn_on_half_levels = data_alloc.zero_field( - grid, dims.EdgeDim, dims.KDim, extend={dims.KDim: 1} + tangential_wind = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + tangential_wind_on_half_levels = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + vn_on_half_levels = self.data_alloc.zero_field( + dims.EdgeDim, dims.KDim, extend={dims.KDim: 1} ) - horizontal_kinetic_energy_at_edges_on_model_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim + horizontal_kinetic_energy_at_edges_on_model_levels = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim ) - contravariant_correction_at_edges_on_model_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim + contravariant_correction_at_edges_on_model_levels = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim ) - vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - w = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - rbf_vec_coeff_e = data_alloc.random_field(grid, dims.EdgeDim, dims.E2C2EDim) - wgtfac_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - ddxn_z_full = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - ddxt_z_full = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - inv_dual_edge_length = data_alloc.random_field(grid, dims.EdgeDim) - inv_primal_edge_length = data_alloc.random_field(grid, dims.EdgeDim) - tangent_orientation = data_alloc.random_field(grid, dims.EdgeDim) - wgtfacq_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - c_intp = data_alloc.random_field(grid, dims.VertexDim, dims.V2CDim) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + w = self.data_alloc.random_field(dims.CellDim, dims.KDim) + rbf_vec_coeff_e = self.data_alloc.random_field(dims.EdgeDim, dims.E2C2EDim) + wgtfac_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + ddxn_z_full = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + ddxt_z_full = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + inv_dual_edge_length = self.data_alloc.random_field(dims.EdgeDim) + inv_primal_edge_length = self.data_alloc.random_field(dims.EdgeDim) + tangent_orientation = self.data_alloc.random_field(dims.EdgeDim) + wgtfacq_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + c_intp = self.data_alloc.random_field(dims.VertexDim, dims.V2CDim) nlev = grid.num_levels nflatlev = 5 # value is set to reflect the MCH ch1 experiment. Changing this value will change the expected runtime diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_divergence_of_fluxes_of_rho_and_theta.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_divergence_of_fluxes_of_rho_and_theta.py index d009a74d8d..981050a15b 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_divergence_of_fluxes_of_rho_and_theta.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_divergence_of_fluxes_of_rho_and_theta.py @@ -5,13 +5,13 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np import pytest -import icon4py.model.common.utils.data_allocation as data_alloc from icon4py.model.atmosphere.dycore.stencils.compute_divergence_of_fluxes_of_rho_and_theta import ( compute_divergence_of_fluxes_of_rho_and_theta, ) @@ -22,12 +22,12 @@ def compute_divergence_of_fluxes_of_rho_and_theta_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], geofac_div: np.ndarray, mass_flux_at_edges_on_model_levels: np.ndarray, theta_v_flux_at_edges_on_model_levels: np.ndarray, ) -> tuple[np.ndarray, np.ndarray]: - c2e = connectivities[dims.C2EDim] + c2e = connectivities[dims.C2E] geofac_div = np.expand_dims(geofac_div, axis=-1) divergence_of_mass_wp = np.sum(geofac_div * mass_flux_at_edges_on_model_levels[c2e], axis=1) @@ -41,14 +41,15 @@ class TestComputeDivergenceConnectivityOfFluxesOfRhoAndTheta(stencil_tests.Stenc PROGRAM = compute_divergence_of_fluxes_of_rho_and_theta OUTPUTS = ("z_flxdiv_mass", "z_flxdiv_theta") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, geofac_div: np.ndarray, mass_fl_e: np.ndarray, z_theta_v_fl_e: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) z_flxdiv_mass, z_flxdiv_theta = compute_divergence_of_fluxes_of_rho_and_theta_numpy( connectivities, geofac_div, @@ -57,13 +58,13 @@ def reference( ) return dict(z_flxdiv_mass=z_flxdiv_mass, z_flxdiv_theta=z_flxdiv_theta) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim, dtype=ta.wpfloat) - z_theta_v_fl_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) - z_flxdiv_theta = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - mass_fl_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) - z_flxdiv_mass = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) + geofac_div = self.data_alloc.random_field(dims.CellDim, dims.C2EDim, dtype=ta.wpfloat) + z_theta_v_fl_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + z_flxdiv_theta = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + mass_fl_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + z_flxdiv_mass = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) return dict( geofac_div=geofac_div, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_dwdz_for_divergence_damping.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_dwdz_for_divergence_damping.py index a14d93aca1..f91a426f30 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_dwdz_for_divergence_damping.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_dwdz_for_divergence_damping.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -17,12 +18,11 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_dwdz_for_divergence_damping_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], inv_ddqz_z_full: np.ndarray, w: np.ndarray, w_concorr_c: np.ndarray, @@ -33,31 +33,34 @@ def compute_dwdz_for_divergence_damping_numpy( return z_dwdz_dd -class TestComputeDwdzForDivergenceDamping(StencilTest): +class TestComputeDwdzForDivergenceDamping(stencil_tests.StencilTest): PROGRAM = _compute_dwdz_for_divergence_damping OUTPUTS = ("out",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, inv_ddqz_z_full: np.ndarray, w: np.ndarray, w_concorr_c: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) z_dwdz_dd = compute_dwdz_for_divergence_damping_numpy( connectivities, inv_ddqz_z_full=inv_ddqz_z_full, w=w, w_concorr_c=w_concorr_c ) return dict(out=z_dwdz_dd) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, Any]: - inv_ddqz_z_full = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - w = random_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=wpfloat) - w_concorr_c = random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=vpfloat + inv_ddqz_z_full = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + w = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=wpfloat ) - z_dwdz_dd = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + w_concorr_c = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=vpfloat + ) + z_dwdz_dd = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) return dict( inv_ddqz_z_full=inv_ddqz_z_full, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_exner_from_rhotheta.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_exner_from_rhotheta.py index 582a4ef0d3..a2a025f7a2 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_exner_from_rhotheta.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_exner_from_rhotheta.py @@ -17,17 +17,16 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestComputeExnerFromRhotheta(StencilTest): +class TestComputeExnerFromRhotheta(stencil_tests.StencilTest): PROGRAM = _compute_exner_from_rhotheta OUTPUTS = ("out",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, rho: np.ndarray, exner: np.ndarray, rd_o_cvd: float, @@ -38,13 +37,13 @@ def reference( exner = np.exp(rd_o_cvd * np.log(rd_o_p0ref * rho * theta_v)) return dict(out=(theta_v, exner)) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, Any]: rd_o_cvd = wpfloat("10.0") rd_o_p0ref = wpfloat("20.0") - rho = random_field(grid, dims.CellDim, dims.KDim, low=1, high=2, dtype=wpfloat) - exner = random_field(grid, dims.CellDim, dims.KDim, low=1, high=2, dtype=wpfloat) - theta_v = zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) + rho = self.data_alloc.random_field(dims.CellDim, dims.KDim, low=1, high=2, dtype=wpfloat) + exner = self.data_alloc.random_field(dims.CellDim, dims.KDim, low=1, high=2, dtype=wpfloat) + theta_v = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat) return dict( rho=rho, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_explicit_part_for_rho_and_exner.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_explicit_part_for_rho_and_exner.py index 347299fd46..fdb389c070 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_explicit_part_for_rho_and_exner.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_explicit_part_for_rho_and_exner.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -17,12 +18,11 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_explicit_part_for_rho_and_exner_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], rho_nnow: np.ndarray, inv_ddqz_z_full: np.ndarray, z_flxdiv_mass: np.ndarray, @@ -51,13 +51,13 @@ def compute_explicit_part_for_rho_and_exner_numpy( return (z_rho_expl, z_exner_expl) -class TestComputeExplicitPartForRhoAndExner(StencilTest): +class TestComputeExplicitPartForRhoAndExner(stencil_tests.StencilTest): PROGRAM = compute_explicit_part_for_rho_and_exner OUTPUTS = ("z_rho_expl", "z_exner_expl") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, rho_nnow: np.ndarray, inv_ddqz_z_full: np.ndarray, z_flxdiv_mass: np.ndarray, @@ -70,6 +70,7 @@ def reference( dtime: float, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) (z_rho_expl, z_exner_expl) = compute_explicit_part_for_rho_and_exner_numpy( connectivities, rho_nnow=rho_nnow, @@ -85,25 +86,25 @@ def reference( ) return dict(z_rho_expl=z_rho_expl, z_exner_expl=z_exner_expl) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: dtime = ta.wpfloat("1.0") - rho_nnow = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - inv_ddqz_z_full = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - z_flxdiv_mass = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - z_contr_w_fl_l = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=ta.wpfloat + rho_nnow = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + inv_ddqz_z_full = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_flxdiv_mass = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_contr_w_fl_l = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=ta.wpfloat ) - exner_pr = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - z_beta = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - z_flxdiv_theta = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - theta_v_ic = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=ta.wpfloat + exner_pr = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + z_beta = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_flxdiv_theta = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + theta_v_ic = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=ta.wpfloat ) - ddt_exner_phy = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) + ddt_exner_phy = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) - z_rho_expl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - z_exner_expl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) + z_rho_expl = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + z_exner_expl = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) return dict( z_rho_expl=z_rho_expl, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_explicit_vertical_wind_from_advection_and_vertical_wind_density.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_explicit_vertical_wind_from_advection_and_vertical_wind_density.py index efd3ba9a10..9c3a58b5c1 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_explicit_vertical_wind_from_advection_and_vertical_wind_density.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_explicit_vertical_wind_from_advection_and_vertical_wind_density.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -17,12 +18,11 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_explicit_vertical_wind_from_advection_and_vertical_wind_density_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], w_nnow: np.ndarray, ddt_w_adv_ntl1: np.ndarray, ddt_w_adv_ntl2: np.ndarray, @@ -43,13 +43,13 @@ def compute_explicit_vertical_wind_from_advection_and_vertical_wind_density_nump return (z_w_expl, z_contr_w_fl_l) -class TestComputeExplicitVerticalWindFromAdvectionAndVerticalWindDensity(StencilTest): +class TestComputeExplicitVerticalWindFromAdvectionAndVerticalWindDensity(stencil_tests.StencilTest): PROGRAM = compute_explicit_vertical_wind_from_advection_and_vertical_wind_density OUTPUTS = ("z_w_expl", "z_contr_w_fl_l") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, w_nnow: np.ndarray, ddt_w_adv_ntl1: np.ndarray, ddt_w_adv_ntl2: np.ndarray, @@ -63,6 +63,7 @@ def reference( cpd: float, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) ( z_w_expl, z_contr_w_fl_l, @@ -82,17 +83,17 @@ def reference( ) return dict(z_w_expl=z_w_expl, z_contr_w_fl_l=z_contr_w_fl_l) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - w_nnow = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - ddt_w_adv_ntl1 = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - ddt_w_adv_ntl2 = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - z_th_ddz_exner_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - z_w_expl = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - rho_ic = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - w_concorr_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - vwind_expl_wgt = data_alloc.random_field(grid, dims.CellDim, dtype=ta.wpfloat) - z_contr_w_fl_l = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) + w_nnow = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + ddt_w_adv_ntl1 = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + ddt_w_adv_ntl2 = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_th_ddz_exner_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_w_expl = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + rho_ic = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + w_concorr_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + vwind_expl_wgt = self.data_alloc.random_field(dims.CellDim, dtype=ta.wpfloat) + z_contr_w_fl_l = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) dtime = ta.wpfloat("5.0") wgt_nnow_vel = ta.wpfloat("8.0") wgt_nnew_vel = ta.wpfloat("9.0") diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_explicit_vertical_wind_speed_and_vertical_wind_times_density.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_explicit_vertical_wind_speed_and_vertical_wind_times_density.py index bdaf78fbc3..198a865f28 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_explicit_vertical_wind_speed_and_vertical_wind_times_density.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_explicit_vertical_wind_speed_and_vertical_wind_times_density.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -18,12 +19,11 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_explicit_vertical_wind_speed_and_vertical_wind_times_density_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], w_nnow: np.ndarray, ddt_w_adv_ntl1: np.ndarray, z_th_ddz_exner_c: np.ndarray, @@ -39,13 +39,13 @@ def compute_explicit_vertical_wind_speed_and_vertical_wind_times_density_numpy( return (z_w_expl, z_contr_w_fl_l) -class TestComputeExplicitVerticalWindSpeedAndVerticalWindTimesDensity(StencilTest): +class TestComputeExplicitVerticalWindSpeedAndVerticalWindTimesDensity(stencil_tests.StencilTest): PROGRAM = compute_explicit_vertical_wind_speed_and_vertical_wind_times_density OUTPUTS = ("z_w_expl", "z_contr_w_fl_l") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, w_nnow: np.ndarray, ddt_w_adv_ntl1: np.ndarray, z_th_ddz_exner_c: np.ndarray, @@ -56,6 +56,7 @@ def reference( cpd: float, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) ( z_w_expl, z_contr_w_fl_l, @@ -72,16 +73,16 @@ def reference( ) return dict(z_w_expl=z_w_expl, z_contr_w_fl_l=z_contr_w_fl_l) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - w_nnow = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - ddt_w_adv_ntl1 = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - z_th_ddz_exner_c = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - z_w_expl = zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - rho_ic = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - w_concorr_c = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - vwind_expl_wgt = random_field(grid, dims.CellDim, dtype=wpfloat) - z_contr_w_fl_l = zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) + w_nnow = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + ddt_w_adv_ntl1 = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + z_th_ddz_exner_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + z_w_expl = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat) + rho_ic = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + w_concorr_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + vwind_expl_wgt = self.data_alloc.random_field(dims.CellDim, dtype=wpfloat) + z_contr_w_fl_l = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat) dtime = wpfloat("5.0") cpd = wpfloat("10.0") diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_graddiv2_of_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_graddiv2_of_vn.py index 3f9eca2afc..b563848b45 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_graddiv2_of_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_graddiv2_of_vn.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -16,16 +17,15 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_graddiv2_of_vn_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], geofac_grdiv: np.ndarray, z_graddiv_vn: np.ndarray, ) -> np.ndarray: - e2c2eO = connectivities[dims.E2C2EODim] + e2c2eO = connectivities[dims.E2C2EO] geofac_grdiv = np.expand_dims(geofac_grdiv, axis=-1) z_graddiv2_vn = np.sum( np.where((e2c2eO != -1)[:, :, np.newaxis], z_graddiv_vn[e2c2eO] * geofac_grdiv, 0), @@ -35,25 +35,26 @@ def compute_graddiv2_of_vn_numpy( @pytest.mark.embedded_remap_error -class TestComputeGraddiv2OfVn(StencilTest): +class TestComputeGraddiv2OfVn(stencil_tests.StencilTest): PROGRAM = compute_graddiv2_of_vn OUTPUTS = ("z_graddiv2_vn",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, geofac_grdiv: np.ndarray, z_graddiv_vn: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) z_graddiv2_vn = compute_graddiv2_of_vn_numpy(connectivities, geofac_grdiv, z_graddiv_vn) return dict(z_graddiv2_vn=z_graddiv2_vn) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - z_graddiv_vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - geofac_grdiv = random_field(grid, dims.EdgeDim, dims.E2C2EODim, dtype=wpfloat) - z_graddiv2_vn = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) + z_graddiv_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + geofac_grdiv = self.data_alloc.random_field(dims.EdgeDim, dims.E2C2EODim, dtype=wpfloat) + z_graddiv2_vn = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) return dict( geofac_grdiv=geofac_grdiv, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_advection_of_rho_and_theta.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_advection_of_rho_and_theta.py index 1b8d733e2b..45e9f726a2 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_advection_of_rho_and_theta.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_advection_of_rho_and_theta.py @@ -5,13 +5,13 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np import pytest -import icon4py.model.common.utils.data_allocation as data_alloc from icon4py.model.atmosphere.dycore.stencils.compute_horizontal_advection_of_rho_and_theta import ( _compute_horizontal_advection_of_rho_and_theta, ) @@ -23,13 +23,13 @@ # TODO(): copied from `test_mo_math_gradients_grad_green_gauss_cell_dsl_numpy`. delete that test? def mo_math_gradients_grad_green_gauss_cell_dsl_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], p_ccpr1: np.ndarray, p_ccpr2: np.ndarray, geofac_grg_x: np.ndarray, geofac_grg_y: np.ndarray, ) -> tuple[np.ndarray, ...]: - c2e2cO = connectivities[dims.C2E2CODim] + c2e2cO = connectivities[dims.C2E2CO] geofac_grg_x = np.expand_dims(geofac_grg_x, axis=-1) p_grad_1_u = np.sum( np.where((c2e2cO != -1)[:, :, np.newaxis], geofac_grg_x * p_ccpr1[c2e2cO], 0), axis=1 @@ -95,7 +95,7 @@ def compute_btraj_numpy( def sten_16_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], p_vn: np.ndarray, rho_ref_me: np.ndarray, theta_ref_me: np.ndarray, @@ -109,7 +109,7 @@ def sten_16_numpy( z_rth_pr_2: np.ndarray, **kwargs: Any, ) -> tuple[np.ndarray, np.ndarray]: - e2c = connectivities[dims.E2CDim] + e2c = connectivities[dims.E2C] z_rth_pr_1_e2c = z_rth_pr_1[e2c] z_rth_pr_2_e2c = z_rth_pr_2[e2c] z_grad_rth_1_e2c = z_grad_rth_1[e2c] @@ -145,7 +145,7 @@ def sten_16_numpy( def compute_horizontal_advection_of_rho_and_theta_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], p_vn: np.ndarray, p_vt: np.ndarray, pos_on_tplane_e_1: np.ndarray, @@ -211,9 +211,9 @@ class TestComputeHorizontalAvectionOfRhoAndTheta(stencil_tests.StencilTest): PROGRAM = _compute_horizontal_advection_of_rho_and_theta OUTPUTS = ("out",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_vn: np.ndarray, p_vt: np.ndarray, pos_on_tplane_e_1: np.ndarray, @@ -231,6 +231,7 @@ def reference( geofac_grg_y: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) z_rho_e, z_theta_v_e = compute_horizontal_advection_of_rho_and_theta_numpy( connectivities, p_vn, @@ -251,44 +252,44 @@ def reference( ) return dict(out=(z_rho_e, z_theta_v_e)) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data( self, grid: base.Grid ) -> dict[str, gtx.Field | state_utils.ScalarType | gtx.Domain | tuple[gtx.Field, ...]]: - p_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) - p_vt = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) - pos_on_tplane_e_1 = data_alloc.random_field( - grid, dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat + p_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + p_vt = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + pos_on_tplane_e_1 = self.data_alloc.random_field( + dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat ) - pos_on_tplane_e_2 = data_alloc.random_field( - grid, dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat + pos_on_tplane_e_2 = self.data_alloc.random_field( + dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat ) - primal_normal_cell_1 = data_alloc.random_field( - grid, dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat + primal_normal_cell_1 = self.data_alloc.random_field( + dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat ) - dual_normal_cell_1 = data_alloc.random_field( - grid, dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat + dual_normal_cell_1 = self.data_alloc.random_field( + dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat ) - primal_normal_cell_2 = data_alloc.random_field( - grid, dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat + primal_normal_cell_2 = self.data_alloc.random_field( + dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat ) - dual_normal_cell_2 = data_alloc.random_field( - grid, dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat + dual_normal_cell_2 = self.data_alloc.random_field( + dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat ) p_dthalf = 2.0 - rho_ref_me = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) - theta_ref_me = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) - perturbed_rho_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat + rho_ref_me = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + theta_ref_me = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + perturbed_rho_at_cells_on_model_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim, dtype=ta.vpfloat ) - perturbed_theta_v_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat + perturbed_theta_v_at_cells_on_model_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim, dtype=ta.vpfloat ) - geofac_grg_x = data_alloc.random_field(grid, dims.CellDim, dims.C2E2CODim, dtype=ta.wpfloat) - geofac_grg_y = data_alloc.random_field(grid, dims.CellDim, dims.C2E2CODim, dtype=ta.wpfloat) - z_rho_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) - z_theta_v_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + geofac_grg_x = self.data_alloc.random_field(dims.CellDim, dims.C2E2CODim, dtype=ta.wpfloat) + geofac_grg_y = self.data_alloc.random_field(dims.CellDim, dims.C2E2CODim, dtype=ta.wpfloat) + z_rho_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + z_theta_v_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) return dict( p_vn=p_vn, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_advection_term_for_vertical_velocity.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_advection_term_for_vertical_velocity.py index 0c47cf433c..971306dadf 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_advection_term_for_vertical_velocity.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_advection_term_for_vertical_velocity.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -18,12 +19,11 @@ from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_horizontal_advection_term_for_vertical_velocity_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], vn_ie: np.ndarray, inv_dual_edge_length: np.ndarray, w: np.ndarray, @@ -36,8 +36,8 @@ def compute_horizontal_advection_term_for_vertical_velocity_numpy( inv_primal_edge_length = np.expand_dims(inv_primal_edge_length, axis=-1) tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) - w_e2c = w[connectivities[dims.E2CDim]] - z_w_v_e2v = z_w_v[connectivities[dims.E2VDim]] + w_e2c = w[connectivities[dims.E2C]] + z_w_v_e2v = z_w_v[connectivities[dims.E2V]] red_w = w_e2c[:, 0] - w_e2c[:, 1] red_z_w_v = z_w_v_e2v[:, 0] - z_w_v_e2v[:, 1] @@ -49,13 +49,13 @@ def compute_horizontal_advection_term_for_vertical_velocity_numpy( return z_v_grad_w -class TestComputeHorizontalAdvectionTermForVerticalVelocity(StencilTest): +class TestComputeHorizontalAdvectionTermForVerticalVelocity(stencil_tests.StencilTest): PROGRAM = compute_horizontal_advection_term_for_vertical_velocity OUTPUTS = ("z_v_grad_w",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, vn_ie: np.ndarray, inv_dual_edge_length: np.ndarray, w: np.ndarray, @@ -68,6 +68,7 @@ def reference( horizontal_end: int, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) z_v_grad_w[horizontal_start:horizontal_end, :] = ( compute_horizontal_advection_term_for_vertical_velocity_numpy( connectivities, @@ -82,16 +83,16 @@ def reference( )[horizontal_start:horizontal_end, :] return dict(z_v_grad_w=z_v_grad_w) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - vn_ie = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - inv_dual_edge_length = random_field(grid, dims.EdgeDim, dtype=wpfloat) - w = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - z_vt_ie = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - inv_primal_edge_length = random_field(grid, dims.EdgeDim, dtype=wpfloat) - tangent_orientation = random_field(grid, dims.EdgeDim, dtype=wpfloat) - z_w_v = random_field(grid, dims.VertexDim, dims.KDim, dtype=vpfloat) - z_v_grad_w = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) + vn_ie = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + inv_dual_edge_length = self.data_alloc.random_field(dims.EdgeDim, dtype=wpfloat) + w = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + z_vt_ie = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + inv_primal_edge_length = self.data_alloc.random_field(dims.EdgeDim, dtype=wpfloat) + tangent_orientation = self.data_alloc.random_field(dims.EdgeDim, dtype=wpfloat) + z_w_v = self.data_alloc.random_field(dims.VertexDim, dims.KDim, dtype=vpfloat) + z_v_grad_w = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) edge_domain = h_grid.domain(dims.EdgeDim) horizontal_start = grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_7)) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates.py index 09ade5f64b..16ff2042c6 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -18,18 +19,17 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], inv_dual_edge_length: np.ndarray, z_exner_ex_pr: np.ndarray, ) -> np.ndarray: inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) - z_exner_ex_pr_e2c = z_exner_ex_pr[connectivities[dims.E2CDim]] + z_exner_ex_pr_e2c = z_exner_ex_pr[connectivities[dims.E2C]] z_exner_ex_weighted = z_exner_ex_pr_e2c[:, 1] - z_exner_ex_pr_e2c[:, 0] z_gradh_exner = inv_dual_edge_length * z_exner_ex_weighted @@ -37,27 +37,28 @@ def compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates_numpy( @pytest.mark.skip_value_error -class TestComputeHorizontalGradientOfExnerPressureForFlatCoordinates(StencilTest): +class TestComputeHorizontalGradientOfExnerPressureForFlatCoordinates(stencil_tests.StencilTest): PROGRAM = compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates OUTPUTS = ("z_gradh_exner",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, inv_dual_edge_length: np.ndarray, z_exner_ex_pr: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) z_gradh_exner = compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates_numpy( connectivities, inv_dual_edge_length, z_exner_ex_pr ) return dict(z_gradh_exner=z_gradh_exner) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - inv_dual_edge_length = random_field(grid, dims.EdgeDim, dtype=wpfloat) - z_exner_ex_pr = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - z_gradh_exner = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) + inv_dual_edge_length = self.data_alloc.random_field(dims.EdgeDim, dtype=wpfloat) + z_exner_ex_pr = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + z_gradh_exner = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) return dict( inv_dual_edge_length=inv_dual_edge_length, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_multiple_levels.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_multiple_levels.py index 53555a19e6..965900e73b 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_multiple_levels.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_multiple_levels.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -17,12 +18,11 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_horizontal_gradient_of_exner_pressure_for_multiple_levels_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], inv_dual_edge_length: np.ndarray, z_exner_ex_pr: np.ndarray, zdiff_gradp: np.ndarray, @@ -43,7 +43,7 @@ def _apply_index_field( ] return indexed - e2c = connectivities[dims.E2CDim] + e2c = connectivities[dims.E2C] inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, -1) full_shape = ikoffset.shape @@ -65,13 +65,13 @@ def at_neighbor(i: int) -> np.ndarray: @pytest.mark.skip_value_error @pytest.mark.uses_as_offset -class TestComputeHorizontalGradientOfExnerPressureForMultipleLevels(StencilTest): +class TestComputeHorizontalGradientOfExnerPressureForMultipleLevels(stencil_tests.StencilTest): PROGRAM = compute_horizontal_gradient_of_exner_pressure_for_multiple_levels OUTPUTS = ("z_gradh_exner",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, inv_dual_edge_length: np.ndarray, z_exner_ex_pr: np.ndarray, zdiff_gradp: np.ndarray, @@ -80,6 +80,7 @@ def reference( z_dexner_dz_c_2: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) z_gradh_exner = compute_horizontal_gradient_of_exner_pressure_for_multiple_levels_numpy( connectivities, inv_dual_edge_length, @@ -91,12 +92,14 @@ def reference( ) return dict(z_gradh_exner=z_gradh_exner) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - inv_dual_edge_length = random_field(grid, dims.EdgeDim, dtype=ta.wpfloat) - z_exner_ex_pr = random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - zdiff_gradp = random_field(grid, dims.EdgeDim, dims.E2CDim, dims.KDim, dtype=ta.vpfloat) - ikoffset = zero_field(grid, dims.EdgeDim, dims.E2CDim, dims.KDim, dtype=gtx.int32) + inv_dual_edge_length = self.data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat) + z_exner_ex_pr = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + zdiff_gradp = self.data_alloc.random_field( + dims.EdgeDim, dims.E2CDim, dims.KDim, dtype=ta.vpfloat + ) + ikoffset = self.data_alloc.zero_field(dims.EdgeDim, dims.E2CDim, dims.KDim, dtype=gtx.int32) rng = np.random.default_rng() for k in range(grid.num_levels): # construct offsets that reach all k-levels except the last (because we are using the entries of this field with `+1`) @@ -106,9 +109,9 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala size=(ikoffset.shape[0], ikoffset.shape[1]), ) - z_dexner_dz_c_1 = random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - z_dexner_dz_c_2 = random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - z_gradh_exner = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + z_dexner_dz_c_1 = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_dexner_dz_c_2 = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_gradh_exner = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) return dict( inv_dual_edge_length=inv_dual_edge_length, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates.py index 0369fc2b9c..d7b679e3d0 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -18,19 +19,18 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], inv_dual_edge_length: np.ndarray, z_exner_ex_pr: np.ndarray, ddxn_z_full: np.ndarray, c_lin_e: np.ndarray, z_dexner_dz_c_1: np.ndarray, ) -> np.ndarray: - e2c = connectivities[dims.E2CDim] + e2c = connectivities[dims.E2C] inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) c_lin_e = np.expand_dims(c_lin_e, axis=-1) @@ -44,13 +44,13 @@ def compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates_numpy( @pytest.mark.skip_value_error -class TestComputeHorizontalGradientOfExnerPressureForNonflatCoordinates(StencilTest): +class TestComputeHorizontalGradientOfExnerPressureForNonflatCoordinates(stencil_tests.StencilTest): PROGRAM = compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates OUTPUTS = ("z_gradh_exner",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, inv_dual_edge_length: np.ndarray, z_exner_ex_pr: np.ndarray, ddxn_z_full: np.ndarray, @@ -58,6 +58,7 @@ def reference( z_dexner_dz_c_1: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) z_gradh_exner = compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates_numpy( connectivities, inv_dual_edge_length, @@ -68,14 +69,14 @@ def reference( ) return dict(z_gradh_exner=z_gradh_exner) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - inv_dual_edge_length = random_field(grid, dims.EdgeDim, dtype=wpfloat) - z_exner_ex_pr = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - ddxn_z_full = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - c_lin_e = random_field(grid, dims.EdgeDim, dims.E2CDim, dtype=wpfloat) - z_dexner_dz_c_1 = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - z_gradh_exner = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) + inv_dual_edge_length = self.data_alloc.random_field(dims.EdgeDim, dtype=wpfloat) + z_exner_ex_pr = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + ddxn_z_full = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + c_lin_e = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim, dtype=wpfloat) + z_dexner_dz_c_1 = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + z_gradh_exner = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) return dict( inv_dual_edge_length=inv_dual_edge_length, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_kinetic_energy.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_kinetic_energy.py index a81376ec23..87a1cbf62f 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_kinetic_energy.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_kinetic_energy.py @@ -18,8 +18,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_horizontal_kinetic_energy_numpy(vn: np.ndarray, vt: np.ndarray) -> tuple: @@ -29,13 +28,13 @@ def compute_horizontal_kinetic_energy_numpy(vn: np.ndarray, vt: np.ndarray) -> t return vn_ie, z_vt_ie, z_kin_hor_e -class TestComputeHorizontalKineticEnergy(StencilTest): +class TestComputeHorizontalKineticEnergy(stencil_tests.StencilTest): PROGRAM = compute_horizontal_kinetic_energy OUTPUTS = ("vn_ie", "z_vt_ie", "z_kin_hor_e") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, vn: np.ndarray, vt: np.ndarray, **kwargs: Any, @@ -43,14 +42,14 @@ def reference( vn_ie, z_vt_ie, z_kin_hor_e = compute_horizontal_kinetic_energy_numpy(vn, vt) return dict(vn_ie=vn_ie, z_vt_ie=z_vt_ie, z_kin_hor_e=z_kin_hor_e) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - vt = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + vt = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) - vn_ie = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - z_vt_ie = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - z_kin_hor_e = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) + vn_ie = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + z_vt_ie = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + z_kin_hor_e = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) return dict( vn=vn, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_velocity_quantities_and_fluxes.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_velocity_quantities_and_fluxes.py index ed59ee339d..5e302d48d1 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_velocity_quantities_and_fluxes.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_velocity_quantities_and_fluxes.py @@ -5,13 +5,13 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np import pytest -import icon4py.model.common.utils.data_allocation as data_alloc from icon4py.model.atmosphere.dycore.stencils.compute_horizontal_velocity_quantities import ( compute_horizontal_velocity_quantities_and_fluxes, ) @@ -32,7 +32,7 @@ def compute_vt_vn_on_half_levels_and_kinetic_energy_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], vn: np.ndarray, tangential_wind: np.ndarray, vn_on_half_levels: np.ndarray, @@ -113,9 +113,9 @@ class TestComputeHorizontalVelocityQuantitiesAndFluxes(stencil_tests.StencilTest ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, spatially_averaged_vn: np.ndarray, horizontal_gradient_of_normal_wind_divergence: np.ndarray, tangential_wind: np.ndarray, @@ -143,6 +143,7 @@ def reference( vertical_end: gtx.int32, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) k = np.arange(vertical_end)[np.newaxis, :] k_nlev = k[:, :-1] @@ -273,37 +274,37 @@ def reference( contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - spatially_averaged_vn = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - horizontal_gradient_of_normal_wind_divergence = data_alloc.zero_field( - grid, dims.EdgeDim, dims.KDim + spatially_averaged_vn = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + horizontal_gradient_of_normal_wind_divergence = self.data_alloc.zero_field( + dims.EdgeDim, dims.KDim ) - tangential_wind = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - mass_flux_at_edges_on_model_levels = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - theta_v_flux_at_edges_on_model_levels = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - tangential_wind_on_half_levels = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - vn_on_half_levels = data_alloc.zero_field( - grid, dims.EdgeDim, dims.KDim, extend={dims.KDim: 1} + tangential_wind = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + mass_flux_at_edges_on_model_levels = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + theta_v_flux_at_edges_on_model_levels = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + tangential_wind_on_half_levels = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim) + vn_on_half_levels = self.data_alloc.zero_field( + dims.EdgeDim, dims.KDim, extend={dims.KDim: 1} ) - horizontal_kinetic_energy_at_edges_on_model_levels = data_alloc.zero_field( - grid, dims.EdgeDim, dims.KDim + horizontal_kinetic_energy_at_edges_on_model_levels = self.data_alloc.zero_field( + dims.EdgeDim, dims.KDim ) - contravariant_correction_at_edges_on_model_levels = data_alloc.zero_field( - grid, dims.EdgeDim, dims.KDim + contravariant_correction_at_edges_on_model_levels = self.data_alloc.zero_field( + dims.EdgeDim, dims.KDim ) - vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - wgtfac_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - wgtfacq_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - e_flx_avg = data_alloc.random_field(grid, dims.EdgeDim, dims.E2C2EODim) - geofac_grdiv = data_alloc.random_field(grid, dims.EdgeDim, dims.E2C2EODim) - rbf_vec_coeff_e = data_alloc.random_field(grid, dims.EdgeDim, dims.E2C2EDim) - rho_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - theta_v_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - ddqz_z_full_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - ddxn_z_full = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - ddxt_z_full = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + wgtfac_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + wgtfacq_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + e_flx_avg = self.data_alloc.random_field(dims.EdgeDim, dims.E2C2EODim) + geofac_grdiv = self.data_alloc.random_field(dims.EdgeDim, dims.E2C2EODim) + rbf_vec_coeff_e = self.data_alloc.random_field(dims.EdgeDim, dims.E2C2EDim) + rho_at_edges_on_model_levels = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + theta_v_at_edges_on_model_levels = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + ddqz_z_full_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + ddxn_z_full = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + ddxt_z_full = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) nflatlev = 5 # value is set to reflect the MCH ch1 experiment. Changing this value will change the expected runtime diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py index 7eed5b8929..eaf52b88e5 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -17,12 +18,11 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest +from icon4py.model.testing import stencil_tests def compute_hydrostatic_correction_term_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], theta_v: np.ndarray, ikoffset: np.ndarray, zdiff_gradp: np.ndarray, @@ -48,7 +48,7 @@ def _apply_index_field( ] return indexed, indexed_p1 - e2c = connectivities[dims.E2CDim] + e2c = connectivities[dims.E2C] full_shape = ikoffset.shape inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, -1) @@ -89,26 +89,26 @@ def _apply_index_field( @pytest.mark.uses_as_offset @pytest.mark.continuous_benchmarking -class TestComputeHydrostaticCorrectionTerm(StencilTest): +class TestComputeHydrostaticCorrectionTerm(stencil_tests.StencilTest): OUTPUTS = ("z_hydro_corr",) PROGRAM = compute_hydrostatic_correction_term STATIC_PARAMS = { - StandardStaticVariants.NONE: (), - StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", "vertical_start", "vertical_end", ), - StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( "vertical_start", "vertical_end", ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, theta_v: np.ndarray, ikoffset: np.ndarray, zdiff_gradp: np.ndarray, @@ -118,6 +118,7 @@ def reference( grav_o_cpd: float, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) z_hydro_corr = compute_hydrostatic_correction_term_numpy( connectivities, theta_v, @@ -130,11 +131,9 @@ def reference( ) return dict(z_hydro_corr=z_hydro_corr) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - ikoffset = data_alloc.zero_field( - grid, dims.EdgeDim, dims.E2CDim, dims.KDim, dtype=gtx.int32 - ) + ikoffset = self.data_alloc.zero_field(dims.EdgeDim, dims.E2CDim, dims.KDim, dtype=gtx.int32) rng = np.random.default_rng() for k in range(grid.num_levels): # construct offsets that reach all k-levels except the last (because we are using the entries of this field with `+1`) @@ -144,16 +143,16 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala size=(ikoffset.shape[0], ikoffset.shape[1]), ) - theta_v = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - zdiff_gradp = data_alloc.random_field( - grid, dims.EdgeDim, dims.E2CDim, dims.KDim, dtype=ta.vpfloat + theta_v = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + zdiff_gradp = self.data_alloc.random_field( + dims.EdgeDim, dims.E2CDim, dims.KDim, dtype=ta.vpfloat ) - theta_v_ic = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - inv_ddqz_z_full = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - inv_dual_edge_length = data_alloc.random_field(grid, dims.EdgeDim, dtype=ta.wpfloat) + theta_v_ic = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + inv_ddqz_z_full = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + inv_dual_edge_length = self.data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat) grav_o_cpd = ta.wpfloat("10.0") - z_hydro_corr = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + z_hydro_corr = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) z_hydro_corr = gtx.constructors.zeros( domain={ diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_mass_flux.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_mass_flux.py index 330d0410a0..5dbd968ff0 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_mass_flux.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_mass_flux.py @@ -16,8 +16,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_mass_flux_numpy( @@ -32,13 +31,13 @@ def compute_mass_flux_numpy( return mass_fl_e, z_theta_v_fl_e -class TestComputeMassFlux(StencilTest): +class TestComputeMassFlux(stencil_tests.StencilTest): PROGRAM = compute_mass_flux OUTPUTS = ("mass_fl_e", "z_theta_v_fl_e") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_rho_e: np.ndarray, z_vn_avg: np.ndarray, ddqz_z_full_e: np.ndarray, @@ -54,14 +53,14 @@ def reference( return dict(mass_fl_e=mass_fl_e, z_theta_v_fl_e=z_theta_v_fl_e) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - z_rho_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - z_vn_avg = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - ddqz_z_full_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - mass_fl_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - z_theta_v_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - z_theta_v_fl_e = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + z_rho_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + z_vn_avg = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + ddqz_z_full_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + mass_fl_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + z_theta_v_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + z_theta_v_fl_e = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) return dict( z_rho_e=z_rho_e, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py index 2cd2a83c8f..3a43b8f3d4 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py @@ -17,8 +17,7 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils.data_allocation import random_field, random_mask, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_maximum_cfl_and_clip_contravariant_vertical_velocity_numpy( @@ -43,13 +42,13 @@ def compute_maximum_cfl_and_clip_contravariant_vertical_velocity_numpy( return cfl_clipping, vcfl, z_w_con_c -class TestComputeMaximumCflAndClipContravariantVerticalVelocity(StencilTest): +class TestComputeMaximumCflAndClipContravariantVerticalVelocity(stencil_tests.StencilTest): PROGRAM = compute_maximum_cfl_and_clip_contravariant_vertical_velocity OUTPUTS = ("cfl_clipping", "vcfl", "z_w_con_c") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, ddqz_z_half: np.ndarray, z_w_con_c: np.ndarray, cfl_w_limit: ta.wpfloat, @@ -70,12 +69,12 @@ def reference( z_w_con_c=z_w_con_c, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - ddqz_z_half = random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - z_w_con_c = random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - cfl_clipping = random_mask(grid, dims.CellDim, dims.KDim, dtype=bool) - vcfl = zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) + ddqz_z_half = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_w_con_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + cfl_clipping = self.data_alloc.random_mask(dims.CellDim, dims.KDim, dtype=bool) + vcfl = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) cfl_w_limit = ta.vpfloat("5.0") dtime = ta.wpfloat("9.0") diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbation_of_rho_and_theta.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbation_of_rho_and_theta.py index 79d910c5b9..9ab18a1d48 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbation_of_rho_and_theta.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbation_of_rho_and_theta.py @@ -18,8 +18,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_perturbation_of_rho_and_theta_numpy( @@ -33,13 +32,13 @@ def compute_perturbation_of_rho_and_theta_numpy( return (z_rth_pr_1, z_rth_pr_2) -class TestComputePerturbationOfRhoAndTheta(StencilTest): +class TestComputePerturbationOfRhoAndTheta(stencil_tests.StencilTest): PROGRAM = compute_perturbation_of_rho_and_theta OUTPUTS = ("z_rth_pr_1", "z_rth_pr_2") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, rho: np.ndarray, rho_ref_mc: np.ndarray, theta_v: np.ndarray, @@ -51,14 +50,14 @@ def reference( ) return dict(z_rth_pr_1=z_rth_pr_1, z_rth_pr_2=z_rth_pr_2) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - rho = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - rho_ref_mc = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - theta_v = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - theta_ref_mc = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - z_rth_pr_1 = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - z_rth_pr_2 = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + rho = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + rho_ref_mc = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + theta_v = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + theta_ref_mc = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + z_rth_pr_1 = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) + z_rth_pr_2 = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) return dict( rho=rho, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers.py index e224318650..d916f0afec 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers.py @@ -18,8 +18,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers_numpy( @@ -39,13 +38,13 @@ def compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers_numpy( return rho_ic, z_rth_pr_1, z_rth_pr_2 -class TestComputePerturbationOfRhoAndThetaAndRhoInterfaceCellCenters(StencilTest): +class TestComputePerturbationOfRhoAndThetaAndRhoInterfaceCellCenters(stencil_tests.StencilTest): PROGRAM = compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers OUTPUTS = ("rho_ic", "z_rth_pr_1", "z_rth_pr_2") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, wgtfac_c: np.ndarray, rho: np.ndarray, rho_ref_mc: np.ndarray, @@ -66,16 +65,16 @@ def reference( ) return dict(rho_ic=rho_ic, z_rth_pr_1=z_rth_pr_1, z_rth_pr_2=z_rth_pr_2) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - wgtfac_c = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - rho = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - rho_ref_mc = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - theta_v = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - theta_ref_mc = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - rho_ic = zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - z_rth_pr_1 = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - z_rth_pr_2 = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + wgtfac_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + rho = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + rho_ref_mc = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + theta_v = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + theta_ref_mc = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + rho_ic = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat) + z_rth_pr_1 = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) + z_rth_pr_2 = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) return dict( wgtfac_c=wgtfac_c, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py index 9611d9bdc1..35e2ca2c69 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py @@ -18,7 +18,8 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -31,7 +32,6 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests from .test_compute_approx_of_2nd_vertical_derivative_of_exner import ( @@ -104,9 +104,9 @@ class TestComputePerturbedQuantitiesAndInterpolation(stencil_tests.StencilTest): ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, current_rho: np.ndarray, reference_rho_at_cells_on_model_levels: np.ndarray, current_theta_v: np.ndarray, @@ -144,6 +144,7 @@ def reference( end_cell_halo_level_2: gtx.int32, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) vert_idx = np.arange(kwargs["surface_level"]) cell = np.arange(end_cell_halo_level_2) horz_idx = cell[:, np.newaxis] @@ -347,65 +348,63 @@ def reference( d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - reference_rho_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim + reference_rho_at_cells_on_model_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim ) - reference_theta_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim + reference_theta_at_cells_on_model_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim ) - wgtfacq_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) - perturbed_rho_at_cells_on_model_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim + wgtfacq_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, extend={dims.KDim: 1}) + perturbed_rho_at_cells_on_model_levels = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + perturbed_theta_v_at_cells_on_model_levels = self.data_alloc.zero_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - perturbed_theta_v_at_cells_on_model_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + perturbed_theta_v_at_cells_on_half_levels = self.data_alloc.zero_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - perturbed_theta_v_at_cells_on_half_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + reference_theta_at_cells_on_half_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - reference_theta_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + d2dexdz2_fac1_mc = self.data_alloc.random_field(dims.CellDim, dims.KDim) + d2dexdz2_fac2_mc = self.data_alloc.random_field(dims.CellDim, dims.KDim) + wgtfac_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, extend={dims.KDim: 1}) + exner_w_explicit_weight_parameter = self.data_alloc.random_field(dims.CellDim) + perturbed_exner_at_cells_on_model_levels = self.data_alloc.zero_field( + dims.CellDim, dims.KDim ) - d2dexdz2_fac1_mc = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - d2dexdz2_fac2_mc = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - wgtfac_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) - exner_w_explicit_weight_parameter = data_alloc.random_field(grid, dims.CellDim) - perturbed_exner_at_cells_on_model_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim + ddz_of_reference_exner_at_cells_on_half_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim ) - ddz_of_reference_exner_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim + ddqz_z_half = self.data_alloc.random_field(dims.CellDim, dims.KDim) + nonhydro_buoy_at_cells_on_half_levels = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + rho_at_cells_on_half_levels = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + exner_at_cells_on_half_levels = self.data_alloc.zero_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - ddqz_z_half = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - nonhydro_buoy_at_cells_on_half_levels = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - rho_at_cells_on_half_levels = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - exner_at_cells_on_half_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + time_extrapolation_parameter_for_exner = self.data_alloc.random_field( + dims.CellDim, dims.KDim ) - time_extrapolation_parameter_for_exner = data_alloc.random_field( - grid, dims.CellDim, dims.KDim + current_exner = self.data_alloc.random_field(dims.CellDim, dims.KDim) + reference_exner_at_cells_on_model_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim ) - current_exner = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - reference_exner_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim + temporal_extrapolation_of_perturbed_exner = self.data_alloc.zero_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - temporal_extrapolation_of_perturbed_exner = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( + self.data_alloc.zero_field(dims.CellDim, dims.KDim) ) - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim + d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( + self.data_alloc.zero_field(dims.CellDim, dims.KDim) ) - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim + theta_v_at_cells_on_half_levels = self.data_alloc.zero_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - theta_v_at_cells_on_half_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} - ) - inv_ddqz_z_full = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - current_rho = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - current_theta_v = data_alloc.random_field(grid, dims.CellDim, dims.KDim) + inv_ddqz_z_full = self.data_alloc.random_field(dims.CellDim, dims.KDim) + current_rho = self.data_alloc.random_field(dims.CellDim, dims.KDim) + current_theta_v = self.data_alloc.random_field(dims.CellDim, dims.KDim) igradp_method = horzpres_discr_type.TAYLOR_HYDRO diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_results_for_thermodynamic_variables.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_results_for_thermodynamic_variables.py index 2efbdbddd9..ff3dce9086 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_results_for_thermodynamic_variables.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_results_for_thermodynamic_variables.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any, Final +from collections.abc import Mapping +from typing import Any, Final, cast import gt4py.next as gtx import numpy as np @@ -17,15 +18,14 @@ from icon4py.model.common import constants, dimension as dims, type_alias as ta from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests dycore_consts: Final = constants.PhysicsConstants() def compute_results_for_thermodynamic_variables_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], z_rho_expl: np.ndarray, vwind_impl_wgt: np.ndarray, inv_ddqz_z_full: np.ndarray, @@ -62,13 +62,13 @@ def compute_results_for_thermodynamic_variables_numpy( return rho_new, exner_new, theta_v_new -class TestComputeResultsForThermodynamicVariables(StencilTest): +class TestComputeResultsForThermodynamicVariables(stencil_tests.StencilTest): PROGRAM = compute_results_for_thermodynamic_variables OUTPUTS = ("rho_new", "exner_new", "theta_v_new") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_rho_expl: np.ndarray, vwind_impl_wgt: np.ndarray, inv_ddqz_z_full: np.ndarray, @@ -84,6 +84,7 @@ def reference( dtime: float, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) (rho_new, exner_new, theta_v_new) = compute_results_for_thermodynamic_variables_numpy( connectivities, z_rho_expl=z_rho_expl, @@ -102,29 +103,29 @@ def reference( ) return dict(rho_new=rho_new, exner_new=exner_new, theta_v_new=theta_v_new) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - z_rho_expl = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - vwind_impl_wgt = data_alloc.random_field(grid, dims.CellDim, dtype=ta.wpfloat) - inv_ddqz_z_full = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - rho_ic = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=ta.wpfloat + z_rho_expl = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + vwind_impl_wgt = self.data_alloc.random_field(dims.CellDim, dtype=ta.wpfloat) + inv_ddqz_z_full = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + rho_ic = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=ta.wpfloat ) - w = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=ta.wpfloat + w = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=ta.wpfloat ) - z_exner_expl = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - exner_ref_mc = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - z_alpha = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=ta.vpfloat + z_exner_expl = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + exner_ref_mc = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_alpha = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=ta.vpfloat ) - z_beta = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - rho_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - theta_v_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - exner_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - rho_new = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - exner_new = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - theta_v_new = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) + z_beta = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + rho_now = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + theta_v_now = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + exner_now = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + rho_new = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + exner_new = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + theta_v_new = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) dtime = ta.wpfloat("5.0") return dict( diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_rho_virtual_potential_temperatures_and_pressure_gradient.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_rho_virtual_potential_temperatures_and_pressure_gradient.py index 6f71a1f98a..8479287b53 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_rho_virtual_potential_temperatures_and_pressure_gradient.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_rho_virtual_potential_temperatures_and_pressure_gradient.py @@ -17,8 +17,7 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_rho_virtual_potential_temperatures_and_pressure_gradient_numpy( @@ -75,13 +74,13 @@ def compute_rho_virtual_potential_temperatures_and_pressure_gradient_numpy( return (rho_ic, z_theta_v_pr_ic, theta_v_ic, z_th_ddz_exner_c) -class TestComputeRhoVirtualPotentialTemperaturesAndPressureGradient(StencilTest): +class TestComputeRhoVirtualPotentialTemperaturesAndPressureGradient(stencil_tests.StencilTest): PROGRAM = compute_rho_virtual_potential_temperatures_and_pressure_gradient OUTPUTS = ("rho_ic", "z_theta_v_pr_ic", "theta_v_ic", "z_th_ddz_exner_c") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, w: np.ndarray, w_concorr_c: np.ndarray, ddqz_z_half: np.ndarray, @@ -128,27 +127,27 @@ def reference( z_th_ddz_exner_c=z_th_ddz_exner_c, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: dtime = ta.wpfloat("1.0") wgt_nnow_rth = ta.wpfloat("2.0") wgt_nnew_rth = ta.wpfloat("3.0") - w = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - w_concorr_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - ddqz_z_half = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - rho_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - rho_var = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - theta_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - theta_var = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - wgtfac_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - theta_ref_mc = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - vwind_expl_wgt = data_alloc.random_field(grid, dims.CellDim, dtype=ta.wpfloat) - exner_pr = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - d_exner_dz_ref_ic = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - rho_ic = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - z_theta_v_pr_ic = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - theta_v_ic = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - z_th_ddz_exner_c = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) + w = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + w_concorr_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + ddqz_z_half = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + rho_now = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + rho_var = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + theta_now = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + theta_var = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + wgtfac_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + theta_ref_mc = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + vwind_expl_wgt = self.data_alloc.random_field(dims.CellDim, dtype=ta.wpfloat) + exner_pr = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + d_exner_dz_ref_ic = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + rho_ic = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + z_theta_v_pr_ic = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + theta_v_ic = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + z_th_ddz_exner_c = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) return dict( w=w, w_concorr_c=w_concorr_c, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_solver_coefficients_matrix.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_solver_coefficients_matrix.py index da06853a0a..a4b5243d40 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_solver_coefficients_matrix.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_solver_coefficients_matrix.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -17,12 +18,11 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_solver_coefficients_matrix_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], exner_nnow: np.ndarray, rho_nnow: np.ndarray, theta_v_nnow: np.ndarray, @@ -40,13 +40,13 @@ def compute_solver_coefficients_matrix_numpy( return (z_beta, z_alpha) -class TestComputeSolverCoefficientsMatrix(StencilTest): +class TestComputeSolverCoefficientsMatrix(stencil_tests.StencilTest): PROGRAM = compute_solver_coefficients_matrix OUTPUTS = ("z_beta", "z_alpha") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, exner_nnow: np.ndarray, rho_nnow: np.ndarray, theta_v_nnow: np.ndarray, @@ -59,6 +59,7 @@ def reference( cvd: ta.wpfloat, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) (z_beta, z_alpha) = compute_solver_coefficients_matrix_numpy( connectivities, exner_nnow=exner_nnow, @@ -74,17 +75,17 @@ def reference( ) return dict(z_beta=z_beta, z_alpha=z_alpha) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - exner_nnow = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - rho_nnow = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - theta_v_nnow = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - inv_ddqz_z_full = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - vwind_impl_wgt = data_alloc.random_field(grid, dims.CellDim, dtype=ta.wpfloat) - theta_v_ic = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - rho_ic = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - z_alpha = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - z_beta = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) + exner_nnow = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + rho_nnow = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + theta_v_nnow = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + inv_ddqz_z_full = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + vwind_impl_wgt = self.data_alloc.random_field(dims.CellDim, dtype=ta.wpfloat) + theta_v_ic = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + rho_ic = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + z_alpha = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_beta = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) dtime = ta.wpfloat("10.0") rd = ta.wpfloat("5.0") cvd = ta.wpfloat("3.0") diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_tangential_wind.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_tangential_wind.py index fbcc397990..3309827784 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_tangential_wind.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_tangential_wind.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -16,39 +17,41 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_tangential_wind_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], vn: np.ndarray, rbf_vec_coeff_e: np.ndarray + connectivities: Mapping[gtx.FieldOffset, np.ndarray], + vn: np.ndarray, + rbf_vec_coeff_e: np.ndarray, ) -> np.ndarray: rbf_vec_coeff_e = np.expand_dims(rbf_vec_coeff_e, axis=-1) - e2c2e = connectivities[dims.E2C2EDim] + e2c2e = connectivities[dims.E2C2E] vt = np.sum(np.where((e2c2e != -1)[:, :, np.newaxis], vn[e2c2e] * rbf_vec_coeff_e, 0), axis=1) return vt @pytest.mark.embedded_remap_error -class TestComputeTangentialWind(StencilTest): +class TestComputeTangentialWind(stencil_tests.StencilTest): PROGRAM = compute_tangential_wind OUTPUTS = ("vt",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, vn: np.ndarray, rbf_vec_coeff_e: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) vt = compute_tangential_wind_numpy(connectivities, vn, rbf_vec_coeff_e) return dict(vt=vt) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - rbf_vec_coeff_e = random_field(grid, dims.EdgeDim, dims.E2C2EDim, dtype=wpfloat) - vt = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + rbf_vec_coeff_e = self.data_alloc.random_field(dims.EdgeDim, dims.E2C2EDim, dtype=wpfloat) + vt = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) return dict( vn=vn, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py index ec566da678..734367b11f 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py @@ -6,7 +6,8 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -22,7 +23,6 @@ ) from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -31,7 +31,7 @@ def compute_theta_rho_face_value_by_miura_scheme_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], vn: np.ndarray, tangential_wind: np.ndarray, pos_on_tplane_e_x: np.ndarray, @@ -51,7 +51,7 @@ def compute_theta_rho_face_value_by_miura_scheme_numpy( perturbed_theta_v_at_cells_on_model_levels: np.ndarray, **kwargs: Any, ) -> tuple[np.ndarray, np.ndarray]: - e2c = connectivities[dims.E2CDim] + e2c = connectivities[dims.E2C] pos_on_tplane_e_x = pos_on_tplane_e_x.reshape(e2c.shape) pos_on_tplane_e_y = pos_on_tplane_e_y.reshape(e2c.shape) primal_normal_cell_x = primal_normal_cell_x.reshape(e2c.shape) @@ -160,9 +160,9 @@ class TestComputeThetaRhoPressureGradientAndUpdateVn(stencil_tests.StencilTest): ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, rho_at_edges_on_model_levels: np.ndarray, theta_v_at_edges_on_model_levels: np.ndarray, horizontal_pressure_gradient: np.ndarray, @@ -210,6 +210,7 @@ def reference( vertical_end: gtx.int32, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) vert_idx = np.arange(vertical_end) horz_idx = np.arange(horizontal_end)[:, np.newaxis] default_shape = perturbed_rho_at_cells_on_model_levels.shape @@ -220,7 +221,7 @@ def reference( ddy_perturbed_theta_v = np.zeros(default_shape) # Compute Green-Gauss gradients for rho and theta - c2e2cO = connectivities[dims.C2E2CODim] + c2e2cO = connectivities[dims.C2E2CO] geofac_grg_x = np.expand_dims(geofac_grg_x, axis=-1) ddx_perturbed_rho = np.sum( @@ -301,7 +302,7 @@ def reference( ) # Remaining computations at edge points - e2c = connectivities[dims.E2CDim] + e2c = connectivities[dims.E2C] temporal_extrapolation_of_perturbed_exner_at_edges = ( temporal_extrapolation_of_perturbed_exner[e2c] ) @@ -431,69 +432,67 @@ def at_neighbor(i: int) -> np.ndarray: next_vn=next_vn, ) - @pytest.fixture( + @stencil_tests.input_data_fixture( params=[ {"is_iau_active": value} for value in [True, False] ], # True for testing, False for benchmarking ids=lambda param: f"is_iau_active[{param['is_iau_active']}]", scope="class", ) - def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: - geofac_grg_x = data_alloc.random_field(grid, dims.CellDim, dims.C2E2CODim) - geofac_grg_y = data_alloc.random_field(grid, dims.CellDim, dims.C2E2CODim) - current_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - tangential_wind = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - pos_on_tplane_e_x = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - pos_on_tplane_e_y = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - primal_normal_cell_x = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - dual_normal_cell_x = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - primal_normal_cell_y = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - dual_normal_cell_y = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - reference_rho_at_edges_on_model_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim + def input_data(self, grid: base.Grid, request: pytest.FixtureRequest) -> dict: + geofac_grg_x = self.data_alloc.random_field(dims.CellDim, dims.C2E2CODim) + geofac_grg_y = self.data_alloc.random_field(dims.CellDim, dims.C2E2CODim) + current_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + tangential_wind = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + pos_on_tplane_e_x = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + pos_on_tplane_e_y = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + primal_normal_cell_x = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + dual_normal_cell_x = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + primal_normal_cell_y = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + dual_normal_cell_y = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + reference_rho_at_edges_on_model_levels = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim ) - reference_theta_at_edges_on_model_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim + reference_theta_at_edges_on_model_levels = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim ) - perturbed_rho_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim + perturbed_rho_at_cells_on_model_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim ) - perturbed_theta_v_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim + perturbed_theta_v_at_cells_on_model_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim ) - ddxn_z_full = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - c_lin_e = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - temporal_extrapolation_of_perturbed_exner = data_alloc.random_field( - grid, dims.CellDim, dims.KDim + ddxn_z_full = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + c_lin_e = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + temporal_extrapolation_of_perturbed_exner = self.data_alloc.random_field( + dims.CellDim, dims.KDim ) - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( + self.data_alloc.random_field(dims.CellDim, dims.KDim) ) d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( - data_alloc.random_field(grid, dims.CellDim, dims.KDim) + self.data_alloc.random_field(dims.CellDim, dims.KDim) ) - hydrostatic_correction_on_lowest_level = data_alloc.random_field(grid, dims.EdgeDim) - zdiff_gradp = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim, dims.KDim) - pg_exdist = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim + hydrostatic_correction_on_lowest_level = self.data_alloc.random_field(dims.EdgeDim) + zdiff_gradp = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim, dims.KDim) + pg_exdist = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim ) # TODO(havogt): should be allocated with a sparse pattern - inv_dual_edge_length = data_alloc.random_field(grid, dims.EdgeDim) - predictor_normal_wind_advective_tendency = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim + inv_dual_edge_length = self.data_alloc.random_field(dims.EdgeDim) + predictor_normal_wind_advective_tendency = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim ) - normal_wind_tendency_due_to_slow_physics_process = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim - ) - normal_wind_iau_increment = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - grf_tend_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - next_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - theta_v_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - horizontal_pressure_gradient = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - rho_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - - ikoffset = data_alloc.zero_field( - grid, dims.EdgeDim, dims.E2CDim, dims.KDim, dtype=gtx.int32 + normal_wind_tendency_due_to_slow_physics_process = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim ) + normal_wind_iau_increment = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + grf_tend_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + next_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + theta_v_at_edges_on_model_levels = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + horizontal_pressure_gradient = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + rho_at_edges_on_model_levels = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + + ikoffset = self.data_alloc.zero_field(dims.EdgeDim, dims.E2CDim, dims.KDim, dtype=gtx.int32) rng = np.random.default_rng() k_levels = grid.num_levels diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_virtual_potential_temperatures_and_pressure_gradient.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_virtual_potential_temperatures_and_pressure_gradient.py index 4ef0d53930..831425cd74 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_virtual_potential_temperatures_and_pressure_gradient.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_virtual_potential_temperatures_and_pressure_gradient.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -18,12 +19,11 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_virtual_potential_temperatures_and_pressure_gradient_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], wgtfac_c: np.ndarray, z_rth_pr_2: np.ndarray, theta_v: np.ndarray, @@ -55,13 +55,13 @@ def compute_virtual_potential_temperatures_and_pressure_gradient_numpy( ) -class TestComputeVirtualPotentialTemperaturesAndPressureGradient(StencilTest): +class TestComputeVirtualPotentialTemperaturesAndPressureGradient(stencil_tests.StencilTest): PROGRAM = compute_virtual_potential_temperatures_and_pressure_gradient OUTPUTS = ("z_theta_v_pr_ic", "theta_v_ic", "z_th_ddz_exner_c") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, wgtfac_c: np.ndarray, z_rth_pr_2: np.ndarray, theta_v: np.ndarray, @@ -71,6 +71,7 @@ def reference( ddqz_z_half: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) ( z_theta_v_pr_ic, theta_v_ic, @@ -92,18 +93,18 @@ def reference( z_th_ddz_exner_c=z_th_ddz_exner_c, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - wgtfac_c = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - z_rth_pr_2 = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - theta_v = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - vwind_expl_wgt = random_field(grid, dims.CellDim, dtype=wpfloat) - exner_pr = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - d_exner_dz_ref_ic = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - ddqz_z_half = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - z_theta_v_pr_ic = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - theta_v_ic = zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - z_th_ddz_exner_c = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + wgtfac_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + z_rth_pr_2 = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + theta_v = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + vwind_expl_wgt = self.data_alloc.random_field(dims.CellDim, dtype=wpfloat) + exner_pr = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + d_exner_dz_ref_ic = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + ddqz_z_half = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + z_theta_v_pr_ic = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) + theta_v_ic = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat) + z_th_ddz_exner_c = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) return dict( wgtfac_c=wgtfac_c, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_vn_on_lateral_boundary.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_vn_on_lateral_boundary.py index 50aec37834..483dfe5dda 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_vn_on_lateral_boundary.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_vn_on_lateral_boundary.py @@ -17,17 +17,16 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestComputeVnOnLateralBoundary(StencilTest): +class TestComputeVnOnLateralBoundary(stencil_tests.StencilTest): PROGRAM = compute_vn_on_lateral_boundary OUTPUTS = ("vn_new",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, grf_tend_vn: np.ndarray, vn_now: np.ndarray, dtime: ta.wpfloat, @@ -36,11 +35,11 @@ def reference( vn_new = vn_now + dtime * grf_tend_vn return dict(vn_new=vn_new) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - grf_tend_vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) - vn_now = random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) - vn_new = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + grf_tend_vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + vn_now = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + vn_new = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) dtime = ta.wpfloat("6.0") return dict( diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_copy_cell_kdim_field_to_vp.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_copy_cell_kdim_field_to_vp.py index e1408308a9..8cd43204cb 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_copy_cell_kdim_field_to_vp.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_copy_cell_kdim_field_to_vp.py @@ -18,8 +18,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def copy_cell_kdim_field_to_vp_numpy(field: np.ndarray) -> np.ndarray: @@ -27,21 +26,19 @@ def copy_cell_kdim_field_to_vp_numpy(field: np.ndarray) -> np.ndarray: return field_copy -class TestCopyCellKdimFieldToVp(StencilTest): +class TestCopyCellKdimFieldToVp(stencil_tests.StencilTest): PROGRAM = copy_cell_kdim_field_to_vp OUTPUTS = ("field_copy",) - @staticmethod - def reference( - connectivities: dict[gtx.Dimension, np.ndarray], field: np.ndarray, **kwargs: Any - ) -> dict: + @stencil_tests.static_reference + def reference(grid: base.Grid, field: np.ndarray, **kwargs: Any) -> dict: field_copy = copy_cell_kdim_field_to_vp_numpy(field) return dict(field_copy=field_copy) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - field = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - field_copy = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + field = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + field_copy = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) return dict( field=field, field_copy=field_copy, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_at_top.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_at_top.py index 74c9b6845d..c9ff1406e3 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_at_top.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_at_top.py @@ -16,8 +16,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def extrapolate_at_top_numpy(wgtfacq_e: np.ndarray, vn: np.ndarray) -> np.ndarray: @@ -37,13 +36,13 @@ def extrapolate_at_top_numpy(wgtfacq_e: np.ndarray, vn: np.ndarray) -> np.ndarra return vn_ie -class TestExtrapolateAtTop(StencilTest): +class TestExtrapolateAtTop(stencil_tests.StencilTest): PROGRAM = extrapolate_at_top OUTPUTS = ("vn_ie",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, wgtfacq_e: np.ndarray, vn: np.ndarray, **kwargs: Any, @@ -51,11 +50,13 @@ def reference( vn_ie = extrapolate_at_top_numpy(wgtfacq_e, vn) return dict(vn_ie=vn_ie) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - wgtfacq_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - vn_ie = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat, extend={dims.KDim: 1}) + wgtfacq_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + vn_ie = self.data_alloc.zero_field( + dims.EdgeDim, dims.KDim, dtype=vpfloat, extend={dims.KDim: 1} + ) return dict( wgtfacq_e=wgtfacq_e, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_temporally_exner_pressure.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_temporally_exner_pressure.py index a87d532ff4..aaf796b70b 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_temporally_exner_pressure.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_temporally_exner_pressure.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -18,12 +19,11 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def extrapolate_temporally_exner_pressure_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], exner: np.ndarray, exner_ref_mc: np.ndarray, exner_pr: np.ndarray, @@ -34,19 +34,20 @@ def extrapolate_temporally_exner_pressure_numpy( return (z_exner_ex_pr, exner_pr) -class TestExtrapolateTemporallyExnerPressure(StencilTest): +class TestExtrapolateTemporallyExnerPressure(stencil_tests.StencilTest): PROGRAM = extrapolate_temporally_exner_pressure OUTPUTS = ("z_exner_ex_pr", "exner_pr") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, exner: np.ndarray, exner_ref_mc: np.ndarray, exner_pr: np.ndarray, exner_exfac: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) (z_exner_ex_pr, exner_pr) = extrapolate_temporally_exner_pressure_numpy( connectivities, exner=exner, @@ -57,13 +58,13 @@ def reference( return dict(z_exner_ex_pr=z_exner_ex_pr, exner_pr=exner_pr) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - exner = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - exner_ref_mc = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - exner_pr = zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - exner_exfac = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - z_exner_ex_pr = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + exner = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + exner_ref_mc = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + exner_pr = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat) + exner_exfac = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + z_exner_ex_pr = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) return dict( exner_exfac=exner_exfac, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py index a0c2e0f902..bb1ad6f27f 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py @@ -18,40 +18,39 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils.data_allocation import zero_field -from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest +from icon4py.model.testing import stencil_tests @pytest.mark.continuous_benchmarking -class TestInitCellKdimFieldWithZeroWp(StencilTest): +class TestInitCellKdimFieldWithZeroWp(stencil_tests.StencilTest): PROGRAM = init_cell_kdim_field_with_zero_wp OUTPUTS = ("field_with_zero_wp",) STATIC_PARAMS = { - StandardStaticVariants.NONE: (), - StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", "vertical_start", "vertical_end", ), - StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( "vertical_start", "vertical_end", ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, field_with_zero_wp: np.ndarray, **kwargs: Any, ) -> dict: field_with_zero_wp = np.zeros_like(field_with_zero_wp) return dict(field_with_zero_wp=field_with_zero_wp) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - field_with_zero_wp = zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) + field_with_zero_wp = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat) return dict( field_with_zero_wp=field_with_zero_wp, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_two_edge_kdim_fields_with_zero_wp.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_two_edge_kdim_fields_with_zero_wp.py index 77c5de7055..05a7d65eb2 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_two_edge_kdim_fields_with_zero_wp.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_two_edge_kdim_fields_with_zero_wp.py @@ -18,17 +18,16 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils.data_allocation import zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestInitTwoEdgeKdimFieldsWithZeroWp(StencilTest): +class TestInitTwoEdgeKdimFieldsWithZeroWp(stencil_tests.StencilTest): PROGRAM = init_two_edge_kdim_fields_with_zero_wp OUTPUTS = ("edge_kdim_field_with_zero_wp_1", "edge_kdim_field_with_zero_wp_2") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, edge_kdim_field_with_zero_wp_1: np.ndarray, edge_kdim_field_with_zero_wp_2: np.ndarray, **kwargs: Any, @@ -40,10 +39,14 @@ def reference( edge_kdim_field_with_zero_wp_2=edge_kdim_field_with_zero_wp_2, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - edge_kdim_field_with_zero_wp_1 = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - edge_kdim_field_with_zero_wp_2 = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + edge_kdim_field_with_zero_wp_1 = self.data_alloc.zero_field( + dims.EdgeDim, dims.KDim, dtype=wpfloat + ) + edge_kdim_field_with_zero_wp_2 = self.data_alloc.zero_field( + dims.EdgeDim, dims.KDim, dtype=wpfloat + ) return dict( edge_kdim_field_with_zero_wp_1=edge_kdim_field_with_zero_wp_1, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_cell_field_to_half_levels_vp.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_cell_field_to_half_levels_vp.py index 4cc48fd101..e87563b9df 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_cell_field_to_half_levels_vp.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_cell_field_to_half_levels_vp.py @@ -18,7 +18,6 @@ interpolate_cell_field_to_half_levels_vp, ) from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils.data_allocation import random_field, zero_field from icon4py.model.testing import stencil_tests @@ -38,9 +37,9 @@ class TestInterpolateToHalfLevelsVp(stencil_tests.StencilTest): PROGRAM = interpolate_cell_field_to_half_levels_vp OUTPUTS = ("interpolation_to_half_levels_vp",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, wgtfac_c: np.ndarray, interpolant: np.ndarray, **kwargs: Any, @@ -50,12 +49,12 @@ def reference( ) return dict(interpolation_to_half_levels_vp=interpolation_to_half_levels_vp) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - interpolant = random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - wgtfac_c = random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - interpolation_to_half_levels_vp = zero_field( - grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat + interpolant = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + wgtfac_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + interpolation_to_half_levels_vp = self.data_alloc.zero_field( + dims.CellDim, dims.KDim, dtype=ta.vpfloat ) return dict( diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_contravariant_vertical_velocity_to_full_levels.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_contravariant_vertical_velocity_to_full_levels.py index 12f5fd8506..7d969f3d03 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_contravariant_vertical_velocity_to_full_levels.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_contravariant_vertical_velocity_to_full_levels.py @@ -18,8 +18,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def interpolate_contravariant_vertical_velocity_to_full_levels_numpy( @@ -29,24 +28,22 @@ def interpolate_contravariant_vertical_velocity_to_full_levels_numpy( return z_w_con_c_full -class TestInterpolateContravariantVerticalVelocityToFullLevels(StencilTest): +class TestInterpolateContravariantVerticalVelocityToFullLevels(stencil_tests.StencilTest): PROGRAM = interpolate_contravariant_vertical_velocity_to_full_levels OUTPUTS = ("z_w_con_c_full",) - @staticmethod - def reference( - connectivities: dict[gtx.Dimension, np.ndarray], z_w_con_c: np.ndarray, **kwargs: Any - ) -> dict: + @stencil_tests.static_reference + def reference(grid: base.Grid, z_w_con_c: np.ndarray, **kwargs: Any) -> dict: z_w_con_c_full = interpolate_contravariant_vertical_velocity_to_full_levels_numpy(z_w_con_c) return dict(z_w_con_c_full=z_w_con_c_full) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - z_w_con_c = random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=vpfloat + z_w_con_c = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=vpfloat ) - z_w_con_c_full = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + z_w_con_c_full = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) return dict( z_w_con_c=z_w_con_c, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py index c36fa68bfa..0c53614e96 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py @@ -30,7 +30,6 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -59,9 +58,9 @@ class TestInterpolateRhoThetaVToHalfLevelsAndComputePressureBuoyancyAcceleration ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, rho_at_cells_on_half_levels: np.ndarray, perturbed_theta_v_at_cells_on_half_levels: np.ndarray, theta_v_at_cells_on_half_levels: np.ndarray, @@ -182,36 +181,36 @@ def reference( nonhydro_buoy_at_cells_on_half_levels=nonhydro_buoy_at_cells_on_half_levels, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - w = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - contravariant_correction_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim - ) - ddqz_z_half = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - current_rho = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - next_rho = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - current_theta_v = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - next_theta_v = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - wgtfac_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - reference_theta_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim - ) - exner_w_explicit_weight_parameter = data_alloc.random_field(grid, dims.CellDim) - perturbed_exner_at_cells_on_model_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim - ) - ddz_of_reference_exner_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim - ) - rho_at_cells_on_half_levels = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - perturbed_theta_v_at_cells_on_half_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} - ) - theta_v_at_cells_on_half_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} - ) - nonhydro_buoy_at_cells_on_half_levels = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + w = self.data_alloc.random_field(dims.CellDim, dims.KDim) + contravariant_correction_at_cells_on_half_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim + ) + ddqz_z_half = self.data_alloc.random_field(dims.CellDim, dims.KDim) + current_rho = self.data_alloc.random_field(dims.CellDim, dims.KDim) + next_rho = self.data_alloc.random_field(dims.CellDim, dims.KDim) + current_theta_v = self.data_alloc.random_field(dims.CellDim, dims.KDim) + next_theta_v = self.data_alloc.random_field(dims.CellDim, dims.KDim) + wgtfac_c = self.data_alloc.random_field(dims.CellDim, dims.KDim) + reference_theta_at_cells_on_model_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim + ) + exner_w_explicit_weight_parameter = self.data_alloc.random_field(dims.CellDim) + perturbed_exner_at_cells_on_model_levels = self.data_alloc.zero_field( + dims.CellDim, dims.KDim + ) + ddz_of_reference_exner_at_cells_on_half_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim + ) + rho_at_cells_on_half_levels = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + perturbed_theta_v_at_cells_on_half_levels = self.data_alloc.zero_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} + ) + theta_v_at_cells_on_half_levels = self.data_alloc.zero_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} + ) + nonhydro_buoy_at_cells_on_half_levels = self.data_alloc.zero_field(dims.CellDim, dims.KDim) dtime = 0.9 rhotheta_explicit_weight_parameter = 0.25 diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_to_cell_center.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_to_cell_center.py index 8f018db9ca..e11895ca45 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_to_cell_center.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_to_cell_center.py @@ -5,13 +5,13 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np import pytest -import icon4py.model.common.utils.data_allocation as data_alloc from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base from icon4py.model.common.interpolation.stencils.interpolate_to_cell_center import ( @@ -22,13 +22,13 @@ def interpolate_to_cell_center_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], interpolant: np.ndarray, e_bln_c_s: np.ndarray, **kwargs: Any, ) -> np.ndarray: e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) - c2e = connectivities[dims.C2EDim] + c2e = connectivities[dims.C2E] interpolation = np.sum( interpolant[c2e] * e_bln_c_s, @@ -41,21 +41,22 @@ class TestInterpolateToCellCenter(stencil_tests.StencilTest): PROGRAM = interpolate_to_cell_center OUTPUTS = ("interpolation",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, interpolant: np.ndarray, e_bln_c_s: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) interpolation = interpolate_to_cell_center_numpy(connectivities, interpolant, e_bln_c_s) return dict(interpolation=interpolation) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - interpolant = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) - e_bln_c_s = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim, dtype=ta.wpfloat) - interpolation = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) + interpolant = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + e_bln_c_s = self.data_alloc.random_field(dims.CellDim, dims.C2EDim, dtype=ta.wpfloat) + interpolation = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) return dict( interpolant=interpolant, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_to_surface.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_to_surface.py index 2dc775909d..a79d2974c2 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_to_surface.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_to_surface.py @@ -16,8 +16,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def interpolate_to_surface_numpy( @@ -32,13 +31,13 @@ def interpolate_to_surface_numpy( return interpolation_to_surface -class TestInterpolateToSurface(StencilTest): +class TestInterpolateToSurface(stencil_tests.StencilTest): PROGRAM = interpolate_to_surface OUTPUTS = ("interpolation_to_surface",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, interpolant: np.ndarray, wgtfacq_c: np.ndarray, interpolation_to_surface: np.ndarray, @@ -51,11 +50,13 @@ def reference( ) return dict(interpolation_to_surface=interpolation_to_surface) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - interpolant = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - wgtfacq_c = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - interpolation_to_surface = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + interpolant = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + wgtfacq_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + interpolation_to_surface = self.data_alloc.zero_field( + dims.CellDim, dims.KDim, dtype=vpfloat + ) return dict( interpolant=interpolant, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py index b998025020..ac60673207 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py @@ -18,17 +18,16 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestInterpolateVnAndVtToIeAndComputeEkinOnEdges(StencilTest): +class TestInterpolateVnAndVtToIeAndComputeEkinOnEdges(stencil_tests.StencilTest): PROGRAM = interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges OUTPUTS = ("vn_ie", "z_vt_ie", "z_kin_hor_e") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, wgtfac_e: np.ndarray, vn: np.ndarray, vt: np.ndarray, @@ -45,14 +44,14 @@ def reference( z_kin_hor_e[:, 0] = 0 return dict(vn_ie=vn_ie, z_vt_ie=z_vt_ie, z_kin_hor_e=z_kin_hor_e) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - wgtfac_e = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - vt = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - vn_ie = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - z_vt_ie = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - z_kin_hor_e = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) + wgtfac_e = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + vt = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + vn_ie = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + z_vt_ie = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + z_kin_hor_e = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) return dict( wgtfac_e=wgtfac_e, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges.py index b2a19e3566..44765370e3 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges.py @@ -18,7 +18,6 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field from icon4py.model.testing import stencil_tests @@ -57,9 +56,9 @@ class TestInterpolateVnToHalfLevelsAndComputeKineticEnergyOnEdges(stencil_tests. PROGRAM = interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges OUTPUTS = ("vn_ie", "z_kin_hor_e") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, wgtfac_e: np.ndarray, vn: np.ndarray, vt: np.ndarray, @@ -84,14 +83,14 @@ def reference( z_kin_hor_e=z_kin_hor_e, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - wgtfac_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - vt = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) + wgtfac_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + vn = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + vt = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) - vn_ie = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - z_kin_hor_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) + vn_ie = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + z_kin_hor_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) return dict( wgtfac_e=wgtfac_e, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vt_to_interface_edges.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vt_to_interface_edges.py index 4fc1660fb6..6e494f89bd 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vt_to_interface_edges.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vt_to_interface_edges.py @@ -17,8 +17,7 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def interpolate_vt_to_interface_edges_numpy( @@ -30,13 +29,13 @@ def interpolate_vt_to_interface_edges_numpy( return z_vt_ie -class TestInterpolateVtToInterfaceEdges(StencilTest): +class TestInterpolateVtToInterfaceEdges(stencil_tests.StencilTest): PROGRAM = interpolate_vt_to_interface_edges OUTPUTS = ("z_vt_ie",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, wgtfac_e: np.ndarray, vt: np.ndarray, z_vt_ie: np.ndarray, @@ -50,12 +49,12 @@ def reference( z_vt_ie[subset] = interpolate_vt_to_interface_edges_numpy(wgtfac_e, vt)[subset] return dict(z_vt_ie=z_vt_ie) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - wgtfac_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) - vt = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + wgtfac_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + vt = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) - z_vt_ie = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + z_vt_ie = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) return dict( wgtfac_e=wgtfac_e, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py index 3ca6d13109..28d3dcc29d 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -18,30 +19,30 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], p_cell_in: np.ndarray, c_intp: np.ndarray + connectivities: Mapping[gtx.FieldOffset, np.ndarray], p_cell_in: np.ndarray, c_intp: np.ndarray ) -> np.ndarray: - v2c = connectivities[dims.V2CDim] + v2c = connectivities[dims.V2C] c_intp = np.expand_dims(c_intp, axis=-1) p_vert_out = np.sum(np.where((v2c != -1)[:, :, np.newaxis], p_cell_in[v2c] * c_intp, 0), axis=1) return p_vert_out -class TestMoIconInterpolationScalarCells2vertsScalarRiDsl(StencilTest): +class TestMoIconInterpolationScalarCells2vertsScalarRiDsl(stencil_tests.StencilTest): PROGRAM = mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl OUTPUTS = ("p_vert_out",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_cell_in: np.ndarray, c_intp: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) p_vert_out = mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl_numpy( connectivities, p_cell_in, c_intp ) @@ -49,11 +50,11 @@ def reference( p_vert_out=p_vert_out, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - p_cell_in = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - c_intp = random_field(grid, dims.VertexDim, dims.V2CDim, dtype=wpfloat) - p_vert_out = zero_field(grid, dims.VertexDim, dims.KDim, dtype=vpfloat) + p_cell_in = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + c_intp = self.data_alloc.random_field(dims.VertexDim, dims.V2CDim, dtype=wpfloat) + p_vert_out = self.data_alloc.zero_field(dims.VertexDim, dims.KDim, dtype=vpfloat) return dict( p_cell_in=p_cell_in, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_mo_math_divrot_rot_vertex_ri_dsl.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_mo_math_divrot_rot_vertex_ri_dsl.py index b819687c83..fa3d2dd6c6 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_mo_math_divrot_rot_vertex_ri_dsl.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_mo_math_divrot_rot_vertex_ri_dsl.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -18,38 +19,38 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def mo_math_divrot_rot_vertex_ri_dsl_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], vec_e: np.ndarray, geofac_rot: np.ndarray + connectivities: Mapping[gtx.FieldOffset, np.ndarray], vec_e: np.ndarray, geofac_rot: np.ndarray ) -> np.ndarray: - v2e = connectivities[dims.V2EDim] + v2e = connectivities[dims.V2E] geofac_rot = np.expand_dims(geofac_rot, axis=-1) rot_vec = np.sum(np.where((v2e != -1)[:, :, np.newaxis], vec_e[v2e] * geofac_rot, 0), axis=1) return rot_vec -class TestMoMathDivrotRotVertexRiDsl(StencilTest): +class TestMoMathDivrotRotVertexRiDsl(stencil_tests.StencilTest): PROGRAM = mo_math_divrot_rot_vertex_ri_dsl OUTPUTS = ("rot_vec",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, vec_e: np.ndarray, geofac_rot: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) rot_vec = mo_math_divrot_rot_vertex_ri_dsl_numpy(connectivities, vec_e, geofac_rot) return dict(rot_vec=rot_vec) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - vec_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - geofac_rot = random_field(grid, dims.VertexDim, dims.V2EDim, dtype=wpfloat) - rot_vec = zero_field(grid, dims.VertexDim, dims.KDim, dtype=vpfloat) + vec_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + geofac_rot = self.data_alloc.random_field(dims.VertexDim, dims.V2EDim, dtype=wpfloat) + rot_vec = self.data_alloc.zero_field(dims.VertexDim, dims.KDim, dtype=vpfloat) return dict( vec_e=vec_e, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_set_lower_boundary_condition_for_w_and_contravariant_correction.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_set_lower_boundary_condition_for_w_and_contravariant_correction.py index af727d761e..7a5dfa0c62 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_set_lower_boundary_condition_for_w_and_contravariant_correction.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_set_lower_boundary_condition_for_w_and_contravariant_correction.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -18,12 +19,11 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def set_lower_boundary_condition_for_w_and_contravariant_correction_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], w_concorr_c: np.ndarray, z_contr_w_fl_l: np.ndarray, ) -> tuple[np.ndarray, np.ndarray]: @@ -32,17 +32,18 @@ def set_lower_boundary_condition_for_w_and_contravariant_correction_numpy( return (w_nnew, z_contr_w_fl_l) -class TestInitLowerBoundaryConditionForWAndContravariantCorrection(StencilTest): +class TestInitLowerBoundaryConditionForWAndContravariantCorrection(stencil_tests.StencilTest): PROGRAM = set_lower_boundary_condition_for_w_and_contravariant_correction OUTPUTS = ("w_nnew", "z_contr_w_fl_l") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, w_concorr_c: np.ndarray, z_contr_w_fl_l: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) ( w_nnew, z_contr_w_fl_l, @@ -51,11 +52,11 @@ def reference( ) return dict(w_nnew=w_nnew, z_contr_w_fl_l=z_contr_w_fl_l) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - w_concorr_c = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - z_contr_w_fl_l = zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - w_nnew = zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) + w_concorr_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + z_contr_w_fl_l = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat) + w_nnew = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat) return dict( w_nnew=w_nnew, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_set_theta_v_prime_ic_at_lower_boundary.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_set_theta_v_prime_ic_at_lower_boundary.py index 122f1913ef..1a3f694204 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_set_theta_v_prime_ic_at_lower_boundary.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_set_theta_v_prime_ic_at_lower_boundary.py @@ -18,8 +18,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests from .test_interpolate_to_surface import interpolate_to_surface_numpy @@ -40,13 +39,13 @@ def set_theta_v_prime_ic_at_lower_boundary_numpy( return (z_theta_v_pr_ic, theta_v_ic) -class TestInitThetaVPrimeIcAtLowerBoundary(StencilTest): +class TestInitThetaVPrimeIcAtLowerBoundary(stencil_tests.StencilTest): PROGRAM = set_theta_v_prime_ic_at_lower_boundary OUTPUTS = ("z_theta_v_pr_ic", "theta_v_ic") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, wgtfacq_c: np.ndarray, z_rth_pr: np.ndarray, theta_ref_ic: np.ndarray, @@ -63,13 +62,13 @@ def reference( ) return dict(z_theta_v_pr_ic=z_theta_v_pr_ic, theta_v_ic=theta_v_ic) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - wgtfacq_c = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - z_rth_pr = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - theta_ref_ic = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - z_theta_v_pr_ic = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - theta_v_ic = zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) + wgtfacq_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + z_rth_pr = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + theta_ref_ic = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + z_theta_v_pr_ic = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + theta_v_ic = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat) return dict( wgtfacq_c=wgtfacq_c, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_solve_tridiagonal_matrix_for_w_back_substitution.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_solve_tridiagonal_matrix_for_w_back_substitution.py index 076e84b767..e519e03dad 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_solve_tridiagonal_matrix_for_w_back_substitution.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_solve_tridiagonal_matrix_for_w_back_substitution.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -18,12 +19,11 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def solve_tridiagonal_matrix_for_w_back_substitution_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], z_q: np.ndarray, w: np.ndarray, ) -> np.ndarray: @@ -38,24 +38,25 @@ def solve_tridiagonal_matrix_for_w_back_substitution_numpy( return w_new -class TestSolveTridiagonalMatrixForWBackSubstitution(StencilTest): +class TestSolveTridiagonalMatrixForWBackSubstitution(stencil_tests.StencilTest): PROGRAM = solve_tridiagonal_matrix_for_w_back_substitution OUTPUTS = ("w",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_q: np.ndarray, w: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) w_new = solve_tridiagonal_matrix_for_w_back_substitution_numpy(connectivities, z_q=z_q, w=w) return dict(w=w_new) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - z_q = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - w = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) + z_q = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + w = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) h_start = 0 h_end = gtx.int32(grid.num_cells) v_start = 1 diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_solve_tridiagonal_matrix_for_w_forward_sweep.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_solve_tridiagonal_matrix_for_w_forward_sweep.py index 7bd38f9619..de2e20adfd 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_solve_tridiagonal_matrix_for_w_forward_sweep.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_solve_tridiagonal_matrix_for_w_forward_sweep.py @@ -15,10 +15,9 @@ solve_tridiagonal_matrix_for_w_forward_sweep, ) from icon4py.model.common import dimension as dims, type_alias as ta -from icon4py.model.common.grid import base as base_grid +from icon4py.model.common.grid import base, base as base_grid from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def solve_tridiagonal_matrix_for_w_forward_sweep_numpy( @@ -57,13 +56,13 @@ def solve_tridiagonal_matrix_for_w_forward_sweep_numpy( return z_q, w -class TestSolveTridiagonalMatrixForWForwardSweep(StencilTest): +class TestSolveTridiagonalMatrixForWForwardSweep(stencil_tests.StencilTest): PROGRAM = solve_tridiagonal_matrix_for_w_forward_sweep OUTPUTS = ("w", "z_q") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, vwind_impl_wgt: np.ndarray, theta_v_ic: np.ndarray, ddqz_z_half: np.ndarray, @@ -92,23 +91,23 @@ def reference( ) return dict(z_q=z_q_ref, w=w_ref) - @pytest.fixture - def input_data(self, grid: base_grid.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - vwind_impl_wgt = data_alloc.random_field(grid, dims.CellDim, dtype=ta.wpfloat) - theta_v_ic = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - ddqz_z_half = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - z_alpha = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=ta.vpfloat + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: + vwind_impl_wgt = self.data_alloc.random_field(dims.CellDim, dtype=ta.wpfloat) + theta_v_ic = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + ddqz_z_half = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_alpha = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=ta.vpfloat ) - z_beta = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) - z_exner_expl = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - z_w_expl = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=ta.wpfloat + z_beta = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_exner_expl = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + z_w_expl = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=ta.wpfloat ) - z_q = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) + z_q = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat) # z_q first level should always be initialized to zero when solve_tridiagonal_matrix_for_w_forward_sweep is called z_q.asnumpy()[:, 0] = 0.0 - w = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) + w = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) # w first level should always be initialized to zero when solve_tridiagonal_matrix_for_w_forward_sweep is called w.asnumpy()[:, 0] = 0.0 diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_spatially_average_flux_or_velocity.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_spatially_average_flux_or_velocity.py index d9d8ae062c..bed1b5002e 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_spatially_average_flux_or_velocity.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_spatially_average_flux_or_velocity.py @@ -6,7 +6,8 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -19,16 +20,15 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def spatially_average_flux_or_velocity_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], e_flx_avg: np.ndarray, flux_or_velocity: np.ndarray, ) -> np.ndarray: - e2c2eO = connectivities[dims.E2C2EODim] + e2c2eO = connectivities[dims.E2C2EO] e_flx_avg = np.expand_dims(e_flx_avg, axis=-1) spatially_averaged_flux_or_velocity = np.sum(flux_or_velocity[e2c2eO] * e_flx_avg, axis=1) @@ -36,29 +36,30 @@ def spatially_average_flux_or_velocity_numpy( @pytest.mark.embedded_remap_error -class TestSpatiallyAverageFluxOrVelocity(StencilTest): +class TestSpatiallyAverageFluxOrVelocity(stencil_tests.StencilTest): PROGRAM = spatially_average_flux_or_velocity OUTPUTS = ("spatially_averaged_flux_or_velocity",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, e_flx_avg: np.ndarray, flux_or_velocity: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) spatially_averaged_flux_or_velocity = spatially_average_flux_or_velocity_numpy( connectivities, e_flx_avg, flux_or_velocity ) return dict(spatially_averaged_flux_or_velocity=spatially_averaged_flux_or_velocity) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - e_flx_avg = random_field(grid, dims.EdgeDim, dims.E2C2EODim, dtype=wpfloat) - flux_or_velocity = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - spatially_averaged_flux_or_velocity = zero_field( - grid, dims.EdgeDim, dims.KDim, dtype=wpfloat + e_flx_avg = self.data_alloc.random_field(dims.EdgeDim, dims.E2C2EODim, dtype=wpfloat) + flux_or_velocity = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + spatially_averaged_flux_or_velocity = self.data_alloc.zero_field( + dims.EdgeDim, dims.KDim, dtype=wpfloat ) return dict( diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_density_exner_wind.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_density_exner_wind.py index 612a940991..170ab6777b 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_density_exner_wind.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_density_exner_wind.py @@ -17,17 +17,16 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestUpdateDensityExnerWind(StencilTest): +class TestUpdateDensityExnerWind(stencil_tests.StencilTest): PROGRAM = update_density_exner_wind OUTPUTS = ("rho_new", "exner_new", "w_new") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, rho_now: np.ndarray, grf_tend_rho: np.ndarray, theta_v_now: np.ndarray, @@ -42,18 +41,18 @@ def reference( w_new = w_now + dtime * grf_tend_w return dict(rho_new=rho_new, exner_new=exner_new, w_new=w_new) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - rho_now = random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - grf_tend_rho = random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - theta_v_now = random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - grf_tend_thv = random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - w_now = random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - grf_tend_w = random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) + rho_now = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + grf_tend_rho = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + theta_v_now = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + grf_tend_thv = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + w_now = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + grf_tend_w = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) dtime = ta.wpfloat("5.0") - rho_new = zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - exner_new = zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - w_new = zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) + rho_new = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + exner_new = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + w_new = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) return dict( rho_now=rho_now, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_dynamical_exner_time_increment.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_dynamical_exner_time_increment.py index ffcb582b55..6c00cf78be 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_dynamical_exner_time_increment.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_dynamical_exner_time_increment.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -18,12 +19,11 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def update_dynamical_exner_time_increment_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], exner: np.ndarray, ddt_exner_phy: np.ndarray, exner_dyn_incr: np.ndarray, @@ -34,13 +34,13 @@ def update_dynamical_exner_time_increment_numpy( return exner_dyn_incr -class TestUpdateDynamicalExnerTimeIncrement(StencilTest): +class TestUpdateDynamicalExnerTimeIncrement(stencil_tests.StencilTest): PROGRAM = update_dynamical_exner_time_increment OUTPUTS = ("exner_dyn_incr",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, exner: np.ndarray, ddt_exner_phy: np.ndarray, exner_dyn_incr: np.ndarray, @@ -48,6 +48,7 @@ def reference( dtime: float, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) exner_dyn_incr = update_dynamical_exner_time_increment_numpy( connectivities, exner, @@ -58,12 +59,12 @@ def reference( ) return dict(exner_dyn_incr=exner_dyn_incr) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: ndyn_substeps_var, dtime = wpfloat("10.0"), wpfloat("12.0") - exner = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - ddt_exner_phy = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - exner_dyn_incr = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + exner = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + ddt_exner_phy = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + exner_dyn_incr = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) return dict( exner=exner, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py index 5a0e5422da..bca9055e28 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py @@ -18,31 +18,30 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest +from icon4py.model.testing import stencil_tests @pytest.mark.continuous_benchmarking -class TestUpdateMassFluxWeighted(StencilTest): +class TestUpdateMassFluxWeighted(stencil_tests.StencilTest): PROGRAM = update_mass_flux_weighted OUTPUTS = ("mass_flx_ic",) STATIC_PARAMS = { - StandardStaticVariants.NONE: (), - StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", "vertical_start", "vertical_end", ), - StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( "vertical_start", "vertical_end", ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, rho_ic: np.ndarray, vwind_expl_wgt: np.ndarray, vwind_impl_wgt: np.ndarray, @@ -60,16 +59,16 @@ def reference( ) return dict(mass_flx_ic=mass_flx_ic) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: r_nsubsteps = wpfloat("10.0") - rho_ic = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - vwind_expl_wgt = random_field(grid, dims.CellDim, dtype=wpfloat) - vwind_impl_wgt = random_field(grid, dims.CellDim, dtype=wpfloat) - w_now = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - w_new = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - w_concorr_c = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - mass_flx_ic = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) + rho_ic = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + vwind_expl_wgt = self.data_alloc.random_field(dims.CellDim, dtype=wpfloat) + vwind_impl_wgt = self.data_alloc.random_field(dims.CellDim, dtype=wpfloat) + w_now = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + w_new = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + w_concorr_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + mass_flx_ic = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) return dict( rho_ic=rho_ic, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_volume_flux.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_volume_flux.py index 64cd7fab4f..e42ad12daf 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_volume_flux.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_volume_flux.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -15,12 +16,11 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def update_mass_volume_flux_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], z_contr_w_fl_l: np.ndarray, rho_ic: np.ndarray, vwind_impl_wgt: np.ndarray, @@ -36,16 +36,16 @@ def update_mass_volume_flux_numpy( return (mass_flx_ic, vol_flx_ic) -class TestUpdateMassVolumeFlux(StencilTest): +class TestUpdateMassVolumeFlux(stencil_tests.StencilTest): PROGRAM = update_mass_volume_flux OUTPUTS = ( "mass_flx_ic", "vol_flx_ic", ) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, z_contr_w_fl_l: np.ndarray, rho_ic: np.ndarray, vwind_impl_wgt: np.ndarray, @@ -55,6 +55,7 @@ def reference( r_nsubsteps: float, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) (mass_flx_ic, vol_flx_ic) = update_mass_volume_flux_numpy( connectivities, z_contr_w_fl_l=z_contr_w_fl_l, @@ -67,14 +68,14 @@ def reference( ) return dict(mass_flx_ic=mass_flx_ic, vol_flx_ic=vol_flx_ic) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - z_contr_w_fl_l = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - rho_ic = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - vwind_impl_wgt = data_alloc.random_field(grid, dims.CellDim, dtype=ta.wpfloat) - w = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - mass_flx_ic = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - vol_flx_ic = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) + z_contr_w_fl_l = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + rho_ic = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + vwind_impl_wgt = self.data_alloc.random_field(dims.CellDim, dtype=ta.wpfloat) + w = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + mass_flx_ic = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + vol_flx_ic = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) r_nsubsteps = 7.0 return dict( diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_theta_and_exner_in_halo.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_theta_and_exner_in_halo.py index 8f01a98f00..86fce3d110 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_theta_and_exner_in_halo.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_theta_and_exner_in_halo.py @@ -18,23 +18,22 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests dycore_consts: Final = constants.PhysicsConstants() -class TestUpdateThetaV(StencilTest): +class TestUpdateThetaV(stencil_tests.StencilTest): PROGRAM = update_theta_and_exner_in_halo OUTPUTS = ( "theta_v_new", "exner_new", ) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, mask_prog_halo_c: np.ndarray, rho_now: np.ndarray, rho_new: np.ndarray, @@ -63,15 +62,15 @@ def reference( ) return dict(theta_v_new=theta_v_new, exner_new=exner_new) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - mask_prog_halo_c = data_alloc.random_mask(grid, dims.CellDim) - rho_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - rho_new = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - theta_v_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - theta_v_new = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - exner_now = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - exner_new = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) + mask_prog_halo_c = self.data_alloc.random_mask(dims.CellDim) + rho_now = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + rho_new = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + theta_v_now = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + theta_v_new = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + exner_now = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + exner_new = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) return dict( mask_prog_halo_c=mask_prog_halo_c, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_wind.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_wind.py index 32f31fe859..7c34c41a15 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_wind.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_wind.py @@ -16,17 +16,16 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestUpdateWind(StencilTest): +class TestUpdateWind(stencil_tests.StencilTest): PROGRAM = update_wind OUTPUTS = ("w_new",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, w_now: np.ndarray, grf_tend_w: np.ndarray, dtime: float, @@ -35,12 +34,12 @@ def reference( w_new = w_now + dtime * grf_tend_w return dict(w_new=w_new) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: dtime = wpfloat("10.0") - w_now = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - grf_tend_w = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - w_new = zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) + w_now = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + grf_tend_w = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=wpfloat) + w_new = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat) return dict( w_now=w_now, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py index 4d871fcfbf..1df006ad68 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -17,7 +18,6 @@ from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests from .test_add_analysis_increments_from_data_assimilation import ( @@ -91,9 +91,9 @@ class TestVerticallyImplicitSolverAtCorrectorStep(stencil_tests.StencilTest): ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, next_w: np.ndarray, next_rho: np.ndarray, next_exner: np.ndarray, @@ -139,6 +139,7 @@ def reference( kstart_moist: int, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) horizontal_start = kwargs["start_cell_index_nudging"] horizontal_end = kwargs["end_cell_index_local"] n_lev = kwargs["vertical_end_index_model_surface"] - 1 @@ -396,7 +397,7 @@ def reference( exner_dynamical_increment=exner_dynamical_increment, ) - @pytest.fixture( + @stencil_tests.input_data_fixture( params=[ {"at_first_substep": afs, "at_last_substep": als, "lprep_adv": la, "is_iau_active": ia} for afs, als, la, ia in [ @@ -414,60 +415,60 @@ def reference( ), ) def input_data( - self, request: pytest.FixtureRequest, grid: base.Grid + self, grid: base.Grid, request: pytest.FixtureRequest ) -> dict[str, gtx.Field | state_utils.ScalarType]: - geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) - mass_flux_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - theta_v_flux_at_edges_on_model_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim + geofac_div = self.data_alloc.random_field(dims.CellDim, dims.C2EDim) + mass_flux_at_edges_on_model_levels = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + theta_v_flux_at_edges_on_model_levels = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim ) - current_w = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - predictor_vertical_wind_advective_tendency = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + current_w = self.data_alloc.random_field(dims.CellDim, dims.KDim) + predictor_vertical_wind_advective_tendency = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - corrector_vertical_wind_advective_tendency = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + corrector_vertical_wind_advective_tendency = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - nonhydro_buoy_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim + nonhydro_buoy_at_cells_on_half_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim ) - rho_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, low=1.0e-5 + rho_at_cells_on_half_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, low=1.0e-5 ) - contravariant_correction_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + contravariant_correction_at_cells_on_half_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - exner_w_explicit_weight_parameter = data_alloc.random_field(grid, dims.CellDim) - current_exner = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - current_rho = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - current_theta_v = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - inv_ddqz_z_full = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - exner_w_implicit_weight_parameter = data_alloc.random_field(grid, dims.CellDim) - theta_v_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, low=1.0e-5 + exner_w_explicit_weight_parameter = self.data_alloc.random_field(dims.CellDim) + current_exner = self.data_alloc.random_field(dims.CellDim, dims.KDim, low=1.0e-5) + current_rho = self.data_alloc.random_field(dims.CellDim, dims.KDim, low=1.0e-5) + current_theta_v = self.data_alloc.random_field(dims.CellDim, dims.KDim, low=1.0e-5) + inv_ddqz_z_full = self.data_alloc.random_field(dims.CellDim, dims.KDim, low=1.0e-5) + exner_w_implicit_weight_parameter = self.data_alloc.random_field(dims.CellDim) + theta_v_at_cells_on_half_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, low=1.0e-5 ) - perturbed_exner_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim + perturbed_exner_at_cells_on_model_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim ) - exner_tendency_due_to_slow_physics = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - rho_iau_increment = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - exner_iau_increment = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - ddqz_z_half = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - rayleigh_damping_factor = data_alloc.random_field(grid, dims.KDim) - reference_exner_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, low=1.0e-5 + exner_tendency_due_to_slow_physics = self.data_alloc.random_field(dims.CellDim, dims.KDim) + rho_iau_increment = self.data_alloc.random_field(dims.CellDim, dims.KDim) + exner_iau_increment = self.data_alloc.random_field(dims.CellDim, dims.KDim) + ddqz_z_half = self.data_alloc.random_field(dims.CellDim, dims.KDim, low=1.0e-5) + rayleigh_damping_factor = self.data_alloc.random_field(dims.KDim) + reference_exner_at_cells_on_model_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim, low=1.0e-5 ) - next_w = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) - next_rho = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) - next_exner = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) - next_theta_v = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) - exner_dynamical_increment = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - dynamical_vertical_mass_flux_at_cells_on_half_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + next_w = self.data_alloc.zero_field(dims.CellDim, dims.KDim, extend={dims.KDim: 1}) + next_rho = self.data_alloc.constant_field(1.0e-5, dims.CellDim, dims.KDim) + next_exner = self.data_alloc.constant_field(1.0e-5, dims.CellDim, dims.KDim) + next_theta_v = self.data_alloc.constant_field(1.0e-5, dims.CellDim, dims.KDim) + exner_dynamical_increment = self.data_alloc.random_field(dims.CellDim, dims.KDim) + dynamical_vertical_mass_flux_at_cells_on_half_levels = self.data_alloc.zero_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - dynamical_vertical_volumetric_flux_at_cells_on_half_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + dynamical_vertical_volumetric_flux_at_cells_on_half_levels = self.data_alloc.zero_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) lprep_adv = request.param["lprep_adv"] diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py index cc48d17c80..34da38a921 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -17,7 +18,6 @@ from icon4py.model.common import constants, dimension as dims, type_alias as ta from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests from .test_add_analysis_increments_from_data_assimilation import ( @@ -95,9 +95,9 @@ class TestVerticallyImplicitSolverAtPredictorStep(stencil_tests.StencilTest): ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, contravariant_correction_at_cells_on_half_levels: np.ndarray, next_w: np.ndarray, next_rho: np.ndarray, @@ -141,6 +141,7 @@ def reference( flat_level_index_plus1: int, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) start_cell_index_nudging = kwargs["start_cell_index_nudging"] end_cell_index_local = kwargs["end_cell_index_local"] start_cell_index_lateral_lvl3 = kwargs["start_cell_index_lateral_lvl3"] @@ -394,7 +395,7 @@ def reference( exner_dynamical_increment=exner_dynamical_increment, ) - @pytest.fixture( + @stencil_tests.input_data_fixture( params=[ {"at_first_substep": afs, "is_iau_active": ia, "divdamp_type": dvdmpt} for afs, ia, dvdmpt in [ @@ -408,61 +409,61 @@ def reference( ), ) def input_data( - self, request: pytest.FixtureRequest, grid: base.Grid + self, grid: base.Grid, request: pytest.FixtureRequest ) -> dict[str, gtx.Field | state_utils.ScalarType]: - geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) - mass_flux_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - theta_v_flux_at_edges_on_model_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim + geofac_div = self.data_alloc.random_field(dims.CellDim, dims.C2EDim) + mass_flux_at_edges_on_model_levels = self.data_alloc.random_field(dims.EdgeDim, dims.KDim) + theta_v_flux_at_edges_on_model_levels = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim ) - predictor_vertical_wind_advective_tendency = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + predictor_vertical_wind_advective_tendency = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - nonhydro_buoy_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim + nonhydro_buoy_at_cells_on_half_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim ) - rho_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, low=1.0e-5 + rho_at_cells_on_half_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, low=1.0e-5 ) - contravariant_correction_at_cells_on_half_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + contravariant_correction_at_cells_on_half_levels = self.data_alloc.zero_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - contravariant_correction_at_edges_on_model_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim + contravariant_correction_at_edges_on_model_levels = self.data_alloc.random_field( + dims.EdgeDim, dims.KDim ) - exner_w_explicit_weight_parameter = data_alloc.random_field(grid, dims.CellDim) - current_exner = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - current_rho = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - current_theta_v = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - current_w = data_alloc.random_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) - inv_ddqz_z_full = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - exner_w_implicit_weight_parameter = data_alloc.random_field(grid, dims.CellDim) - theta_v_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, low=1.0e-5 + exner_w_explicit_weight_parameter = self.data_alloc.random_field(dims.CellDim) + current_exner = self.data_alloc.random_field(dims.CellDim, dims.KDim, low=1.0e-5) + current_rho = self.data_alloc.random_field(dims.CellDim, dims.KDim, low=1.0e-5) + current_theta_v = self.data_alloc.random_field(dims.CellDim, dims.KDim, low=1.0e-5) + current_w = self.data_alloc.random_field(dims.CellDim, dims.KDim, extend={dims.KDim: 1}) + inv_ddqz_z_full = self.data_alloc.random_field(dims.CellDim, dims.KDim, low=1.0e-5) + exner_w_implicit_weight_parameter = self.data_alloc.random_field(dims.CellDim) + theta_v_at_cells_on_half_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, low=1.0e-5 ) - perturbed_exner_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim + perturbed_exner_at_cells_on_model_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim ) - exner_tendency_due_to_slow_physics = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - rho_iau_increment = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - exner_iau_increment = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - ddqz_z_half = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - rayleigh_damping_factor = data_alloc.random_field(grid, dims.KDim) - reference_exner_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, low=1.0e-5 + exner_tendency_due_to_slow_physics = self.data_alloc.random_field(dims.CellDim, dims.KDim) + rho_iau_increment = self.data_alloc.random_field(dims.CellDim, dims.KDim) + exner_iau_increment = self.data_alloc.random_field(dims.CellDim, dims.KDim) + ddqz_z_half = self.data_alloc.random_field(dims.CellDim, dims.KDim, low=1.0e-5) + rayleigh_damping_factor = self.data_alloc.random_field(dims.KDim) + reference_exner_at_cells_on_model_levels = self.data_alloc.random_field( + dims.CellDim, dims.KDim, low=1.0e-5 ) - e_bln_c_s = data_alloc.random_field( - grid, dims.CellDim, dims.C2EDim, low=1.0e-5, high=0.99999 + e_bln_c_s = self.data_alloc.random_field( + dims.CellDim, dims.C2EDim, low=1.0e-5, high=0.99999 ) - wgtfac_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5, high=0.99999) - wgtfacq_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5, high=0.99999) + wgtfac_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, low=1.0e-5, high=0.99999) + wgtfacq_c = self.data_alloc.random_field(dims.CellDim, dims.KDim, low=1.0e-5, high=0.99999) - next_w = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) - next_rho = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) - next_exner = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) - next_theta_v = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) - dwdz_at_cells_on_model_levels = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - exner_dynamical_increment = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + next_w = self.data_alloc.zero_field(dims.CellDim, dims.KDim, extend={dims.KDim: 1}) + next_rho = self.data_alloc.constant_field(1.0e-5, dims.CellDim, dims.KDim) + next_exner = self.data_alloc.constant_field(1.0e-5, dims.CellDim, dims.KDim) + next_theta_v = self.data_alloc.constant_field(1.0e-5, dims.CellDim, dims.KDim) + dwdz_at_cells_on_model_levels = self.data_alloc.zero_field(dims.CellDim, dims.KDim) + exner_dynamical_increment = self.data_alloc.zero_field(dims.CellDim, dims.KDim) is_iau_active = request.param["is_iau_active"] at_first_substep = request.param["at_first_substep"] diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_T_from_internal_energy.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_T_from_internal_energy.py index 1dddcb000c..f149a1413e 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_T_from_internal_energy.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_T_from_internal_energy.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.thermo import T_from_internal_energy from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestTFromInternalEnergy(StencilTest): +class TestTFromInternalEnergy(stencil_tests.StencilTest): PROGRAM = T_from_internal_energy OUTPUTS = ("temperature",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, u: np.ndarray, qv: np.ndarray, qliq: np.ndarray, @@ -32,20 +32,20 @@ def reference( ) -> dict: return dict(temperature=np.full(u.shape, 255.75599999999997)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - u=data_alloc.constant_field( - grid, 38265357.270336017, dims.CellDim, dims.KDim, dtype=wpfloat + u=self.data_alloc.constant_field( + 38265357.270336017, dims.CellDim, dims.KDim, dtype=wpfloat ), - qv=data_alloc.constant_field(grid, 0.00122576, dims.CellDim, dims.KDim, dtype=wpfloat), - qliq=data_alloc.constant_field( - grid, 1.63837e-20, dims.CellDim, dims.KDim, dtype=wpfloat + qv=self.data_alloc.constant_field(0.00122576, dims.CellDim, dims.KDim, dtype=wpfloat), + qliq=self.data_alloc.constant_field( + 1.63837e-20, dims.CellDim, dims.KDim, dtype=wpfloat ), - qice=data_alloc.constant_field( - grid, 1.09462e-08, dims.CellDim, dims.KDim, dtype=wpfloat + qice=self.data_alloc.constant_field( + 1.09462e-08, dims.CellDim, dims.KDim, dtype=wpfloat ), - rho=data_alloc.constant_field(grid, 0.83444, dims.CellDim, dims.KDim, dtype=wpfloat), - dz=data_alloc.constant_field(grid, 249.569, dims.CellDim, dims.KDim, dtype=wpfloat), - temperature=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(0.83444, dims.CellDim, dims.KDim, dtype=wpfloat), + dz=self.data_alloc.constant_field(249.569, dims.CellDim, dims.KDim, dtype=wpfloat), + temperature=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_graupel.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_graupel.py index c362dd7ece..d08f9b9401 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_graupel.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_graupel.py @@ -10,27 +10,27 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import cloud_to_graupel from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestCloudToGraupel(StencilTest): +class TestCloudToGraupel(stencil_tests.StencilTest): PROGRAM = cloud_to_graupel OUTPUTS = ("riming_graupel_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, t: np.ndarray, rho: np.ndarray, qc: np.ndarray, qg: np.ndarray, **kwargs + grid: base.Grid, t: np.ndarray, rho: np.ndarray, qc: np.ndarray, qg: np.ndarray, **kwargs ) -> dict: return dict(riming_graupel_rate=np.full(t.shape, 2.7054723496793982e-10)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 256.983, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 0.909677, dims.CellDim, dims.KDim, dtype=wpfloat), - qc=data_alloc.constant_field(grid, 8.60101e-06, dims.CellDim, dims.KDim, dtype=wpfloat), - qg=data_alloc.constant_field(grid, 4.11575e-06, dims.CellDim, dims.KDim, dtype=wpfloat), - riming_graupel_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(256.983, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(0.909677, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.constant_field(8.60101e-06, dims.CellDim, dims.KDim, dtype=wpfloat), + qg=self.data_alloc.constant_field(4.11575e-06, dims.CellDim, dims.KDim, dtype=wpfloat), + riming_graupel_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_graupel_default.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_graupel_default.py index 9000f55ab3..9707b5cae2 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_graupel_default.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_graupel_default.py @@ -6,31 +6,30 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause import numpy as np -import pytest from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import cloud_to_graupel from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestCloudToGraupelDefault(StencilTest): +class TestCloudToGraupelDefault(stencil_tests.StencilTest): PROGRAM = cloud_to_graupel OUTPUTS = ("riming_graupel_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, t: np.ndarray, rho: np.ndarray, qc: np.ndarray, qg: np.ndarray, **kwargs + grid: base.Grid, t: np.ndarray, rho: np.ndarray, qc: np.ndarray, qg: np.ndarray, **kwargs ) -> dict: return dict(riming_graupel_rate=np.full(t.shape, 0.0)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 281.787, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 1.24783, dims.CellDim, dims.KDim, dtype=wpfloat), - qc=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - qg=data_alloc.constant_field(grid, 1.03636e-25, dims.CellDim, dims.KDim, dtype=wpfloat), - riming_graupel_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(281.787, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.24783, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), + qg=self.data_alloc.constant_field(1.03636e-25, dims.CellDim, dims.KDim, dtype=wpfloat), + riming_graupel_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_rain.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_rain.py index 4a6b064888..277d063fe2 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_rain.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_rain.py @@ -10,27 +10,27 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import cloud_to_rain from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestCloudToRain(StencilTest): +class TestCloudToRain(stencil_tests.StencilTest): PROGRAM = cloud_to_rain OUTPUTS = ("conversion_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, t: np.ndarray, qc: np.ndarray, qr: np.ndarray, nc: np.ndarray, **kwargs + grid: base.Grid, t: np.ndarray, qc: np.ndarray, qr: np.ndarray, nc: np.ndarray, **kwargs ) -> dict: return dict(conversion_rate=np.full(t.shape, 0.0045484481075162512)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 267.25, dims.CellDim, dims.KDim, dtype=wpfloat), - qc=data_alloc.constant_field(grid, 5.52921e-05, dims.CellDim, dims.KDim, dtype=wpfloat), - qr=data_alloc.constant_field(grid, 2.01511e-12, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(267.25, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.constant_field(5.52921e-05, dims.CellDim, dims.KDim, dtype=wpfloat), + qr=self.data_alloc.constant_field(2.01511e-12, dims.CellDim, dims.KDim, dtype=wpfloat), nc=100.0, - conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + conversion_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_rain_default.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_rain_default.py index 820355c8c2..3fdfd0a98c 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_rain_default.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_rain_default.py @@ -10,27 +10,27 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import cloud_to_rain from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestCloudToRainDefault(StencilTest): +class TestCloudToRainDefault(stencil_tests.StencilTest): PROGRAM = cloud_to_rain OUTPUTS = ("conversion_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, t: np.ndarray, qc: np.ndarray, qr: np.ndarray, nc: np.ndarray, **kwargs + grid: base.Grid, t: np.ndarray, qc: np.ndarray, qr: np.ndarray, nc: np.ndarray, **kwargs ) -> dict: return dict(conversion_rate=np.full(t.shape, 0.0)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 281.787, dims.CellDim, dims.KDim, dtype=wpfloat), - qc=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - qr=data_alloc.constant_field(grid, 52312e-7, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(281.787, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), + qr=self.data_alloc.constant_field(52312e-7, dims.CellDim, dims.KDim, dtype=wpfloat), nc=100.0, - conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + conversion_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_snow.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_snow.py index 90a9f18d34..cfc173f62f 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_snow.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_snow.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import cloud_to_snow from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestCloudToSnow(StencilTest): +class TestCloudToSnow(stencil_tests.StencilTest): PROGRAM = cloud_to_snow OUTPUTS = ("riming_snow_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, qc: np.ndarray, qs: np.ndarray, @@ -31,13 +31,13 @@ def reference( ) -> dict: return dict(riming_snow_rate=np.full(t.shape, 9.5431874564438999e-10)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 256.571, dims.CellDim, dims.KDim, dtype=wpfloat), - qc=data_alloc.constant_field(grid, 3.31476e-05, dims.CellDim, dims.KDim, dtype=wpfloat), - qs=data_alloc.constant_field(grid, 7.47365e-06, dims.CellDim, dims.KDim, dtype=wpfloat), - ns=data_alloc.constant_field(grid, 3.37707e07, dims.CellDim, dims.KDim, dtype=wpfloat), - lam=data_alloc.constant_field(grid, 8989.78, dims.CellDim, dims.KDim, dtype=wpfloat), - riming_snow_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(256.571, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.constant_field(3.31476e-05, dims.CellDim, dims.KDim, dtype=wpfloat), + qs=self.data_alloc.constant_field(7.47365e-06, dims.CellDim, dims.KDim, dtype=wpfloat), + ns=self.data_alloc.constant_field(3.37707e07, dims.CellDim, dims.KDim, dtype=wpfloat), + lam=self.data_alloc.constant_field(8989.78, dims.CellDim, dims.KDim, dtype=wpfloat), + riming_snow_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_snow_default.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_snow_default.py index 701a03e949..93fad432b2 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_snow_default.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_snow_default.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import cloud_to_snow from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestCloudToSnowDefault(StencilTest): +class TestCloudToSnowDefault(stencil_tests.StencilTest): PROGRAM = cloud_to_snow OUTPUTS = ("riming_snow_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, qc: np.ndarray, qs: np.ndarray, @@ -31,13 +31,13 @@ def reference( ) -> dict: return dict(riming_snow_rate=np.full(t.shape, 0.0)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 281.787, dims.CellDim, dims.KDim, dtype=wpfloat), - qc=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - qs=data_alloc.constant_field(grid, 3.63983e-40, dims.CellDim, dims.KDim, dtype=wpfloat), - ns=data_alloc.constant_field(grid, 800000.0, dims.CellDim, dims.KDim, dtype=wpfloat), - lam=data_alloc.constant_field(grid, 1.0e10, dims.CellDim, dims.KDim, dtype=wpfloat), - riming_snow_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(281.787, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), + qs=self.data_alloc.constant_field(3.63983e-40, dims.CellDim, dims.KDim, dtype=wpfloat), + ns=self.data_alloc.constant_field(800000.0, dims.CellDim, dims.KDim, dtype=wpfloat), + lam=self.data_alloc.constant_field(1.0e10, dims.CellDim, dims.KDim, dtype=wpfloat), + riming_snow_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_x_ice.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_x_ice.py index ae897cab44..8ae42ece86 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_x_ice.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_x_ice.py @@ -10,29 +10,29 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import cloud_x_ice from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestCloudXIce(StencilTest): +class TestCloudXIce(stencil_tests.StencilTest): PROGRAM = cloud_x_ice OUTPUTS = ("freezing_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, t: np.ndarray, qc: np.ndarray, qi: np.ndarray, dt: wpfloat, **kwargs + grid: base.Grid, t: np.ndarray, qc: np.ndarray, qi: np.ndarray, dt: wpfloat, **kwargs ) -> dict: return dict(freezing_rate=np.full(t.shape, -1.5008166666666666e-08)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field( - grid, 274.15, dims.CellDim, dims.KDim, dtype=wpfloat + t=self.data_alloc.constant_field( + 274.15, dims.CellDim, dims.KDim, dtype=wpfloat ), # tmelt + 1.0 - qc=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - qi=data_alloc.constant_field(grid, 4.50245e-07, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), + qi=self.data_alloc.constant_field(4.50245e-07, dims.CellDim, dims.KDim, dtype=wpfloat), dt=30.0, - freezing_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + freezing_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_x_ice_default.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_x_ice_default.py index 2a09f91f90..2c50fcd5bd 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_x_ice_default.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_x_ice_default.py @@ -10,27 +10,27 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import cloud_x_ice from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestCloudXIceDefault(StencilTest): +class TestCloudXIceDefault(stencil_tests.StencilTest): PROGRAM = cloud_x_ice OUTPUTS = ("freezing_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, t: np.ndarray, qc: np.ndarray, qi: np.ndarray, dt: wpfloat, **kwargs + grid: base.Grid, t: np.ndarray, qc: np.ndarray, qi: np.ndarray, dt: wpfloat, **kwargs ) -> dict: return dict(freezing_rate=np.full(t.shape, 0.0)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 256.835, dims.CellDim, dims.KDim, dtype=wpfloat), - qc=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - qi=data_alloc.constant_field(grid, 4.50245e-7, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(256.835, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), + qi=self.data_alloc.constant_field(4.50245e-7, dims.CellDim, dims.KDim, dtype=wpfloat), dt=30.0, - freezing_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + freezing_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_deposition_auto_conversion.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_deposition_auto_conversion.py index 377c940ed3..45358dd179 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_deposition_auto_conversion.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_deposition_auto_conversion.py @@ -13,28 +13,30 @@ deposition_auto_conversion, ) from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestDepositionAutoConversion(StencilTest): +class TestDepositionAutoConversion(stencil_tests.StencilTest): PROGRAM = deposition_auto_conversion OUTPUTS = ("conversion_rate",) - @staticmethod - def reference(grid, qi: np.ndarray, m_ice: np.ndarray, ice_dep: np.ndarray, **kwargs) -> dict: + @stencil_tests.static_reference + def reference( + grid: base.Grid, qi: np.ndarray, m_ice: np.ndarray, ice_dep: np.ndarray, **kwargs + ) -> dict: return dict(conversion_rate=np.full(qi.shape, 6.6430804299795412e-08)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - qi=data_alloc.constant_field( - grid, 2.02422e-2 + GraupelCt.qmin, dims.CellDim, dims.KDim, dtype=wpfloat + qi=self.data_alloc.constant_field( + 2.02422e-2 + GraupelCt.qmin, dims.CellDim, dims.KDim, dtype=wpfloat ), - m_ice=data_alloc.constant_field(grid, 1.0e-12, dims.CellDim, dims.KDim, dtype=wpfloat), - ice_dep=data_alloc.constant_field( - grid, 2.06276e-05, dims.CellDim, dims.KDim, dtype=wpfloat + m_ice=self.data_alloc.constant_field(1.0e-12, dims.CellDim, dims.KDim, dtype=wpfloat), + ice_dep=self.data_alloc.constant_field( + 2.06276e-05, dims.CellDim, dims.KDim, dtype=wpfloat ), - conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + conversion_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_deposition_auto_conversion_default.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_deposition_auto_conversion_default.py index 971bc9c76b..23c00208ee 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_deposition_auto_conversion_default.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_deposition_auto_conversion_default.py @@ -12,26 +12,28 @@ deposition_auto_conversion, ) from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestDepositionAutoConversionDefault(StencilTest): +class TestDepositionAutoConversionDefault(stencil_tests.StencilTest): PROGRAM = deposition_auto_conversion OUTPUTS = ("conversion_rate",) - @staticmethod - def reference(grid, qi: np.ndarray, m_ice: np.ndarray, ice_dep: np.ndarray, **kwargs) -> dict: + @stencil_tests.static_reference + def reference( + grid: base.Grid, qi: np.ndarray, m_ice: np.ndarray, ice_dep: np.ndarray, **kwargs + ) -> dict: return dict(conversion_rate=np.full(qi.shape, 0.0)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - qi=data_alloc.constant_field(grid, 2.02422e-23, dims.CellDim, dims.KDim, dtype=wpfloat), - m_ice=data_alloc.constant_field(grid, 1.0e-12, dims.CellDim, dims.KDim, dtype=wpfloat), - ice_dep=data_alloc.constant_field( - grid, -2.06276e-05, dims.CellDim, dims.KDim, dtype=wpfloat + qi=self.data_alloc.constant_field(2.02422e-23, dims.CellDim, dims.KDim, dtype=wpfloat), + m_ice=self.data_alloc.constant_field(1.0e-12, dims.CellDim, dims.KDim, dtype=wpfloat), + ice_dep=self.data_alloc.constant_field( + -2.06276e-05, dims.CellDim, dims.KDim, dtype=wpfloat ), - conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + conversion_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_deposition_factor.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_deposition_factor.py index d89b8452a9..192dc298f9 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_deposition_factor.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_deposition_factor.py @@ -10,25 +10,23 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.properties import deposition_factor from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestDepositionFactor(StencilTest): +class TestDepositionFactor(stencil_tests.StencilTest): PROGRAM = deposition_factor OUTPUTS = ("deposition_rate",) - @staticmethod - def reference(grid, t: np.ndarray, qvsi: np.ndarray, **kwargs) -> dict: + @stencil_tests.static_reference + def reference(grid: base.Grid, t: np.ndarray, qvsi: np.ndarray, **kwargs) -> dict: return dict(deposition_rate=np.full(t.shape, 1.3234329478493952e-05)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 272.731, dims.CellDim, dims.KDim, dtype=wpfloat), - qvsi=data_alloc.constant_field( - grid, 0.00416891, dims.CellDim, dims.KDim, dtype=wpfloat - ), - deposition_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(272.731, dims.CellDim, dims.KDim, dtype=wpfloat), + qvsi=self.data_alloc.constant_field(0.00416891, dims.CellDim, dims.KDim, dtype=wpfloat), + deposition_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_dqsatdT_rho.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_dqsatdT_rho.py index 65522326e2..7d1386a6d8 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_dqsatdT_rho.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_dqsatdT_rho.py @@ -10,23 +10,23 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.thermo import dqsatdT_rho from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestQsatRho(StencilTest): +class TestQsatRho(stencil_tests.StencilTest): PROGRAM = dqsatdT_rho OUTPUTS = ("derivative",) - @staticmethod - def reference(grid, qs: np.ndarray, t: np.ndarray, **kwargs) -> dict: + @stencil_tests.static_reference + def reference(grid: base.Grid, qs: np.ndarray, t: np.ndarray, **kwargs) -> dict: return dict(derivative=np.full(t.shape, 0.00030825070286492049)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - qs=data_alloc.constant_field(grid, 0.00448941, dims.CellDim, dims.KDim, dtype=wpfloat), - t=data_alloc.constant_field(grid, 273.909, dims.CellDim, dims.KDim, dtype=wpfloat), - derivative=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + qs=self.data_alloc.constant_field(0.00448941, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(273.909, dims.CellDim, dims.KDim, dtype=wpfloat), + derivative=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_fall_speed.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_fall_speed.py index 162ad4c301..cdb152824e 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_fall_speed.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_fall_speed.py @@ -11,27 +11,32 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.common.constants import Idx from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.properties import fall_speed from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestFallSpeed(StencilTest): +class TestFallSpeed(stencil_tests.StencilTest): PROGRAM = fall_speed OUTPUTS = ("speed",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, density: np.ndarray, prefactor: wpfloat, offset: wpfloat, exponent: wpfloat, **kwargs + grid: base.Grid, + density: np.ndarray, + prefactor: wpfloat, + offset: wpfloat, + exponent: wpfloat, + **kwargs, ) -> dict: return dict(speed=np.full(density.shape, 0.67882452435647411)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - density=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + density=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), prefactor=Idx.prefactor_r, offset=Idx.offset_r, exponent=Idx.exponent_r, - speed=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + speed=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_graupel_to_rain.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_graupel_to_rain.py index 6f52329dc9..7667e0f7e6 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_graupel_to_rain.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_graupel_to_rain.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import graupel_to_rain from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestGraupelToRain(StencilTest): +class TestGraupelToRain(stencil_tests.StencilTest): PROGRAM = graupel_to_rain OUTPUTS = ("rain_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, p: np.ndarray, rho: np.ndarray, @@ -31,15 +31,15 @@ def reference( ) -> dict: return dict(rain_rate=np.full(t.shape, 5.9748142538569357e-13)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 280.156, dims.CellDim, dims.KDim, dtype=wpfloat), - p=data_alloc.constant_field(grid, 98889.4, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 1.22804, dims.CellDim, dims.KDim, dtype=wpfloat), - dvsw0=data_alloc.constant_field( - grid, -0.00167867, dims.CellDim, dims.KDim, dtype=wpfloat + t=self.data_alloc.constant_field(280.156, dims.CellDim, dims.KDim, dtype=wpfloat), + p=self.data_alloc.constant_field(98889.4, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.22804, dims.CellDim, dims.KDim, dtype=wpfloat), + dvsw0=self.data_alloc.constant_field( + -0.00167867, dims.CellDim, dims.KDim, dtype=wpfloat ), - qg=data_alloc.constant_field(grid, 1.53968e-15, dims.CellDim, dims.KDim, dtype=wpfloat), - rain_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + qg=self.data_alloc.constant_field(1.53968e-15, dims.CellDim, dims.KDim, dtype=wpfloat), + rain_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_graupel_to_rain_default.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_graupel_to_rain_default.py index bd8b18b74e..99401b4bfd 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_graupel_to_rain_default.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_graupel_to_rain_default.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import graupel_to_rain from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestGraupelToRainDefault(StencilTest): +class TestGraupelToRainDefault(stencil_tests.StencilTest): PROGRAM = graupel_to_rain OUTPUTS = ("rain_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, p: np.ndarray, rho: np.ndarray, @@ -31,15 +31,15 @@ def reference( ) -> dict: return dict(rain_rate=np.full(t.shape, 0.0)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 280.156, dims.CellDim, dims.KDim, dtype=wpfloat), - p=data_alloc.constant_field(grid, 98889.4, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 1.22804, dims.CellDim, dims.KDim, dtype=wpfloat), - dvsw0=data_alloc.constant_field( - grid, -0.00167867, dims.CellDim, dims.KDim, dtype=wpfloat + t=self.data_alloc.constant_field(280.156, dims.CellDim, dims.KDim, dtype=wpfloat), + p=self.data_alloc.constant_field(98889.4, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.22804, dims.CellDim, dims.KDim, dtype=wpfloat), + dvsw0=self.data_alloc.constant_field( + -0.00167867, dims.CellDim, dims.KDim, dtype=wpfloat ), - qg=data_alloc.constant_field(grid, 1.53968e-17, dims.CellDim, dims.KDim, dtype=wpfloat), - rain_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + qg=self.data_alloc.constant_field(1.53968e-17, dims.CellDim, dims.KDim, dtype=wpfloat), + rain_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_deposition_nucleation.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_deposition_nucleation.py index 374b5c738e..9aea953e26 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_deposition_nucleation.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_deposition_nucleation.py @@ -12,18 +12,18 @@ ice_deposition_nucleation, ) from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestIceDepositionNucleation(StencilTest): +class TestIceDepositionNucleation(stencil_tests.StencilTest): PROGRAM = ice_deposition_nucleation OUTPUTS = ("vapor_deposition_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, qc: np.ndarray, qi: np.ndarray, @@ -34,16 +34,16 @@ def reference( ) -> dict: return dict(vapor_deposition_rate=np.full(t.shape, 1.6836299999999999e-13)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 160.9, dims.CellDim, dims.KDim, dtype=wpfloat), - qc=data_alloc.constant_field(grid, 1.0e-2, dims.CellDim, dims.KDim, dtype=wpfloat), - qi=data_alloc.constant_field(grid, 2.02422e-23, dims.CellDim, dims.KDim, dtype=wpfloat), - ni=data_alloc.constant_field(grid, 5.05089, dims.CellDim, dims.KDim, dtype=wpfloat), - dvsi=data_alloc.constant_field(grid, 0.0001, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(160.9, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.constant_field(1.0e-2, dims.CellDim, dims.KDim, dtype=wpfloat), + qi=self.data_alloc.constant_field(2.02422e-23, dims.CellDim, dims.KDim, dtype=wpfloat), + ni=self.data_alloc.constant_field(5.05089, dims.CellDim, dims.KDim, dtype=wpfloat), + dvsi=self.data_alloc.constant_field(0.0001, dims.CellDim, dims.KDim, dtype=wpfloat), dt=30.0, - vapor_deposition_rate=data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, dtype=wpfloat + vapor_deposition_rate=self.data_alloc.zero_field( + dims.CellDim, dims.KDim, dtype=wpfloat ), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_deposition_nucleation_default.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_deposition_nucleation_default.py index b18697e349..d449d4fee0 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_deposition_nucleation_default.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_deposition_nucleation_default.py @@ -12,18 +12,18 @@ ice_deposition_nucleation, ) from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestIceDepositionNucleationDefault(StencilTest): +class TestIceDepositionNucleationDefault(stencil_tests.StencilTest): PROGRAM = ice_deposition_nucleation OUTPUTS = ("vapor_deposition_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, qc: np.ndarray, qi: np.ndarray, @@ -34,18 +34,18 @@ def reference( ) -> dict: return dict(vapor_deposition_rate=np.full(t.shape, 0.0)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 272.731, dims.CellDim, dims.KDim, dtype=wpfloat), - qc=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - qi=data_alloc.constant_field(grid, 2.02422e-23, dims.CellDim, dims.KDim, dtype=wpfloat), - ni=data_alloc.constant_field(grid, 5.05089, dims.CellDim, dims.KDim, dtype=wpfloat), - dvsi=data_alloc.constant_field( - grid, -0.000618828, dims.CellDim, dims.KDim, dtype=wpfloat + t=self.data_alloc.constant_field(272.731, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), + qi=self.data_alloc.constant_field(2.02422e-23, dims.CellDim, dims.KDim, dtype=wpfloat), + ni=self.data_alloc.constant_field(5.05089, dims.CellDim, dims.KDim, dtype=wpfloat), + dvsi=self.data_alloc.constant_field( + -0.000618828, dims.CellDim, dims.KDim, dtype=wpfloat ), dt=30.0, - vapor_deposition_rate=data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, dtype=wpfloat + vapor_deposition_rate=self.data_alloc.zero_field( + dims.CellDim, dims.KDim, dtype=wpfloat ), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_mass.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_mass.py index bb9bdf1a45..718cb1fd31 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_mass.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_mass.py @@ -10,23 +10,23 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.properties import ice_mass from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestIceNumber(StencilTest): +class TestIceNumber(stencil_tests.StencilTest): PROGRAM = ice_mass OUTPUTS = ("mass",) - @staticmethod - def reference(grid, qi: np.ndarray, ni: np.ndarray, **kwargs) -> dict: + @stencil_tests.static_reference + def reference(grid: base.Grid, qi: np.ndarray, ni: np.ndarray, **kwargs) -> dict: return dict(mass=np.full(qi.shape, 1.0e-12)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - qi=data_alloc.constant_field(grid, 2.02422e-23, dims.CellDim, dims.KDim, dtype=wpfloat), - ni=data_alloc.constant_field(grid, 5.05089, dims.CellDim, dims.KDim, dtype=wpfloat), - mass=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + qi=self.data_alloc.constant_field(2.02422e-23, dims.CellDim, dims.KDim, dtype=wpfloat), + ni=self.data_alloc.constant_field(5.05089, dims.CellDim, dims.KDim, dtype=wpfloat), + mass=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_number.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_number.py index 05b5fa336c..4176092cf3 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_number.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_number.py @@ -10,23 +10,23 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.properties import ice_number from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestIceNumber(StencilTest): +class TestIceNumber(stencil_tests.StencilTest): PROGRAM = ice_number OUTPUTS = ("number",) - @staticmethod - def reference(grid, t: np.ndarray, rho: np.ndarray, **kwargs) -> dict: + @stencil_tests.static_reference + def reference(grid: base.Grid, t: np.ndarray, rho: np.ndarray, **kwargs) -> dict: return dict(number=np.full(t.shape, 5.0507995893464388)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 272.731, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 1.12442, dims.CellDim, dims.KDim, dtype=wpfloat), - number=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(272.731, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.12442, dims.CellDim, dims.KDim, dtype=wpfloat), + number=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_sticking.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_sticking.py index 52532333b9..5e722a5787 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_sticking.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_sticking.py @@ -10,22 +10,22 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.properties import ice_sticking from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestIceSticking(StencilTest): +class TestIceSticking(stencil_tests.StencilTest): PROGRAM = ice_sticking OUTPUTS = ("sticking_factor",) - @staticmethod - def reference(grid, t: np.ndarray, **kwargs) -> dict: + @stencil_tests.static_reference + def reference(grid: base.Grid, t: np.ndarray, **kwargs) -> dict: return dict(sticking_factor=np.full(t.shape, 0.8697930232044021)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 271.6, dims.CellDim, dims.KDim, dtype=wpfloat), - sticking_factor=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(271.6, dims.CellDim, dims.KDim, dtype=wpfloat), + sticking_factor=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_to_graupel.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_to_graupel.py index 4f2bff9e32..1a8bc05d68 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_to_graupel.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_to_graupel.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import ice_to_graupel from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestIceToGraupel(StencilTest): +class TestIceToGraupel(stencil_tests.StencilTest): PROGRAM = ice_to_graupel OUTPUTS = ("aggregation",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, rho: np.ndarray, qr: np.ndarray, qg: np.ndarray, @@ -31,15 +31,15 @@ def reference( ) -> dict: return dict(aggregation=np.full(rho.shape, 7.1049436957697864e-19)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - rho=data_alloc.constant_field(grid, 1.04848, dims.CellDim, dims.KDim, dtype=wpfloat), - qr=data_alloc.constant_field(grid, 6.00408e-13, dims.CellDim, dims.KDim, dtype=wpfloat), - qg=data_alloc.constant_field(grid, 1.19022e-18, dims.CellDim, dims.KDim, dtype=wpfloat), - qi=data_alloc.constant_field(grid, 1.9584e-08, dims.CellDim, dims.KDim, dtype=wpfloat), - sticking_eff=data_alloc.constant_field( - grid, 1.9584e-08, dims.CellDim, dims.KDim, dtype=wpfloat + rho=self.data_alloc.constant_field(1.04848, dims.CellDim, dims.KDim, dtype=wpfloat), + qr=self.data_alloc.constant_field(6.00408e-13, dims.CellDim, dims.KDim, dtype=wpfloat), + qg=self.data_alloc.constant_field(1.19022e-18, dims.CellDim, dims.KDim, dtype=wpfloat), + qi=self.data_alloc.constant_field(1.9584e-08, dims.CellDim, dims.KDim, dtype=wpfloat), + sticking_eff=self.data_alloc.constant_field( + 1.9584e-08, dims.CellDim, dims.KDim, dtype=wpfloat ), - aggregation=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + aggregation=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_to_graupel_default.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_to_graupel_default.py index 4f2bff9e32..1a8bc05d68 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_to_graupel_default.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_to_graupel_default.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import ice_to_graupel from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestIceToGraupel(StencilTest): +class TestIceToGraupel(stencil_tests.StencilTest): PROGRAM = ice_to_graupel OUTPUTS = ("aggregation",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, rho: np.ndarray, qr: np.ndarray, qg: np.ndarray, @@ -31,15 +31,15 @@ def reference( ) -> dict: return dict(aggregation=np.full(rho.shape, 7.1049436957697864e-19)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - rho=data_alloc.constant_field(grid, 1.04848, dims.CellDim, dims.KDim, dtype=wpfloat), - qr=data_alloc.constant_field(grid, 6.00408e-13, dims.CellDim, dims.KDim, dtype=wpfloat), - qg=data_alloc.constant_field(grid, 1.19022e-18, dims.CellDim, dims.KDim, dtype=wpfloat), - qi=data_alloc.constant_field(grid, 1.9584e-08, dims.CellDim, dims.KDim, dtype=wpfloat), - sticking_eff=data_alloc.constant_field( - grid, 1.9584e-08, dims.CellDim, dims.KDim, dtype=wpfloat + rho=self.data_alloc.constant_field(1.04848, dims.CellDim, dims.KDim, dtype=wpfloat), + qr=self.data_alloc.constant_field(6.00408e-13, dims.CellDim, dims.KDim, dtype=wpfloat), + qg=self.data_alloc.constant_field(1.19022e-18, dims.CellDim, dims.KDim, dtype=wpfloat), + qi=self.data_alloc.constant_field(1.9584e-08, dims.CellDim, dims.KDim, dtype=wpfloat), + sticking_eff=self.data_alloc.constant_field( + 1.9584e-08, dims.CellDim, dims.KDim, dtype=wpfloat ), - aggregation=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + aggregation=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_to_snow.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_to_snow.py index 7072fc1c8c..2d453a86fe 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_to_snow.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_to_snow.py @@ -10,29 +10,34 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import ice_to_snow from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestIceToSnow(StencilTest): +class TestIceToSnow(stencil_tests.StencilTest): PROGRAM = ice_to_snow OUTPUTS = ("conversion_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, qi: np.ndarray, ns: np.ndarray, lam: np.ndarray, sticking_eff: np.ndarray, **kwargs + grid: base.Grid, + qi: np.ndarray, + ns: np.ndarray, + lam: np.ndarray, + sticking_eff: np.ndarray, + **kwargs, ) -> dict: return dict(conversion_rate=np.full(qi.shape, 3.3262745200740486e-11)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - qi=data_alloc.constant_field(grid, 6.43223e-08, dims.CellDim, dims.KDim, dtype=wpfloat), - ns=data_alloc.constant_field(grid, 1.93157e07, dims.CellDim, dims.KDim, dtype=wpfloat), - lam=data_alloc.constant_field(grid, 10576.8, dims.CellDim, dims.KDim, dtype=wpfloat), - sticking_eff=data_alloc.constant_field( - grid, 0.511825, dims.CellDim, dims.KDim, dtype=wpfloat + qi=self.data_alloc.constant_field(6.43223e-08, dims.CellDim, dims.KDim, dtype=wpfloat), + ns=self.data_alloc.constant_field(1.93157e07, dims.CellDim, dims.KDim, dtype=wpfloat), + lam=self.data_alloc.constant_field(10576.8, dims.CellDim, dims.KDim, dtype=wpfloat), + sticking_eff=self.data_alloc.constant_field( + 0.511825, dims.CellDim, dims.KDim, dtype=wpfloat ), - conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + conversion_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_to_snow_default.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_to_snow_default.py index bf6ad46db8..5b92b8c347 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_to_snow_default.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_ice_to_snow_default.py @@ -10,29 +10,34 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import ice_to_snow from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestIceToSnowDefault(StencilTest): +class TestIceToSnowDefault(stencil_tests.StencilTest): PROGRAM = ice_to_snow OUTPUTS = ("conversion_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, qi: np.ndarray, ns: np.ndarray, lam: np.ndarray, sticking_eff: np.ndarray, **kwargs + grid: base.Grid, + qi: np.ndarray, + ns: np.ndarray, + lam: np.ndarray, + sticking_eff: np.ndarray, + **kwargs, ) -> dict: return dict(conversion_rate=np.full(qi.shape, 0.0)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - qi=data_alloc.constant_field(grid, 7.95122e-25, dims.CellDim, dims.KDim, dtype=wpfloat), - ns=data_alloc.constant_field(grid, 2.23336e07, dims.CellDim, dims.KDim, dtype=wpfloat), - lam=data_alloc.constant_field(grid, 61911.1, dims.CellDim, dims.KDim, dtype=wpfloat), - sticking_eff=data_alloc.constant_field( - grid, 0.241568, dims.CellDim, dims.KDim, dtype=wpfloat + qi=self.data_alloc.constant_field(7.95122e-25, dims.CellDim, dims.KDim, dtype=wpfloat), + ns=self.data_alloc.constant_field(2.23336e07, dims.CellDim, dims.KDim, dtype=wpfloat), + lam=self.data_alloc.constant_field(61911.1, dims.CellDim, dims.KDim, dtype=wpfloat), + sticking_eff=self.data_alloc.constant_field( + 0.241568, dims.CellDim, dims.KDim, dtype=wpfloat ), - conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + conversion_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_internal_energy.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_internal_energy.py index 3da5d59413..183230f526 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_internal_energy.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_internal_energy.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.thermo import internal_energy from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestInternalEnergy(StencilTest): +class TestInternalEnergy(stencil_tests.StencilTest): PROGRAM = internal_energy OUTPUTS = ("energy",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, qv: np.ndarray, qliq: np.ndarray, @@ -32,18 +32,18 @@ def reference( ) -> dict: return dict(energy=np.full(t.shape, 38265357.270336017)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 255.756, dims.CellDim, dims.KDim, dtype=wpfloat), - qv=data_alloc.constant_field(grid, 0.00122576, dims.CellDim, dims.KDim, dtype=wpfloat), - qliq=data_alloc.constant_field( - grid, 1.63837e-20, dims.CellDim, dims.KDim, dtype=wpfloat + t=self.data_alloc.constant_field(255.756, dims.CellDim, dims.KDim, dtype=wpfloat), + qv=self.data_alloc.constant_field(0.00122576, dims.CellDim, dims.KDim, dtype=wpfloat), + qliq=self.data_alloc.constant_field( + 1.63837e-20, dims.CellDim, dims.KDim, dtype=wpfloat ), - qice=data_alloc.constant_field( - grid, 1.09462e-08, dims.CellDim, dims.KDim, dtype=wpfloat + qice=self.data_alloc.constant_field( + 1.09462e-08, dims.CellDim, dims.KDim, dtype=wpfloat ), - rho=data_alloc.constant_field(grid, 0.83444, dims.CellDim, dims.KDim, dtype=wpfloat), - dz=data_alloc.constant_field(grid, 249.569, dims.CellDim, dims.KDim, dtype=wpfloat), - energy=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(0.83444, dims.CellDim, dims.KDim, dtype=wpfloat), + dz=self.data_alloc.constant_field(249.569, dims.CellDim, dims.KDim, dtype=wpfloat), + energy=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_qsat_ice_rho.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_qsat_ice_rho.py index e29dc9911a..241e233dd3 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_qsat_ice_rho.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_qsat_ice_rho.py @@ -10,23 +10,23 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.thermo import qsat_ice_rho from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestQsatIceRho(StencilTest): +class TestQsatIceRho(stencil_tests.StencilTest): PROGRAM = qsat_ice_rho OUTPUTS = ("pressure",) - @staticmethod - def reference(grid, t: np.ndarray, rho: np.ndarray, **kwargs) -> dict: + @stencil_tests.static_reference + def reference(grid: base.Grid, t: np.ndarray, rho: np.ndarray, **kwargs) -> dict: return dict(pressure=np.full(t.shape, 0.0074981245870634101)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 281.787, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 1.24783, dims.CellDim, dims.KDim, dtype=wpfloat), - pressure=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(281.787, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.24783, dims.CellDim, dims.KDim, dtype=wpfloat), + pressure=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_qsat_rho.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_qsat_rho.py index 8e29622ff0..1dde875e5d 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_qsat_rho.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_qsat_rho.py @@ -10,23 +10,23 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.thermo import qsat_rho from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestQsatRho(StencilTest): +class TestQsatRho(stencil_tests.StencilTest): PROGRAM = qsat_rho OUTPUTS = ("pressure",) - @staticmethod - def reference(grid, t: np.ndarray, rho: np.ndarray, **kwargs) -> dict: + @stencil_tests.static_reference + def reference(grid: base.Grid, t: np.ndarray, rho: np.ndarray, **kwargs) -> dict: return dict(pressure=np.full(t.shape, 0.0069027592942577506)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 281.787, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 1.24783, dims.CellDim, dims.KDim, dtype=wpfloat), - pressure=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(281.787, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.24783, dims.CellDim, dims.KDim, dtype=wpfloat), + pressure=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_qsat_rho_tmelt.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_qsat_rho_tmelt.py index 6e13021be7..a747819a84 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_qsat_rho_tmelt.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_qsat_rho_tmelt.py @@ -10,22 +10,22 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.thermo import qsat_rho_tmelt from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestQsatRho(StencilTest): +class TestQsatRho(stencil_tests.StencilTest): PROGRAM = qsat_rho_tmelt OUTPUTS = ("pressure",) - @staticmethod - def reference(grid, rho: np.ndarray, **kwargs) -> dict: + @stencil_tests.static_reference + def reference(grid: base.Grid, rho: np.ndarray, **kwargs) -> dict: return dict(pressure=np.full(rho.shape, 0.0038828182695875113)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - rho=data_alloc.constant_field(grid, 1.24783, dims.CellDim, dims.KDim, dtype=wpfloat), - pressure=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.24783, dims.CellDim, dims.KDim, dtype=wpfloat), + pressure=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_graupel_1.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_graupel_1.py index bc1c75ce0d..139cbba786 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_graupel_1.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_graupel_1.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import rain_to_graupel from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestRainToGraupel1(StencilTest): +class TestRainToGraupel1(stencil_tests.StencilTest): PROGRAM = rain_to_graupel OUTPUTS = ("conversion_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, rho: np.ndarray, qc: np.ndarray, @@ -35,17 +35,17 @@ def reference( ) -> dict: return dict(conversion_rate=np.full(t.shape, 5.1570340525841922e-17)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 258.542, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 0.956089, dims.CellDim, dims.KDim, dtype=wpfloat), - qc=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - qr=data_alloc.constant_field(grid, 3.01332e-11, dims.CellDim, dims.KDim, dtype=wpfloat), - qi=data_alloc.constant_field(grid, 5.57166e-06, dims.CellDim, dims.KDim, dtype=wpfloat), - qs=data_alloc.constant_field(grid, 3.55432e-05, dims.CellDim, dims.KDim, dtype=wpfloat), - mi=data_alloc.constant_field(grid, 1.0e-9, dims.CellDim, dims.KDim, dtype=wpfloat), - dvsw=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(258.542, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(0.956089, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), + qr=self.data_alloc.constant_field(3.01332e-11, dims.CellDim, dims.KDim, dtype=wpfloat), + qi=self.data_alloc.constant_field(5.57166e-06, dims.CellDim, dims.KDim, dtype=wpfloat), + qs=self.data_alloc.constant_field(3.55432e-05, dims.CellDim, dims.KDim, dtype=wpfloat), + mi=self.data_alloc.constant_field(1.0e-9, dims.CellDim, dims.KDim, dtype=wpfloat), + dvsw=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), dt=30.0, - conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + conversion_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_graupel_2.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_graupel_2.py index 1a26ea9b2e..6318c365bd 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_graupel_2.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_graupel_2.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import rain_to_graupel from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestRainToGraupel1(StencilTest): +class TestRainToGraupel1(stencil_tests.StencilTest): PROGRAM = rain_to_graupel OUTPUTS = ("conversion_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, rho: np.ndarray, qc: np.ndarray, @@ -35,17 +35,17 @@ def reference( ) -> dict: return dict(conversion_rate=np.full(t.shape, 1.0044914238516472e-12)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 230.542, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 0.956089, dims.CellDim, dims.KDim, dtype=wpfloat), - qc=data_alloc.constant_field(grid, 8.6157e-5, dims.CellDim, dims.KDim, dtype=wpfloat), - qr=data_alloc.constant_field(grid, 3.01332e-11, dims.CellDim, dims.KDim, dtype=wpfloat), - qi=data_alloc.constant_field(grid, 5.57166e-06, dims.CellDim, dims.KDim, dtype=wpfloat), - qs=data_alloc.constant_field(grid, 3.55432e-05, dims.CellDim, dims.KDim, dtype=wpfloat), - mi=data_alloc.constant_field(grid, 1.0e-9, dims.CellDim, dims.KDim, dtype=wpfloat), - dvsw=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(230.542, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(0.956089, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.constant_field(8.6157e-5, dims.CellDim, dims.KDim, dtype=wpfloat), + qr=self.data_alloc.constant_field(3.01332e-11, dims.CellDim, dims.KDim, dtype=wpfloat), + qi=self.data_alloc.constant_field(5.57166e-06, dims.CellDim, dims.KDim, dtype=wpfloat), + qs=self.data_alloc.constant_field(3.55432e-05, dims.CellDim, dims.KDim, dtype=wpfloat), + mi=self.data_alloc.constant_field(1.0e-9, dims.CellDim, dims.KDim, dtype=wpfloat), + dvsw=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), dt=30.0, - conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + conversion_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_graupel_3.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_graupel_3.py index 0d6a3dc0a5..a0a42aa909 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_graupel_3.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_graupel_3.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import rain_to_graupel from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestRainToGraupel3(StencilTest): +class TestRainToGraupel3(stencil_tests.StencilTest): PROGRAM = rain_to_graupel OUTPUTS = ("conversion_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, rho: np.ndarray, qc: np.ndarray, @@ -35,17 +35,17 @@ def reference( ) -> dict: return dict(conversion_rate=np.full(t.shape, 5.1423851647153399e-17)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 258.542, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 0.956089, dims.CellDim, dims.KDim, dtype=wpfloat), - qc=data_alloc.constant_field(grid, 8.6157e-5, dims.CellDim, dims.KDim, dtype=wpfloat), - qr=data_alloc.constant_field(grid, 3.01332e-11, dims.CellDim, dims.KDim, dtype=wpfloat), - qi=data_alloc.constant_field(grid, 5.57166e-06, dims.CellDim, dims.KDim, dtype=wpfloat), - qs=data_alloc.constant_field(grid, 3.55432e-05, dims.CellDim, dims.KDim, dtype=wpfloat), - mi=data_alloc.constant_field(grid, 1.0e-9, dims.CellDim, dims.KDim, dtype=wpfloat), - dvsw=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(258.542, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(0.956089, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.constant_field(8.6157e-5, dims.CellDim, dims.KDim, dtype=wpfloat), + qr=self.data_alloc.constant_field(3.01332e-11, dims.CellDim, dims.KDim, dtype=wpfloat), + qi=self.data_alloc.constant_field(5.57166e-06, dims.CellDim, dims.KDim, dtype=wpfloat), + qs=self.data_alloc.constant_field(3.55432e-05, dims.CellDim, dims.KDim, dtype=wpfloat), + mi=self.data_alloc.constant_field(1.0e-9, dims.CellDim, dims.KDim, dtype=wpfloat), + dvsw=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), dt=30.0, - conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + conversion_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_graupel_default.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_graupel_default.py index 00ece8df68..a39288bc5f 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_graupel_default.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_graupel_default.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import rain_to_graupel from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestRainToGraupelDefault(StencilTest): +class TestRainToGraupelDefault(stencil_tests.StencilTest): PROGRAM = rain_to_graupel OUTPUTS = ("conversion_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, rho: np.ndarray, qc: np.ndarray, @@ -35,19 +35,19 @@ def reference( ) -> dict: return dict(conversion_rate=np.full(t.shape, 0.0)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 272.731, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 1.12442, dims.CellDim, dims.KDim, dtype=wpfloat), - qc=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - qr=data_alloc.constant_field(grid, 1.34006e-17, dims.CellDim, dims.KDim, dtype=wpfloat), - qi=data_alloc.constant_field(grid, 2.02422e-23, dims.CellDim, dims.KDim, dtype=wpfloat), - qs=data_alloc.constant_field(grid, 1.02627e-19, dims.CellDim, dims.KDim, dtype=wpfloat), - mi=data_alloc.constant_field(grid, 1.0e-12, dims.CellDim, dims.KDim, dtype=wpfloat), - dvsw=data_alloc.constant_field( - grid, -0.000635669, dims.CellDim, dims.KDim, dtype=wpfloat + t=self.data_alloc.constant_field(272.731, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.12442, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), + qr=self.data_alloc.constant_field(1.34006e-17, dims.CellDim, dims.KDim, dtype=wpfloat), + qi=self.data_alloc.constant_field(2.02422e-23, dims.CellDim, dims.KDim, dtype=wpfloat), + qs=self.data_alloc.constant_field(1.02627e-19, dims.CellDim, dims.KDim, dtype=wpfloat), + mi=self.data_alloc.constant_field(1.0e-12, dims.CellDim, dims.KDim, dtype=wpfloat), + dvsw=self.data_alloc.constant_field( + -0.000635669, dims.CellDim, dims.KDim, dtype=wpfloat ), dt=30.0, - conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + conversion_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_vapor.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_vapor.py index c5528bc37a..44c35230a6 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_vapor.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_vapor.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import rain_to_vapor from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestRainToVapor(StencilTest): +class TestRainToVapor(stencil_tests.StencilTest): PROGRAM = rain_to_vapor OUTPUTS = ("conversion_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, rho: np.ndarray, qc: np.ndarray, @@ -32,14 +32,14 @@ def reference( ) -> dict: return dict(conversion_rate=np.full(t.shape, 2.8556697055499901e-19)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 258.542, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 0.956089, dims.CellDim, dims.KDim, dtype=wpfloat), - qc=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - qr=data_alloc.constant_field(grid, 3.01332e-11, dims.CellDim, dims.KDim, dtype=wpfloat), - dvsw=data_alloc.constant_field(grid, -1.0e-10, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(258.542, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(0.956089, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), + qr=self.data_alloc.constant_field(3.01332e-11, dims.CellDim, dims.KDim, dtype=wpfloat), + dvsw=self.data_alloc.constant_field(-1.0e-10, dims.CellDim, dims.KDim, dtype=wpfloat), dt=30.0, - conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + conversion_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_vapor_default.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_vapor_default.py index fb681877af..268dea89ce 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_vapor_default.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_vapor_default.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import rain_to_vapor from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestRainToVaporDefault(StencilTest): +class TestRainToVaporDefault(stencil_tests.StencilTest): PROGRAM = rain_to_vapor OUTPUTS = ("conversion_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, rho: np.ndarray, qc: np.ndarray, @@ -32,14 +32,14 @@ def reference( ) -> dict: return dict(conversion_rate=np.full(t.shape, 0.0)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 258.542, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 0.956089, dims.CellDim, dims.KDim, dtype=wpfloat), - qc=data_alloc.constant_field(grid, 8.6157e-5, dims.CellDim, dims.KDim, dtype=wpfloat), - qr=data_alloc.constant_field(grid, 3.01332e-11, dims.CellDim, dims.KDim, dtype=wpfloat), - dvsw=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(258.542, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(0.956089, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.constant_field(8.6157e-5, dims.CellDim, dims.KDim, dtype=wpfloat), + qr=self.data_alloc.constant_field(3.01332e-11, dims.CellDim, dims.KDim, dtype=wpfloat), + dvsw=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), dt=30.0, - conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + conversion_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_sat_pres_ice.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_sat_pres_ice.py index d56e9a0a01..310971b4d7 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_sat_pres_ice.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_sat_pres_ice.py @@ -10,22 +10,22 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.thermo import sat_pres_ice from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestSatPresIce(StencilTest): +class TestSatPresIce(stencil_tests.StencilTest): PROGRAM = sat_pres_ice OUTPUTS = ("pressure",) - @staticmethod - def reference(grid, t: np.ndarray, **kwargs) -> dict: + @stencil_tests.static_reference + def reference(grid: base.Grid, t: np.ndarray, **kwargs) -> dict: return dict(pressure=np.full(t.shape, 1216.7746246067475)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 281.787, dims.CellDim, dims.KDim, dtype=wpfloat), - pressure=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(281.787, dims.CellDim, dims.KDim, dtype=wpfloat), + pressure=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_sat_pres_water.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_sat_pres_water.py index 6bde29be18..2347a958e8 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_sat_pres_water.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_sat_pres_water.py @@ -10,22 +10,22 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.thermo import sat_pres_water from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestSatPresWater(StencilTest): +class TestSatPresWater(stencil_tests.StencilTest): PROGRAM = sat_pres_water OUTPUTS = ("pressure",) - @staticmethod - def reference(grid, t: np.ndarray, **kwargs) -> dict: + @stencil_tests.static_reference + def reference(grid: base.Grid, t: np.ndarray, **kwargs) -> dict: return dict(pressure=np.full(t.shape, 1120.1604149806028)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 281.787, dims.CellDim, dims.KDim, dtype=wpfloat), - pressure=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(281.787, dims.CellDim, dims.KDim, dtype=wpfloat), + pressure=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py index 69c06c9e0b..1c9cf4bb54 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py @@ -18,21 +18,20 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests if TYPE_CHECKING: - from icon4py.model.common.grid import base as base_grid + from icon4py.model.common.grid import base, base as base_grid -class TestSaturationAdjustment(StencilTest): +class TestSaturationAdjustment(stencil_tests.StencilTest): PROGRAM = saturation_adjustment OUTPUTS = ("te_out", "qve_out", "qce_out") - @staticmethod + @stencil_tests.static_reference def reference( - grid: base_grid.Grid, + grid: base.Grid, te: np.ndarray, **kwargs, ) -> dict: @@ -42,38 +41,38 @@ def reference( qce_out=np.full(te.shape, 9.5724552280369163e-007), ) - @pytest.fixture - def input_data(self, grid: base_grid.Grid) -> dict: + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid) -> dict: return dict( - te=data_alloc.constant_field( - grid, 273.90911754406039, dims.CellDim, dims.KDim, dtype=wpfloat + te=self.data_alloc.constant_field( + 273.90911754406039, dims.CellDim, dims.KDim, dtype=wpfloat ), q_in=Q( - v=data_alloc.constant_field( - grid, 4.4913424511676030e-003, dims.CellDim, dims.KDim, dtype=wpfloat + v=self.data_alloc.constant_field( + 4.4913424511676030e-003, dims.CellDim, dims.KDim, dtype=wpfloat ), - c=data_alloc.constant_field( - grid, 6.0066941654987605e-013, dims.CellDim, dims.KDim, dtype=wpfloat + c=self.data_alloc.constant_field( + 6.0066941654987605e-013, dims.CellDim, dims.KDim, dtype=wpfloat ), - r=data_alloc.constant_field( - grid, 2.5939378002267028e-004, dims.CellDim, dims.KDim, dtype=wpfloat + r=self.data_alloc.constant_field( + 2.5939378002267028e-004, dims.CellDim, dims.KDim, dtype=wpfloat ), - s=data_alloc.constant_field( - grid, 3.582312533881839e-06, dims.CellDim, dims.KDim, dtype=wpfloat + s=self.data_alloc.constant_field( + 3.582312533881839e-06, dims.CellDim, dims.KDim, dtype=wpfloat ), - i=data_alloc.constant_field( - grid, 3.582312533881839e-06, dims.CellDim, dims.KDim, dtype=wpfloat + i=self.data_alloc.constant_field( + 3.582312533881839e-06, dims.CellDim, dims.KDim, dtype=wpfloat ), - g=data_alloc.constant_field( - grid, 3.582312533881839e-06, dims.CellDim, dims.KDim, dtype=wpfloat + g=self.data_alloc.constant_field( + 3.582312533881839e-06, dims.CellDim, dims.KDim, dtype=wpfloat ), ), - rho=data_alloc.constant_field( - grid, 1.1371657035251757, dims.CellDim, dims.KDim, dtype=wpfloat + rho=self.data_alloc.constant_field( + 1.1371657035251757, dims.CellDim, dims.KDim, dtype=wpfloat ), - te_out=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - qve_out=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - qce_out=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + te_out=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), + qve_out=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), + qce_out=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), horizontal_start=0, horizontal_end=grid.num_cells, vertical_start=0, diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_lambda.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_lambda.py index 3c30334f90..62a15a85f0 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_lambda.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_lambda.py @@ -11,26 +11,28 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.common.constants import GraupelCt from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.properties import snow_lambda from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestSnowLambda(StencilTest): +class TestSnowLambda(stencil_tests.StencilTest): PROGRAM = snow_lambda OUTPUTS = ("riming_snow_rate",) - @staticmethod - def reference(grid, rho: np.ndarray, qs: np.ndarray, ns: np.ndarray, **kwargs) -> dict: + @stencil_tests.static_reference + def reference( + grid: base.Grid, rho: np.ndarray, qs: np.ndarray, ns: np.ndarray, **kwargs + ) -> dict: return dict(riming_snow_rate=np.full(rho.shape, 1.0e10)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - rho=data_alloc.constant_field(grid, 1.12204, dims.CellDim, dims.KDim, dtype=wpfloat), - qs=data_alloc.constant_field( - grid, GraupelCt.qmin, dims.CellDim, dims.KDim, dtype=wpfloat + rho=self.data_alloc.constant_field(1.12204, dims.CellDim, dims.KDim, dtype=wpfloat), + qs=self.data_alloc.constant_field( + GraupelCt.qmin, dims.CellDim, dims.KDim, dtype=wpfloat ), - ns=data_alloc.constant_field(grid, 1.76669e07, dims.CellDim, dims.KDim, dtype=wpfloat), - riming_snow_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + ns=self.data_alloc.constant_field(1.76669e07, dims.CellDim, dims.KDim, dtype=wpfloat), + riming_snow_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_number.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_number.py index 0887fae8c9..a6d6f4b795 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_number.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_number.py @@ -10,24 +10,26 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.properties import snow_number from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestSnowNumber(StencilTest): +class TestSnowNumber(stencil_tests.StencilTest): PROGRAM = snow_number OUTPUTS = ("number",) - @staticmethod - def reference(grid, t: np.ndarray, rho: np.ndarray, qs: np.ndarray, **kwargs) -> dict: + @stencil_tests.static_reference + def reference( + grid: base.Grid, t: np.ndarray, rho: np.ndarray, qs: np.ndarray, **kwargs + ) -> dict: return dict(number=np.full(t.shape, 3813750.0)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 276.302, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 1.17797, dims.CellDim, dims.KDim, dtype=wpfloat), - qs=data_alloc.constant_field(grid, 8.28451e-4, dims.CellDim, dims.KDim, dtype=wpfloat), - number=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(276.302, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.17797, dims.CellDim, dims.KDim, dtype=wpfloat), + qs=self.data_alloc.constant_field(8.28451e-4, dims.CellDim, dims.KDim, dtype=wpfloat), + number=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_number_default.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_number_default.py index 86c5a7aacd..23a20e2657 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_number_default.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_number_default.py @@ -10,24 +10,26 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.properties import snow_number from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestSnowNumberDefault(StencilTest): +class TestSnowNumberDefault(stencil_tests.StencilTest): PROGRAM = snow_number OUTPUTS = ("number",) - @staticmethod - def reference(grid, t: np.ndarray, rho: np.ndarray, qs: np.ndarray, **kwargs) -> dict: + @stencil_tests.static_reference + def reference( + grid: base.Grid, t: np.ndarray, rho: np.ndarray, qs: np.ndarray, **kwargs + ) -> dict: return dict(number=np.full(t.shape, 8.00e5)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 276.302, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 1.17797, dims.CellDim, dims.KDim, dtype=wpfloat), - qs=data_alloc.constant_field(grid, 8.28451e-24, dims.CellDim, dims.KDim, dtype=wpfloat), - number=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(276.302, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.17797, dims.CellDim, dims.KDim, dtype=wpfloat), + qs=self.data_alloc.constant_field(8.28451e-24, dims.CellDim, dims.KDim, dtype=wpfloat), + number=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_to_graupel.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_to_graupel.py index e69fc05ae6..2eb2c0b9d3 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_to_graupel.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_to_graupel.py @@ -10,27 +10,27 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import snow_to_graupel from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestSnowToGraupel(StencilTest): +class TestSnowToGraupel(stencil_tests.StencilTest): PROGRAM = snow_to_graupel OUTPUTS = ("conversion_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, t: np.ndarray, rho: np.ndarray, qc: np.ndarray, qs: np.ndarray, **kwargs + grid: base.Grid, t: np.ndarray, rho: np.ndarray, qc: np.ndarray, qs: np.ndarray, **kwargs ) -> dict: return dict(conversion_rate=np.full(t.shape, 6.2696154545048011e-10)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 265.85, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 1.04848, dims.CellDim, dims.KDim, dtype=wpfloat), - qc=data_alloc.constant_field(grid, 7.02792e-5, dims.CellDim, dims.KDim, dtype=wpfloat), - qs=data_alloc.constant_field(grid, 4.44664e-7, dims.CellDim, dims.KDim, dtype=wpfloat), - conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(265.85, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.04848, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.constant_field(7.02792e-5, dims.CellDim, dims.KDim, dtype=wpfloat), + qs=self.data_alloc.constant_field(4.44664e-7, dims.CellDim, dims.KDim, dtype=wpfloat), + conversion_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_to_graupel_default.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_to_graupel_default.py index 1614d1aa14..ea1beb6e4a 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_to_graupel_default.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_to_graupel_default.py @@ -10,27 +10,27 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import snow_to_graupel from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestSnowToGraupelDefault(StencilTest): +class TestSnowToGraupelDefault(stencil_tests.StencilTest): PROGRAM = snow_to_graupel OUTPUTS = ("conversion_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, t: np.ndarray, rho: np.ndarray, qc: np.ndarray, qs: np.ndarray, **kwargs + grid: base.Grid, t: np.ndarray, rho: np.ndarray, qc: np.ndarray, qs: np.ndarray, **kwargs ) -> dict: return dict(conversion_rate=np.full(t.shape, 0.0)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 281.157, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 0.93171, dims.CellDim, dims.KDim, dtype=wpfloat), - qc=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - qs=data_alloc.constant_field(grid, 4.34854e-5, dims.CellDim, dims.KDim, dtype=wpfloat), - conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(281.157, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(0.93171, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), + qs=self.data_alloc.constant_field(4.34854e-5, dims.CellDim, dims.KDim, dtype=wpfloat), + conversion_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_to_rain.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_to_rain.py index 4b25992492..1ba293b1d7 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_to_rain.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_to_rain.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import snow_to_rain from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestSnowToRainDefault(StencilTest): +class TestSnowToRainDefault(stencil_tests.StencilTest): PROGRAM = snow_to_rain OUTPUTS = ("conversion_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, p: np.ndarray, rho: np.ndarray, @@ -31,15 +31,15 @@ def reference( ) -> dict: return dict(conversion_rate=np.full(t.shape, 3.7268547760462804e-07)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 275.83, dims.CellDim, dims.KDim, dtype=wpfloat), - p=data_alloc.constant_field(grid, 80134.5, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 1.04892, dims.CellDim, dims.KDim, dtype=wpfloat), - dvsw0=data_alloc.constant_field( - grid, 0.00258631, dims.CellDim, dims.KDim, dtype=wpfloat + t=self.data_alloc.constant_field(275.83, dims.CellDim, dims.KDim, dtype=wpfloat), + p=self.data_alloc.constant_field(80134.5, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.04892, dims.CellDim, dims.KDim, dtype=wpfloat), + dvsw0=self.data_alloc.constant_field( + 0.00258631, dims.CellDim, dims.KDim, dtype=wpfloat ), - qs=data_alloc.constant_field(grid, 1.47687e-6, dims.CellDim, dims.KDim, dtype=wpfloat), - conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + qs=self.data_alloc.constant_field(1.47687e-6, dims.CellDim, dims.KDim, dtype=wpfloat), + conversion_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_to_rain_default.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_to_rain_default.py index b3a6b6ff6c..60903ee3fc 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_to_rain_default.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_to_rain_default.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import snow_to_rain from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestSnowToRainDefault(StencilTest): +class TestSnowToRainDefault(stencil_tests.StencilTest): PROGRAM = snow_to_rain OUTPUTS = ("conversion_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, p: np.ndarray, rho: np.ndarray, @@ -31,15 +31,15 @@ def reference( ) -> dict: return dict(conversion_rate=np.full(t.shape, 0.0)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 265.83, dims.CellDim, dims.KDim, dtype=wpfloat), - p=data_alloc.constant_field(grid, 80134.5, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 1.04892, dims.CellDim, dims.KDim, dtype=wpfloat), - dvsw0=data_alloc.constant_field( - grid, -0.00258631, dims.CellDim, dims.KDim, dtype=wpfloat + t=self.data_alloc.constant_field(265.83, dims.CellDim, dims.KDim, dtype=wpfloat), + p=self.data_alloc.constant_field(80134.5, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.04892, dims.CellDim, dims.KDim, dtype=wpfloat), + dvsw0=self.data_alloc.constant_field( + -0.00258631, dims.CellDim, dims.KDim, dtype=wpfloat ), - qs=data_alloc.constant_field(grid, 1.47687e-6, dims.CellDim, dims.KDim, dtype=wpfloat), - conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + qs=self.data_alloc.constant_field(1.47687e-6, dims.CellDim, dims.KDim, dtype=wpfloat), + conversion_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_graupel.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_graupel.py index 318c53c90e..60d27b0da0 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_graupel.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_graupel.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import vapor_x_graupel from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestVaporXGraupel(StencilTest): +class TestVaporXGraupel(stencil_tests.StencilTest): PROGRAM = vapor_x_graupel OUTPUTS = ("exchange_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, p: np.ndarray, rho: np.ndarray, @@ -34,22 +34,22 @@ def reference( ) -> dict: return dict(exchange_rate=np.full(t.shape, -6.8498666666666675e-13)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 278.026, dims.CellDim, dims.KDim, dtype=wpfloat), - p=data_alloc.constant_field(grid, 95987.1, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 1.20041, dims.CellDim, dims.KDim, dtype=wpfloat), - qg=data_alloc.constant_field(grid, 2.05496e-11, dims.CellDim, dims.KDim, dtype=wpfloat), - dvsw=data_alloc.constant_field( - grid, -0.00234674, dims.CellDim, dims.KDim, dtype=wpfloat + t=self.data_alloc.constant_field(278.026, dims.CellDim, dims.KDim, dtype=wpfloat), + p=self.data_alloc.constant_field(95987.1, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.20041, dims.CellDim, dims.KDim, dtype=wpfloat), + qg=self.data_alloc.constant_field(2.05496e-11, dims.CellDim, dims.KDim, dtype=wpfloat), + dvsw=self.data_alloc.constant_field( + -0.00234674, dims.CellDim, dims.KDim, dtype=wpfloat ), - dvsi=data_alloc.constant_field( - grid, -0.00261576, dims.CellDim, dims.KDim, dtype=wpfloat + dvsi=self.data_alloc.constant_field( + -0.00261576, dims.CellDim, dims.KDim, dtype=wpfloat ), - dvsw0=data_alloc.constant_field( - grid, -0.00076851, dims.CellDim, dims.KDim, dtype=wpfloat + dvsw0=self.data_alloc.constant_field( + -0.00076851, dims.CellDim, dims.KDim, dtype=wpfloat ), dt=30.0, - exchange_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + exchange_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_graupel_default.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_graupel_default.py index 8420bd7063..82129c1e46 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_graupel_default.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_graupel_default.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import vapor_x_graupel from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestVaporXGraupelDefault(StencilTest): +class TestVaporXGraupelDefault(stencil_tests.StencilTest): PROGRAM = vapor_x_graupel OUTPUTS = ("exchange_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, p: np.ndarray, rho: np.ndarray, @@ -34,24 +34,22 @@ def reference( ) -> dict: return dict(exchange_rate=np.full(t.shape, 0.0)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 278.026, dims.CellDim, dims.KDim, dtype=wpfloat), - p=data_alloc.constant_field(grid, 95987.1, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 1.20041, dims.CellDim, dims.KDim, dtype=wpfloat), - qg=data_alloc.constant_field( - grid, 2.056496e-16, dims.CellDim, dims.KDim, dtype=wpfloat + t=self.data_alloc.constant_field(278.026, dims.CellDim, dims.KDim, dtype=wpfloat), + p=self.data_alloc.constant_field(95987.1, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.20041, dims.CellDim, dims.KDim, dtype=wpfloat), + qg=self.data_alloc.constant_field(2.056496e-16, dims.CellDim, dims.KDim, dtype=wpfloat), + dvsw=self.data_alloc.constant_field( + -0.00234674, dims.CellDim, dims.KDim, dtype=wpfloat ), - dvsw=data_alloc.constant_field( - grid, -0.00234674, dims.CellDim, dims.KDim, dtype=wpfloat + dvsi=self.data_alloc.constant_field( + -0.00261576, dims.CellDim, dims.KDim, dtype=wpfloat ), - dvsi=data_alloc.constant_field( - grid, -0.00261576, dims.CellDim, dims.KDim, dtype=wpfloat - ), - dvsw0=data_alloc.constant_field( - grid, -0.00076851, dims.CellDim, dims.KDim, dtype=wpfloat + dvsw0=self.data_alloc.constant_field( + -0.00076851, dims.CellDim, dims.KDim, dtype=wpfloat ), dt=30.0, - exchange_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + exchange_rate=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_ice.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_ice.py index e76dc730df..fbee9c910e 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_ice.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_ice.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import vapor_x_ice from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestVaporXIceDefault(StencilTest): +class TestVaporXIceDefault(stencil_tests.StencilTest): PROGRAM = vapor_x_ice OUTPUTS = ("vapor_deposition_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, qi: np.ndarray, mi: np.ndarray, eta: np.ndarray, @@ -32,20 +32,18 @@ def reference( ) -> dict: return dict(vapor_deposition_rate=np.full(qi.shape, 2.2106162342610385e-09)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - qi=data_alloc.constant_field(grid, 9.53048e-07, dims.CellDim, dims.KDim, dtype=wpfloat), - mi=data_alloc.constant_field(grid, 1.0e-9, dims.CellDim, dims.KDim, dtype=wpfloat), - eta=data_alloc.constant_field( - grid, 1.90278e-05, dims.CellDim, dims.KDim, dtype=wpfloat + qi=self.data_alloc.constant_field(9.53048e-07, dims.CellDim, dims.KDim, dtype=wpfloat), + mi=self.data_alloc.constant_field(1.0e-9, dims.CellDim, dims.KDim, dtype=wpfloat), + eta=self.data_alloc.constant_field(1.90278e-05, dims.CellDim, dims.KDim, dtype=wpfloat), + dvsi=self.data_alloc.constant_field( + 0.000120375, dims.CellDim, dims.KDim, dtype=wpfloat ), - dvsi=data_alloc.constant_field( - grid, 0.000120375, dims.CellDim, dims.KDim, dtype=wpfloat - ), - rho=data_alloc.constant_field(grid, 1.19691, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.19691, dims.CellDim, dims.KDim, dtype=wpfloat), dt=30.0, - vapor_deposition_rate=data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, dtype=wpfloat + vapor_deposition_rate=self.data_alloc.zero_field( + dims.CellDim, dims.KDim, dtype=wpfloat ), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_ice_default.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_ice_default.py index e45648728a..6fbed57762 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_ice_default.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_ice_default.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import vapor_x_ice from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestVaporXIceDefault(StencilTest): +class TestVaporXIceDefault(stencil_tests.StencilTest): PROGRAM = vapor_x_ice OUTPUTS = ("vapor_deposition_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, qi: np.ndarray, mi: np.ndarray, eta: np.ndarray, @@ -32,20 +32,18 @@ def reference( ) -> dict: return dict(vapor_deposition_rate=np.full(qi.shape, 0.0)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - qi=data_alloc.constant_field(grid, 2.02422e-23, dims.CellDim, dims.KDim, dtype=wpfloat), - mi=data_alloc.constant_field(grid, 1.0e-12, dims.CellDim, dims.KDim, dtype=wpfloat), - eta=data_alloc.constant_field( - grid, 1.32343e-05, dims.CellDim, dims.KDim, dtype=wpfloat + qi=self.data_alloc.constant_field(2.02422e-23, dims.CellDim, dims.KDim, dtype=wpfloat), + mi=self.data_alloc.constant_field(1.0e-12, dims.CellDim, dims.KDim, dtype=wpfloat), + eta=self.data_alloc.constant_field(1.32343e-05, dims.CellDim, dims.KDim, dtype=wpfloat), + dvsi=self.data_alloc.constant_field( + -0.000618828, dims.CellDim, dims.KDim, dtype=wpfloat ), - dvsi=data_alloc.constant_field( - grid, -0.000618828, dims.CellDim, dims.KDim, dtype=wpfloat - ), - rho=data_alloc.constant_field(grid, 1.19691, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.19691, dims.CellDim, dims.KDim, dtype=wpfloat), dt=30.0, - vapor_deposition_rate=data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, dtype=wpfloat + vapor_deposition_rate=self.data_alloc.zero_field( + dims.CellDim, dims.KDim, dtype=wpfloat ), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_snow.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_snow.py index bb074808aa..69e7c578ee 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_snow.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_snow.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import vapor_x_snow from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestVaporXSnow(StencilTest): +class TestVaporXSnow(stencil_tests.StencilTest): PROGRAM = vapor_x_snow OUTPUTS = ("vapor_deposition_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, p: np.ndarray, rho: np.ndarray, @@ -38,28 +38,28 @@ def reference( ) -> dict: return dict(vapor_deposition_rate=np.full(t.shape, -8.6584296264775935e-13)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 278.748, dims.CellDim, dims.KDim, dtype=wpfloat), - p=data_alloc.constant_field(grid, 95995.5, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 1.19691, dims.CellDim, dims.KDim, dtype=wpfloat), - qs=data_alloc.constant_field(grid, 1.25653e-10, dims.CellDim, dims.KDim, dtype=wpfloat), - ns=data_alloc.constant_field(grid, 800000.0, dims.CellDim, dims.KDim, dtype=wpfloat), - lam=data_alloc.constant_field(grid, 1.0e10, dims.CellDim, dims.KDim, dtype=wpfloat), - eta=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - ice_dep=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - dvsw=data_alloc.constant_field( - grid, -0.00196781, dims.CellDim, dims.KDim, dtype=wpfloat + t=self.data_alloc.constant_field(278.748, dims.CellDim, dims.KDim, dtype=wpfloat), + p=self.data_alloc.constant_field(95995.5, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.19691, dims.CellDim, dims.KDim, dtype=wpfloat), + qs=self.data_alloc.constant_field(1.25653e-10, dims.CellDim, dims.KDim, dtype=wpfloat), + ns=self.data_alloc.constant_field(800000.0, dims.CellDim, dims.KDim, dtype=wpfloat), + lam=self.data_alloc.constant_field(1.0e10, dims.CellDim, dims.KDim, dtype=wpfloat), + eta=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), + ice_dep=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), + dvsw=self.data_alloc.constant_field( + -0.00196781, dims.CellDim, dims.KDim, dtype=wpfloat ), - dvsi=data_alloc.constant_field( - grid, -0.00229367, dims.CellDim, dims.KDim, dtype=wpfloat + dvsi=self.data_alloc.constant_field( + -0.00229367, dims.CellDim, dims.KDim, dtype=wpfloat ), - dvsw0=data_alloc.constant_field( - grid, -0.000110022, dims.CellDim, dims.KDim, dtype=wpfloat + dvsw0=self.data_alloc.constant_field( + -0.000110022, dims.CellDim, dims.KDim, dtype=wpfloat ), dt=30.0, - vapor_deposition_rate=data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, dtype=wpfloat + vapor_deposition_rate=self.data_alloc.zero_field( + dims.CellDim, dims.KDim, dtype=wpfloat ), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_snow_default.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_snow_default.py index 26b02e3dee..14589e0582 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_snow_default.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vapor_x_snow_default.py @@ -10,18 +10,18 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import vapor_x_snow from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestVaporXSnowDefault(StencilTest): +class TestVaporXSnowDefault(stencil_tests.StencilTest): PROGRAM = vapor_x_snow OUTPUTS = ("vapor_deposition_rate",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, + grid: base.Grid, t: np.ndarray, p: np.ndarray, rho: np.ndarray, @@ -38,28 +38,28 @@ def reference( ) -> dict: return dict(vapor_deposition_rate=np.full(t.shape, 0.0)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - t=data_alloc.constant_field(grid, 278.748, dims.CellDim, dims.KDim, dtype=wpfloat), - p=data_alloc.constant_field(grid, 95995.5, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 1.19691, dims.CellDim, dims.KDim, dtype=wpfloat), - qs=data_alloc.constant_field(grid, 1.25653e-20, dims.CellDim, dims.KDim, dtype=wpfloat), - ns=data_alloc.constant_field(grid, 800000.0, dims.CellDim, dims.KDim, dtype=wpfloat), - lam=data_alloc.constant_field(grid, 1.0e10, dims.CellDim, dims.KDim, dtype=wpfloat), - eta=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - ice_dep=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - dvsw=data_alloc.constant_field( - grid, -0.00196781, dims.CellDim, dims.KDim, dtype=wpfloat + t=self.data_alloc.constant_field(278.748, dims.CellDim, dims.KDim, dtype=wpfloat), + p=self.data_alloc.constant_field(95995.5, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(1.19691, dims.CellDim, dims.KDim, dtype=wpfloat), + qs=self.data_alloc.constant_field(1.25653e-20, dims.CellDim, dims.KDim, dtype=wpfloat), + ns=self.data_alloc.constant_field(800000.0, dims.CellDim, dims.KDim, dtype=wpfloat), + lam=self.data_alloc.constant_field(1.0e10, dims.CellDim, dims.KDim, dtype=wpfloat), + eta=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), + ice_dep=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), + dvsw=self.data_alloc.constant_field( + -0.00196781, dims.CellDim, dims.KDim, dtype=wpfloat ), - dvsi=data_alloc.constant_field( - grid, -0.00229367, dims.CellDim, dims.KDim, dtype=wpfloat + dvsi=self.data_alloc.constant_field( + -0.00229367, dims.CellDim, dims.KDim, dtype=wpfloat ), - dvsw0=data_alloc.constant_field( - grid, -0.000110022, dims.CellDim, dims.KDim, dtype=wpfloat + dvsw0=self.data_alloc.constant_field( + -0.000110022, dims.CellDim, dims.KDim, dtype=wpfloat ), dt=30.0, - vapor_deposition_rate=data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, dtype=wpfloat + vapor_deposition_rate=self.data_alloc.zero_field( + dims.CellDim, dims.KDim, dtype=wpfloat ), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vel_scale_factor_ice.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vel_scale_factor_ice.py index 8b9f9795e8..6aed43f79d 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vel_scale_factor_ice.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vel_scale_factor_ice.py @@ -12,22 +12,22 @@ vel_scale_factor_ice, ) from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestVelScaleFactorIce(StencilTest): +class TestVelScaleFactorIce(stencil_tests.StencilTest): PROGRAM = vel_scale_factor_ice OUTPUTS = ("scale_factor",) - @staticmethod - def reference(grid, xrho: np.ndarray, **kwargs) -> dict: + @stencil_tests.static_reference + def reference(grid: base.Grid, xrho: np.ndarray, **kwargs) -> dict: return dict(scale_factor=np.full(xrho.shape, 1.1158596098981044)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - xrho=data_alloc.constant_field(grid, 1.17873, dims.CellDim, dims.KDim, dtype=wpfloat), - scale_factor=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + xrho=self.data_alloc.constant_field(1.17873, dims.CellDim, dims.KDim, dtype=wpfloat), + scale_factor=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vel_scale_factor_snow.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vel_scale_factor_snow.py index 18e7048ab7..46e88e6727 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vel_scale_factor_snow.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_vel_scale_factor_snow.py @@ -12,27 +12,27 @@ vel_scale_factor_snow, ) from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests -class TestVelScaleFactorSnow(StencilTest): +class TestVelScaleFactorSnow(stencil_tests.StencilTest): PROGRAM = vel_scale_factor_snow OUTPUTS = ("scale_factor",) - @staticmethod + @stencil_tests.static_reference def reference( - grid, xrho: np.ndarray, rho: np.ndarray, t: np.ndarray, qs: np.ndarray, **kwargs + grid: base.Grid, xrho: np.ndarray, rho: np.ndarray, t: np.ndarray, qs: np.ndarray, **kwargs ) -> dict: return dict(scale_factor=np.full(xrho.shape, 0.06633230453931642)) - @pytest.fixture - def input_data(self, grid): + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid): return dict( - xrho=data_alloc.constant_field(grid, 1.17787, dims.CellDim, dims.KDim, dtype=wpfloat), - rho=data_alloc.constant_field(grid, 0.882961, dims.CellDim, dims.KDim, dtype=wpfloat), - t=data_alloc.constant_field(grid, 257.101, dims.CellDim, dims.KDim, dtype=wpfloat), - qs=data_alloc.constant_field(grid, 5.78761e-06, dims.CellDim, dims.KDim, dtype=wpfloat), - scale_factor=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + xrho=self.data_alloc.constant_field(1.17787, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=self.data_alloc.constant_field(0.882961, dims.CellDim, dims.KDim, dtype=wpfloat), + t=self.data_alloc.constant_field(257.101, dims.CellDim, dims.KDim, dtype=wpfloat), + qs=self.data_alloc.constant_field(5.78761e-06, dims.CellDim, dims.KDim, dtype=wpfloat), + scale_factor=self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=wpfloat), ) diff --git a/model/common/src/icon4py/model/common/grid/base.py b/model/common/src/icon4py/model/common/grid/base.py index c981b7a8dd..b9fe62b522 100644 --- a/model/common/src/icon4py/model/common/grid/base.py +++ b/model/common/src/icon4py/model/common/grid/base.py @@ -5,6 +5,10 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import collections.abc import dataclasses import enum import functools @@ -102,6 +106,8 @@ class Grid: def __post_init__(self): # TODO(havogt): replace `Koff[k]` by `KDim + k` syntax and remove the following line. + assert isinstance(self.connectivities, collections.abc.MutableMapping) + assert isinstance(dims.Koff.value, str) self.connectivities[dims.Koff.value] = dims.KDim @functools.cached_property @@ -155,7 +161,9 @@ def limited_area(self) -> bool: def get_connectivity(self, offset: str | gtx.FieldOffset) -> gtx_common.NeighborTable: """Get the connectivity by its name.""" if isinstance(offset, gtx.FieldOffset): - offset = offset.value + offset_value = offset.value + assert isinstance(offset_value, str) + offset = offset_value if offset not in self.connectivities: raise exceptions.MissingConnectivityError( f"Missing connectivity for offset {offset} in grid {self.id}." diff --git a/model/common/tests/common/interpolation/stencil_tests/test_cell_2_edge_interpolation.py b/model/common/tests/common/interpolation/stencil_tests/test_cell_2_edge_interpolation.py index eb8ad994c6..c21dbc1a89 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_cell_2_edge_interpolation.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_cell_2_edge_interpolation.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -16,7 +17,6 @@ from icon4py.model.common.interpolation.stencils.cell_2_edge_interpolation import ( cell_2_edge_interpolation, ) -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -25,14 +25,15 @@ class TestCell2EdgeInterpolation(stencil_tests.StencilTest): PROGRAM = cell_2_edge_interpolation OUTPUTS = ("out_field",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, in_field: np.ndarray, coeff: np.ndarray, **kwargs: Any, ) -> dict: - e2c = connectivities[dims.E2CDim] + connectivities = stencil_tests.connectivities_asnumpy(grid) + e2c = connectivities[dims.E2C] coeff_ = np.expand_dims(coeff, axis=-1) out_field = np.sum(in_field[e2c] * coeff_, axis=1) @@ -40,11 +41,11 @@ def reference( out_field=out_field, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - in_field = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - coeff = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat) - out_field = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + in_field = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + coeff = self.data_alloc.random_field(dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat) + out_field = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) return dict( in_field=in_field, diff --git a/model/common/tests/common/interpolation/stencil_tests/test_compute_cell_2_vertex_interpolation.py b/model/common/tests/common/interpolation/stencil_tests/test_compute_cell_2_vertex_interpolation.py index b68124bbda..c67f6419f2 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_compute_cell_2_vertex_interpolation.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_compute_cell_2_vertex_interpolation.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -17,7 +18,6 @@ from icon4py.model.common.interpolation.stencils.compute_cell_2_vertex_interpolation import ( compute_cell_2_vertex_interpolation, ) -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -26,14 +26,15 @@ class TestComputeCells2VertsInterpolation(stencil_tests.StencilTest): PROGRAM = compute_cell_2_vertex_interpolation OUTPUTS = ("vert_out",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, cell_in: np.ndarray, c_int: np.ndarray, **kwargs: Any, ) -> dict: - v2c = connectivities[dims.V2CDim] + connectivities = stencil_tests.connectivities_asnumpy(grid) + v2c = connectivities[dims.V2C] c_int = np.expand_dims(c_int, axis=-1) out_field = np.sum(cell_in[v2c] * c_int, axis=1) @@ -41,11 +42,11 @@ def reference( vert_out=out_field, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - cell_in = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=types.wpfloat) - c_int = data_alloc.random_field(grid, dims.VertexDim, dims.V2CDim, dtype=types.wpfloat) - vert_out = data_alloc.zero_field(grid, dims.VertexDim, dims.KDim, dtype=types.wpfloat) + cell_in = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=types.wpfloat) + c_int = self.data_alloc.random_field(dims.VertexDim, dims.V2CDim, dtype=types.wpfloat) + vert_out = self.data_alloc.zero_field(dims.VertexDim, dims.KDim, dtype=types.wpfloat) return dict( cell_in=cell_in, diff --git a/model/common/tests/common/interpolation/stencil_tests/test_diagnose_pressure.py b/model/common/tests/common/interpolation/stencil_tests/test_diagnose_pressure.py index 2f0471e4fe..09c30da153 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_diagnose_pressure.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_diagnose_pressure.py @@ -16,7 +16,6 @@ diagnose_pressure, ) from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -24,9 +23,9 @@ class TestDiagnosePressure(stencil_tests.StencilTest): PROGRAM = diagnose_pressure OUTPUTS = ("pressure", "pressure_ifc") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, surface_pressure: np.ndarray, virtual_temperature: np.ndarray, ddqz_z_full: np.ndarray, @@ -52,17 +51,17 @@ def reference( pressure_ifc=pressure_ifc, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - ddqz_z_full = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, low=1.0, dtype=ta.wpfloat + ddqz_z_full = self.data_alloc.random_field( + dims.CellDim, dims.KDim, low=1.0, dtype=ta.wpfloat ) - virtual_temperature = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, low=1.0e-2, dtype=ta.wpfloat + virtual_temperature = self.data_alloc.random_field( + dims.CellDim, dims.KDim, low=1.0e-2, dtype=ta.wpfloat ) - surface_pressure = data_alloc.random_field(grid, dims.CellDim, low=1.0, dtype=ta.wpfloat) - pressure = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - pressure_ifc = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) + surface_pressure = self.data_alloc.random_field(dims.CellDim, low=1.0, dtype=ta.wpfloat) + pressure = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + pressure_ifc = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) return dict( ddqz_z_full=ddqz_z_full, diff --git a/model/common/tests/common/interpolation/stencil_tests/test_diagnose_surface_pressure.py b/model/common/tests/common/interpolation/stencil_tests/test_diagnose_surface_pressure.py index 55e0967f6a..b814e87bbc 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_diagnose_surface_pressure.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_diagnose_surface_pressure.py @@ -16,7 +16,6 @@ diagnose_surface_pressure, ) from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -24,9 +23,9 @@ class TestDiagnoseSurfacePressure(stencil_tests.StencilTest): PROGRAM = diagnose_surface_pressure OUTPUTS = ("surface_pressure",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, exner: np.ndarray, virtual_temperature: np.ndarray, ddqz_z_full: np.ndarray, @@ -47,18 +46,18 @@ def reference( surface_pressure=surface_pressure, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: low = 1.0e-2 - exner = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-6, dtype=ta.wpfloat) - virtual_temperature = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, low=low, dtype=ta.wpfloat + exner = self.data_alloc.random_field(dims.CellDim, dims.KDim, low=1.0e-6, dtype=ta.wpfloat) + virtual_temperature = self.data_alloc.random_field( + dims.CellDim, dims.KDim, low=low, dtype=ta.wpfloat ) - ddqz_z_full = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, low=low, dtype=ta.wpfloat + ddqz_z_full = self.data_alloc.random_field( + dims.CellDim, dims.KDim, low=low, dtype=ta.wpfloat ) - surface_pressure = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, extend={dims.KDim: 1} + surface_pressure = self.data_alloc.zero_field( + dims.CellDim, dims.KDim, dtype=ta.wpfloat, extend={dims.KDim: 1} ) return dict( diff --git a/model/common/tests/common/interpolation/stencil_tests/test_diagnose_temperature.py b/model/common/tests/common/interpolation/stencil_tests/test_diagnose_temperature.py index 1e03d6b648..311c25b8ab 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_diagnose_temperature.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_diagnose_temperature.py @@ -16,7 +16,6 @@ diagnose_virtual_temperature_and_temperature, ) from icon4py.model.common.grid import base -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -24,9 +23,9 @@ class TestDiagnoseTemperature(stencil_tests.StencilTest): PROGRAM = diagnose_virtual_temperature_and_temperature OUTPUTS = ("virtual_temperature", "temperature") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, qv: np.ndarray, qc: np.ndarray, qi: np.ndarray, @@ -45,34 +44,34 @@ def reference( temperature=temperature, ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - theta_v = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, low=1.0e-4, high=1.0, dtype=ta.wpfloat + theta_v = self.data_alloc.random_field( + dims.CellDim, dims.KDim, low=1.0e-4, high=1.0, dtype=ta.wpfloat ) - exner = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, low=1.0e-4, high=1.0, dtype=ta.wpfloat + exner = self.data_alloc.random_field( + dims.CellDim, dims.KDim, low=1.0e-4, high=1.0, dtype=ta.wpfloat ) - qv = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, low=0.0, high=1.0, dtype=ta.wpfloat + qv = self.data_alloc.random_field( + dims.CellDim, dims.KDim, low=0.0, high=1.0, dtype=ta.wpfloat ) - qc = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, low=0.0, high=1.0, dtype=ta.wpfloat + qc = self.data_alloc.random_field( + dims.CellDim, dims.KDim, low=0.0, high=1.0, dtype=ta.wpfloat ) - qi = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, low=0.0, high=1.0, dtype=ta.wpfloat + qi = self.data_alloc.random_field( + dims.CellDim, dims.KDim, low=0.0, high=1.0, dtype=ta.wpfloat ) - qr = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, low=0.0, high=1.0, dtype=ta.wpfloat + qr = self.data_alloc.random_field( + dims.CellDim, dims.KDim, low=0.0, high=1.0, dtype=ta.wpfloat ) - qs = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, low=0.0, high=1.0, dtype=ta.wpfloat + qs = self.data_alloc.random_field( + dims.CellDim, dims.KDim, low=0.0, high=1.0, dtype=ta.wpfloat ) - qg = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, low=0.0, high=1.0, dtype=ta.wpfloat + qg = self.data_alloc.random_field( + dims.CellDim, dims.KDim, low=0.0, high=1.0, dtype=ta.wpfloat ) - virtual_temperature = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - temperature = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) + virtual_temperature = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + temperature = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) return dict( qv=qv, diff --git a/model/common/tests/common/interpolation/stencil_tests/test_edge_2_cell_vector_rbf_interpolation.py b/model/common/tests/common/interpolation/stencil_tests/test_edge_2_cell_vector_rbf_interpolation.py index a4dea61823..3f53373348 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_edge_2_cell_vector_rbf_interpolation.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_edge_2_cell_vector_rbf_interpolation.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -16,7 +17,6 @@ from icon4py.model.common.interpolation.stencils.edge_2_cell_vector_rbf_interpolation import ( edge_2_cell_vector_rbf_interpolation, ) -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -25,15 +25,16 @@ class TestEdge2CellVectorRBFInterpolation(stencil_tests.StencilTest): PROGRAM = edge_2_cell_vector_rbf_interpolation OUTPUTS = ("p_u_out", "p_v_out") - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_e_in: np.ndarray, ptr_coeff_1: np.ndarray, ptr_coeff_2: np.ndarray, **kwargs: Any, ) -> dict: - c2e2c2e = connectivities[dims.C2E2C2EDim] + connectivities = stencil_tests.connectivities_asnumpy(grid) + c2e2c2e = connectivities[dims.C2E2C2E] ptr_coeff_1 = np.expand_dims(ptr_coeff_1, axis=-1) ptr_coeff_2 = np.expand_dims(ptr_coeff_2, axis=-1) p_u_out = np.sum(p_e_in[c2e2c2e] * ptr_coeff_1, axis=1) @@ -41,13 +42,13 @@ def reference( return dict(p_v_out=p_v_out, p_u_out=p_u_out) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - p_e_in = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) - ptr_coeff_1 = data_alloc.random_field(grid, dims.CellDim, dims.C2E2C2EDim, dtype=ta.wpfloat) - ptr_coeff_2 = data_alloc.random_field(grid, dims.CellDim, dims.C2E2C2EDim, dtype=ta.wpfloat) - p_v_out = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) - p_u_out = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) + p_e_in = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) + ptr_coeff_1 = self.data_alloc.random_field(dims.CellDim, dims.C2E2C2EDim, dtype=ta.wpfloat) + ptr_coeff_2 = self.data_alloc.random_field(dims.CellDim, dims.C2E2C2EDim, dtype=ta.wpfloat) + p_v_out = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) + p_u_out = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=ta.wpfloat) return dict( p_e_in=p_e_in, diff --git a/model/common/tests/common/interpolation/stencil_tests/test_interpolate_edge_field_to_half_levels_vp.py b/model/common/tests/common/interpolation/stencil_tests/test_interpolate_edge_field_to_half_levels_vp.py index 027b0ab492..b6588b695c 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_interpolate_edge_field_to_half_levels_vp.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_interpolate_edge_field_to_half_levels_vp.py @@ -11,14 +11,13 @@ import numpy as np import pytest -import icon4py.model.testing.stencil_tests as test_helpers from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import base as base_grid +from icon4py.model.common.grid import base, base as base_grid from icon4py.model.common.interpolation.stencils.interpolate_edge_field_to_half_levels_vp import ( interpolate_edge_field_to_half_levels_vp, ) from icon4py.model.common.type_alias import vpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field +from icon4py.model.testing import stencil_tests def interpolate_edge_field_to_half_levels_vp_numpy( @@ -33,24 +32,26 @@ def interpolate_edge_field_to_half_levels_vp_numpy( return interpolation_to_half_levels_vp -class TestInterpolateToHalfLevelsVp(test_helpers.StencilTest): +class TestInterpolateToHalfLevelsVp(stencil_tests.StencilTest): PROGRAM = interpolate_edge_field_to_half_levels_vp OUTPUTS = ("interpolation_to_half_levels_vp",) - @staticmethod + @stencil_tests.static_reference def reference( - grid: base_grid.Grid, wgtfac_e: np.ndarray, interpolant: np.ndarray, **kwargs: Any + grid: base.Grid, wgtfac_e: np.ndarray, interpolant: np.ndarray, **kwargs: Any ) -> dict: interpolation_to_half_levels_vp = interpolate_edge_field_to_half_levels_vp_numpy( wgtfac_e=wgtfac_e, interpolant=interpolant ) return dict(interpolation_to_half_levels_vp=interpolation_to_half_levels_vp) - @pytest.fixture - def input_data(self, grid: base_grid.Grid) -> dict: - interpolant = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - wgtfac_e = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - interpolation_to_half_levels_vp = zero_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid) -> dict: + interpolant = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + wgtfac_e = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) + interpolation_to_half_levels_vp = self.data_alloc.zero_field( + dims.EdgeDim, dims.KDim, dtype=vpfloat + ) return dict( wgtfac_e=wgtfac_e, diff --git a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index 07104cbd89..c525d725d5 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -17,31 +18,30 @@ mo_intp_rbf_rbf_vec_interpol_vertex, ) from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest +from icon4py.model.testing import stencil_tests @pytest.mark.continuous_benchmarking -class TestMoIntpRbfRbfVecInterpolVertex(StencilTest): +class TestMoIntpRbfRbfVecInterpolVertex(stencil_tests.StencilTest): PROGRAM = mo_intp_rbf_rbf_vec_interpol_vertex OUTPUTS = ("p_u_out", "p_v_out") STATIC_PARAMS = { - StandardStaticVariants.NONE: (), - StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", "vertical_start", "vertical_end", ), - StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( "vertical_start", "vertical_end", ), } - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, p_e_in: np.ndarray, ptr_coeff_1: np.ndarray, ptr_coeff_2: np.ndarray, @@ -49,7 +49,8 @@ def reference( horizontal_end: int, **kwargs: Any, ) -> dict[str, np.ndarray]: - v2e = connectivities[dims.V2EDim] + connectivities = stencil_tests.connectivities_asnumpy(grid) + v2e = connectivities[dims.V2E] ptr_coeff_1 = np.expand_dims(ptr_coeff_1, axis=-1) p_u_out = np.sum( np.where(np.expand_dims(v2e, axis=-1) >= 0, p_e_in[v2e] * ptr_coeff_1, 0.0), axis=1 @@ -70,13 +71,13 @@ def reference( return dict(p_v_out=p_v_final_out, p_u_out=p_u_final_out) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - p_e_in = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) - ptr_coeff_1 = data_alloc.random_field(grid, dims.VertexDim, dims.V2EDim, dtype=wpfloat) - ptr_coeff_2 = data_alloc.random_field(grid, dims.VertexDim, dims.V2EDim, dtype=wpfloat) - p_v_out = data_alloc.zero_field(grid, dims.VertexDim, dims.KDim, dtype=wpfloat) - p_u_out = data_alloc.zero_field(grid, dims.VertexDim, dims.KDim, dtype=wpfloat) + p_e_in = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + ptr_coeff_1 = self.data_alloc.random_field(dims.VertexDim, dims.V2EDim, dtype=wpfloat) + ptr_coeff_2 = self.data_alloc.random_field(dims.VertexDim, dims.V2EDim, dtype=wpfloat) + p_v_out = self.data_alloc.zero_field(dims.VertexDim, dims.KDim, dtype=wpfloat) + p_u_out = self.data_alloc.zero_field(dims.VertexDim, dims.KDim, dtype=wpfloat) vertex_domain = h_grid.domain(dims.VertexDim) horizontal_start = grid.start_index(vertex_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)) diff --git a/model/common/tests/common/math/stencil_tests/test_cell_horizontal_gradients_by_green_gauss_method.py b/model/common/tests/common/math/stencil_tests/test_cell_horizontal_gradients_by_green_gauss_method.py index 6c4ae32a04..d63a4c0d9b 100644 --- a/model/common/tests/common/math/stencil_tests/test_cell_horizontal_gradients_by_green_gauss_method.py +++ b/model/common/tests/common/math/stencil_tests/test_cell_horizontal_gradients_by_green_gauss_method.py @@ -5,7 +5,11 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -18,17 +22,16 @@ ) from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def cell_horizontal_gradients_by_green_gauss_method_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], + connectivities: Mapping[gtx.FieldOffset, np.ndarray], scalar_field: np.ndarray, geofac_grg_x: np.ndarray, geofac_grg_y: np.ndarray, ) -> tuple[np.ndarray, ...]: - c2e2cO = connectivities[dims.C2E2CODim] + c2e2cO = connectivities[dims.C2E2CO] geofac_grg_x = np.expand_dims(geofac_grg_x, axis=-1) p_grad_1_u = np.sum( np.where((c2e2cO != -1)[:, :, np.newaxis], geofac_grg_x * scalar_field[c2e2cO], 0), axis=1 @@ -44,18 +47,19 @@ def cell_horizontal_gradients_by_green_gauss_method_numpy( @pytest.mark.embedded_remap_error -class TestMoMathGradientsGradGreenGaussCellDsl(StencilTest): +class TestMoMathGradientsGradGreenGaussCellDsl(stencil_tests.StencilTest): PROGRAM = cell_horizontal_gradients_by_green_gauss_method OUTPUTS = ("out",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, scalar_field: np.ndarray, geofac_grg_x: np.ndarray, geofac_grg_y: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) ( p_grad_1_u, p_grad_1_v, @@ -66,15 +70,15 @@ def reference( out=(p_grad_1_u, p_grad_1_v), ) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data( self, grid: base.Grid ) -> dict[str, gtx.Field | state_utils.ScalarType | gtx.Domain | tuple[gtx.Field, ...]]: - scalar_field = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - geofac_grg_x = random_field(grid, dims.CellDim, dims.C2E2CODim, dtype=wpfloat) - geofac_grg_y = random_field(grid, dims.CellDim, dims.C2E2CODim, dtype=wpfloat) - p_grad_1_u = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - p_grad_1_v = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + scalar_field = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + geofac_grg_x = self.data_alloc.random_field(dims.CellDim, dims.C2E2CODim, dtype=wpfloat) + geofac_grg_y = self.data_alloc.random_field(dims.CellDim, dims.C2E2CODim, dtype=wpfloat) + p_grad_1_u = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) + p_grad_1_v = self.data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) return dict( scalar_field=scalar_field, diff --git a/model/common/tests/common/math/stencil_tests/test_compute_first_vertical_derivative.py b/model/common/tests/common/math/stencil_tests/test_compute_first_vertical_derivative.py index 5cd46b15ff..061a7f3905 100644 --- a/model/common/tests/common/math/stencil_tests/test_compute_first_vertical_derivative.py +++ b/model/common/tests/common/math/stencil_tests/test_compute_first_vertical_derivative.py @@ -16,8 +16,7 @@ from icon4py.model.common.math.derivative import compute_first_vertical_derivative_at_cells from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import stencil_tests def compute_first_vertical_derivative_numpy( @@ -27,13 +26,13 @@ def compute_first_vertical_derivative_numpy( return first_vertical_derivative -class TestComputeFirstVerticalDerivative(StencilTest): +class TestComputeFirstVerticalDerivative(stencil_tests.StencilTest): PROGRAM = compute_first_vertical_derivative_at_cells OUTPUTS = ("first_vertical_derivative",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, cell_kdim_field: np.ndarray, inv_ddqz_z_full: np.ndarray, **kwargs: Any, @@ -43,13 +42,15 @@ def reference( ) return dict(first_vertical_derivative=first_vertical_derivative) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - cell_kdim_field = random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=vpfloat + cell_kdim_field = self.data_alloc.random_field( + dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=vpfloat + ) + inv_ddqz_z_full = self.data_alloc.random_field(dims.CellDim, dims.KDim, dtype=vpfloat) + first_vertical_derivative = self.data_alloc.zero_field( + dims.CellDim, dims.KDim, dtype=vpfloat ) - inv_ddqz_z_full = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - first_vertical_derivative = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) return dict( cell_kdim_field=cell_kdim_field, diff --git a/model/common/tests/common/math/unit_tests/test_helpers.py b/model/common/tests/common/math/unit_tests/test_helpers.py index e1de50194d..79130fe92e 100644 --- a/model/common/tests/common/math/unit_tests/test_helpers.py +++ b/model/common/tests/common/math/unit_tests/test_helpers.py @@ -16,13 +16,14 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, simple from icon4py.model.common.math import helpers -from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests from icon4py.model.testing.fixtures.datatest import backend, backend_like from icon4py.model.testing.fixtures.stencil_tests import grid, grid_manager def test_cross_product(backend: gtx_typing.Backend) -> None: + from icon4py.model.common.utils import data_allocation as data_alloc + mesh = simple.simple_grid(allocator=backend) x1 = data_alloc.random_field(mesh, dims.EdgeDim, allocator=backend) y1 = data_alloc.random_field(mesh, dims.EdgeDim, allocator=backend) @@ -56,9 +57,9 @@ class TestAverageTwoVerticalLevelsDownwardsOnEdges(stencil_tests.StencilTest): ), ) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, input_field: np.ndarray, **kwargs: Any, ) -> dict: @@ -66,10 +67,10 @@ def reference( average = 0.5 * (input_field + offset) return dict(average=average) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - input_field = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, extend={dims.KDim: 1}) - result = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, extend={dims.KDim: 1}) + input_field = self.data_alloc.zero_field(dims.EdgeDim, dims.KDim, extend={dims.KDim: 1}) + result = self.data_alloc.random_field(dims.EdgeDim, dims.KDim, extend={dims.KDim: 1}) return dict( input_field=input_field, average=result, @@ -90,9 +91,9 @@ class TestAverageTwoVerticalLevelsDownwardsOnCells(stencil_tests.StencilTest): ), ) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, input_field: np.ndarray, **kwargs: Any, ) -> dict: @@ -100,10 +101,10 @@ def reference( res = 0.5 * (input_field + np.roll(input_field, shift=-1, axis=1))[:, : shp[1] - 1] return dict(average=res) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - input_field = data_alloc.random_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) - result = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + input_field = self.data_alloc.random_field(dims.CellDim, dims.KDim, extend={dims.KDim: 1}) + result = self.data_alloc.zero_field(dims.CellDim, dims.KDim) return dict( input_field=input_field, average=result, diff --git a/model/common/tests/common/math/unit_tests/test_operators.py b/model/common/tests/common/math/unit_tests/test_operators.py index b487d1a50b..f16fb5eeac 100644 --- a/model/common/tests/common/math/unit_tests/test_operators.py +++ b/model/common/tests/common/math/unit_tests/test_operators.py @@ -5,7 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any +from collections.abc import Mapping +from typing import Any, cast import gt4py.next as gtx import numpy as np @@ -15,33 +16,32 @@ from icon4py.model.common.grid import base, base as base_grid from icon4py.model.common.math.stencils.compute_nabla2_on_cell import compute_nabla2_on_cell from icon4py.model.common.math.stencils.compute_nabla2_on_cell_k import compute_nabla2_on_cell_k -from icon4py.model.common.utils.data_allocation import constant_field, zero_field -from icon4py.model.testing import reference_funcs +from icon4py.model.testing import reference_funcs, stencil_tests from icon4py.model.testing.fixtures.datatest import backend_like from icon4py.model.testing.fixtures.stencil_tests import grid, grid_manager -from icon4py.model.testing.stencil_tests import StencilTest @pytest.mark.embedded_remap_error -class TestNabla2OnCell(StencilTest): +class TestNabla2OnCell(stencil_tests.StencilTest): PROGRAM = compute_nabla2_on_cell OUTPUTS = ("nabla2_psi_c",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, psi_c: np.ndarray, geofac_n2s: np.ndarray, **kwargs: Any, ) -> dict[str, np.ndarray]: + connectivities = stencil_tests.connectivities_asnumpy(grid) nabla2_psi_c_np = reference_funcs.nabla2_on_cell_numpy(connectivities, psi_c, geofac_n2s) return dict(nabla2_psi_c=nabla2_psi_c_np) - @pytest.fixture - def input_data(self, grid: base_grid.Grid) -> dict: - psi_c = constant_field(grid, 1.0, dims.CellDim) - geofac_n2s = constant_field(grid, 2.0, dims.CellDim, dims.C2E2CODim) - nabla2_psi_c = zero_field(grid, dims.CellDim) + @stencil_tests.input_data_fixture + def input_data(self, grid: base.Grid) -> dict: + psi_c = self.data_alloc.constant_field(1.0, dims.CellDim) + geofac_n2s = self.data_alloc.constant_field(2.0, dims.CellDim, dims.C2E2CODim) + nabla2_psi_c = self.data_alloc.zero_field(dims.CellDim) return dict( psi_c=psi_c, geofac_n2s=geofac_n2s, @@ -52,25 +52,26 @@ def input_data(self, grid: base_grid.Grid) -> dict: @pytest.mark.embedded_remap_error -class TestNabla2OnCellK(StencilTest): +class TestNabla2OnCellK(stencil_tests.StencilTest): PROGRAM = compute_nabla2_on_cell_k OUTPUTS = ("nabla2_psi_c",) - @staticmethod + @stencil_tests.static_reference def reference( - connectivities: dict[gtx.Dimension, np.ndarray], + grid: base.Grid, psi_c: np.ndarray, geofac_n2s: np.ndarray, **kwargs: Any, ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) nabla2_psi_c_np = reference_funcs.nabla2_on_cell_k_numpy(connectivities, psi_c, geofac_n2s) return dict(nabla2_psi_c=nabla2_psi_c_np) - @pytest.fixture + @stencil_tests.input_data_fixture def input_data(self, grid: base.Grid) -> dict: - psi_c = constant_field(grid, 1.0, dims.CellDim, dims.KDim) - geofac_n2s = constant_field(grid, 2.0, dims.CellDim, dims.C2E2CODim) - nabla2_psi_c = zero_field(grid, dims.CellDim, dims.KDim) + psi_c = self.data_alloc.constant_field(1.0, dims.CellDim, dims.KDim) + geofac_n2s = self.data_alloc.constant_field(2.0, dims.CellDim, dims.C2E2CODim) + nabla2_psi_c = self.data_alloc.zero_field(dims.CellDim, dims.KDim) return dict( psi_c=psi_c, geofac_n2s=geofac_n2s, diff --git a/model/testing/src/icon4py/model/testing/reference_funcs.py b/model/testing/src/icon4py/model/testing/reference_funcs.py index af278c05bd..17ac2cf4f6 100644 --- a/model/testing/src/icon4py/model/testing/reference_funcs.py +++ b/model/testing/src/icon4py/model/testing/reference_funcs.py @@ -6,6 +6,8 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from collections.abc import Mapping + import gt4py.next as gtx import numpy as np @@ -31,17 +33,17 @@ def enhanced_smagorinski_factor_numpy( def nabla2_on_cell_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], psi_c: np.ndarray, geofac_n2s: np.ndarray + connectivities: Mapping[gtx.FieldOffset, np.ndarray], psi_c: np.ndarray, geofac_n2s: np.ndarray ) -> np.ndarray: - c2e2cO = connectivities[dims.C2E2CODim] + c2e2cO = connectivities[dims.C2E2CO] nabla2_psi_c = np.sum(np.where((c2e2cO != -1), psi_c[c2e2cO] * geofac_n2s, 0), axis=1) return nabla2_psi_c def nabla2_on_cell_k_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], psi_c: np.ndarray, geofac_n2s: np.ndarray + connectivities: Mapping[gtx.FieldOffset, np.ndarray], psi_c: np.ndarray, geofac_n2s: np.ndarray ) -> np.ndarray: - c2e2cO = connectivities[dims.C2E2CODim] + c2e2cO = connectivities[dims.C2E2CO] geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) nabla2_psi_c = np.sum( np.where((c2e2cO != -1)[:, :, np.newaxis], psi_c[c2e2cO] * geofac_n2s, 0), axis=1 diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index d4b5d63a6b..060a537c52 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -10,69 +10,229 @@ import contextlib import dataclasses +import functools +import inspect import os -from collections.abc import Callable, Generator, Mapping, Sequence -from typing import Any, ClassVar, Final +import types +from collections.abc import Callable, Generator, Iterator, Mapping, Sequence +from typing import TYPE_CHECKING, Any, ClassVar, Final, TypeAlias, cast import gt4py.next as gtx import numpy as np import pytest from gt4py import eve -from gt4py.next import ( - constructors, - named_collections as gtx_named_collections, - typing as gtx_typing, -) +from gt4py.next import common as gtx_common, typing as gtx_typing # TODO(havogt): import will disappear after FieldOperators support `.compile` from gt4py.next.ffront.decorator import FieldOperator from gt4py.next.instrumentation import hooks as gtx_hooks, metrics as gtx_metrics -from icon4py.model.common import model_backends, model_options +from icon4py.model.common import model_backends, model_options, type_alias as ta from icon4py.model.common.grid import base -from icon4py.model.common.utils import device_utils +from icon4py.model.common.utils import data_allocation, device_utils from icon4py.model.testing import test_utils -def allocate_data( - allocator: gtx_typing.Allocator | None, - input_data: dict[ - str, Any - ], # `Field`s or collection of `Field`s are re-allocated, the rest is passed through -) -> dict[str, Any]: - def _allocate_field(f: gtx.Field) -> gtx.Field: - return constructors.as_field(domain=f.domain, data=f.ndarray, allocator=allocator) +if TYPE_CHECKING: + import numpy.typing as npt - input_data = { - k: gtx_named_collections.tree_map_named_collection(_allocate_field)(v) - if not gtx.is_scalar_type(v) and k != "domain" - else v - for k, v in input_data.items() - } - return input_data +_STENCIL_REFERENCE_MARKER: Final = "__stencil_test_reference__" +_INPUT_DATA_FIXTURE_MARKER: Final = "__stencil_test_input_fixture__" -@dataclasses.dataclass(frozen=True) -class Output: - name: str - refslice: tuple[slice, ...] = dataclasses.field(default_factory=lambda: (slice(None),)) - gtslice: tuple[slice, ...] = dataclasses.field(default_factory=lambda: (slice(None),)) +def _static_reference(func: types.FunctionType | staticmethod) -> staticmethod: + """Decorator to mark the `reference` method of a `StencilTest` suite.""" + if not isinstance(func, (types.FunctionType, staticmethod)): + raise TypeError( + f"The 'reference' function must be a regular function or staticmethod but got {type(func)}." + ) + if func.__name__ != "reference": + raise ValueError( + f"The 'reference' method must be named 'reference' but got '{func.__name__}'." + ) + func_params = tuple(inspect.signature(func).parameters.keys()) + if func_params[0] != "grid": + raise ValueError( + f"The 'reference' method signature must be 'reference(grid, ...)' but got" + f" '{func.__name__}{func_params}'." + ) + if not isinstance(func, staticmethod): + func = staticmethod(func) + + setattr(func, _STENCIL_REFERENCE_MARKER, True) + + return func + + +def _input_data_fixture( + func: types.FunctionType | None = None, **kwargs: Any +) -> types.FunctionType | Callable[[types.FunctionType], types.FunctionType]: + """ + Decorator to mark the `input_data` method of a `StencilTest` suite as a pytest fixture. + + Perform some checks on the decorated function and forward all the keyword + arguments to `pytest.fixture` for parametrization and scoping (default: "class"). + """ + if func is None: + return functools.partial(input_data_fixture, **kwargs) + + if not isinstance(func, types.FunctionType): + raise TypeError(f"The 'input_data' method must be a regular function but got {type(func)}.") + if func.__name__ != "input_data": + raise ValueError( + f"The 'input_data' method must be named 'input_data' but got '{func.__name__}'." + ) + func_params = tuple(inspect.signature(func).parameters.keys()) + if func_params[:2] != ("self", "grid"): + raise ValueError( + f"The 'input_data' method signature must be 'input_data(self, grid, ...)' but got" + f" '{func.__name__}{func_params}'." + ) + + # This allows us to check that the `input_data` fixture does not call any `data_allocation` + # functions directly and thus it only uses the `self.data_alloc` wrapper, which ensures + # that the backend and grid are properly bound. However, it might be a bit too strict, + # since it means that the `data_allocation` module cannot be imported in the global scope + # of the test module (it can be still imported in the local scope of other functions). + # We can remove it in the future if it causes many issues. + cv = inspect.getclosurevars(func) + if any(ref is data_allocation for ref in [*cv.globals.values(), *cv.nonlocals.values()]): + raise TypeError( + "The 'input_data_fixture' should not call 'data_allocation' functions directly. " + "Use `self.data_alloc` inside the fixture to access data allocation functions instead." + ) + + kwargs.setdefault("scope", "class") + fixt = pytest.fixture(**kwargs)(func) + setattr(fixt, _INPUT_DATA_FIXTURE_MARKER, True) + + return fixt -@dataclasses.dataclass -class _ConnectivityConceptFixer: +if TYPE_CHECKING: + static_reference: TypeAlias = staticmethod + input_data_fixture: Final = pytest.fixture +else: + static_reference = _static_reference + input_data_fixture = _input_data_fixture + + +@dataclasses.dataclass(frozen=True) +class DataAllocationWrapper: """ - This works around a misuse of dimensions as an identifier for connectivities. - Since GT4Py might change the way the mesh is represented, we could - keep this for a while, otherwise we need to touch all StencilTests. + This wrapper mimics the 'icon4py.model.common.utils.data_allocation' functions, + but with 'backend' and `grid` bound in the respective functions. """ - _grid: base.Grid + grid: base.Grid + allocator: gtx_typing.Allocator | None + + def constant_field( + self, + value: float, + *dims: gtx.Dimension, + dtype: npt.DTypeLike | None = ta.wpfloat, + ) -> gtx.Field: + return data_allocation.constant_field( + self.grid, value, *dims, dtype=dtype, allocator=self.allocator + ) + + def index_field( + self, + dim: gtx.Dimension, + extend: dict[gtx.Dimension, int] | None = None, + dtype: npt.DTypeLike = gtx.int32, + ) -> gtx.Field: + return data_allocation.index_field( + grid=self.grid, dim=dim, extend=extend, dtype=dtype, allocator=self.allocator + ) + + def random_field( + self, + *dims: gtx.Dimension, + low: float = -1.0, + high: float = 1.0, + dtype: npt.DTypeLike | None = None, + extend: dict[gtx.Dimension, int] | None = None, + ) -> gtx.Field: + return data_allocation.random_field( + self.grid, + *dims, + low=low, + high=high, + dtype=dtype, + allocator=self.allocator, + extend=extend, + ) + + def random_mask( + self, + *dims: gtx.Dimension, + dtype: npt.DTypeLike | None = None, + extend: dict[gtx.Dimension, int] | None = None, + ) -> gtx.Field: + return data_allocation.random_mask( + self.grid, *dims, dtype=dtype, allocator=self.allocator, extend=extend + ) + + def random_sign( + self, + *dims: gtx.Dimension, + dtype: npt.DTypeLike | None = None, + extend: dict[gtx.Dimension, int] | None = None, + ) -> gtx.Field: + return data_allocation.random_sign( + self.grid, *dims, dtype=dtype, allocator=self.allocator, extend=extend + ) + + def zero_field( + self, + *dims: gtx.Dimension, + dtype: npt.DTypeLike | None = ta.wpfloat, + extend: dict[gtx.Dimension, int] | None = None, + ) -> gtx.Field: + return data_allocation.zero_field( + self.grid, *dims, dtype=dtype, allocator=self.allocator, extend=extend + ) + + +class NumPyGridConnectivitiesView(Mapping[str | gtx.FieldOffset, np.ndarray]): + """View on the grid connectivities that allows to access them as numpy arrays.""" + + def __init__(self, grid: base.Grid): + self.grid = grid - def __getitem__(self, dim: gtx.Dimension | str) -> np.ndarray: - if isinstance(dim, gtx.Dimension): - dim = dim.value - return self._grid.get_connectivity(dim).asnumpy() + def __getitem__(self, key: str | gtx.FieldOffset) -> np.ndarray: + connectivity = self.grid.get_connectivity(key) + if gtx_common.is_neighbor_table(connectivity): + return connectivity.asnumpy() + else: + raise TypeError(f"Connectivity '{key}' is not a neighbor table.") + + def __iter__(self) -> Iterator[str | gtx.FieldOffset]: + return ( + key + for key, connectivity in self.grid.connectivities.items() + if gtx_common.is_neighbor_table(connectivity) + ) + + def __len__(self) -> int: + return sum( + 1 + for connectivity in self.grid.connectivities.values() + if gtx_common.is_neighbor_table(connectivity) + ) + + +def connectivities_asnumpy(grid: base.Grid) -> Mapping[gtx.FieldOffset, np.ndarray]: + return cast(Mapping[gtx.FieldOffset, np.ndarray], NumPyGridConnectivitiesView(grid)) + + +@dataclasses.dataclass(frozen=True) +class Output: + name: str + refslice: tuple[slice, ...] = dataclasses.field(default_factory=lambda: (slice(None),)) + gtslice: tuple[slice, ...] = dataclasses.field(default_factory=lambda: (slice(None),)) class StandardStaticVariants(eve.StrEnum): @@ -85,10 +245,16 @@ def test_and_benchmark( self: StencilTest, benchmark: Any, # should be `pytest_benchmark.fixture.BenchmarkFixture` but pytest_benchmark is not typed grid: base.Grid, - _properly_allocated_input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], - _configured_program: Callable[..., None], + input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], + configured_program: Callable[..., None], request: pytest.FixtureRequest, ) -> None: + """ + Test and benchmark the stencil program. + + Note that it is defined as a standalone function and then attached to the `StencilTest` + subclasses in order to use a meaningful name for the test in pytest output. + """ skip_stenciltest_verification = request.config.getoption( "skip_stenciltest_verification" ) # skip verification if `--skip-stenciltest-verification` CLI option is set @@ -96,19 +262,12 @@ def test_and_benchmark( if not skip_stenciltest_verification: reference_outputs = self.reference( - _ConnectivityConceptFixer( - grid # TODO(havogt): pass as keyword argument (needs fixes in some tests) - ), - **{ - k: v.asnumpy() if isinstance(v, gtx.Field) else v - for k, v in _properly_allocated_input_data.items() - }, + grid=grid, + **{k: v.asnumpy() if isinstance(v, gtx.Field) else v for k, v in input_data.items()}, ) - _configured_program(**_properly_allocated_input_data, offset_provider=grid.connectivities) - self._verify_stencil_test( - input_data=_properly_allocated_input_data, reference_outputs=reference_outputs - ) + configured_program(**input_data, offset_provider=grid.connectivities) + self.verify_data(input_data=input_data, reference_outputs=reference_outputs) if not skip_stenciltest_benchmark: warmup_rounds = int(os.getenv("ICON4PY_STENCIL_TEST_WARMUP_ROUNDS", "1")) @@ -116,9 +275,9 @@ def test_and_benchmark( # Use of `pedantic` to explicitly control warmup rounds and iterations benchmark.pedantic( - _configured_program, + configured_program, args=(), - kwargs=dict(**_properly_allocated_input_data, offset_provider=grid.connectivities), + kwargs=dict(**input_data, offset_provider=grid.connectivities), rounds=int( os.getenv("ICON4PY_STENCIL_TEST_BENCHMARK_ROUNDS", "3") ), # 30 iterations in total should be stable enough @@ -148,18 +307,18 @@ def _get_metrics_id_program_callback( gtx_hooks.program_call_context.register( _get_metrics_id_program_callback, name=METRICS_KEY_EXTRACTOR ) - _configured_program( - **_properly_allocated_input_data, offset_provider=grid.connectivities - ) + configured_program(**input_data, offset_provider=grid.connectivities) gtx_hooks.program_call_context.remove(METRICS_KEY_EXTRACTOR) - assert metrics_key is not None, "Metrics key could not be recovered during run." - assert metrics_key.startswith( - _configured_program.__name__ - ), f"Metrics key ({metrics_key}) does not start with the program name ({_configured_program.__name__})" - - assert ( - len(_configured_program._compiled_programs.compiled_programs) == 1 - ), "Multiple compiled programs found, cannot extract metrics." + + if metrics_key is None: + raise RuntimeError("Metrics key could not be recovered during run.") + if not metrics_key.startswith(configured_program.__name__): + raise RuntimeError( + f"Metrics key ({metrics_key}) does not start with the program name ({configured_program.__name__})" + ) + if len(configured_program._compiled_programs.compiled_programs) != 1: + raise RuntimeError("Multiple compiled programs found, cannot extract metrics.") + metrics_data = gtx_metrics.sources compute_samples = metrics_data[metrics_key].metrics["compute"].samples # exclude: @@ -170,9 +329,10 @@ def _get_metrics_id_program_callback( initial_program_iterations_to_skip = warmup_rounds * iterations + ( 2 if skip_stenciltest_verification else 3 ) - assert ( - len(compute_samples) > initial_program_iterations_to_skip - ), "Not enough samples collected to compute metrics." + + if len(compute_samples) <= initial_program_iterations_to_skip: + raise RuntimeError("Not enough samples collected to compute metrics.") + benchmark.extra_info["gtx_metrics"] = compute_samples[ initial_program_iterations_to_skip: ] @@ -203,9 +363,12 @@ class StencilTest: STATIC_PARAMS: ClassVar[dict[str, Sequence[str]] | None] = None reference: ClassVar[Callable[..., Mapping[str, np.ndarray | tuple[np.ndarray, ...]]]] + input_data: ClassVar[Callable[..., dict[str, Any]]] + + data_alloc: DataAllocationWrapper @pytest.fixture - def _configured_program( + def configured_program( self, backend_like: model_backends.BackendLike, static_variant: Sequence[str], @@ -235,19 +398,24 @@ def _configured_program( test_func = device_utils.synchronized_function(program, allocator=backend) return test_func - @pytest.fixture - def _properly_allocated_input_data( - self, - input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], - backend_like: model_backends.BackendLike, - ) -> dict[str, Any]: - # TODO(havogt): this is a workaround, - # because in the `input_data` fixture provided by the user - # it does not allocate for the correct device. - allocator = model_backends.get_allocator(backend_like) - return allocate_data(allocator=allocator, input_data=input_data) - - def _verify_stencil_test( + @pytest.fixture(autouse=True, scope="class") + def _instance_setup_fixture( + self, backend_like: model_backends.BackendLike, grid: base.Grid + ) -> Generator[None, None, None]: + """ + Convenience fixture to provide data allocation functions with backend and grid already bound. + """ + self.data_alloc_wrapper = DataAllocationWrapper( + grid=grid, allocator=model_backends.get_allocator(backend_like) + ) + try: + self.data_alloc = self.data_alloc_wrapper + yield + + finally: + del self.data_alloc + + def verify_data( self, input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], reference_outputs: Mapping[str, np.ndarray | tuple[np.ndarray, ...]], @@ -296,12 +464,30 @@ def static_variant(request: pytest.FixtureRequest) -> Sequence[str]: _, variant = request.param return () if variant is None else variant - def __init_subclass__(cls, **kwargs: Any) -> None: - super().__init_subclass__(**kwargs) + def __init_subclass__(cls, *args: Any, **kwargs: Any) -> None: + super().__init_subclass__(*args, **kwargs) + + # Check the conventions for `reference` and `input_data` methods + if not hasattr(cls, "reference"): + raise TypeError( + f"{cls.__name__} StencilTest subclass does not implement a 'reference' method." + ) + if not getattr(cls.__dict__["reference"], _STENCIL_REFERENCE_MARKER, False): + raise RuntimeError( + f"The 'reference' method of {cls.__name__} must be decorated with '@static_reference'." + ) + if not hasattr(cls, "input_data"): + raise TypeError( + f"{cls.__name__} StencilTest subclass does not implement an 'input_data' method." + ) + if not getattr(cls.__dict__["input_data"], _INPUT_DATA_FIXTURE_MARKER, False): + raise RuntimeError( + f"The 'input_data' method of {cls.__name__} must be decorated with '@input_data_fixture'." + ) setattr(cls, f"test_{cls.__name__}", test_and_benchmark) - # decorate `static_variant` with parametrized fixtures, since the + # Decorate `static_variant` with parametrized fixtures, since the # parametrization is only available in the concrete subclass definition if cls.STATIC_PARAMS is None: # not parametrized, return an empty tuple