diff --git a/ci/dace.yml b/ci/dace.yml index 62b26e83cf..e96243d918 100644 --- a/ci/dace.yml +++ b/ci/dace.yml @@ -22,12 +22,11 @@ test_model_stencils_aarch64: stage: test script: - pip install dace==$DACE_VERSION - - tox -r -e run_model_tests -c model/ -- --backend=$BACKEND $DACE_ORCHESTRATION $COMPONENT --verbose + - tox -r -e run_model_tests -c model/ -- --backend=$BACKEND $COMPONENT --verbose parallel: matrix: - COMPONENT: [atmosphere/diffusion/tests/diffusion_tests] BACKEND: [dace_cpu_noopt] - DACE_ORCHESTRATION: ['--dace-orchestration=True', ''] test_model_datatests_x86_64: extends: [.test_model_datatests, .test_template_x86_64] test_model_datatests_aarch64: diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py index 7c7f56159e..72e2606599 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -65,7 +65,7 @@ from icon4py.model.common.utils import gt4py_field_allocation as field_alloc -from icon4py.model.common.orchestration import decorator as orchestration +from icon4py.model.common.orchestration import decorator as dace_orchestration """ @@ -364,9 +364,11 @@ def __init__( edge_params: grid_states.EdgeParams, cell_params: grid_states.CellParams, backend: backend.Backend, + orchestration: bool = False, exchange: decomposition.ExchangeRuntime = decomposition.SingleNodeExchange(), ): self._backend = backend + self._orchestration = orchestration self._exchange = exchange self.config = config self._params = params @@ -461,7 +463,7 @@ def __init__( self._determine_horizontal_domains() - self.compile_time_connectivities = orchestration.build_compile_time_connectivities( + self.compile_time_connectivities = dace_orchestration.build_compile_time_connectivities( self._grid.offset_providers ) @@ -621,7 +623,7 @@ def _sync_cell_fields(self, prognostic_state): ) log.debug("communication of prognostic cell fields: theta, w, exner - done") - @orchestration.orchestrate + @dace_orchestration.orchestrate def _do_diffusion_step( self, diagnostic_state: diffusion_states.DiffusionDiagnosticState, @@ -929,6 +931,6 @@ def orchestration_uid(self) -> str: if isinstance(self.__dict__[name], gtx.ffront.decorator.Program) ], ] - return orchestration.generate_orchestration_uid( + return dace_orchestration.generate_orchestration_uid( self, members_to_disregard=members_to_disregard ) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py index 6d2c3b0866..0078758c15 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py @@ -13,7 +13,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import KDim, VertexDim from icon4py.model.common.math.smagorinsky import _en_smag_fac_for_zero_nshift -from icon4py.model.common.settings import backend @gtx.field_operator @@ -46,7 +45,7 @@ def _init_zero_v_k() -> gtx.Field[[dims.VertexDim, dims.KDim], float]: return broadcast(0.0, (VertexDim, KDim)) -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED, backend=backend) +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def init_zero_v_k(field: gtx.Field[[dims.VertexDim, dims.KDim], float]): _init_zero_v_k(out=field) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py index 8fb0513a22..fe4994caca 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py @@ -22,7 +22,6 @@ _update_theta_and_exner, ) from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -58,7 +57,7 @@ def _apply_diffusion_to_theta_and_exner( return theta_v, exner -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def apply_diffusion_to_theta_and_exner( kh_smag_e: fa.EdgeKField[vpfloat], inv_dual_edge_length: fa.EdgeField[wpfloat], diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py index 6c0930b4a8..dc734fd447 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py @@ -11,7 +11,6 @@ from gt4py.next.ffront.fbuiltins import astype from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -32,7 +31,7 @@ def _apply_nabla2_and_nabla4_global_to_vn( return vn_wp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def apply_nabla2_and_nabla4_global_to_vn( area_edge: fa.EdgeField[wpfloat], kh_smag_e: fa.EdgeKField[vpfloat], diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_to_vn.py index dd25d46ad9..64f6afeca5 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_to_vn.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_to_vn.py @@ -11,7 +11,6 @@ from gt4py.next.ffront.fbuiltins import astype, broadcast, maximum from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -38,7 +37,7 @@ def _apply_nabla2_and_nabla4_to_vn( return vn_wp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def apply_nabla2_and_nabla4_to_vn( area_edge: fa.EdgeField[wpfloat], kh_smag_e: fa.EdgeKField[vpfloat], diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py index 988a286188..2561244878 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py @@ -10,7 +10,6 @@ from gt4py.next.ffront.decorator import field_operator, program from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat @@ -25,7 +24,7 @@ def _apply_nabla2_to_vn_in_lateral_boundary( return vn_wp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def apply_nabla2_to_vn_in_lateral_boundary( z_nabla2_e: fa.EdgeKField[wpfloat], area_edge: fa.EdgeField[wpfloat], diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w.py index bfdbc34a8b..3bf0129e42 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import C2E2CO, C2E2CODim -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -32,7 +31,7 @@ def _apply_nabla2_to_w( return w_wp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def apply_nabla2_to_w( area: fa.CellField[wpfloat], z_nabla2_c: fa.CellKField[vpfloat], diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w_in_upper_damping_layer.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w_in_upper_damping_layer.py index fd60a69fb7..a15ac68405 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w_in_upper_damping_layer.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w_in_upper_damping_layer.py @@ -11,7 +11,6 @@ from gt4py.next.ffront.fbuiltins import astype, broadcast from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -29,7 +28,7 @@ def _apply_nabla2_to_w_in_upper_damping_layer( return w_wp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def apply_nabla2_to_w_in_upper_damping_layer( w: fa.CellKField[wpfloat], diff_multfac_n2w: fa.KField[wpfloat], diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py index b15f2402a6..5aa03e693e 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py @@ -12,7 +12,6 @@ from icon4py.model.common import field_type_aliases as fa from icon4py.model.common.dimension import Koff -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -33,7 +32,7 @@ def _calculate_diagnostics_for_turbulence( return astype((div_ic_wp, hdef_ic_wp), vpfloat) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def calculate_diagnostics_for_turbulence( div: fa.CellKField[vpfloat], kh_c: fa.CellKField[vpfloat], diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_horizontal_gradients_for_turbulence.py index 6e6f81fbba..5404bda7c3 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_horizontal_gradients_for_turbulence.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import C2E2CO, C2E2CODim -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -27,7 +26,7 @@ def _calculate_horizontal_gradients_for_turbulence( return astype((dwdx_wp, dwdy_wp), vpfloat) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def calculate_horizontal_gradients_for_turbulence( w: fa.CellKField[wpfloat], geofac_grg_x: gtx.Field[gtx.Dims[dims.CellDim, C2E2CODim], wpfloat], diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_w.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_w.py index 3fcedf378d..ce87f8715b 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_w.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_w.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import C2E2CO, C2E2CODim -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -24,7 +23,7 @@ def _calculate_nabla2_for_w( return astype(z_nabla2_c_wp, vpfloat) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def calculate_nabla2_for_w( w: fa.CellKField[wpfloat], geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, C2E2CODim], wpfloat], diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_z.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_z.py index 3c22185b33..eb7f9f3116 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_z.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_z.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import E2C -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -28,7 +27,7 @@ def _calculate_nabla2_for_z( return z_nabla2_e_wp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def calculate_nabla2_for_z( kh_smag_e: fa.EdgeKField[vpfloat], inv_dual_edge_length: fa.EdgeField[wpfloat], diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py index 875985cf65..f05bb9b8c2 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import C2CE, C2E, C2EDim -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -25,7 +24,7 @@ def _calculate_nabla2_of_theta( return astype(z_temp_wp, vpfloat) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def calculate_nabla2_of_theta( z_nabla2_e: fa.EdgeKField[wpfloat], geofac_div: gtx.Field[gtx.Dims[dims.CEDim], wpfloat], diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py index 73d0714bb5..6efdea24cc 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import E2C2V, E2ECV, ECVDim -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -56,7 +55,7 @@ def _calculate_nabla4( return astype(z_nabla4_e2_wp, vpfloat) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def calculate_nabla4( u_vert: fa.VertexKField[vpfloat], v_vert: fa.VertexKField[vpfloat], diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py index a97b1409c2..f8de6677bc 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import E2C, E2CDim -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat @@ -25,7 +24,7 @@ def _enhance_diffusion_coefficient_for_grid_point_cold_pools( return kh_smag_e_vp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def enhance_diffusion_coefficient_for_grid_point_cold_pools( kh_smag_e: fa.EdgeKField[vpfloat], enh_diffu_3d: fa.CellKField[vpfloat], diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_field_for_grid_point_cold_pools_enhancement.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_field_for_grid_point_cold_pools_enhancement.py index 97255f0290..e3100ba7bc 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_field_for_grid_point_cold_pools_enhancement.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_field_for_grid_point_cold_pools_enhancement.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import C2E2C, C2E2CDim -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -38,7 +37,7 @@ def _temporary_field_for_grid_point_cold_pools_enhancement( return enh_diffu_3d_vp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def temporary_field_for_grid_point_cold_pools_enhancement( theta_v: fa.CellKField[wpfloat], theta_ref_mc: fa.CellKField[vpfloat], diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py index f6f5727bf5..051689f74d 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import C2CE, C2E, C2EDim -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -31,7 +30,7 @@ def _temporary_fields_for_turbulence_diagnostics( return astype((kh_c_wp, div_wp), vpfloat) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def temporary_fields_for_turbulence_diagnostics( kh_smag_ec: fa.EdgeKField[vpfloat], vn: fa.EdgeKField[wpfloat], 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 e72d356c77..186e5c08fe 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 @@ -28,7 +28,10 @@ def reference(grid, kh_smag_e: np.array, enh_diffu_3d: np.array, **kwargs) -> np e2c = grid.connectivities[dims.E2CDim] kh_smag_e = np.maximum( kh_smag_e, - np.max(np.where((e2c != -1)[:, :, np.newaxis], enh_diffu_3d[e2c], -math.inf), axis=1), + np.max( + np.where((e2c != -1)[:, :, np.newaxis], enh_diffu_3d[e2c], -math.inf), + axis=1, + ), ) return dict(kh_smag_e=kh_smag_e) diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/mpi_tests/test_parallel_diffusion.py b/model/atmosphere/diffusion/tests/diffusion_tests/mpi_tests/test_parallel_diffusion.py index b1889adb95..85aaf10c47 100644 --- a/model/atmosphere/diffusion/tests/diffusion_tests/mpi_tests/test_parallel_diffusion.py +++ b/model/atmosphere/diffusion/tests/diffusion_tests/mpi_tests/test_parallel_diffusion.py @@ -9,7 +9,7 @@ import pytest from icon4py.model.atmosphere.diffusion import diffusion as diffusion_, diffusion_states -from icon4py.model.common import dimension as dims, settings +from icon4py.model.common import dimension as dims from icon4py.model.common.decomposition import definitions from icon4py.model.common.grid import vertical as v_grid from icon4py.model.common.test_utils import datatest_utils, helpers, parallel_helpers @@ -20,7 +20,8 @@ @pytest.mark.mpi @pytest.mark.parametrize("experiment", [datatest_utils.REGIONAL_EXPERIMENT]) @pytest.mark.parametrize("ndyn_substeps", [2]) -@pytest.mark.parametrize("linit", [True, False]) +@pytest.mark.parametrize("linit", [([True, False])]) +@pytest.mark.parametrize("orchestration", [([True, False])]) def test_parallel_diffusion( experiment, step_date_init, @@ -40,7 +41,10 @@ def test_parallel_diffusion( damping_height, caplog, backend, + orchestration, ): + if orchestration and ("dace" not in backend.name.lower()): + raise pytest.skip("This test is only executed for `dace backends.") caplog.set_level("INFO") parallel_helpers.check_comm_size(processor_props) print( @@ -107,6 +111,7 @@ def test_parallel_diffusion( cell_params=cell_geometry, exchange=exchange, backend=backend, + orchestration=orchestration, ) print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ") @@ -168,9 +173,8 @@ def test_parallel_diffusion_multiple_steps( caplog, backend, ): - if settings.dace_orchestration is None: - raise pytest.skip("This test is only executed for `--dace-orchestration=True`.") - + if "dace" not in backend.name.lower(): + raise pytest.skip("This test is only executed for `dace backends.") ###################################################################### # Diffusion initialization ###################################################################### @@ -232,7 +236,6 @@ def test_parallel_diffusion_multiple_steps( ###################################################################### # DaCe NON-Orchestrated Backend ###################################################################### - settings.dace_orchestration = None diffusion = diffusion_.Diffusion( grid=icon_grid, @@ -247,6 +250,7 @@ def test_parallel_diffusion_multiple_steps( cell_params=cell_geometry, backend=backend, exchange=exchange, + orchestration=False, ) print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ") @@ -281,7 +285,6 @@ def test_parallel_diffusion_multiple_steps( ###################################################################### # DaCe Orchestrated Backend ###################################################################### - settings.dace_orchestration = True exchange = definitions.create_exchange(processor_props, decomposition_info) diffusion = diffusion_.Diffusion( @@ -297,6 +300,7 @@ def test_parallel_diffusion_multiple_steps( cell_params=cell_geometry, exchange=exchange, backend=backend, + orchestration=True, ) print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ") diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py index 0250ca085f..6437de5a2d 100644 --- a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py +++ b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py @@ -5,17 +5,16 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +import numpy as np import pytest import icon4py.model.common.dimension as dims import icon4py.model.common.grid.states as grid_states from icon4py.model.atmosphere.diffusion import diffusion, diffusion_states, diffusion_utils -from icon4py.model.common import settings from icon4py.model.common.grid import ( geometry_attributes as geometry_meta, vertical as v_grid, ) -from icon4py.model.common.settings import backend, xp from icon4py.model.common.test_utils import ( datatest_utils as dt_utils, grid_utils, @@ -161,7 +160,7 @@ def test_smagorinski_factor_diffusion_type_5(experiment): params = diffusion.DiffusionParams(construct_diffusion_config(experiment, ndyn_substeps=5)) assert len(params.smagorinski_factor) == len(params.smagorinski_height) assert len(params.smagorinski_factor) == 4 - assert xp.all(params.smagorinski_factor >= xp.zeros(len(params.smagorinski_factor))) + assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor))) @pytest.mark.datatest @@ -272,7 +271,7 @@ def test_diffusion_init( def _verify_init_values_against_savepoint( - savepoint: sb.IconDiffusionInitSavepoint, diffusion_granule: diffusion.Diffusion + savepoint: sb.IconDiffusionInitSavepoint, diffusion_granule: diffusion.Diffusion, backend ): dtime = savepoint.get_metadata("dtime")["dtime"] @@ -374,10 +373,11 @@ def test_verify_diffusion_init_against_savepoint( interpolation_state, edge_params, cell_params, + orchestration=True, backend=backend, ) - _verify_init_values_against_savepoint(savepoint_diffusion_init, diffusion_granule) + _verify_init_values_against_savepoint(savepoint_diffusion_init, diffusion_granule, backend) @pytest.mark.datatest @@ -388,7 +388,7 @@ def test_verify_diffusion_init_against_savepoint( (dt_utils.GLOBAL_EXPERIMENT, "2000-01-01T00:00:02.000", "2000-01-01T00:00:02.000"), ], ) -@pytest.mark.parametrize("ndyn_substeps", (2,)) +@pytest.mark.parametrize("ndyn_substeps, orchestration", [(2, [True, False])]) def test_run_diffusion_single_step( savepoint_diffusion_init, savepoint_diffusion_exit, @@ -401,7 +401,10 @@ def test_run_diffusion_single_step( damping_height, ndyn_substeps, backend, + orchestration, ): + if orchestration and ("dace" not in backend.name.lower()): + raise pytest.skip("This test is only executed for `dace backends.") grid = get_grid_for_experiment(experiment, backend) cell_geometry = get_cell_geometry_for_experiment(experiment, backend) edge_geometry = get_edge_geometry_for_experiment(experiment, backend) @@ -462,6 +465,7 @@ def test_run_diffusion_single_step( edge_params=edge_geometry, cell_params=cell_geometry, backend=backend, + orchestration=orchestration, ) verify_diffusion_fields(config, diagnostic_state, prognostic_state, savepoint_diffusion_init) assert savepoint_diffusion_init.fac_bdydiff_v() == diffusion_granule.fac_bdydiff_v @@ -498,9 +502,8 @@ def test_run_diffusion_multiple_steps( backend, icon_grid, ): - if settings.dace_orchestration is None: - raise pytest.skip("This test is only executed for `--dace-orchestration=True`.") - + if "dace" not in backend.name.lower(): + raise pytest.skip("This test is only executed for `dace backends.") ###################################################################### # Diffusion initialization ###################################################################### @@ -547,7 +550,6 @@ def test_run_diffusion_multiple_steps( ###################################################################### # DaCe NON-Orchestrated Backend ###################################################################### - settings.dace_orchestration = None diagnostic_state_dace_non_orch = diffusion_states.DiffusionDiagnosticState( hdef_ic=savepoint_diffusion_init.hdef_ic(), @@ -566,6 +568,7 @@ def test_run_diffusion_multiple_steps( interpolation_state=interpolation_state, edge_params=edge_geometry, cell_params=cell_geometry, + orchestration=False, backend=backend, ) @@ -579,7 +582,6 @@ def test_run_diffusion_multiple_steps( ###################################################################### # DaCe Orchestrated Backend ###################################################################### - settings.dace_orchestration = True diagnostic_state_dace_orch = diffusion_states.DiffusionDiagnosticState( hdef_ic=savepoint_diffusion_init.hdef_ic(), @@ -599,6 +601,7 @@ def test_run_diffusion_multiple_steps( edge_params=edge_geometry, cell_params=cell_geometry, backend=backend, + orchestration=True, ) for _ in range(3): @@ -621,7 +624,7 @@ def test_run_diffusion_multiple_steps( @pytest.mark.datatest @pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT]) -@pytest.mark.parametrize("linit", [True]) +@pytest.mark.parametrize("linit, orchestration", [(True, [True, False])]) def test_run_diffusion_initial_step( experiment, linit, @@ -634,7 +637,10 @@ def test_run_diffusion_initial_step( interpolation_savepoint, metrics_savepoint, backend, + orchestration, ): + if orchestration and ("dace" not in backend.name.lower()): + raise pytest.skip("This test is only executed for `dace backends.") grid = get_grid_for_experiment(experiment, backend) cell_geometry = get_cell_geometry_for_experiment(experiment, backend) edge_geometry = get_edge_geometry_for_experiment(experiment, backend) @@ -691,6 +697,7 @@ def test_run_diffusion_initial_step( edge_params=edge_geometry, cell_params=cell_geometry, backend=backend, + orchestration=orchestration, ) assert savepoint_diffusion_init.fac_bdydiff_v() == diffusion_granule.fac_bdydiff_v diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion_utils.py b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion_utils.py index 7634462cfe..75dca66da5 100644 --- a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion_utils.py +++ b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion_utils.py @@ -12,7 +12,6 @@ from icon4py.model.atmosphere.diffusion import diffusion, diffusion_utils from icon4py.model.common import dimension as dims from icon4py.model.common.grid import simple as simple_grid -from icon4py.model.common.settings import backend from icon4py.model.common.test_utils import helpers from .utils import construct_diffusion_config, diff_multfac_vn_numpy, smag_limit_numpy @@ -99,14 +98,14 @@ def test_diff_multfac_vn_smag_limit_for_loop_run_with_k4_substeps(backend): def test_init_zero_vertex_k(backend): grid = simple_grid.SimpleGrid() f = helpers.random_field(grid, dims.VertexDim, dims.KDim) - diffusion_utils.init_zero_v_k(f, offset_provider={}) + diffusion_utils.init_zero_v_k.with_backend(backend)(f, offset_provider={}) assert np.allclose(0.0, f.asnumpy()) @pytest.mark.datatest @pytest.mark.parametrize("linit", [True]) def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( - savepoint_diffusion_init, experiment, icon_grid, linit, ndyn_substeps + savepoint_diffusion_init, experiment, icon_grid, linit, ndyn_substeps, backend ): savepoint = savepoint_diffusion_init config = construct_diffusion_config(experiment, ndyn_substeps=ndyn_substeps) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_utils.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_utils.py index 7d8cef2901..69ce0f95cb 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_utils.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_utils.py @@ -14,7 +14,6 @@ from icon4py.model.common import field_type_aliases as fa from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat @@ -23,7 +22,7 @@ def _scale_k(field: fa.KField[float], factor: float) -> fa.KField[float]: return field * factor -@gtx.program(backend=backend) +@gtx.program def scale_k(field: fa.KField[float], factor: float, scaled_field: fa.KField[float]): _scale_k(field, factor, out=scaled_field) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/add_interpolated_horizontal_advection_of_w.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/add_interpolated_horizontal_advection_of_w.py index 1d873c5ded..2819494f84 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/add_interpolated_horizontal_advection_of_w.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/add_interpolated_horizontal_advection_of_w.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import C2CE, C2E, C2EDim -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -28,7 +27,7 @@ def _add_interpolated_horizontal_advection_of_w( return astype(ddt_w_adv_wp, vpfloat) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def add_interpolated_horizontal_advection_of_w( e_bln_c_s: gtx.Field[gtx.Dims[dims.CEDim], wpfloat], z_v_grad_w: fa.EdgeKField[vpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_advective_vertical_wind_tendency.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_advective_vertical_wind_tendency.py index 4a8b1ce227..1770361d68 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_advective_vertical_wind_tendency.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_advective_vertical_wind_tendency.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import Koff -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -35,7 +34,7 @@ def _compute_advective_vertical_wind_tendency( return astype(ddt_w_adv_wp, vpfloat) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def compute_advective_vertical_wind_tendency( z_w_con_c: fa.CellKField[vpfloat], w: fa.CellKField[wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_airmass.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_airmass.py index afca5afbad..c2bbf153b7 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_airmass.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_airmass.py @@ -9,7 +9,6 @@ from gt4py.next.ffront.decorator import field_operator, program from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat @@ -22,7 +21,7 @@ def _compute_airmass( return rho_in * ddqz_z_full_in * deepatmo_t1mc_in -@program(grid_type=gtx.GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_airmass( rho_in: fa.CellKField[wpfloat], ddqz_z_full_in: fa.CellKField[wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_contravariant_correction.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_contravariant_correction.py index 6b3fba85a3..df5ab91c09 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_contravariant_correction.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_contravariant_correction.py @@ -11,7 +11,6 @@ from gt4py.next.ffront.fbuiltins import astype from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -29,7 +28,7 @@ def _compute_contravariant_correction( return astype(z_w_concorr_me_wp, vpfloat) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def compute_contravariant_correction( vn: fa.EdgeKField[wpfloat], ddxn_z_full: fa.EdgeKField[vpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_contravariant_correction_of_w.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_contravariant_correction_of_w.py index 1d5e5a93c9..2e3d14a1f9 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_contravariant_correction_of_w.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_contravariant_correction_of_w.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import C2CE, C2E, C2EDim, Koff -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -43,7 +42,7 @@ def _compute_contravariant_correction_of_w( return w_concorr_c_vp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def compute_contravariant_correction_of_w( e_bln_c_s: gtx.Field[gtx.Dims[dims.CEDim], wpfloat], z_w_concorr_me: fa.EdgeKField[vpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_contravariant_correction_of_w_for_lower_boundary.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_contravariant_correction_of_w_for_lower_boundary.py index b6620e07b0..2dbc74114a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_contravariant_correction_of_w_for_lower_boundary.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_contravariant_correction_of_w_for_lower_boundary.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import C2CE, C2E, C2EDim, Koff -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -45,7 +44,7 @@ def _compute_contravariant_correction_of_w_for_lower_boundary( return w_concorr_c_vp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def compute_contravariant_correction_of_w_for_lower_boundary( e_bln_c_s: gtx.Field[gtx.Dims[dims.CEDim], wpfloat], z_w_concorr_me: fa.EdgeKField[vpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_explicit_part_for_rho_and_exner.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_explicit_part_for_rho_and_exner.py index 050c20f843..7032c3cc5c 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_explicit_part_for_rho_and_exner.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_explicit_part_for_rho_and_exner.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import Koff -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -51,7 +50,7 @@ def _compute_explicit_part_for_rho_and_exner( return z_rho_expl_wp, z_exner_expl_wp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def compute_explicit_part_for_rho_and_exner( z_rho_expl: fa.CellKField[wpfloat], z_exner_expl: fa.CellKField[wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_explicit_vertical_wind_from_advection_and_vertical_wind_density.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_explicit_vertical_wind_from_advection_and_vertical_wind_density.py index 1a9d835b48..dc351c2e60 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_explicit_vertical_wind_from_advection_and_vertical_wind_density.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_explicit_vertical_wind_from_advection_and_vertical_wind_density.py @@ -11,7 +11,6 @@ from gt4py.next.ffront.fbuiltins import astype from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -43,7 +42,7 @@ def _compute_explicit_vertical_wind_from_advection_and_vertical_wind_density( return z_w_expl_wp, z_contr_w_fl_l_wp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def compute_explicit_vertical_wind_from_advection_and_vertical_wind_density( z_w_expl: fa.CellKField[wpfloat], w_nnow: fa.CellKField[wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_explicit_vertical_wind_speed_and_vertical_wind_times_density.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_explicit_vertical_wind_speed_and_vertical_wind_times_density.py index f16801dfdc..9ccfa0ae5e 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_explicit_vertical_wind_speed_and_vertical_wind_times_density.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_explicit_vertical_wind_speed_and_vertical_wind_times_density.py @@ -11,7 +11,6 @@ from gt4py.next.ffront.fbuiltins import astype from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -36,7 +35,7 @@ def _compute_explicit_vertical_wind_speed_and_vertical_wind_times_density( return z_w_expl_wp, z_contr_w_fl_l_wp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def compute_explicit_vertical_wind_speed_and_vertical_wind_times_density( z_w_expl: fa.CellKField[wpfloat], w_nnow: fa.CellKField[wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_first_vertical_derivative.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_first_vertical_derivative.py index 495c48ee35..01e7b7b43e 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_first_vertical_derivative.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_first_vertical_derivative.py @@ -11,7 +11,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import Koff -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat @@ -25,7 +24,7 @@ def _compute_first_vertical_derivative( return z_dexner_dz_c_1 -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def compute_first_vertical_derivative( z_exner_ic: fa.CellKField[vpfloat], inv_ddqz_z_full: fa.CellKField[vpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_advection_of_rho_and_theta.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_advection_of_rho_and_theta.py index 22da6a6090..ae2852cb6d 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_advection_of_rho_and_theta.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_advection_of_rho_and_theta.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import E2C, E2EC -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -172,7 +171,7 @@ def _compute_horizontal_advection_of_rho_and_theta( return z_rho_e, z_theta_v_e -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def compute_horizontal_advection_of_rho_and_theta( p_vn: fa.EdgeKField[wpfloat], p_vt: fa.EdgeKField[vpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_kinetic_energy.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_kinetic_energy.py index eaa2181b79..af39b7e49e 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_kinetic_energy.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_kinetic_energy.py @@ -11,7 +11,6 @@ from gt4py.next.ffront.fbuiltins import astype from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -31,7 +30,7 @@ def _compute_horizontal_kinetic_energy( return astype(vn_ie_wp, vpfloat), z_vt_ie_vp, astype(z_kin_hor_e_wp, vpfloat) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def compute_horizontal_kinetic_energy( vn: fa.EdgeKField[wpfloat], vt: fa.EdgeKField[vpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py index 6da2afa19b..f6a94d2a13 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py @@ -16,7 +16,6 @@ ) from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -58,7 +57,7 @@ def _compute_maximum_cfl_and_clip_contravariant_vertical_velocity( return cfl_clipping, vcfl_vp, astype(z_w_con_c_wp, vpfloat) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def compute_maximum_cfl_and_clip_contravariant_vertical_velocity( ddqz_z_half: fa.CellKField[vpfloat], z_w_con_c: fa.CellKField[vpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers.py index af7978c0a8..733116fb69 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers.py @@ -15,7 +15,6 @@ ) from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import Koff -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -41,7 +40,7 @@ def _compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers( return rho_ic, z_rth_pr_1, z_rth_pr_2 -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers( wgtfac_c: fa.CellKField[vpfloat], rho: fa.CellKField[wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_solver_coefficients_matrix.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_solver_coefficients_matrix.py index 2a64bad0a9..f69bb3bd42 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_solver_coefficients_matrix.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_solver_coefficients_matrix.py @@ -11,7 +11,6 @@ from gt4py.next.ffront.fbuiltins import astype from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -36,7 +35,7 @@ def _compute_solver_coefficients_matrix( return astype((z_beta_wp, z_alpha_wp), vpfloat) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def compute_solver_coefficients_matrix( z_beta: fa.CellKField[vpfloat], exner_nnow: fa.CellKField[wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_virtual_potential_temperatures_and_pressure_gradient.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_virtual_potential_temperatures_and_pressure_gradient.py index 36d1403d46..1343c5c462 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_virtual_potential_temperatures_and_pressure_gradient.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_virtual_potential_temperatures_and_pressure_gradient.py @@ -15,7 +15,6 @@ ) from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import Koff -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -44,7 +43,7 @@ def _compute_virtual_potential_temperatures_and_pressure_gradient( return z_theta_v_pr_ic_vp, theta_v_ic_wp, astype(z_th_ddz_exner_c_wp, vpfloat) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def compute_virtual_potential_temperatures_and_pressure_gradient( wgtfac_c: fa.CellKField[vpfloat], z_rth_pr_2: fa.CellKField[vpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/correct_contravariant_vertical_velocity.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/correct_contravariant_vertical_velocity.py index 033e2acf8d..601bbf0de9 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/correct_contravariant_vertical_velocity.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/correct_contravariant_vertical_velocity.py @@ -10,7 +10,6 @@ from gt4py.next.ffront.decorator import field_operator, program from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat @@ -24,7 +23,7 @@ def _correct_contravariant_vertical_velocity( return z_w_con_c_vp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def correct_contravariant_vertical_velocity( w_concorr_c: fa.CellKField[vpfloat], z_w_con_c: fa.CellKField[vpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_at_top.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_at_top.py index b03e6ffab9..afe7b0692b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_at_top.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_at_top.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import Koff -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -33,7 +32,7 @@ def _extrapolate_at_top( return astype(vn_ie_wp, vpfloat) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def extrapolate_at_top( wgtfacq_e: fa.EdgeKField[vpfloat], vn: fa.EdgeKField[wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_temporally_exner_pressure.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_temporally_exner_pressure.py index b88a6bb0a8..5c6a1cdf60 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_temporally_exner_pressure.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_temporally_exner_pressure.py @@ -11,7 +11,6 @@ from gt4py.next.ffront.fbuiltins import astype from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -32,7 +31,7 @@ def _extrapolate_temporally_exner_pressure( return astype(z_exner_ex_pr_wp, vpfloat), exner_pr_wp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def extrapolate_temporally_exner_pressure( exner_exfac: fa.CellKField[vpfloat], exner: fa.CellKField[wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_solve_nonhydro_stencil_39_40.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_solve_nonhydro_stencil_39_40.py index 88248ecad5..884c183f83 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_solve_nonhydro_stencil_39_40.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_solve_nonhydro_stencil_39_40.py @@ -16,7 +16,6 @@ _compute_contravariant_correction_of_w_for_lower_boundary, ) from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -40,7 +39,7 @@ def _fused_solve_nonhydro_stencil_39_40( return w_concorr_c -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def fused_solve_nonhydro_stencil_39_40( e_bln_c_s: gtx.Field[gtx.Dims[dims.CEDim], wpfloat], z_w_concorr_me: fa.EdgeKField[vpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_15_to_18.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_15_to_18.py index 6c2a60614b..102f23ed39 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_15_to_18.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_15_to_18.py @@ -23,7 +23,6 @@ _interpolate_contravariant_vertical_velocity_to_full_levels, ) from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -153,7 +152,7 @@ def _fused_velocity_advection_stencil_15_to_18( return (z_w_con_c_full, ddt_w_adv) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def fused_velocity_advection_stencil_15_to_18( z_w_con_c: fa.CellKField[vpfloat], w: fa.CellKField[wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_19_to_20.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_19_to_20.py index 3c57cc6928..142dcfd836 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_19_to_20.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_19_to_20.py @@ -20,7 +20,6 @@ _mo_math_divrot_rot_vertex_ri_dsl, ) from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -93,7 +92,7 @@ def _fused_velocity_advection_stencil_19_to_20( return ddt_vn_apc -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def fused_velocity_advection_stencil_19_to_20( vn: fa.EdgeKField[wpfloat], geofac_rot: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_1_to_7.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_1_to_7.py index 5cee4cb9e0..e6cd92bbbd 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_1_to_7.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_1_to_7.py @@ -33,7 +33,6 @@ _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl, ) from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -408,7 +407,7 @@ def _fused_velocity_advection_stencil_1_to_7_restricted( )[1] -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def fused_velocity_advection_stencil_1_to_7( vn: fa.EdgeKField[wpfloat], rbf_vec_coeff_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EDim], wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_8_to_13.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_8_to_13.py index 63bf7a2446..a671a3b97f 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_8_to_13.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_8_to_13.py @@ -26,7 +26,6 @@ _interpolate_to_half_levels_vp, ) from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -199,7 +198,7 @@ def _fused_velocity_advection_stencil_8_to_13_restricted( )[2] -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def fused_velocity_advection_stencil_8_to_13( z_kin_hor_e: fa.EdgeKField[vpfloat], e_bln_c_s: gtx.Field[gtx.Dims[dims.CEDim], wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_8_to_14.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_8_to_14.py index d608161766..7691573a54 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_8_to_14.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_8_to_14.py @@ -29,7 +29,6 @@ _interpolate_to_half_levels_vp, ) from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -111,7 +110,7 @@ def _fused_velocity_advection_stencil_8_to_14( return z_ekinh, cfl_clipping, pre_levelmask, vcfl, z_w_con_c -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def fused_velocity_advection_stencil_8_to_14( z_kin_hor_e: fa.EdgeKField[vpfloat], e_bln_c_s: gtx.Field[gtx.Dims[dims.CEDim], wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_cell_kdim_field_with_zero_vp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_cell_kdim_field_with_zero_vp.py index f78bc47404..4788e0f0e7 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_cell_kdim_field_with_zero_vp.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_cell_kdim_field_with_zero_vp.py @@ -11,7 +11,6 @@ from gt4py.next.ffront.fbuiltins import broadcast from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat @@ -21,7 +20,7 @@ def _init_cell_kdim_field_with_zero_vp() -> fa.CellKField[vpfloat]: return broadcast(vpfloat("0.0"), (dims.CellDim, dims.KDim)) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def init_cell_kdim_field_with_zero_vp( field_with_zero_vp: fa.CellKField[vpfloat], horizontal_start: gtx.int32, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_two_cell_kdim_fields_index_with_zero_vp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_two_cell_kdim_fields_index_with_zero_vp.py index 0e9141621c..78b9e266bd 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_two_cell_kdim_fields_index_with_zero_vp.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_two_cell_kdim_fields_index_with_zero_vp.py @@ -14,7 +14,6 @@ _init_cell_kdim_field_with_zero_vp, ) from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat @@ -39,7 +38,7 @@ def _init_two_cell_kdim_fields_index_with_zero_vp( return field_index_with_zero_1, field_index_with_zero_2 -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def init_two_cell_kdim_fields_index_with_zero_vp( field_index_with_zero_1: fa.CellKField[vpfloat], field_index_with_zero_2: fa.CellKField[vpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_to_half_levels_vp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_to_half_levels_vp.py index d3893178d2..ee02f5bb16 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_to_half_levels_vp.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_to_half_levels_vp.py @@ -11,7 +11,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import Koff -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat @@ -27,7 +26,7 @@ def _interpolate_to_half_levels_vp( return interpolation_to_half_levels_vp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def interpolate_to_half_levels_vp( wgtfac_c: fa.CellKField[vpfloat], interpolant: fa.CellKField[vpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_to_surface.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_to_surface.py index cb68a61d27..ec89d20ba9 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_to_surface.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_to_surface.py @@ -11,7 +11,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import Koff -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat @@ -29,7 +28,7 @@ def _interpolate_to_surface( return interpolation_to_surface -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def interpolate_to_surface( wgtfacq_c: fa.CellKField[vpfloat], interpolant: fa.CellKField[vpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py index ab8ac9ee96..9a247262a4 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py @@ -16,7 +16,6 @@ _interpolate_vt_to_interface_edges, ) from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -38,7 +37,7 @@ def _interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges( return vn_ie, z_vt_ie, z_kin_hor_e -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges( wgtfac_e: fa.EdgeKField[vpfloat], vn: fa.EdgeKField[wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/mo_solve_nonhydro_stencil_51.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/mo_solve_nonhydro_stencil_51.py index b7fba3ec9d..78e4238e76 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/mo_solve_nonhydro_stencil_51.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/mo_solve_nonhydro_stencil_51.py @@ -11,7 +11,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import Koff -from icon4py.model.common.settings import backend @field_operator @@ -36,7 +35,7 @@ def _mo_solve_nonhydro_stencil_51( return z_q, w_nnew -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def mo_solve_nonhydro_stencil_51( z_q: fa.CellKField[float], w_nnew: fa.CellKField[float], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/set_lower_boundary_condition_for_w_and_contravariant_correction.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/set_lower_boundary_condition_for_w_and_contravariant_correction.py index 3802a7fffc..11f0655830 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/set_lower_boundary_condition_for_w_and_contravariant_correction.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/set_lower_boundary_condition_for_w_and_contravariant_correction.py @@ -14,7 +14,6 @@ _init_cell_kdim_field_with_zero_wp, ) from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -30,7 +29,7 @@ def _set_lower_boundary_condition_for_w_and_contravariant_correction( return w_nnew_wp, z_contr_w_fl_l_wp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def set_lower_boundary_condition_for_w_and_contravariant_correction( w_nnew: fa.CellKField[wpfloat], z_contr_w_fl_l: fa.CellKField[wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/set_theta_v_prime_ic_at_lower_boundary.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/set_theta_v_prime_ic_at_lower_boundary.py index d0a834022f..f4c1b73779 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/set_theta_v_prime_ic_at_lower_boundary.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/set_theta_v_prime_ic_at_lower_boundary.py @@ -12,7 +12,6 @@ from icon4py.model.atmosphere.dycore.stencils.interpolate_to_surface import _interpolate_to_surface from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -28,7 +27,7 @@ def _set_theta_v_prime_ic_at_lower_boundary( return z_theta_v_pr_ic_vp, astype(theta_v_ic_vp, wpfloat) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def set_theta_v_prime_ic_at_lower_boundary( wgtfacq_c: fa.CellKField[vpfloat], z_rth_pr: fa.CellKField[vpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/update_density_exner_wind.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/update_density_exner_wind.py index 772c1110d0..9592944ba5 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/update_density_exner_wind.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/update_density_exner_wind.py @@ -11,7 +11,6 @@ from icon4py.model.atmosphere.dycore.stencils.update_wind import _update_wind from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat @@ -36,7 +35,7 @@ def _update_density_exner_wind( return rho_new_wp, exner_new_wp, w_new_wp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def update_density_exner_wind( rho_now: fa.CellKField[wpfloat], grf_tend_rho: fa.CellKField[wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/update_wind.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/update_wind.py index 5137c42f66..4942ccfe67 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/update_wind.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/update_wind.py @@ -10,7 +10,6 @@ from gt4py.next.ffront.decorator import field_operator, program from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat @@ -25,7 +24,7 @@ def _update_wind( return w_new_wp -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def update_wind( w_now: fa.CellKField[wpfloat], grf_tend_w: fa.CellKField[wpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py index 7e8408ebfb..a927b04c2d 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py @@ -50,7 +50,6 @@ states as grid_states, vertical as v_grid, ) -from icon4py.model.common.settings import xp from icon4py.model.common.states import prognostic_state as prognostics from icon4py.model.common.utils import gt4py_field_allocation as field_alloc @@ -460,6 +459,7 @@ def run_predictor_step( ) def _update_levmask_from_cfl_clipping(self): + xp = field_alloc.import_array_ns(self._backend) self.levmask = gtx.as_field( domain=(dims.KDim,), data=(xp.any(self.cfl_clipping.ndarray, 0)), dtype=bool ) diff --git a/model/atmosphere/dycore/tests/dycore_tests/test_dycore_utils.py b/model/atmosphere/dycore/tests/dycore_tests/test_dycore_utils.py index d33c839a5f..4050ffc71c 100644 --- a/model/atmosphere/dycore/tests/dycore_tests/test_dycore_utils.py +++ b/model/atmosphere/dycore/tests/dycore_tests/test_dycore_utils.py @@ -11,7 +11,6 @@ from icon4py.model.atmosphere.dycore import dycore_utils from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import simple as simple_grid -from icon4py.model.common.settings import backend from icon4py.model.common.test_utils import helpers @@ -24,7 +23,7 @@ def bdy_divdamp_numpy(coeff: float, field: np.array): return 0.75 / (coeff + constants.DBL_EPS) * np.abs(field) -def test_caclulate_scal_divdamp_order_24(): +def test_caclulate_scal_divdamp_order_24(backend): divdamp_fac_o2 = 3.0 divdamp_order = 24 mean_cell_area = 1000.0 @@ -45,7 +44,7 @@ def test_caclulate_scal_divdamp_order_24(): assert helpers.dallclose(ref, out.asnumpy()) -def test_calculate_scal_divdamp_any_order(): +def test_calculate_scal_divdamp_any_order(backend): divdamp_fac_o2 = 4.2 divdamp_order = 3 mean_cell_area = 1000.0 @@ -65,7 +64,7 @@ def test_calculate_scal_divdamp_any_order(): assert helpers.dallclose(enhanced_factor, out.asnumpy()) -def test_calculate_bdy_divdamp(): +def test_calculate_bdy_divdamp(backend): grid = simple_grid.SimpleGrid() scal_divdamp = helpers.random_field(grid, dims.KDim) out = helpers.zero_field(grid, dims.KDim) @@ -76,7 +75,7 @@ def test_calculate_bdy_divdamp(): assert helpers.dallclose(out.asnumpy(), bdy_divdamp_numpy(coeff, scal_divdamp.asnumpy())) -def test_calculate_divdamp_fields(): +def test_calculate_divdamp_fields(backend): grid = simple_grid.SimpleGrid() divdamp_field = helpers.random_field(grid, dims.KDim) scal_divdamp = helpers.zero_field(grid, dims.KDim) diff --git a/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py index e5850724fd..5f8df30586 100644 --- a/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py @@ -19,7 +19,6 @@ from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import horizontal as h_grid, vertical as v_grid from icon4py.model.common.math import smagorinsky -from icon4py.model.common.settings import backend from icon4py.model.common.test_utils import ( datatest_utils as dt_utils, helpers, @@ -31,9 +30,7 @@ @pytest.mark.datatest def test_validate_divdamp_fields_against_savepoint_values( - grid_savepoint, - savepoint_nonhydro_init, - icon_grid, + grid_savepoint, savepoint_nonhydro_init, icon_grid, backend ): config = solve_nh.NonHydrostaticConfig() divdamp_fac_o2 = 0.032 @@ -150,6 +147,7 @@ def test_nonhydro_predictor_step( ndyn_substeps, at_initial_timestep, caplog, + backend, ): caplog.set_level(logging.DEBUG) config = utils.construct_solve_nh_config(experiment, ndyn_substeps) @@ -537,6 +535,7 @@ def test_nonhydro_corrector_step( ndyn_substeps, at_initial_timestep, caplog, + backend, ): caplog.set_level(logging.DEBUG) config = utils.construct_solve_nh_config(experiment, ndyn_substeps) @@ -744,6 +743,7 @@ def test_run_solve_nonhydro_single_step( savepoint_nonhydro_step_exit, at_initial_timestep, caplog, + backend, ): caplog.set_level(logging.DEBUG) config = utils.construct_solve_nh_config(experiment, ndyn_substeps) @@ -864,6 +864,7 @@ def test_run_solve_nonhydro_multi_step( savepoint_nonhydro_step_exit, experiment, ndyn_substeps, + backend, at_initial_timestep, ): config = utils.construct_solve_nh_config(experiment, ndyn_substeps) diff --git a/model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py b/model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py index 120ed364ab..2120adffe0 100644 --- a/model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py +++ b/model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py @@ -9,6 +9,7 @@ from typing import Final import gt4py.next as gtx +import numpy as np from gt4py.eve.utils import FrozenNamespace from gt4py.next import backend, broadcast from gt4py.next.ffront.fbuiltins import ( @@ -30,7 +31,6 @@ ) from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid, vertical as v_grid -from icon4py.model.common.settings import xp from icon4py.model.common.states import ( diagnostic_state as diagnostics, model, @@ -382,7 +382,7 @@ def run( temperature_list = [self._temperature1, self._temperature2] ncurrent, nnext = 0, 1 for _ in range(self.config.max_iter): - if xp.any( + if np.any( self._newton_iteration_mask.ndarray[ start_cell_nudging:end_cell_local, 0 : self.grid.num_levels ] @@ -428,13 +428,13 @@ def run( nnext = (nnext + 1) % 2 else: break - if xp.any( + if np.any( self._newton_iteration_mask.ndarray[ start_cell_nudging:end_cell_local, 0 : self.grid.num_levels ] ): raise ConvergenceError( - f"Maximum iteration of saturation adjustment ({self.config.max_iter}) is not enough. The max absolute error is {xp.abs(self.new_temperature1.ndarray - self.new_temperature2.ndarray).max()} . Please raise max_iter" + f"Maximum iteration of saturation adjustment ({self.config.max_iter}) is not enough. The max absolute error is {np.abs(self.new_temperature1.ndarray - self.new_temperature2.ndarray).max()} . Please raise max_iter" ) self.update_temperature_qv_qc_tendencies( dtime, diff --git a/model/common/src/icon4py/model/common/decomposition/definitions.py b/model/common/src/icon4py/model/common/decomposition/definitions.py index 9d13a4d7eb..fc77a0b9aa 100644 --- a/model/common/src/icon4py/model/common/decomposition/definitions.py +++ b/model/common/src/icon4py/model/common/decomposition/definitions.py @@ -12,12 +12,13 @@ import logging from dataclasses import dataclass from enum import IntEnum -from typing import Any, Optional, Protocol, Sequence, runtime_checkable +from typing import Any, Optional, Protocol, Sequence, Union, runtime_checkable +import numpy as np from gt4py.next import Dimension from icon4py.model.common import utils -from icon4py.model.common.settings import xp +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc try: @@ -76,7 +77,9 @@ class EntryType(IntEnum): HALO = 2 @utils.chainable - def with_dimension(self, dim: Dimension, global_index: xp.ndarray, owner_mask: xp.ndarray): + def with_dimension( + self, dim: Dimension, global_index: field_alloc.NDArray, owner_mask: field_alloc.NDArray + ): self._global_index[dim] = global_index self._owner_mask[dim] = owner_mask @@ -126,9 +129,15 @@ def local_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL): def _to_local_index(self, dim): data = self._global_index[dim] assert data.ndim == 1 + if isinstance(data, np.ndarray): + import numpy as xp + else: + import cupy as xp + + xp.arange(data.shape[0]) return xp.arange(data.shape[0]) - def owner_mask(self, dim: Dimension) -> xp.ndarray: + def owner_mask(self, dim: Dimension) -> field_alloc.NDArray: return self._owner_mask[dim] def global_index(self, dim: Dimension, entry_type: EntryType = EntryType.ALL): @@ -359,12 +368,12 @@ def get_runtype(with_mpi: bool = False) -> RunType: @functools.singledispatch -def get_processor_properties(runtime: RunType) -> ProcessProperties: +def get_processor_properties(runtime: RunType, comm_id: Union[int, None]) -> ProcessProperties: raise TypeError(f"Cannot define ProcessProperties for ({type(runtime)})") @get_processor_properties.register(SingleNodeRun) -def get_single_node_properties(s: SingleNodeRun) -> ProcessProperties: +def get_single_node_properties(s: SingleNodeRun, comm_id=None) -> ProcessProperties: return SingleNodeProcessProperties() diff --git a/model/common/src/icon4py/model/common/decomposition/mpi_decomposition.py b/model/common/src/icon4py/model/common/decomposition/mpi_decomposition.py index f615f71504..5490c71e09 100644 --- a/model/common/src/icon4py/model/common/decomposition/mpi_decomposition.py +++ b/model/common/src/icon4py/model/common/decomposition/mpi_decomposition.py @@ -19,7 +19,7 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.decomposition import definitions from icon4py.model.common.decomposition.definitions import SingleNodeExchange -from icon4py.model.common.settings import xp +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc try: @@ -200,7 +200,9 @@ def _create_pattern(self, horizontal_dim: Dimension): ) return pattern - def _slice_field_based_on_dim(self, field: Field, dim: definitions.Dimension) -> xp.ndarray: + def _slice_field_based_on_dim( + self, field: Field, dim: definitions.Dimension + ) -> field_alloc.NDArray: """ Slices the field based on the dimension passed in. """ @@ -235,7 +237,7 @@ def exchange(self, dim: definitions.Dimension, *fields: Sequence[Field]): make_field_descriptor( domain_descriptor, f, - arch=Architecture.CPU if isinstance(f, np.ndarray) else Architecture.GPU, + arch=Architecture.CPU if isinstance(f, np.ndarrray) else Architecture.GPU, ) ) for f in sliced_fields diff --git a/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_pressure.py b/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_pressure.py index 95342d7837..1c2510ec16 100644 --- a/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_pressure.py +++ b/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_pressure.py @@ -11,7 +11,6 @@ from gt4py.next.ffront.fbuiltins import exp, sqrt from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta -from icon4py.model.common.settings import backend @scan_operator(axis=dims.KDim, forward=False, init=(0.0, 0.0, True)) @@ -60,7 +59,7 @@ def _diagnose_pressure( return pressure, pressure_ifc -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def diagnose_pressure( ddqz_z_full: fa.CellKField[ta.wpfloat], virtual_temperature: fa.CellKField[ta.wpfloat], diff --git a/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_surface_pressure.py b/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_surface_pressure.py index cf5cc6f781..4bd37c749e 100644 --- a/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_surface_pressure.py +++ b/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_surface_pressure.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta from icon4py.model.common.dimension import Koff -from icon4py.model.common.settings import backend @field_operator @@ -36,7 +35,7 @@ def _diagnose_surface_pressure( return surface_pressure -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def diagnose_surface_pressure( exner: fa.CellKField[ta.wpfloat], virtual_temperature: fa.CellKField[ta.wpfloat], diff --git a/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_temperature.py b/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_temperature.py index 6a7245ecf9..247e23308c 100644 --- a/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_temperature.py +++ b/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_temperature.py @@ -8,7 +8,6 @@ import gt4py.next as gtx from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta -from icon4py.model.common.settings import backend @gtx.field_operator @@ -29,7 +28,7 @@ def _diagnose_virtual_temperature_and_temperature( return virtual_temperature, temperature -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED, backend=backend) +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def diagnose_virtual_temperature_and_temperature( qv: fa.CellKField[ta.wpfloat], # TODO (Chia Rui): This should be changed to a list hydrometeors with mass instead of directly specifying each hydrometeor, as in trHydroMass list in ICON. Otherwise, the input arguments may need to be changed when different microphysics is used. diff --git a/model/common/src/icon4py/model/common/grid/base.py b/model/common/src/icon4py/model/common/grid/base.py index a8051d66ec..86e0472c7d 100644 --- a/model/common/src/icon4py/model/common/grid/base.py +++ b/model/common/src/icon4py/model/common/grid/base.py @@ -17,7 +17,7 @@ from icon4py.model.common import dimension as dims, utils from icon4py.model.common.grid import utils as grid_utils -from icon4py.model.common.settings import xp +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc class MissingConnectivity(ValueError): @@ -113,7 +113,7 @@ def offset_providers(self): return offset_providers @utils.chainable - def with_connectivities(self, connectivity: Dict[gtx.Dimension, np.ndarray]): + def with_connectivities(self, connectivity: Dict[gtx.Dimension, field_alloc.NDArray]): self.connectivities.update({d: k.astype(gtx.int32) for d, k in connectivity.items()}) self.size.update({d: t.shape[1] for d, t in connectivity.items()}) @@ -137,7 +137,7 @@ def _get_offset_provider(self, dim, from_dim, to_dim): dim, self.connectivities[dim].dtype ) return gtx.NeighborTableOffsetProvider( - xp.asarray(self.connectivities[dim]), + self.connectivities[dim], from_dim, to_dim, self.size[dim], diff --git a/model/common/src/icon4py/model/common/grid/geometry.py b/model/common/src/icon4py/model/common/grid/geometry.py index 5ffa4a5840..9d9ba6cc3a 100644 --- a/model/common/src/icon4py/model/common/grid/geometry.py +++ b/model/common/src/icon4py/model/common/grid/geometry.py @@ -8,6 +8,8 @@ import functools from typing import Any, Callable, Literal, Mapping, Optional, Sequence, TypeAlias, TypeVar +# TODO (@halungge ) test on GPU (NEP 18 ?) +import numpy as np from gt4py import next as gtx from gt4py.next import backend as gtx_backend @@ -26,7 +28,6 @@ horizontal as h_grid, icon, ) -from icon4py.model.common.settings import xp from icon4py.model.common.states import factory, model, utils as state_utils @@ -532,7 +533,7 @@ def as_sparse_field( fields = [] for t in data: buffers = list(b.ndarray for b in t) - field = gtx.as_field(target_dims, data=(xp.vstack(buffers).T), dtype=buffers[0].dtype) + field = gtx.as_field(target_dims, data=(np.vstack(buffers).T), dtype=buffers[0].dtype) fields.append(field) return fields @@ -573,7 +574,7 @@ def create_auxiliary_coordinate_arrays_for_orientation( lat = cell_lat.ndarray[e2c_table] lon = cell_lon.ndarray[e2c_table] for i in (0, 1): - boundary_edges = xp.where(e2c_table[:, i] == gm.GridFile.INVALID_INDEX) + boundary_edges = np.where(e2c_table[:, i] == gm.GridFile.INVALID_INDEX) lat[boundary_edges, i] = edge_lat.ndarray[boundary_edges] lon[boundary_edges, i] = edge_lon.ndarray[boundary_edges] diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py index f08c656778..56b616ae4e 100644 --- a/model/common/src/icon4py/model/common/grid/grid_manager.py +++ b/model/common/src/icon4py/model/common/grid/grid_manager.py @@ -14,18 +14,12 @@ import gt4py.next.backend as gtx_backend import numpy as np -import icon4py.model.common.utils as common_utils from icon4py.model.common import dimension as dims, exceptions, type_alias as ta from icon4py.model.common.decomposition import ( definitions as decomposition, ) from icon4py.model.common.grid import base, icon, vertical as v_grid - - -try: - import cupy as xp -except ImportError: - import numpy as xp +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc try: @@ -39,8 +33,6 @@ def __init__(self, *args, **kwargs): raise ModuleNotFoundError("NetCDF4 is not installed.") -NDArray: TypeAlias = Union[np.ndarray, xp.ndarray] - _log = logging.getLogger(__name__) @@ -276,7 +268,7 @@ def int_variable( name: name of the field to read transpose: flag to indicate whether the file should be transposed (for 2d fields) Returns: - np.ndarray: field data + NDArray: field data """ _log.debug(f"reading {name}: transposing = {transpose}") @@ -328,20 +320,20 @@ class IndexTransformation(Protocol): def __call__( self, - array: NDArray, - ) -> NDArray: + array: field_alloc.NDArray, + ) -> field_alloc.NDArray: ... class NoTransformation(IndexTransformation): """Empty implementation of the Protocol. Just return zeros.""" - def __call__(self, array: NDArray): + def __call__(self, array: field_alloc.NDArray): return np.zeros_like(array) class ToZeroBasedIndexTransformation(IndexTransformation): - def __call__(self, array: NDArray): + def __call__(self, array: field_alloc.NDArray): """ Calculate the index offset needed for usage with python. @@ -407,7 +399,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): def __call__(self, backend: Optional[gtx_backend.Backend], limited_area=True): if not self._reader: self.open() - on_gpu = common_utils.gt4py_field_allocation.is_cupy_device(backend) + on_gpu = field_alloc.is_cupy_device(backend) self._grid = self._construct_grid(on_gpu=on_gpu, limited_area=limited_area) self._refinement = self._read_grid_refinement_fields(backend) self._coordinates = self._read_coordinates(backend) @@ -493,8 +485,8 @@ def _read_geometry_fields(self, backend: Optional[gtx_backend.Backend]): def _read_start_end_indices( self, ) -> tuple[ - dict[dims.Dimension : np.ndarray], - dict[dims.Dimension : np.ndarray], + dict[dims.Dimension : field_alloc.NDArray], + dict[dims.Dimension : field_alloc.NDArray], dict[dims.Dimension : gtx.int32], ]: """ " @@ -549,14 +541,14 @@ def _read_grid_refinement_fields( self, backend: gtx_backend.Backend, decomposition_info: Optional[decomposition.DecompositionInfo] = None, - ) -> tuple[dict[dims.Dimension : NDArray]]: + ) -> tuple[dict[dims.Dimension : field_alloc.NDArray]]: """ Reads the refinement control fields from the grid file. Refinement control contains the classification of each entry in a field to predefined horizontal grid zones as for example the distance to the boundaries, see [refinement.py](refinement.py) """ - xp = common_utils.gt4py_field_allocation.import_array_ns(backend) + xp = field_alloc.import_array_ns(backend) refinement_control_names = { dims.CellDim: GridRefinementName.CONTROL_CELLS, dims.EdgeDim: GridRefinementName.CONTROL_EDGES, @@ -706,7 +698,9 @@ def _update_size_for_1d_sparse_dims(grid): ) -def _construct_diamond_vertices(e2v: NDArray, c2v: NDArray, e2c: NDArray) -> NDArray: +def _construct_diamond_vertices( + e2v: field_alloc.NDArray, c2v: field_alloc.NDArray, e2c: field_alloc.NDArray +) -> field_alloc.NDArray: r""" Construct the connectivity table for the vertices of a diamond in the ICON triangular grid. @@ -745,7 +739,9 @@ def _construct_diamond_vertices(e2v: NDArray, c2v: NDArray, e2c: NDArray) -> NDA return np.hstack((e2v, far_indices)) -def _construct_diamond_edges(e2c: NDArray, c2e: NDArray) -> NDArray: +def _construct_diamond_edges( + e2c: field_alloc.NDArray, c2e: field_alloc.NDArray +) -> field_alloc.NDArray: r""" Construct the connectivity table for the edges of a diamond in the ICON triangular grid. @@ -784,7 +780,9 @@ def _construct_diamond_edges(e2c: NDArray, c2e: NDArray) -> NDArray: return e2c2e -def _construct_triangle_edges(c2e2c: NDArray, c2e: NDArray) -> NDArray: +def _construct_triangle_edges( + c2e2c: field_alloc.NDArray, c2e: field_alloc.NDArray +) -> field_alloc.NDArray: r"""Compute the connectivity from a central cell to all neighboring edges of its cell neighbors. ----e3---- ----e7---- @@ -813,7 +811,7 @@ def _construct_triangle_edges(c2e2c: NDArray, c2e: NDArray) -> NDArray: return table -def _construct_butterfly_cells(c2e2c: NDArray) -> NDArray: +def _construct_butterfly_cells(c2e2c: field_alloc.NDArray) -> field_alloc.NDArray: r"""Compute the connectivity from a central cell to all neighboring cells of its cell neighbors. / \ / \ diff --git a/model/common/src/icon4py/model/common/grid/icon.py b/model/common/src/icon4py/model/common/grid/icon.py index 860e0fd7a1..86518b6619 100644 --- a/model/common/src/icon4py/model/common/grid/icon.py +++ b/model/common/src/icon4py/model/common/grid/icon.py @@ -127,7 +127,10 @@ def __eq__(self, other: "IconGrid"): @utils.chainable def with_start_end_indices( - self, dim: gtx.Dimension, start_indices: np.ndarray, end_indices: np.ndarray + self, + dim: gtx.Dimension, + start_indices: np.ndarray, + end_indices: np.ndarray, ): log.debug(f"Using start_indices {dim} {start_indices}, end_indices {dim} {end_indices}") self._start_indices[dim] = start_indices.astype(gtx.int32) diff --git a/model/common/src/icon4py/model/common/grid/refinement.py b/model/common/src/icon4py/model/common/grid/refinement.py index 52c6c446f9..8cd740545f 100644 --- a/model/common/src/icon4py/model/common/grid/refinement.py +++ b/model/common/src/icon4py/model/common/grid/refinement.py @@ -9,11 +9,12 @@ import logging from typing import Final +import numpy as np from gt4py import next as gtx import icon4py.model.common.grid.horizontal as h_grid from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc """ @@ -79,21 +80,23 @@ def is_ordered(self) -> bool: return self.value not in _UNORDERED[self.dim] -def is_unordered_field(field: xp.ndarray, dim: dims.Dimension) -> xp.ndarray: - assert field.dtype in (xp.int32, xp.int64), f"not an integer type {field.dtype}" - return xp.where( - field == _UNORDERED[dim][0], True, xp.where(field == _UNORDERED[dim][1], True, False) +def is_unordered_field(field: field_alloc.NDArray, dim: dims.Dimension) -> field_alloc.NDArray: + assert field.dtype in (gtx.int32, gtx.int64), f"not an integer type {field.dtype}" + return np.where( + field == _UNORDERED[dim][0], True, np.where(field == _UNORDERED[dim][1], True, False) ) -def convert_to_unnested_refinement_values(field: xp.ndarray, dim: dims.Dimension) -> xp.ndarray: +def convert_to_unnested_refinement_values( + field: field_alloc.NDArray, dim: dims.Dimension +) -> field_alloc.NDArray: """Convenience function that converts the grid refinement value from a coarser parent grid to the canonical values used in an unnested setup. The nested values are used for example in the radiation grids. """ - assert field.dtype in (xp.int32, xp.int64), f"not an integer type {field.dtype}" - return xp.where(field == _UNORDERED[dim][1], 0, xp.where(field < 0, -field, field)) + assert field.dtype in (gtx.int32, gtx.int64), f"not an integer type {field.dtype}" + return np.where(field == _UNORDERED[dim][1], 0, np.where(field < 0, -field, field)) def refine_control_value(dim: gtx.Dimension, zone: h_grid.Zone) -> RefinementValue: diff --git a/model/common/src/icon4py/model/common/grid/simple.py b/model/common/src/icon4py/model/common/grid/simple.py index 3e5931049f..32e23acea5 100644 --- a/model/common/src/icon4py/model/common/grid/simple.py +++ b/model/common/src/icon4py/model/common/grid/simple.py @@ -33,6 +33,7 @@ # | 15c \ | 16c \ | 17c \ # 0v 1v 2v 0v from icon4py.model.common.grid.vertical import VerticalGridConfig +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc @dataclasses.dataclass @@ -446,7 +447,7 @@ def num_edges(self) -> int: return self.config.num_edges @property - def diamond_table(self) -> np.ndarray: + def diamond_table(self) -> field_alloc.NDArray: return SimpleGridData.e2c2v_table @property diff --git a/model/common/src/icon4py/model/common/grid/utils.py b/model/common/src/icon4py/model/common/grid/utils.py index 424bd51aca..19821e81b7 100644 --- a/model/common/src/icon4py/model/common/grid/utils.py +++ b/model/common/src/icon4py/model/common/grid/utils.py @@ -6,10 +6,9 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np from gt4py.next import Dimension, NeighborTableOffsetProvider -from icon4py.model.common.settings import xp - def neighbortable_offset_provider_for_1d_sparse_fields( old_shape: tuple[int, int], @@ -17,7 +16,7 @@ def neighbortable_offset_provider_for_1d_sparse_fields( neighbor_axis: Dimension, has_skip_values: bool, ): - table = xp.arange(old_shape[0] * old_shape[1], dtype=gtx.int32).reshape(old_shape) + table = np.arange(old_shape[0] * old_shape[1], dtype=gtx.int32).reshape(old_shape) assert ( table.dtype == gtx.int32 ), 'Neighbor table\'s ("{}" to "{}") data type for 1d sparse fields must be gtx.int32. Instead it\'s "{}"'.format( diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py index feedd66526..5c0b4951fd 100644 --- a/model/common/src/icon4py/model/common/grid/vertical.py +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -14,10 +14,11 @@ from typing import Final import gt4py.next as gtx +import numpy as np import icon4py.model.common.states.metadata as data from icon4py.model.common import dimension as dims, exceptions, field_type_aliases as fa -from icon4py.model.common.settings import xp +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc log = logging.getLogger(__name__) @@ -246,30 +247,32 @@ def size(self, dim: gtx.Dimension) -> int: @classmethod def _determine_start_level_of_moist_physics( - cls, vct_a: xp.ndarray, top_moist_threshold: float, nshift_total: int = 0 + cls, vct_a: field_alloc.NDArray, top_moist_threshold: float, nshift_total: int = 0 ) -> gtx.int32: n_levels = vct_a.shape[0] interface_height = 0.5 * (vct_a[: n_levels - 1 - nshift_total] + vct_a[1 + nshift_total :]) - return gtx.int32(xp.min(xp.where(interface_height < top_moist_threshold)[0]).item()) + return gtx.int32(np.min(np.where(interface_height < top_moist_threshold)[0]).item()) @classmethod - def _determine_damping_height_index(cls, vct_a: xp.ndarray, damping_height: float) -> gtx.int32: + def _determine_damping_height_index( + cls, vct_a: field_alloc.NDArray, damping_height: float + ) -> gtx.int32: assert damping_height >= 0.0, "Damping height must be positive." return ( 0 if damping_height > vct_a[0] - else gtx.int32(xp.argmax(xp.where(vct_a >= damping_height)[0]).item()) + else gtx.int32(np.argmax(np.where(vct_a >= damping_height)[0]).item()) ) @classmethod def _determine_end_index_of_flat_layers( - cls, vct_a: xp.ndarray, flat_height: float + cls, vct_a: field_alloc.NDArray, flat_height: float ) -> gtx.int32: assert flat_height >= 0.0, "Flat surface height must be positive." return ( 0 if flat_height > vct_a[0] - else gtx.int32(xp.max(xp.where(vct_a >= flat_height)[0]).item()) + else gtx.int32(np.max(np.where(vct_a >= flat_height)[0]).item()) ) @@ -293,8 +296,8 @@ def _read_vct_a_and_vct_b_from_file( Returns: one dimensional vct_a and vct_b arrays. """ num_levels_plus_one = num_levels + 1 - vct_a = xp.zeros(num_levels_plus_one, dtype=float) - vct_b = xp.zeros(num_levels_plus_one, dtype=float) + vct_a = np.zeros(num_levels_plus_one, dtype=float) + vct_b = np.zeros(num_levels_plus_one, dtype=float) try: with open(file_path, "r") as vertical_grid_file: # skip the first line that contains titles @@ -358,12 +361,12 @@ def _compute_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KF """ num_levels_plus_one = vertical_config.num_levels + 1 if vertical_config.lowest_layer_thickness > 0.01: - vct_a_exponential_factor = xp.log( + vct_a_exponential_factor = np.log( vertical_config.lowest_layer_thickness / vertical_config.model_top_height - ) / xp.log( + ) / np.log( 2.0 / math.pi - * xp.arccos( + * np.arccos( float(vertical_config.num_levels - 1) ** vertical_config.stretch_factor / float(vertical_config.num_levels) ** vertical_config.stretch_factor ) @@ -374,8 +377,8 @@ def _compute_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KF * ( 2.0 / math.pi - * xp.arccos( - xp.arange(vertical_config.num_levels + 1, dtype=float) + * np.arccos( + np.arange(vertical_config.num_levels + 1, dtype=float) ** vertical_config.stretch_factor / float(vertical_config.num_levels) ** vertical_config.stretch_factor ) @@ -389,10 +392,10 @@ def _compute_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KF < 0.5 * vertical_config.top_height_limit_for_maximal_layer_thickness ): layer_thickness = vct_a[: vertical_config.num_levels] - vct_a[1:] - lowest_level_exceeding_limit = xp.max( - xp.where(layer_thickness > vertical_config.maximal_layer_thickness) + lowest_level_exceeding_limit = np.max( + np.where(layer_thickness > vertical_config.maximal_layer_thickness) ) - modified_vct_a = xp.zeros(num_levels_plus_one, dtype=float) + modified_vct_a = np.zeros(num_levels_plus_one, dtype=float) lowest_level_unmodified_thickness = 0 shifted_levels = 0 for k in range(vertical_config.num_levels - 1, -1, -1): @@ -400,7 +403,7 @@ def _compute_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KF modified_vct_a[k + 1] < vertical_config.top_height_limit_for_maximal_layer_thickness ): - modified_vct_a[k] = modified_vct_a[k + 1] + xp.minimum( + modified_vct_a[k] = modified_vct_a[k + 1] + np.minimum( vertical_config.maximal_layer_thickness, layer_thickness[k] ) elif lowest_level_unmodified_thickness == 0: @@ -431,7 +434,7 @@ def _compute_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KF for k in range(vertical_config.num_levels - 1, -1, -1): if vct_a[k + 1] < vertical_config.top_height_limit_for_maximal_layer_thickness: - vct_a[k] = vct_a[k + 1] + xp.minimum( + vct_a[k] = vct_a[k + 1] + np.minimum( vertical_config.maximal_layer_thickness, layer_thickness[k] ) else: @@ -454,7 +457,7 @@ def _compute_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KF ): modified_vct_a[k] = vct_a[k] else: - modified_layer_thickness = xp.minimum( + modified_layer_thickness = np.minimum( 1.025 * (vct_a[k] - vct_a[k + 1]), 1.025 * ( @@ -467,7 +470,7 @@ def _compute_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KF ) * (modified_vct_a[k + 1] - modified_vct_a[k + 2]), ) - modified_vct_a[k] = xp.minimum( + modified_vct_a[k] = np.minimum( vct_a[k], modified_vct_a[k + 1] + modified_layer_thickness ) if modified_vct_a[0] == vct_a[0]: @@ -484,12 +487,12 @@ def _compute_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KF else: vct_a = ( vertical_config.model_top_height - * (float(vertical_config.num_levels) - xp.arange(num_levels_plus_one, dtype=float)) + * (float(vertical_config.num_levels) - np.arange(num_levels_plus_one, dtype=float)) / float(vertical_config.num_levels) ) - vct_b = xp.exp(-vct_a / 5000.0) + vct_b = np.exp(-vct_a / 5000.0) - if not xp.allclose(vct_a[0], vertical_config.model_top_height): + if not np.allclose(vct_a[0], vertical_config.model_top_height): log.warning( f" Warning. vct_a[0], {vct_a[0]}, is not equal to model top height, {vertical_config.model_top_height}, of vertical configuration. Please consider changing the vertical setting." ) diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py index 14464a00cd..8244bc5880 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py @@ -19,16 +19,16 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.dimension import C2E, V2E from icon4py.model.common.grid import grid_manager as gm -from icon4py.model.common.utils import gt4py_field_allocation as alloc +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc def compute_c_lin_e( - edge_cell_length: alloc.NDArray, - inv_dual_edge_length: alloc.NDArray, - edge_owner_mask: alloc.NDArray, - horizontal_start: gtx.int32, + edge_cell_length: field_alloc.NDArray, + inv_dual_edge_length: field_alloc.NDArray, + edge_owner_mask: field_alloc.NDArray, + horizontal_start: np.int32, array_ns: ModuleType = np, -) -> alloc.NDArray: +) -> field_alloc.NDArray: """ Compute E2C average inverse distance. @@ -91,14 +91,14 @@ def compute_geofac_rot( def compute_geofac_n2s( - dual_edge_length: alloc.NDArray, - geofac_div: alloc.NDArray, - c2e: alloc.NDArray, - e2c: alloc.NDArray, - c2e2c: alloc.NDArray, - horizontal_start: gtx.int32, + dual_edge_length: field_alloc.NDArray, + geofac_div: field_alloc.NDArray, + c2e: field_alloc.NDArray, + e2c: field_alloc.NDArray, + c2e2c: field_alloc.NDArray, + horizontal_start: np.int32, array_ns: ModuleType = np, -) -> alloc.NDArray: +) -> field_alloc.NDArray: """ Compute geometric factor for nabla2-scalar. @@ -147,14 +147,14 @@ def compute_geofac_n2s( def _compute_primal_normal_ec( - primal_normal_cell_x: alloc.NDArray, - primal_normal_cell_y: alloc.NDArray, - owner_mask: alloc.NDArray, - c2e: alloc.NDArray, - e2c: alloc.NDArray, - horizontal_start: gtx.int32, + primal_normal_cell_x: field_alloc.NDArray, + primal_normal_cell_y: field_alloc.NDArray, + owner_mask: field_alloc.NDArray, + c2e: field_alloc.NDArray, + e2c: field_alloc.NDArray, + horizontal_start: np.int32, array_ns: ModuleType = np, -) -> alloc.NDArray: +) -> field_alloc.NDArray: """ Compute primal_normal_ec. @@ -201,15 +201,15 @@ def _compute_primal_normal_ec( def _compute_geofac_grg( - primal_normal_ec: alloc.NDArray, - geofac_div: alloc.NDArray, - c_lin_e: alloc.NDArray, - c2e: alloc.NDArray, - e2c: alloc.NDArray, - c2e2c: alloc.NDArray, - horizontal_start: gtx.int32, + primal_normal_ec: field_alloc.NDArray, + geofac_div: field_alloc.NDArray, + c_lin_e: field_alloc.NDArray, + c2e: field_alloc.NDArray, + e2c: field_alloc.NDArray, + c2e2c: field_alloc.NDArray, + horizontal_start: np.int32, array_ns: ModuleType = np, -) -> tuple[alloc.NDArray, alloc.NDArray]: +) -> tuple[field_alloc.NDArray, field_alloc.NDArray]: """ Compute geometrical factor for Green-Gauss gradient. @@ -262,17 +262,17 @@ def _compute_geofac_grg( def compute_geofac_grg( - primal_normal_cell_x: alloc.NDArray, - primal_normal_cell_y: alloc.NDArray, - owner_mask: alloc.NDArray, - geofac_div: alloc.NDArray, - c_lin_e: alloc.NDArray, - c2e: alloc.NDArray, - e2c: alloc.NDArray, - c2e2c: alloc.NDArray, + primal_normal_cell_x: field_alloc.NDArray, + primal_normal_cell_y: field_alloc.NDArray, + owner_mask: field_alloc.NDArray, + geofac_div: field_alloc.NDArray, + c_lin_e: field_alloc.NDArray, + c2e: field_alloc.NDArray, + e2c: field_alloc.NDArray, + c2e2c: field_alloc.NDArray, horizontal_start: gtx.int32, array_ns: ModuleType = np, -) -> tuple[alloc.NDArray, alloc.NDArray]: +) -> tuple[field_alloc.NDArray, field_alloc.NDArray]: primal_normal_ec = functools.partial(_compute_primal_normal_ec, array_ns=array_ns)( primal_normal_cell_x, primal_normal_cell_y, owner_mask, c2e, e2c, horizontal_start ) @@ -282,15 +282,15 @@ def compute_geofac_grg( def compute_geofac_grdiv( - geofac_div: alloc.NDArray, - inv_dual_edge_length: alloc.NDArray, - owner_mask: alloc.NDArray, - c2e: alloc.NDArray, - e2c: alloc.NDArray, - e2c2e: alloc.NDArray, - horizontal_start: gtx.int32, + geofac_div: field_alloc.NDArray, + inv_dual_edge_length: field_alloc.NDArray, + owner_mask: field_alloc.NDArray, + c2e: field_alloc.NDArray, + e2c: field_alloc.NDArray, + e2c2e: field_alloc.NDArray, + horizontal_start: np.int32, array_ns: ModuleType = np, -) -> alloc.NDArray: +) -> field_alloc.NDArray: """ Compute geometrical factor for gradient of divergence (triangles only). @@ -345,12 +345,12 @@ def compute_geofac_grdiv( def rotate_latlon( - lat: alloc.NDArray, - lon: alloc.NDArray, - pollat: alloc.NDArray, - pollon: alloc.NDArray, + lat: field_alloc.NDArray, + lon: field_alloc.NDArray, + pollat: field_alloc.NDArray, + pollon: field_alloc.NDArray, array_ns: ModuleType = np, -) -> tuple[alloc.NDArray, alloc.NDArray]: +) -> tuple[field_alloc.NDArray, field_alloc.NDArray]: """ (Compute rotation of lattitude and longitude.) @@ -384,13 +384,13 @@ def rotate_latlon( def _weighting_factors( - ytemp: alloc.NDArray, - xtemp: alloc.NDArray, - yloc: alloc.NDArray, - xloc: alloc.NDArray, + ytemp: field_alloc.NDArray, + xtemp: field_alloc.NDArray, + yloc: field_alloc.NDArray, + xloc: field_alloc.NDArray, wgt_loc: ta.wpfloat, array_ns: ModuleType = np, -) -> alloc.NDArray: +) -> field_alloc.NDArray: """ Compute weighting factors. The weighting factors are based on the requirement that sum(w(i)*x(i)) = 0 @@ -456,13 +456,13 @@ def _weighting_factors( def _compute_c_bln_avg( - c2e2c: alloc.NDArray, - lat: alloc.NDArray, - lon: alloc.NDArray, + c2e2c: field_alloc.NDArray, + lat: field_alloc.NDArray, + lon: field_alloc.NDArray, divavg_cntrwgt: ta.wpfloat, - horizontal_start: gtx.int32, + horizontal_start: np.int32, array_ns: ModuleType = np, -) -> alloc.NDArray: +) -> field_alloc.NDArray: """ Compute bilinear cell average weight. @@ -502,15 +502,15 @@ def _compute_c_bln_avg( def _force_mass_conservation_to_c_bln_avg( - c2e2c0: alloc.NDArray, - c_bln_avg: alloc.NDArray, - cell_areas: alloc.NDArray, - cell_owner_mask: alloc.NDArray, + c2e2c0: field_alloc.NDArray, + c_bln_avg: field_alloc.NDArray, + cell_areas: field_alloc.NDArray, + cell_owner_mask: field_alloc.NDArray, divavg_cntrwgt: ta.wpfloat, horizontal_start: gtx.int32, array_ns: ModuleType = np, niter: int = 1000, -) -> alloc.NDArray: +) -> field_alloc.NDArray: """ Iteratively enforce mass conservation to the input field c_bln_avg. @@ -536,7 +536,7 @@ def _force_mass_conservation_to_c_bln_avg( def _compute_local_weights( c_bln_avg, cell_areas, c2e2c0, inverse_neighbor_idx - ) -> alloc.NDArray: + ) -> field_alloc.NDArray: """ Compute the total weight which each local point contributes to the sum. @@ -551,8 +551,10 @@ def _compute_local_weights( return weights def _compute_residual_to_mass_conservation( - owner_mask: alloc.NDArray, local_weight: alloc.NDArray, cell_area: alloc.NDArray - ) -> alloc.NDArray: + owner_mask: field_alloc.NDArray, + local_weight: field_alloc.NDArray, + cell_area: field_alloc.NDArray, + ) -> field_alloc.NDArray: """The local_weight weighted by the area should be 1. We compute how far we are off that weight.""" horizontal_size = local_weight.shape[0] assert horizontal_size == owner_mask.shape[0], "Fields do not have the same shape" @@ -561,12 +563,12 @@ def _compute_residual_to_mass_conservation( return residual def _apply_correction( - c_bln_avg: alloc.NDArray, - residual: alloc.NDArray, - c2e2c0: alloc.NDArray, + c_bln_avg: field_alloc.NDArray, + residual: field_alloc.NDArray, + c2e2c0: field_alloc.NDArray, divavg_cntrwgt: float, - horizontal_start: alloc.NDArray, - ) -> alloc.NDArray: + horizontal_start: gtx.int32, + ) -> field_alloc.NDArray: """Apply correction to local weigths based on the computed residuals.""" maxwgt_loc = divavg_cntrwgt + 0.003 minwgt_loc = divavg_cntrwgt - 0.003 @@ -590,11 +592,11 @@ def _apply_correction( return c_bln_avg def _enforce_mass_conservation( - c_bln_avg: alloc.NDArray, - residual: alloc.NDArray, - owner_mask: alloc.NDArray, + c_bln_avg: field_alloc.NDArray, + residual: field_alloc.NDArray, + owner_mask: field_alloc.NDArray, horizontal_start: gtx.int32, - ) -> alloc.NDArray: + ) -> field_alloc.NDArray: """Enforce the mass conservation condition on the local cells by forcefully subtracting the residual from the central field contribution.""" c_bln_avg[horizontal_start:, 0] = array_ns.where( @@ -637,16 +639,16 @@ def _enforce_mass_conservation( def compute_mass_conserving_bilinear_cell_average_weight( - c2e2c0: alloc.NDArray, - lat: alloc.NDArray, - lon: alloc.NDArray, - cell_areas: alloc.NDArray, - cell_owner_mask: alloc.NDArray, + c2e2c0: field_alloc.NDArray, + lat: field_alloc.NDArray, + lon: field_alloc.NDArray, + cell_areas: field_alloc.NDArray, + cell_owner_mask: field_alloc.NDArray, divavg_cntrwgt: ta.wpfloat, horizontal_start: gtx.int32, horizontal_start_level_3: gtx.int32, array_ns: ModuleType = np, -) -> alloc.NDArray: +) -> field_alloc.NDArray: c_bln_avg = _compute_c_bln_avg( c2e2c0[:, 1:], lat, lon, divavg_cntrwgt, horizontal_start, array_ns ) @@ -816,15 +818,15 @@ def compute_e_flx_avg( def compute_cells_aw_verts( - dual_area: np.ndarray, - edge_vert_length: np.ndarray, - edge_cell_length: np.ndarray, - v2e: np.ndarray, - e2v: np.ndarray, - v2c: np.ndarray, - e2c: np.ndarray, + dual_area: field_alloc.NDArray, + edge_vert_length: field_alloc.NDArray, + edge_cell_length: field_alloc.NDArray, + v2e: field_alloc.NDArray, + e2v: field_alloc.NDArray, + v2c: field_alloc.NDArray, + e2c: field_alloc.NDArray, horizontal_start: gtx.int32, -) -> np.ndarray: +) -> field_alloc.NDArray: """ Compute cells_aw_verts. @@ -877,13 +879,13 @@ def compute_cells_aw_verts( def compute_e_bln_c_s( - c2e: np.ndarray, - cells_lat: np.ndarray, - cells_lon: np.ndarray, - edges_lat: np.ndarray, - edges_lon: np.ndarray, + c2e: field_alloc.NDArray, + cells_lat: field_alloc.NDArray, + cells_lon: field_alloc.NDArray, + edges_lat: field_alloc.NDArray, + edges_lon: field_alloc.NDArray, weighting_factor: float, -) -> np.ndarray: +) -> field_alloc.NDArray: """ Compute e_bln_c_s. diff --git a/model/common/src/icon4py/model/common/interpolation/stencils/cell_2_edge_interpolation.py b/model/common/src/icon4py/model/common/interpolation/stencils/cell_2_edge_interpolation.py index 7848e2466c..8948fddfd4 100644 --- a/model/common/src/icon4py/model/common/interpolation/stencils/cell_2_edge_interpolation.py +++ b/model/common/src/icon4py/model/common/interpolation/stencils/cell_2_edge_interpolation.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta from icon4py.model.common.dimension import E2C, E2CDim -from icon4py.model.common.settings import backend @field_operator @@ -31,7 +30,7 @@ def _cell_2_edge_interpolation( return neighbor_sum(in_field(E2C) * coeff, axis=E2CDim) -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def cell_2_edge_interpolation( in_field: fa.CellKField[ta.wpfloat], coeff: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.wpfloat], diff --git a/model/common/src/icon4py/model/common/interpolation/stencils/compute_cell_2_vertex_interpolation.py b/model/common/src/icon4py/model/common/interpolation/stencils/compute_cell_2_vertex_interpolation.py index 4e77eaf584..c729c08d7f 100644 --- a/model/common/src/icon4py/model/common/interpolation/stencils/compute_cell_2_vertex_interpolation.py +++ b/model/common/src/icon4py/model/common/interpolation/stencils/compute_cell_2_vertex_interpolation.py @@ -8,7 +8,6 @@ import gt4py.next as gtx from gt4py.next import neighbor_sum -import icon4py.model.common.settings as settings import icon4py.model.common.type_alias as types from icon4py.model.common import dimension as dims from icon4py.model.common.dimension import V2C, V2CDim @@ -23,7 +22,7 @@ def _compute_cell_2_vertex_interpolation( return vert_out -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED, backend=settings.backend) +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_cell_2_vertex_interpolation( cell_in: gtx.Field[[dims.CellDim, dims.KDim], types.wpfloat], c_int: gtx.Field[[dims.VertexDim, V2CDim], types.wpfloat], diff --git a/model/common/src/icon4py/model/common/interpolation/stencils/edge_2_cell_vector_rbf_interpolation.py b/model/common/src/icon4py/model/common/interpolation/stencils/edge_2_cell_vector_rbf_interpolation.py index 647f5a5161..0b4141da16 100644 --- a/model/common/src/icon4py/model/common/interpolation/stencils/edge_2_cell_vector_rbf_interpolation.py +++ b/model/common/src/icon4py/model/common/interpolation/stencils/edge_2_cell_vector_rbf_interpolation.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta from icon4py.model.common.dimension import C2E2C2E, C2E2C2EDim -from icon4py.model.common.settings import backend @field_operator @@ -41,7 +40,7 @@ def _edge_2_cell_vector_rbf_interpolation( return p_u_out, p_v_out -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def edge_2_cell_vector_rbf_interpolation( p_e_in: fa.EdgeKField[ta.wpfloat], ptr_coeff_1: gtx.Field[gtx.Dims[dims.CellDim, C2E2C2EDim], ta.wpfloat], diff --git a/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_cell.py b/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_cell.py index 47b9f3cee1..a6afc9bf68 100644 --- a/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_cell.py +++ b/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_cell.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import C2E, C2EDim -from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat @@ -27,7 +26,7 @@ def _mo_intp_rbf_rbf_vec_interpol_cell( return p_u_out, p_v_out -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def mo_intp_rbf_rbf_vec_interpol_cell( p_vn_in: fa.EdgeKField[wpfloat], ptr_coeff_1: gtx.Field[gtx.Dims[dims.CellDim, C2EDim], wpfloat], diff --git a/model/common/src/icon4py/model/common/math/projection.py b/model/common/src/icon4py/model/common/math/projection.py index e29e1c8a5b..a590b254a9 100644 --- a/model/common/src/icon4py/model/common/math/projection.py +++ b/model/common/src/icon4py/model/common/math/projection.py @@ -8,13 +8,15 @@ import numpy as np +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc + def gnomonic_proj( - lon_c: np.ndarray, - lat_c: np.ndarray, - lon: np.ndarray, - lat: np.ndarray, -) -> tuple[np.ndarray, np.ndarray]: + lon_c: field_alloc.NDArray, + lat_c: field_alloc.NDArray, + lon: field_alloc.NDArray, + lat: field_alloc.NDArray, +) -> tuple[field_alloc.NDArray, field_alloc.NDArray]: """ Compute gnomonic projection. diff --git a/model/common/src/icon4py/model/common/metrics/compute_diffusion_metrics.py b/model/common/src/icon4py/model/common/metrics/compute_diffusion_metrics.py index 6f289626ff..d055ddb711 100644 --- a/model/common/src/icon4py/model/common/metrics/compute_diffusion_metrics.py +++ b/model/common/src/icon4py/model/common/metrics/compute_diffusion_metrics.py @@ -8,15 +8,17 @@ import numpy as np +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc + def _compute_nbidx( k_range: range, - z_mc: np.ndarray, - z_mc_off: np.ndarray, - nbidx: np.ndarray, + z_mc: field_alloc.NDArray, + z_mc_off: field_alloc.NDArray, + nbidx: field_alloc.NDArray, jc: int, nlev: int, -) -> np.ndarray: +) -> field_alloc.NDArray: for ind in range(3): jk_start = nlev - 1 for jk in reversed(k_range): @@ -34,12 +36,12 @@ def _compute_nbidx( def _compute_z_vintcoeff( k_range: range, - z_mc: np.ndarray, - z_mc_off: np.ndarray, - z_vintcoeff: np.ndarray, + z_mc: field_alloc.NDArray, + z_mc_off: field_alloc.NDArray, + z_vintcoeff: field_alloc.NDArray, jc: int, nlev: int, -) -> np.ndarray: +) -> field_alloc.NDArray: for ind in range(3): jk_start = nlev - 1 for jk in reversed(k_range): @@ -60,9 +62,9 @@ def _compute_z_vintcoeff( def _compute_ls_params( k_start: list, k_end: list, - z_maxslp_avg: np.ndarray, - z_maxhgtd_avg: np.ndarray, - c_owner_mask: np.ndarray, + z_maxslp_avg: field_alloc.NDArray, + z_maxhgtd_avg: field_alloc.NDArray, + c_owner_mask: field_alloc.NDArray, thslp_zdiffu: float, thhgtd_zdiffu: float, cell_nudging: int, @@ -92,11 +94,11 @@ def _compute_ls_params( def _compute_k_start_end( - z_mc: np.ndarray, - max_nbhgt: np.ndarray, - z_maxslp_avg: np.ndarray, - z_maxhgtd_avg: np.ndarray, - c_owner_mask: np.ndarray, + z_mc: field_alloc.NDArray, + max_nbhgt: field_alloc.NDArray, + z_maxslp_avg: field_alloc.NDArray, + z_maxhgtd_avg: field_alloc.NDArray, + c_owner_mask: field_alloc.NDArray, thslp_zdiffu: float, thhgtd_zdiffu: float, cell_nudging: int, @@ -127,24 +129,24 @@ def _compute_k_start_end( def compute_diffusion_metrics( - z_mc: np.ndarray, - z_mc_off: np.ndarray, - max_nbhgt: np.ndarray, - c_owner_mask: np.ndarray, - nbidx: np.ndarray, - z_vintcoeff: np.ndarray, - z_maxslp_avg: np.ndarray, - z_maxhgtd_avg: np.ndarray, - mask_hdiff: np.ndarray, - zd_diffcoef_dsl: np.ndarray, - zd_intcoef_dsl: np.ndarray, - zd_vertoffset_dsl: np.ndarray, + z_mc: field_alloc.NDArray, + z_mc_off: field_alloc.NDArray, + max_nbhgt: field_alloc.NDArray, + c_owner_mask: field_alloc.NDArray, + nbidx: field_alloc.NDArray, + z_vintcoeff: field_alloc.NDArray, + z_maxslp_avg: field_alloc.NDArray, + z_maxhgtd_avg: field_alloc.NDArray, + mask_hdiff: field_alloc.NDArray, + zd_diffcoef_dsl: field_alloc.NDArray, + zd_intcoef_dsl: field_alloc.NDArray, + zd_vertoffset_dsl: field_alloc.NDArray, thslp_zdiffu: float, thhgtd_zdiffu: float, cell_nudging: int, n_cells: int, nlev: int, -) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: +) -> tuple[field_alloc.NDArray, field_alloc.NDArray, field_alloc.NDArray, field_alloc.NDArray]: k_start, k_end = _compute_k_start_end( z_mc=z_mc, max_nbhgt=max_nbhgt, diff --git a/model/common/src/icon4py/model/common/metrics/compute_vwind_impl_wgt.py b/model/common/src/icon4py/model/common/metrics/compute_vwind_impl_wgt.py index afa0e11041..80ed196b5b 100644 --- a/model/common/src/icon4py/model/common/metrics/compute_vwind_impl_wgt.py +++ b/model/common/src/icon4py/model/common/metrics/compute_vwind_impl_wgt.py @@ -11,8 +11,8 @@ import icon4py.model.common.field_type_aliases as fa from icon4py.model.common.grid import base as grid from icon4py.model.common.metrics.metric_fields import compute_vwind_impl_wgt_partial -from icon4py.model.common.settings import xp from icon4py.model.common.type_alias import wpfloat +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc def compute_vwind_impl_wgt( @@ -29,7 +29,7 @@ def compute_vwind_impl_wgt( experiment: str, vwind_offctr: float, horizontal_start_cell: int, -) -> np.ndarray: +) -> field_alloc.NDArray: compute_vwind_impl_wgt_partial.with_backend(backend)( z_ddxn_z_half_e=z_ddxn_z_half_e, z_ddxt_z_half_e=z_ddxt_z_half_e, @@ -50,8 +50,8 @@ def compute_vwind_impl_wgt( ) vwind_impl_wgt = ( - xp.amin(vwind_impl_wgt_k.ndarray, axis=1) + np.amin(vwind_impl_wgt_k.ndarray, axis=1) if experiment == global_exp - else xp.amax(vwind_impl_wgt_k.ndarray, axis=1) + else np.amax(vwind_impl_wgt_k.ndarray, axis=1) ) return vwind_impl_wgt diff --git a/model/common/src/icon4py/model/common/metrics/compute_wgtfacq.py b/model/common/src/icon4py/model/common/metrics/compute_wgtfacq.py index cd88743772..f2a905db7c 100644 --- a/model/common/src/icon4py/model/common/metrics/compute_wgtfacq.py +++ b/model/common/src/icon4py/model/common/metrics/compute_wgtfacq.py @@ -8,10 +8,12 @@ import numpy as np +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc + def _compute_z1_z2_z3( - z_ifc: np.ndarray, i1: int, i2: int, i3: int, i4: int -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + z_ifc: field_alloc.NDArray, i1: int, i2: int, i3: int, i4: int +) -> tuple[field_alloc.NDArray, field_alloc.NDArray, field_alloc.NDArray]: z1 = 0.5 * (z_ifc[:, i2] - z_ifc[:, i1]) z2 = 0.5 * (z_ifc[:, i2] + z_ifc[:, i3]) - z_ifc[:, i1] z3 = 0.5 * (z_ifc[:, i3] + z_ifc[:, i4]) - z_ifc[:, i1] @@ -19,9 +21,9 @@ def _compute_z1_z2_z3( def compute_wgtfacq_c_dsl( - z_ifc: np.ndarray, + z_ifc: field_alloc.NDArray, nlev: int, -) -> np.ndarray: +) -> field_alloc.NDArray: """ Compute weighting factor for quadratic interpolation to surface. @@ -48,9 +50,9 @@ def compute_wgtfacq_c_dsl( def compute_wgtfacq_e_dsl( e2c, - z_ifc: np.ndarray, - c_lin_e: np.ndarray, - wgtfacq_c_dsl: np.ndarray, + z_ifc: field_alloc.NDArray, + c_lin_e: field_alloc.NDArray, + wgtfacq_c_dsl: field_alloc.NDArray, n_edges: int, nlev: int, ): diff --git a/model/common/src/icon4py/model/common/metrics/compute_zdiff_gradp_dsl.py b/model/common/src/icon4py/model/common/metrics/compute_zdiff_gradp_dsl.py index 2afe9ae78b..3657c7bd38 100644 --- a/model/common/src/icon4py/model/common/metrics/compute_zdiff_gradp_dsl.py +++ b/model/common/src/icon4py/model/common/metrics/compute_zdiff_gradp_dsl.py @@ -8,19 +8,21 @@ import numpy as np +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc + def compute_zdiff_gradp_dsl( e2c, - z_me: np.ndarray, - z_mc: np.ndarray, - z_ifc: np.ndarray, - flat_idx: np.ndarray, - z_aux2: np.ndarray, + z_me: field_alloc.NDArray, + z_mc: field_alloc.NDArray, + z_ifc: field_alloc.NDArray, + flat_idx: field_alloc.NDArray, + z_aux2: field_alloc.NDArray, nlev: int, horizontal_start: int, horizontal_start_1: int, nedges: int, -) -> np.ndarray: +) -> field_alloc.NDArray: zdiff_gradp = np.zeros_like(z_mc[e2c]) zdiff_gradp[horizontal_start:, :, :] = ( np.expand_dims(z_me, axis=1)[horizontal_start:, :, :] - z_mc[e2c][horizontal_start:, :, :] diff --git a/model/common/src/icon4py/model/common/orchestration/decorator.py b/model/common/src/icon4py/model/common/orchestration/decorator.py index 5b303164db..37fcc8be3f 100644 --- a/model/common/src/icon4py/model/common/orchestration/decorator.py +++ b/model/common/src/icon4py/model/common/orchestration/decorator.py @@ -29,7 +29,7 @@ import numpy as np from gt4py._core import definitions as core_defs -from icon4py.model.common import dimension as dims, settings +from icon4py.model.common import dimension as dims from icon4py.model.common.decomposition import definitions as decomposition from icon4py.model.common.grid import icon as icon_grid from icon4py.model.common.orchestration import dtypes as orchestration_dtypes @@ -88,30 +88,27 @@ def orchestrate( """ def _decorator(fuse_func: Callable[P, R]) -> Callable[P, R]: - if settings.dace_orchestration is not None: - orchestrator_cache = {} # Caching + orchestrator_cache = {} # Caching + self_name = next(iter(inspect.signature(fuse_func).parameters)) - if "dace" not in settings.backend.name.lower(): - raise ValueError( - "DaCe Orchestration works only with DaCe backends. Change the backend to a DaCe supported one." - ) - - self_name = next(iter(inspect.signature(fuse_func).parameters)) - - # If not explicitly set by the user, assume the provided callable is a method - # when its first argument is called 'self' - func_is_method = method or (self_name == "self") + # If not explicitly set by the user, assume the provided callable is a method + # when its first argument is called 'self' + func_is_method = method or (self_name == "self") - if not func_is_method: - raise NotImplementedError( - "The orchestration decorator is only for methods -at least for now-." - ) - - # Add DaCe data types annotations for **all args and kwargs** - dace_annotations = to_dace_annotations(fuse_func) + if not func_is_method: + raise NotImplementedError( + "The orchestration decorator is only for methods -at least for now-." + ) - def wrapper(*args, **kwargs): - self = args[0] + def wrapper(*args, **kwargs): + self = args[0] + if self._orchestration: + # Add DaCe data types annotations for **all args and kwargs** + dace_annotations = to_dace_annotations(fuse_func) + if "dace" not in self._backend.name.lower(): + raise ValueError( + "DaCe Orchestration works only with DaCe backends. Change the backend to a DaCe supported one." + ) exchange_obj = None grid = None @@ -146,6 +143,7 @@ def wrapper(*args, **kwargs): cache_item = orchestrator_cache[unique_id] = parse_compile_cache_sdfg( default_build_folder, + self._backend, exchange_obj, fuse_func, compile_time_args_kwargs, @@ -187,13 +185,20 @@ def wrapper(*args, **kwargs): del sdfg_args[self_name] with dace.config.temporary_config(): - configure_dace_temp_env(default_build_folder) + configure_dace_temp_env(default_build_folder, self._backend) return compiled_sdfg(**sdfg_args) + else: + return fuse_func(*args, **kwargs) - return wrapper + # Pytest does not clear the cache between runs in a proper way -pytest.mark.parametrize(...)-. + # This leads to corrupted cache and subsequent errors. + # To avoid this, we provide a way to clear the cache. + def clear_cache(): + orchestrator_cache.clear() - else: - return fuse_func + wrapper.clear_cache = clear_cache + + return wrapper return _decorator(func) if func else _decorator @@ -389,6 +394,7 @@ def dev_type_from_gt4py_to_dace(device_type: core_defs.DeviceType) -> dace.dtype def parse_compile_cache_sdfg( default_build_folder: Path, + backend, exchange_obj: Optional[decomposition.ExchangeRuntime], fuse_func: Callable, compile_time_args_kwargs: dict[str, Any], @@ -399,7 +405,7 @@ def parse_compile_cache_sdfg( cache = {} with dace.config.temporary_config(): - device_type = configure_dace_temp_env(default_build_folder) + device_type = configure_dace_temp_env(default_build_folder, backend) cache["dace_program"] = dace_program = dace.program( auto_optimize=False, @@ -495,7 +501,7 @@ def get_env_bool(env_var_name: str, default: bool = False) -> bool: value = os.getenv(env_var_name, str(default)).lower() return value in ("true", "1") - def configure_dace_temp_env(default_build_folder: Path) -> core_defs.DeviceType: + def configure_dace_temp_env(default_build_folder: Path, backend) -> core_defs.DeviceType: dace.config.Config.set("default_build_folder", value=str(default_build_folder)) dace.config.Config.set( "compiler", "allow_view_arguments", value=True @@ -504,7 +510,7 @@ def configure_dace_temp_env(default_build_folder: Path) -> core_defs.DeviceType: "optimizer", "automatic_simplification", value=False ) # simplifications & optimizations after placing halo exchanges -need a sequential structure of nested sdfgs- dace.config.Config.set("optimizer", "autooptimize", value=False) - device_type = settings.backend.executor.step.translation.device_type + device_type = backend.executor.step.translation.device_type if device_type == core_defs.DeviceType.CPU: device = "cpu" compiler_args = dace.config.Config.get("compiler", "cpu", "args") diff --git a/model/common/src/icon4py/model/common/settings.py b/model/common/src/icon4py/model/common/settings.py deleted file mode 100644 index f696ef7627..0000000000 --- a/model/common/src/icon4py/model/common/settings.py +++ /dev/null @@ -1,18 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause - -from icon4py.model.common.config import Icon4PyConfig - - -config = Icon4PyConfig() -backend = config.gt4py_runner -dace_orchestration = config.icon4py_dace_orchestration -xp = config.array_ns -device = config.device -limited_area = config.limited_area -parallel_run = config.parallel_run diff --git a/model/common/src/icon4py/model/common/states/factory.py b/model/common/src/icon4py/model/common/states/factory.py index b44394cabe..b759e67709 100644 --- a/model/common/src/icon4py/model/common/states/factory.py +++ b/model/common/src/icon4py/model/common/states/factory.py @@ -69,7 +69,6 @@ icon as icon_grid, vertical as v_grid, ) -from icon4py.model.common.settings import xp from icon4py.model.common.states import model, utils as state_utils from icon4py.model.common.utils import gt4py_field_allocation as field_alloc @@ -583,7 +582,7 @@ def _compute( args.update(self._params) results = self._func(**args) ## TODO: can the order of return values be checked? - results = (results,) if isinstance(results, xp.ndarray) else results + results = (results,) if isinstance(results, field_alloc.NDArray) else results self._fields = { k: gtx.as_field(tuple(self._dims), results[i], allocator=backend) for i, k in enumerate(self.fields) diff --git a/model/common/src/icon4py/model/common/states/utils.py b/model/common/src/icon4py/model/common/states/utils.py index b5a6b3f68f..cfad64e579 100644 --- a/model/common/src/icon4py/model/common/states/utils.py +++ b/model/common/src/icon4py/model/common/states/utils.py @@ -11,7 +11,7 @@ import xarray as xa from icon4py.model.common import dimension as dims, type_alias as ta -from icon4py.model.common.settings import xp +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc DimT = TypeVar("DimT", dims.KDim, dims.KHalfDim, dims.CellDim, dims.EdgeDim, dims.VertexDim) @@ -23,9 +23,9 @@ T = TypeVar("T", ta.wpfloat, ta.vpfloat, float, bool, gtx.int32, gtx.int64) GTXFieldType: TypeAlias = gtx.Field[Sequence[gtx.Dims[DimT]], T] -FieldType: TypeAlias = Union[gtx.Field[Sequence[gtx.Dims[DimT]], T], xp.ndarray] +FieldType: TypeAlias = Union[gtx.Field[Sequence[gtx.Dims[DimT]], T], field_alloc.NDArray] def to_data_array(field: FieldType, attrs: dict): - data = field if isinstance(field, xp.ndarray) else field.ndarray + data = field_alloc.as_numpy(field) return xa.DataArray(data, attrs=attrs) diff --git a/model/common/src/icon4py/model/common/test_utils/helpers.py b/model/common/src/icon4py/model/common/test_utils/helpers.py index 5c0b33d478..d021af0455 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -18,7 +18,7 @@ from gt4py.next.ffront.decorator import Program from typing_extensions import Buffer -from icon4py.model.common.settings import xp +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc from ..grid.base import BaseGrid from ..type_alias import wpfloat @@ -73,6 +73,7 @@ def random_mask( *dims: gt_common.Dimension, dtype: Optional[npt.DTypeLike] = None, extend: Optional[dict[gt_common.Dimension, int]] = None, + backend=None, ) -> gt_common.Field: rng = np.random.default_rng() shape = _shape(grid, *dims, extend=extend) @@ -83,7 +84,7 @@ def random_mask( arr = np.reshape(arr, newshape=shape) if dtype: arr = arr.astype(dtype) - return as_field(dims, arr) + return as_field(dims, arr, allocator=backend) def random_field( @@ -93,13 +94,14 @@ def random_field( high: float = 1.0, extend: Optional[dict[gt_common.Dimension, int]] = None, dtype: Optional[npt.DTypeLike] = None, + backend=None, ) -> gt_common.Field: arr = np.random.default_rng().uniform( low=low, high=high, size=_shape(grid, *dims, extend=extend) ) if dtype: arr = arr.astype(dtype) - return as_field(dims, arr) + return as_field(dims, arr, allocator=backend) def zero_field( @@ -107,16 +109,25 @@ def zero_field( *dims: gt_common.Dimension, dtype=wpfloat, extend: Optional[dict[gt_common.Dimension, int]] = None, + backend=None, ) -> gt_common.Field: - return as_field(dims, xp.zeros(shape=_shape(grid, *dims, extend=extend), dtype=dtype)) + if extend is not None: + field_domain = { + dim: (0, grid.size[dim] + extend[dim] if dim in extend.keys() else grid.size[dim]) + for dim in dims + } + else: + field_domain = {dim: (0, grid.size[dim]) for dim in dims} + return constructors.zeros(field_domain, dtype=dtype, allocator=backend) def constant_field( - grid: BaseGrid, value: float, *dims: gt_common.Dimension, dtype=wpfloat + grid: BaseGrid, value: float, *dims: gt_common.Dimension, dtype=wpfloat, backend=None ) -> gt_common.Field: return as_field( dims, value * np.ones(shape=tuple(map(lambda x: grid.size[x], dims)), dtype=dtype), + allocator=backend, ) @@ -126,15 +137,19 @@ def as_1D_sparse_field(field: gt_common.Field, target_dim: gt_common.Dimension) return numpy_to_1D_sparse_field(buffer, target_dim) -def numpy_to_1D_sparse_field(field: np.ndarray, dim: gt_common.Dimension) -> gt_common.Field: +def numpy_to_1D_sparse_field( + field: field_alloc.NDArray, dim: gt_common.Dimension, backend=None +) -> gt_common.Field: """Convert a 2D sparse field to a 1D flattened (Felix-style) sparse field.""" old_shape = field.shape assert len(old_shape) == 2 new_shape = (old_shape[0] * old_shape[1],) - return as_field((dim,), field.reshape(new_shape)) + return as_field((dim,), field.reshape(new_shape), allocator=backend) -def flatten_first_two_dims(*dims: gt_common.Dimension, field: gt_common.Field) -> gt_common.Field: +def flatten_first_two_dims( + *dims: gt_common.Dimension, field: gt_common.Field, backend=None +) -> gt_common.Field: """Convert a n-D sparse field to a (n-1)-D flattened (Felix-style) sparse field.""" buffer = field.ndarray old_shape = buffer.shape @@ -143,7 +158,7 @@ def flatten_first_two_dims(*dims: gt_common.Dimension, field: gt_common.Field) - flattened_shape = (flattened_size,) new_shape = flattened_shape + old_shape[2:] newarray = buffer.reshape(new_shape) - return as_field(dims, newarray) + return as_field(dims, newarray, allocator=backend) def unflatten_first_two_dims(field: gt_common.Field) -> np.array: @@ -180,10 +195,7 @@ class Output: def _test_validation(self, grid, backend, input_data): reference_outputs = self.reference( grid, - **{ - k: v.asnumpy() if isinstance(v, gt_common.Field) else np.array(v) - for k, v in input_data.items() - }, + **{k: v.asnumpy() if isinstance(v, gt_common.Field) else v for k, v in input_data.items()}, ) input_data = allocate_data(backend, input_data) diff --git a/model/common/src/icon4py/model/common/test_utils/pytest_config.py b/model/common/src/icon4py/model/common/test_utils/pytest_config.py index 7d17ecc4d9..ab26efbfcc 100644 --- a/model/common/src/icon4py/model/common/test_utils/pytest_config.py +++ b/model/common/src/icon4py/model/common/test_utils/pytest_config.py @@ -11,7 +11,6 @@ import pytest from gt4py.next import gtfn_cpu, gtfn_gpu, itir_python -import icon4py.model.common.settings as settings from icon4py.model.common.test_utils.datatest_utils import ( GLOBAL_EXPERIMENT, REGIONAL_EXPERIMENT, @@ -75,10 +74,6 @@ def pytest_configure(config): if config.getoption("--backend"): backend = config.getoption("--backend") check_backend_validity(backend) - settings.backend = backends[backend] - - if config.getoption("--dace-orchestration"): - settings.dace_orchestration = True def pytest_addoption(parser): @@ -124,16 +119,6 @@ def pytest_addoption(parser): except ValueError: pass - try: - parser.addoption( - "--dace-orchestration", - action="store", - default=None, - help="Performs DaCe orchestration. Any value will enable it.", - ) - except ValueError: - pass - def pytest_runtest_setup(item): for _ in item.iter_markers(name="datatest"): diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py index d91f264a0d..a5a5d7664c 100644 --- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -10,6 +10,7 @@ import uuid import gt4py.next as gtx +import numpy as np import serialbox import icon4py.model.common.decomposition.definitions as decomposition @@ -18,8 +19,8 @@ import icon4py.model.common.test_utils.helpers as helpers from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, horizontal, icon -from icon4py.model.common.settings import xp from icon4py.model.common.states import prognostic_state +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc log = logging.getLogger(__name__) @@ -45,7 +46,7 @@ def wrapper(self, *args, **kwargs): ) if dims: shp = tuple(self.sizes[d] for d in dims) - return gtx.as_field(dims, xp.zeros(shp)) + return gtx.as_field(dims, np.zeros(shp)) else: return None @@ -57,7 +58,7 @@ def log_meta_info(self): self.log.info(self.savepoint.metainfo) def _get_field(self, name, *dimensions, dtype=float): - buffer = xp.squeeze(self.serializer.read(name, self.savepoint).astype(dtype)) + buffer = np.squeeze(self.serializer.read(name, self.savepoint).astype(dtype)) buffer = self._reduce_to_dim_size(buffer, dimensions) self.log.debug(f"{name} {buffer.shape}") @@ -65,7 +66,7 @@ def _get_field(self, name, *dimensions, dtype=float): def _get_field_component(self, name: str, ntnd: int, dims: tuple[gtx.Dimension, gtx]): buffer = self.serializer.read(name, self.savepoint).astype(float) - buffer = xp.squeeze(buffer)[:, :, ntnd - 1] + buffer = np.squeeze(buffer)[:, :, ntnd - 1] buffer = self._reduce_to_dim_size(buffer, dims) self.log.debug(f"{name} {buffer.shape}") return gtx.as_field(dims, buffer) @@ -318,7 +319,7 @@ def e_owner_mask(self): def f_e(self): return self._get_field("f_e", dims.EdgeDim) - def print_connectivity_info(self, name: str, ar: xp.ndarray): + def print_connectivity_info(self, name: str, ar: field_alloc.NDArray): self.log.debug(f" connectivity {name} {ar.shape}") def c2e(self): @@ -326,12 +327,12 @@ def c2e(self): def _get_connectivity_array(self, name: str, target_dim: gtx.Dimension, reverse: bool = False): if reverse: - connectivity = xp.transpose(self._read_int32(name, offset=1))[ + connectivity = np.transpose(self._read_int32(name, offset=1))[ : self.sizes[target_dim], : ] else: connectivity = self._read_int32(name, offset=1)[: self.sizes[target_dim], :] - connectivity = xp.asarray(connectivity) + connectivity = np.asarray(connectivity) self.log.debug(f" connectivity {name} : {connectivity.shape}") return connectivity @@ -343,7 +344,7 @@ def e2c2e(self): def c2e2c2e(self): if self._c2e2c2e() is None: - return xp.zeros((self.sizes[dims.CellDim], 9), dtype=int) + return np.zeros((self.sizes[dims.CellDim], 9), dtype=int) else: return self._c2e2c2e() @@ -380,7 +381,7 @@ def refin_ctrl(self, dim: gtx.Dimension): field_name = "refin_ctl" return gtx.as_field( (dim,), - xp.squeeze( + np.squeeze( self._read_field_for_dim(field_name, self._read_int32, dim)[: self.num(dim)], 1 ), ) @@ -405,7 +406,7 @@ def _read_field_for_dim(field_name, read_func, dim: gtx.Dimension): def owner_mask(self, dim: gtx.Dimension): field_name = "owner_mask" mask = self._read_field_for_dim(field_name, self._read_bool, dim) - return xp.squeeze(mask) + return np.squeeze(mask) def global_index(self, dim: gtx.Dimension): field_name = "glb_index" @@ -453,8 +454,8 @@ def construct_icon_grid(self, on_gpu: bool) -> icon.IconGrid: ) c2e2c = self.c2e2c() e2c2e = self.e2c2e() - c2e2c0 = xp.column_stack(((range(c2e2c.shape[0])), c2e2c)) - e2c2e0 = xp.column_stack(((range(e2c2e.shape[0])), e2c2e)) + c2e2c0 = np.column_stack((range(c2e2c.shape[0]), c2e2c)) + e2c2e0 = np.column_stack((range(e2c2e.shape[0]), e2c2e)) grid = ( icon.IconGrid(self._grid_id) .with_config(config) @@ -575,7 +576,7 @@ def geofac_grdiv(self): return self._get_field("geofac_grdiv", dims.EdgeDim, dims.E2C2EODim) def geofac_grg(self): - grg = xp.squeeze(self.serializer.read("geofac_grg", self.savepoint)) + grg = np.squeeze(self.serializer.read("geofac_grg", self.savepoint)) num_cells = self.sizes[dims.CellDim] return gtx.as_field((dims.CellDim, dims.C2E2CODim), grg[:num_cells, :, 0]), gtx.as_field( (dims.CellDim, dims.C2E2CODim), grg[:num_cells, :, 1] @@ -603,21 +604,21 @@ def pos_on_tplane_e_y(self): return helpers.as_1D_sparse_field(field[:, 0:2], dims.ECDim) def rbf_vec_coeff_e(self): - buffer = xp.squeeze( + buffer = np.squeeze( self.serializer.read("rbf_vec_coeff_e", self.savepoint).astype(float) ).transpose() return gtx.as_field((dims.EdgeDim, dims.E2C2EDim), buffer) @IconSavepoint.optionally_registered() def rbf_vec_coeff_c1(self): - buffer = xp.squeeze( + buffer = np.squeeze( self.serializer.read("rbf_vec_coeff_c1", self.savepoint).astype(float) ).transpose() return gtx.as_field((dims.CellDim, dims.C2E2C2EDim), buffer) @IconSavepoint.optionally_registered() def rbf_vec_coeff_c2(self): - buffer = xp.squeeze( + buffer = np.squeeze( self.serializer.read("rbf_vec_coeff_c2", self.savepoint).astype(float) ).transpose() return gtx.as_field((dims.CellDim, dims.C2E2C2EDim), buffer) @@ -751,9 +752,9 @@ def wgtfac_e(self): def wgtfacq_e_dsl( self, k_level ): # TODO: @abishekg7 Simplify this after serialized data is fixed - ar = xp.squeeze(self.serializer.read("wgtfacq_e", self.savepoint)) + ar = np.squeeze(self.serializer.read("wgtfacq_e", self.savepoint)) k = k_level - 3 - ar = xp.pad(ar[:, ::-1], ((0, 0), (k, 0)), "constant", constant_values=(0.0,)) + ar = np.pad(ar[:, ::-1], ((0, 0), (k, 0)), "constant", constant_values=(0.0,)) return self._get_field_from_ndarray(ar, dims.EdgeDim, dims.KDim) @IconSavepoint.optionally_registered(dims.CellDim, dims.KDim) @@ -768,17 +769,17 @@ def geopot(self): return self._get_field("geopot", dims.CellDim, dims.KDim) def _read_and_reorder_sparse_field(self, name: str, sparse_size=3): - ser_input = xp.squeeze(self.serializer.read(name, self.savepoint))[:, :, :] + ser_input = np.squeeze(self.serializer.read(name, self.savepoint))[:, :, :] ser_input = self._reduce_to_dim_size(ser_input, (dims.CellDim, dims.C2E2CDim, dims.KDim)) if ser_input.shape[1] != sparse_size: - ser_input = xp.moveaxis(ser_input, 1, -1) + ser_input = np.moveaxis(ser_input, 1, -1) return self._linearize_first_2dims( ser_input, sparse_size=sparse_size, target_dims=(dims.CECDim, dims.KDim) ) def _linearize_first_2dims( - self, data: xp.ndarray, sparse_size: int, target_dims: tuple[gtx.Dimension, ...] + self, data: field_alloc.NDArray, sparse_size: int, target_dims: tuple[gtx.Dimension, ...] ): old_shape = data.shape assert old_shape[1] == sparse_size @@ -789,10 +790,10 @@ def zd_vertoffset(self): return self._read_and_reorder_sparse_field("zd_vertoffset") def zd_vertidx(self): - return xp.squeeze(self.serializer.read("zd_vertidx", self.savepoint)) + return np.squeeze(self.serializer.read("zd_vertidx", self.savepoint)) def zd_indlist(self): - return xp.squeeze(self.serializer.read("zd_indlist", self.savepoint)) + return np.squeeze(self.serializer.read("zd_indlist", self.savepoint)) class LeastSquaresSavepoint(IconSavepoint): @@ -885,16 +886,16 @@ def exner(self): return self._get_field("exner", dims.CellDim, dims.KDim) def diff_multfac_smag(self): - return xp.squeeze(self.serializer.read("diff_multfac_smag", self.savepoint)) + return np.squeeze(self.serializer.read("diff_multfac_smag", self.savepoint)) def enh_smag_fac(self): - return xp.squeeze(self.serializer.read("enh_smag_fac", self.savepoint)) + return np.squeeze(self.serializer.read("enh_smag_fac", self.savepoint)) def smag_limit(self): - return xp.squeeze(self.serializer.read("smag_limit", self.savepoint)) + return np.squeeze(self.serializer.read("smag_limit", self.savepoint)) def diff_multfac_n2w(self): - return xp.squeeze(self.serializer.read("diff_multfac_n2w", self.savepoint)) + return np.squeeze(self.serializer.read("diff_multfac_n2w", self.savepoint)) def nudgezone_diff(self) -> int: return self.serializer.read("nudgezone_diff", self.savepoint)[0] diff --git a/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py b/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py index a31210d50a..8e9ecccd91 100644 --- a/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py +++ b/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py @@ -32,9 +32,7 @@ except ImportError: import numpy as xp - NDArray: TypeAlias = Union[np.ndarray, xp.ndarray] - NDArrayInterface: TypeAlias = Union[np.ndarray, xp.ndarray, gtx.Field] diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py index bc81882061..8413486204 100644 --- a/model/common/tests/grid_tests/test_grid_manager.py +++ b/model/common/tests/grid_tests/test_grid_manager.py @@ -22,6 +22,7 @@ refinement as refin, vertical as v_grid, ) +from icon4py.model.common.grid.grid_manager import GeometryName from icon4py.model.common.test_utils import ( datatest_utils as dt_utils, grid_utils as gridtest_utils, diff --git a/model/common/tests/math_tests/test_helpers.py b/model/common/tests/math_tests/test_helpers.py index 9c864e599d..87e61a9a4d 100644 --- a/model/common/tests/math_tests/test_helpers.py +++ b/model/common/tests/math_tests/test_helpers.py @@ -6,10 +6,12 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause + +import numpy as np + from icon4py.model.common import dimension as dims from icon4py.model.common.grid import simple from icon4py.model.common.math import helpers -from icon4py.model.common.settings import xp from icon4py.model.common.test_utils import helpers as test_helpers @@ -28,9 +30,9 @@ def test_cross_product(backend): helpers.cross_product_on_edges.with_backend(backend)( x1, x2, y1, y2, z1, z2, out=(x, y, z), offset_provider={} ) - a = xp.column_stack((x1.ndarray, y1.ndarray, z1.ndarray)) - b = xp.column_stack((x2.ndarray, y2.ndarray, z2.ndarray)) - c = xp.cross(a, b) + a = np.column_stack((x1.ndarray, y1.ndarray, z1.ndarray)) + b = np.column_stack((x2.ndarray, y2.ndarray, z2.ndarray)) + c = np.cross(a, b) assert test_helpers.dallclose(c[:, 0], x.ndarray) assert test_helpers.dallclose(c[:, 1], y.ndarray) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index a81d3afae2..e4d3108f4e 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -452,7 +452,7 @@ def test_compute_ddxt_z_full( vertical_end = icon_grid.num_levels + 1 cells_aw_verts = interpolation_savepoint.c_intp().asnumpy() z_ifv = zero_field(icon_grid, dims.VertexDim, dims.KDim, extend={dims.KDim: 1}) - compute_cell_2_vertex_interpolation( + compute_cell_2_vertex_interpolation.with_backend(backend)( z_ifc, gtx.as_field((dims.VertexDim, dims.V2CDim), cells_aw_verts), z_ifv, diff --git a/model/driver/src/icon4py/model/driver/test_cases/gauss3d.py b/model/driver/src/icon4py/model/driver/test_cases/gauss3d.py index 726655e2e0..3c7028d14f 100644 --- a/model/driver/src/icon4py/model/driver/test_cases/gauss3d.py +++ b/model/driver/src/icon4py/model/driver/test_cases/gauss3d.py @@ -9,6 +9,7 @@ import pathlib import gt4py.next as gtx +import numpy as np from icon4py.model.atmosphere.diffusion import diffusion_states from icon4py.model.atmosphere.dycore import dycore_states @@ -18,7 +19,6 @@ cell_2_edge_interpolation, edge_2_cell_vector_rbf_interpolation, ) -from icon4py.model.common.settings import xp from icon4py.model.common.states import ( diagnostic_state as diagnostics, prognostic_state as prognostics, @@ -90,22 +90,22 @@ def model_initialization_gauss3d( ) end_cell_end = grid.end_index(cell_domain(h_grid.Zone.END)) - w_numpy = xp.zeros((num_cells, num_levels + 1), dtype=float) - exner_numpy = xp.zeros((num_cells, num_levels), dtype=float) - rho_numpy = xp.zeros((num_cells, num_levels), dtype=float) - temperature_numpy = xp.zeros((num_cells, num_levels), dtype=float) - pressure_numpy = xp.zeros((num_cells, num_levels), dtype=float) - theta_v_numpy = xp.zeros((num_cells, num_levels), dtype=float) - eta_v_numpy = xp.zeros((num_cells, num_levels), dtype=float) + w_numpy = np.zeros((num_cells, num_levels + 1), dtype=float) + exner_numpy = np.zeros((num_cells, num_levels), dtype=float) + rho_numpy = np.zeros((num_cells, num_levels), dtype=float) + temperature_numpy = np.zeros((num_cells, num_levels), dtype=float) + pressure_numpy = np.zeros((num_cells, num_levels), dtype=float) + theta_v_numpy = np.zeros((num_cells, num_levels), dtype=float) + eta_v_numpy = np.zeros((num_cells, num_levels), dtype=float) - mask_array_edge_start_plus1_to_edge_end = xp.ones(num_edges, dtype=bool) + mask_array_edge_start_plus1_to_edge_end = np.ones(num_edges, dtype=bool) mask_array_edge_start_plus1_to_edge_end[0:end_edge_lateral_boundary_level_2] = False - mask = xp.repeat( - xp.expand_dims(mask_array_edge_start_plus1_to_edge_end, axis=-1), + mask = np.repeat( + np.expand_dims(mask_array_edge_start_plus1_to_edge_end, axis=-1), num_levels, axis=1, ) - primal_normal_x = xp.repeat(xp.expand_dims(primal_normal_x, axis=-1), num_levels, axis=1) + primal_normal_x = np.repeat(np.expand_dims(primal_normal_x, axis=-1), num_levels, axis=1) # Define test case parameters # The topography can only be read from serialized data for now, then these @@ -121,7 +121,7 @@ def model_initialization_gauss3d( log.info("Topography can only be read from serialized data for now.") # Horizontal wind field - u = xp.where(mask, nh_u0, 0.0) + u = np.where(mask, nh_u0, 0.0) vn_numpy = u * primal_normal_x log.info("Wind profile assigned.") @@ -129,14 +129,14 @@ def model_initialization_gauss3d( for k_index in range(num_levels - 1, -1, -1): z_help = (nh_brunt_vais / phy_const.GRAV) ** 2 * geopot[:, k_index] # profile of theta is explicitly given - theta_v_numpy[:, k_index] = nh_t0 * xp.exp(z_help) + theta_v_numpy[:, k_index] = nh_t0 * np.exp(z_help) # Lower boundary condition for exner pressure if nh_brunt_vais != 0.0: z_help = (nh_brunt_vais / phy_const.GRAV) ** 2 * geopot[:, num_levels - 1] exner_numpy[:, num_levels - 1] = ( phy_const.GRAV / nh_brunt_vais - ) ** 2 / nh_t0 / phy_const.CPD * (xp.exp(-z_help) - 1.0) + 1.0 + ) ** 2 / nh_t0 / phy_const.CPD * (np.exp(-z_help) - 1.0) + 1.0 else: exner_numpy[:, num_levels - 1] = 1.0 - geopot[:, num_levels - 1] / phy_const.CPD / nh_t0 log.info("Vertical computations completed.") @@ -179,7 +179,7 @@ def model_initialization_gauss3d( virtual_temperature = gtx.as_field((dims.CellDim, dims.KDim), temperature_numpy) pressure = gtx.as_field((dims.CellDim, dims.KDim), pressure_numpy) theta_v = gtx.as_field((dims.CellDim, dims.KDim), theta_v_numpy) - pressure_ifc_numpy = xp.zeros((num_cells, num_levels + 1), dtype=float) + pressure_ifc_numpy = np.zeros((num_cells, num_levels + 1), dtype=float) pressure_ifc_numpy[ :, -1 ] = phy_const.P0REF # set surface pressure to the prescribed value (only used for IC in JABW test case, then actually computed in the dycore) diff --git a/model/driver/src/icon4py/model/driver/test_cases/jablonowski_williamson.py b/model/driver/src/icon4py/model/driver/test_cases/jablonowski_williamson.py index 49b27c44e9..9ec8b1dc17 100644 --- a/model/driver/src/icon4py/model/driver/test_cases/jablonowski_williamson.py +++ b/model/driver/src/icon4py/model/driver/test_cases/jablonowski_williamson.py @@ -10,6 +10,7 @@ import pathlib import gt4py.next as gtx +import numpy as np from icon4py.model.atmosphere.diffusion import diffusion_states from icon4py.model.atmosphere.dycore import dycore_states @@ -19,7 +20,6 @@ cell_2_edge_interpolation, edge_2_cell_vector_rbf_interpolation, ) -from icon4py.model.common.settings import xp from icon4py.model.common.states import ( diagnostic_state as diagnostics, prognostic_state as prognostics, @@ -110,16 +110,16 @@ def model_initialization_jabw( lat_perturbation_center = 2.0 * lon_perturbation_center # latitude of the perturb centre ps_o_p0ref = p_sfc / phy_const.P0REF - w_numpy = xp.zeros((num_cells, num_levels + 1), dtype=float) - exner_numpy = xp.zeros((num_cells, num_levels), dtype=float) - rho_numpy = xp.zeros((num_cells, num_levels), dtype=float) - temperature_numpy = xp.zeros((num_cells, num_levels), dtype=float) - pressure_numpy = xp.zeros((num_cells, num_levels), dtype=float) - theta_v_numpy = xp.zeros((num_cells, num_levels), dtype=float) - eta_v_numpy = xp.zeros((num_cells, num_levels), dtype=float) + w_numpy = np.zeros((num_cells, num_levels + 1), dtype=float) + exner_numpy = np.zeros((num_cells, num_levels), dtype=float) + rho_numpy = np.zeros((num_cells, num_levels), dtype=float) + temperature_numpy = np.zeros((num_cells, num_levels), dtype=float) + pressure_numpy = np.zeros((num_cells, num_levels), dtype=float) + theta_v_numpy = np.zeros((num_cells, num_levels), dtype=float) + eta_v_numpy = np.zeros((num_cells, num_levels), dtype=float) - sin_lat = xp.sin(cell_lat) - cos_lat = xp.cos(cell_lat) + sin_lat = np.sin(cell_lat) + cos_lat = np.cos(cell_lat) fac1 = 1.0 / 6.3 - 2.0 * (sin_lat**6) * (cos_lat**2 + 1.0 / 3.0) fac2 = ( (8.0 / 5.0 * (cos_lat**3) * (sin_lat**2 + 2.0 / 3.0) - 0.25 * math.pi) @@ -128,26 +128,26 @@ def model_initialization_jabw( ) lapse_rate = phy_const.RD * gamma / phy_const.GRAV for k_index in range(num_levels - 1, -1, -1): - eta_old = xp.full(num_cells, fill_value=1.0e-7, dtype=float) + eta_old = np.full(num_cells, fill_value=1.0e-7, dtype=float) log.info(f"In Newton iteration, k = {k_index}") # Newton iteration to determine zeta for _ in range(100): eta_v_numpy[:, k_index] = (eta_old - eta_0) * math.pi * 0.5 - cos_etav = xp.cos(eta_v_numpy[:, k_index]) - sin_etav = xp.sin(eta_v_numpy[:, k_index]) + cos_etav = np.cos(eta_v_numpy[:, k_index]) + sin_etav = np.sin(eta_v_numpy[:, k_index]) temperature_avg = jw_temp0 * (eta_old**lapse_rate) geopot_avg = jw_temp0 * phy_const.GRAV / gamma * (1.0 - eta_old**lapse_rate) - temperature_avg = xp.where( + temperature_avg = np.where( eta_old < eta_t, temperature_avg + dtemp * ((eta_t - eta_old) ** 5), temperature_avg ) - geopot_avg = xp.where( + geopot_avg = np.where( eta_old < eta_t, geopot_avg - phy_const.RD * dtemp * ( - (xp.log(eta_old / eta_t) + 137.0 / 60.0) * (eta_t**5) + (np.log(eta_old / eta_t) + 137.0 / 60.0) * (eta_t**5) - 5.0 * (eta_t**4) * eta_old + 5.0 * (eta_t**3) * (eta_old**2) - 10.0 / 3.0 * (eta_t**2) * (eta_old**3) @@ -168,7 +168,7 @@ def model_initialization_jabw( * jw_u0 / phy_const.RD * sin_etav - * xp.sqrt(cos_etav) + * np.sqrt(cos_etav) * (2.0 * jw_u0 * fac1 * (cos_etav**1.5) + fac2) ) newton_function = geopot_jw - geopot[:, k_index] @@ -240,7 +240,7 @@ def model_initialization_jabw( virutal_temperature = gtx.as_field((dims.CellDim, dims.KDim), temperature_numpy) pressure = gtx.as_field((dims.CellDim, dims.KDim), pressure_numpy) theta_v = gtx.as_field((dims.CellDim, dims.KDim), theta_v_numpy) - pressure_ifc_numpy = xp.zeros((num_cells, num_levels + 1), dtype=float) + pressure_ifc_numpy = np.zeros((num_cells, num_levels + 1), dtype=float) pressure_ifc_numpy[:, -1] = p_sfc pressure_ifc = gtx.as_field((dims.CellDim, dims.KDim), pressure_ifc_numpy) diff --git a/model/driver/src/icon4py/model/driver/test_cases/utils.py b/model/driver/src/icon4py/model/driver/test_cases/utils.py index 40b243c119..b1c556b725 100644 --- a/model/driver/src/icon4py/model/driver/test_cases/utils.py +++ b/model/driver/src/icon4py/model/driver/test_cases/utils.py @@ -6,6 +6,7 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np from icon4py.model.common import ( constants as phy_const, @@ -14,24 +15,18 @@ type_alias as ta, ) from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid -from icon4py.model.common.settings import backend, xp - - -# TODO: this will have to be removed once domain allows for imports -CellDim = dims.CellDim -KDim = dims.KDim def hydrostatic_adjustment_numpy( - wgtfac_c: xp.ndarray, - ddqz_z_half: xp.ndarray, - exner_ref_mc: xp.ndarray, - d_exner_dz_ref_ic: xp.ndarray, - theta_ref_mc: xp.ndarray, - theta_ref_ic: xp.ndarray, - rho: xp.ndarray, - exner: xp.ndarray, - theta_v: xp.ndarray, + wgtfac_c: np.ndarray, + ddqz_z_half: np.ndarray, + exner_ref_mc: np.ndarray, + d_exner_dz_ref_ic: np.ndarray, + theta_ref_mc: np.ndarray, + theta_ref_ic: np.ndarray, + rho: np.ndarray, + exner: np.ndarray, + theta_v: np.ndarray, num_levels: int, ): # virtual temperature @@ -51,7 +46,7 @@ def hydrostatic_adjustment_numpy( ) quadratic_c = -(fac2 * fac3 / ddqz_z_half[:, k + 1] + fac2 * d_exner_dz_ref_ic[:, k + 1]) - exner[:, k] = (quadratic_b + xp.sqrt(quadratic_b**2 + 4.0 * quadratic_a * quadratic_c)) / ( + exner[:, k] = (quadratic_b + np.sqrt(quadratic_b**2 + 4.0 * quadratic_a * quadratic_c)) / ( 2.0 * quadratic_a ) theta_v[:, k] = temp_v[:, k] / exner[:, k] @@ -63,17 +58,17 @@ def hydrostatic_adjustment_numpy( def hydrostatic_adjustment_constant_thetav_numpy( - wgtfac_c: xp.ndarray, - ddqz_z_half: xp.ndarray, - exner_ref_mc: xp.ndarray, - d_exner_dz_ref_ic: xp.ndarray, - theta_ref_mc: xp.ndarray, - theta_ref_ic: xp.ndarray, - rho: xp.ndarray, - exner: xp.ndarray, - theta_v: xp.ndarray, + wgtfac_c: np.ndarray, + ddqz_z_half: np.ndarray, + exner_ref_mc: np.ndarray, + d_exner_dz_ref_ic: np.ndarray, + theta_ref_mc: np.ndarray, + theta_ref_ic: np.ndarray, + rho: np.ndarray, + exner: np.ndarray, + theta_v: np.ndarray, num_levels: int, -) -> tuple[xp.ndarray, xp.ndarray]: +) -> tuple[np.ndarray, np.ndarray]: """ Computes a hydrostatically balanced profile. In constrast to the above hydrostatic_adjustment_numpy, the virtual temperature is kept (assumed) @@ -108,10 +103,10 @@ def zonalwind_2_normalwind_numpy( jw_up: float, lat_perturbation_center: float, lon_perturbation_center: float, - edge_lat: xp.ndarray, - edge_lon: xp.ndarray, - primal_normal_x: xp.ndarray, - eta_v_e: xp.ndarray, + edge_lat: np.ndarray, + edge_lon: np.ndarray, + primal_normal_x: np.ndarray, + eta_v_e: np.ndarray, ): """ Compute normal wind at edge center from vertical eta coordinate (eta_v_e). @@ -130,29 +125,29 @@ def zonalwind_2_normalwind_numpy( """ # TODO (Chia Rui) this function needs a test - mask = xp.ones((grid.num_edges, grid.num_levels), dtype=bool) + mask = np.ones((grid.num_edges, grid.num_levels), dtype=bool) mask[ 0 : grid.end_index(h_grid.domain(dims.EdgeDim)(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)), :, ] = False - edge_lat = xp.repeat(xp.expand_dims(edge_lat, axis=-1), eta_v_e.shape[1], axis=1) - edge_lon = xp.repeat(xp.expand_dims(edge_lon, axis=-1), eta_v_e.shape[1], axis=1) - primal_normal_x = xp.repeat(xp.expand_dims(primal_normal_x, axis=-1), eta_v_e.shape[1], axis=1) - u = xp.where(mask, jw_u0 * (xp.cos(eta_v_e) ** 1.5) * (xp.sin(2.0 * edge_lat) ** 2), 0.0) + edge_lat = np.repeat(np.expand_dims(edge_lat, axis=-1), eta_v_e.shape[1], axis=1) + edge_lon = np.repeat(np.expand_dims(edge_lon, axis=-1), eta_v_e.shape[1], axis=1) + primal_normal_x = np.repeat(np.expand_dims(primal_normal_x, axis=-1), eta_v_e.shape[1], axis=1) + u = np.where(mask, jw_u0 * (np.cos(eta_v_e) ** 1.5) * (np.sin(2.0 * edge_lat) ** 2), 0.0) if jw_up > 1.0e-20: - u = xp.where( + u = np.where( mask, u + jw_up - * xp.exp( + * np.exp( -( ( 10.0 - * xp.arccos( - xp.sin(lat_perturbation_center) * xp.sin(edge_lat) - + xp.cos(lat_perturbation_center) - * xp.cos(edge_lat) - * xp.cos(edge_lon - lon_perturbation_center) + * np.arccos( + np.sin(lat_perturbation_center) * np.sin(edge_lat) + + np.cos(lat_perturbation_center) + * np.cos(edge_lat) + * np.cos(edge_lon - lon_perturbation_center) ) ) ** 2 @@ -188,7 +183,7 @@ def _compute_perturbed_exner( return exner_pr -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED, backend=backend) +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_perturbed_exner( exner: fa.CellKField[ta.wpfloat], exner_ref: fa.CellKField[ta.vpfloat], @@ -203,7 +198,7 @@ def compute_perturbed_exner( exner_ref, out=exner_pr, domain={ - CellDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), }, ) diff --git a/model/driver/tests/initial_condition_tests/test_utils.py b/model/driver/tests/initial_condition_tests/test_utils.py index f675bfe6b2..795c3bf270 100644 --- a/model/driver/tests/initial_condition_tests/test_utils.py +++ b/model/driver/tests/initial_condition_tests/test_utils.py @@ -6,7 +6,8 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from icon4py.model.common.settings import xp +import numpy as np + from icon4py.model.common.test_utils import helpers from icon4py.model.driver.test_cases import utils @@ -20,15 +21,15 @@ def test_hydrostatic_adjustment_numpy(): rho0 = 1.25 exner0 = 0.935 theta_v0 = 293.14 - wgtfac_c = 1.05 * xp.ones((num_cells, num_levels)) - ddqz_z_half = xp.ones((num_cells, num_levels)) - exner_ref_mc = 0.89 * xp.ones((num_cells, num_levels)) - d_exner_dz_ref_ic = 0.0 * xp.ones((num_cells, num_levels)) - theta_ref_mc = 312 * xp.ones((num_cells, num_levels)) - theta_ref_ic = 312 * xp.ones((num_cells, num_levels)) - rho = rho0 * xp.ones((num_cells, num_levels)) - exner = exner0 * xp.ones((num_cells, num_levels)) - theta_v = theta_v0 * xp.ones((num_cells, num_levels)) + wgtfac_c = 1.05 * np.ones((num_cells, num_levels)) + ddqz_z_half = np.ones((num_cells, num_levels)) + exner_ref_mc = 0.89 * np.ones((num_cells, num_levels)) + d_exner_dz_ref_ic = 0.0 * np.ones((num_cells, num_levels)) + theta_ref_mc = 312 * np.ones((num_cells, num_levels)) + theta_ref_ic = 312 * np.ones((num_cells, num_levels)) + rho = rho0 * np.ones((num_cells, num_levels)) + exner = exner0 * np.ones((num_cells, num_levels)) + theta_v = theta_v0 * np.ones((num_cells, num_levels)) # Call the function r_rho, r_exner, r_theta_v = utils.hydrostatic_adjustment_numpy( @@ -50,19 +51,19 @@ def test_hydrostatic_adjustment_numpy(): assert helpers.dallclose( r_rho[:, -1], - rho0 * xp.ones(num_cells), + rho0 * np.ones(num_cells), ) assert helpers.dallclose( r_rho[:, :-1], - 1.0046424441749071 * xp.ones((num_cells, num_levels - 1)), + 1.0046424441749071 * np.ones((num_cells, num_levels - 1)), ) assert helpers.dallclose( r_exner, - exner0 * xp.ones((num_cells, num_levels)), + exner0 * np.ones((num_cells, num_levels)), ) assert helpers.dallclose( r_theta_v, - theta_v0 * xp.ones((num_cells, num_levels)), + theta_v0 * np.ones((num_cells, num_levels)), ) @@ -75,15 +76,15 @@ def test_hydrostatic_adjustment_constant_thetav_numpy(): rho0 = 1.25 exner0 = 0.935 theta_v0 = 293.14 - wgtfac_c = 1.05 * xp.ones((num_cells, num_levels)) - ddqz_z_half = xp.ones((num_cells, num_levels)) - exner_ref_mc = 0.89 * xp.ones((num_cells, num_levels)) - d_exner_dz_ref_ic = 0.0 * xp.ones((num_cells, num_levels)) - theta_ref_mc = 312 * xp.ones((num_cells, num_levels)) - theta_ref_ic = 312 * xp.ones((num_cells, num_levels)) - rho = rho0 * xp.ones((num_cells, num_levels)) - exner = exner0 * xp.ones((num_cells, num_levels)) - theta_v = theta_v0 * xp.ones((num_cells, num_levels)) + wgtfac_c = 1.05 * np.ones((num_cells, num_levels)) + ddqz_z_half = np.ones((num_cells, num_levels)) + exner_ref_mc = 0.89 * np.ones((num_cells, num_levels)) + d_exner_dz_ref_ic = 0.0 * np.ones((num_cells, num_levels)) + theta_ref_mc = 312 * np.ones((num_cells, num_levels)) + theta_ref_ic = 312 * np.ones((num_cells, num_levels)) + rho = rho0 * np.ones((num_cells, num_levels)) + exner = exner0 * np.ones((num_cells, num_levels)) + theta_v = theta_v0 * np.ones((num_cells, num_levels)) # Call the function r_rho, r_exner = utils.hydrostatic_adjustment_constant_thetav_numpy( @@ -104,9 +105,9 @@ def test_hydrostatic_adjustment_constant_thetav_numpy(): assert helpers.dallclose( r_rho, - 1.0046424441749071 * xp.ones((num_cells, num_levels)), + 1.0046424441749071 * np.ones((num_cells, num_levels)), ) assert helpers.dallclose( r_exner, - exner0 * xp.ones((num_cells, num_levels)), + exner0 * np.ones((num_cells, num_levels)), ) diff --git a/tools/src/icon4pytools/py2fgen/cli.py b/tools/src/icon4pytools/py2fgen/cli.py index 6caf15117b..0b6faccb90 100644 --- a/tools/src/icon4pytools/py2fgen/cli.py +++ b/tools/src/icon4pytools/py2fgen/cli.py @@ -9,7 +9,6 @@ import pathlib import click -from icon4py.model.common.config import GT4PyBackend from icon4pytools.common.utils import write_string from icon4pytools.py2fgen.generate import ( @@ -19,6 +18,7 @@ ) from icon4pytools.py2fgen.parsing import parse from icon4pytools.py2fgen.plugin import generate_and_compile_cffi_plugin +from icon4pytools.py2fgen.settings import GT4PyBackend def parse_comma_separated_list(ctx, param, value) -> list[str]: diff --git a/model/common/src/icon4py/model/common/config.py b/tools/src/icon4pytools/py2fgen/settings.py similarity index 89% rename from model/common/src/icon4py/model/common/config.py rename to tools/src/icon4pytools/py2fgen/settings.py index fba138bf14..09f4902b98 100644 --- a/model/common/src/icon4py/model/common/config.py +++ b/tools/src/icon4pytools/py2fgen/settings.py @@ -5,8 +5,6 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause - -# Assuming this code is in a module called icon4py_config.py import dataclasses import os from enum import Enum @@ -21,7 +19,7 @@ try: - import dace + import dace # type: ignore[import-not-found, import-untyped] from gt4py.next.program_processors.runners.dace import ( run_dace_cpu, run_dace_cpu_noopt, @@ -32,7 +30,7 @@ from types import ModuleType from typing import Optional - dace: Optional[ModuleType] = None + dace: Optional[ModuleType] = None # type: ignore[no-redef] # definition needed here class Device(Enum): @@ -71,7 +69,7 @@ def icon4py_dace_orchestration(self): @cached_property def array_ns(self): if self.device == Device.GPU: - import cupy as cp # type: ignore[import-untyped] + import cupy as cp # type: ignore[import-not-found] return cp else: @@ -117,3 +115,10 @@ def limited_area(self): @cached_property def parallel_run(self): return os.environ.get("ICON4PY_PARALLEL", False) + + +config = Icon4PyConfig() +backend = config.gt4py_runner +dace_orchestration = config.icon4py_dace_orchestration +device = config.device +limited_area = config.limited_area diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index 4c6dca71bf..1240533f0a 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -13,7 +13,6 @@ from gt4py.eve.codegen import JinjaTemplate as as_jinja, TemplatedGenerator from gt4py.next import Dimension from gt4py.next.type_system.type_specifications import ScalarKind -from icon4py.model.common.config import GT4PyBackend from icon4pytools.icon4pygen.bindings.codegen.type_conversion import ( BUILTIN_TO_CPP_TYPE, @@ -21,12 +20,15 @@ BUILTIN_TO_NUMPY_TYPE, ) from icon4pytools.py2fgen.plugin import int_array_to_bool_array, unpack, unpack_gpu +from icon4pytools.py2fgen.settings import GT4PyBackend from icon4pytools.py2fgen.utils import flatten_and_get_unique_elts from icon4pytools.py2fgen.wrappers import wrapper_dimension # these arrays are not initialised in global experiments (e.g. ape_r02b04) and are not used # therefore unpacking needs to be skipped as otherwise it will trigger an error. + + UNINITIALISED_ARRAYS = [ "mask_hdiff", "zd_diffcoef", @@ -243,7 +245,8 @@ class PythonWrapperGenerator(TemplatedGenerator): {% if _this_node.backend == 'GPU' %}import cupy as cp {% endif %} from numpy.typing import NDArray from gt4py.next.iterator.embedded import np_as_located_field -from icon4py.model.common.settings import xp +from icon4pytools.py2fgen.settings import config +xp = config.array_ns from icon4py.model.common import dimension as dims {% if _this_node.is_gt4py_program_present %} diff --git a/tools/src/icon4pytools/py2fgen/wrappers/common.py b/tools/src/icon4pytools/py2fgen/wrappers/common.py index eb06af59d0..55e515682d 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/common.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/common.py @@ -10,17 +10,20 @@ import logging from icon4py.model.common import dimension as dims -from icon4py.model.common.decomposition import definitions, mpi_decomposition as mpi +from icon4py.model.common.decomposition import definitions from icon4py.model.common.grid import base, horizontal, icon -from icon4py.model.common.settings import xp +from icon4pytools.py2fgen.settings import config + + +xp = config.array_ns log = logging.getLogger(__name__) def adjust_fortran_indices(inp: xp.ndarray, offset: int) -> xp.ndarray: """For some Fortran arrays we need to subtract 1 to be compatible with Python indexing.""" - return xp.subtract(inp.ndarray, offset) + return xp.subtract(inp, offset) def construct_icon_grid( @@ -61,9 +64,9 @@ def construct_icon_grid( vertex_start_index = adjust_fortran_indices(vertex_starts, offset) edge_start_index = adjust_fortran_indices(edge_starts, offset) - cells_end_index = cell_ends.ndarray - vertex_end_index = vertex_ends.ndarray - edge_end_index = edge_ends.ndarray + cells_end_index = cell_ends + vertex_end_index = vertex_ends + edge_end_index = edge_ends c2e = adjust_fortran_indices(c2e, offset) c2v = adjust_fortran_indices(c2v, offset) @@ -153,9 +156,9 @@ def construct_decomposition( e_glb_index = adjust_fortran_indices(e_glb_index, offset) v_glb_index = adjust_fortran_indices(v_glb_index, offset) - c_owner_mask = c_owner_mask.ndarray[:num_cells] - e_owner_mask = e_owner_mask.ndarray[:num_edges] - v_owner_mask = v_owner_mask.ndarray[:num_vertices] + c_owner_mask = c_owner_mask[:num_cells] + e_owner_mask = e_owner_mask[:num_edges] + v_owner_mask = v_owner_mask[:num_vertices] decomposition_info = ( definitions.DecompositionInfo( @@ -165,7 +168,7 @@ def construct_decomposition( .with_dimension(dims.EdgeDim, e_glb_index, e_owner_mask) .with_dimension(dims.VertexDim, v_glb_index, v_owner_mask) ) - processor_props = mpi.get_multinode_properties(definitions.MultiNodeRun(), comm_id) + processor_props = definitions.get_processor_properties(definitions.MultiNodeRun(), comm_id) exchange = definitions.create_exchange(processor_props, decomposition_info) return processor_props, decomposition_info, exchange diff --git a/tools/src/icon4pytools/py2fgen/wrappers/debug_utils.py b/tools/src/icon4pytools/py2fgen/wrappers/debug_utils.py index 006c328152..fd176d6362 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/debug_utils.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/debug_utils.py @@ -18,11 +18,13 @@ VertexDim, ) from icon4py.model.common.grid.icon import IconGrid -from icon4py.model.common.settings import xp from icon4pytools.common.logger import setup_logger +from icon4pytools.py2fgen.settings import config +xp = config.array_ns + log = setup_logger(__name__) diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion_wrapper.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion_wrapper.py index c3739c42c0..8bb459926b 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion_wrapper.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion_wrapper.py @@ -32,13 +32,12 @@ DiffusionInterpolationState, DiffusionMetricState, ) -from icon4py.model.common import dimension as dims, field_type_aliases as fa, settings +from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.constants import DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO from icon4py.model.common.decomposition import definitions from icon4py.model.common.grid import icon from icon4py.model.common.grid.icon import GlobalGridParams from icon4py.model.common.grid.vertical import VerticalGrid, VerticalGridConfig -from icon4py.model.common.settings import backend, device, parallel_run from icon4py.model.common.states.prognostic_state import PrognosticState from icon4py.model.common.test_utils.helpers import ( as_1D_sparse_field, @@ -47,6 +46,7 @@ from icon4py.model.common.type_alias import wpfloat from icon4pytools.common.logger import setup_logger +from icon4pytools.py2fgen.settings import backend, config as config_settings, device from icon4pytools.py2fgen.wrappers import common as wrapper_common from icon4pytools.py2fgen.wrappers.debug_utils import print_grid_decomp_info from icon4pytools.py2fgen.wrappers.wrapper_dimension import ( @@ -332,21 +332,21 @@ def grid_init_diffusion( global_grid_params = GlobalGridParams(level=global_level, root=global_root) diffusion_wrapper_state["grid"] = wrapper_common.construct_icon_grid( - cell_starts=cell_starts, - cell_ends=cell_ends, - vertex_starts=vertex_starts, - vertex_ends=vertex_ends, - edge_starts=edge_starts, - edge_ends=edge_ends, - c2e=c2e, - e2c=e2c, - c2e2c=c2e2c, - e2c2e=e2c2e, - e2v=e2v, - v2e=v2e, - v2c=v2c, - e2c2v=e2c2v, - c2v=c2v, + cell_starts=cell_starts.ndarray, + cell_ends=cell_ends.ndarray, + vertex_starts=vertex_starts.ndarray, + vertex_ends=vertex_ends.ndarray, + edge_starts=edge_starts.ndarray, + edge_ends=edge_ends.ndarray, + c2e=c2e.ndarray, + e2c=e2c.ndarray, + c2e2c=c2e2c.ndarray, + e2c2e=e2c2e.ndarray, + e2v=e2v.ndarray, + v2e=v2e.ndarray, + v2c=v2c.ndarray, + e2c2v=e2c2v.ndarray, + c2v=c2v.ndarray, grid_id="icon_grid", global_grid_params=global_grid_params, num_vertices=num_vertices, @@ -354,10 +354,11 @@ def grid_init_diffusion( num_edges=num_edges, vertical_size=vertical_size, limited_area=limited_area, - on_gpu=True if settings.device == "GPU" else False, + on_gpu=True if config_settings.device == "GPU" else False, ) - if parallel_run: + if config_settings.parallel_run: + # Set MultiNodeExchange as exchange runtime ( processor_props, decomposition_info, diff --git a/tools/src/icon4pytools/py2fgen/wrappers/dycore_wrapper.py b/tools/src/icon4pytools/py2fgen/wrappers/dycore_wrapper.py index 0240cdce96..6446b3b619 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/dycore_wrapper.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/dycore_wrapper.py @@ -35,7 +35,7 @@ import icon4py.model.common.grid.states as grid_states from gt4py.next import common as gt4py_common from icon4py.model.atmosphere.dycore import dycore_states, solve_nonhydro -from icon4py.model.common import dimension as dims, settings, utils as common_utils +from icon4py.model.common import dimension as dims, utils as common_utils from icon4py.model.common.dimension import ( C2E2CODim, C2EDim, @@ -57,7 +57,6 @@ from icon4py.model.common.grid import icon from icon4py.model.common.grid.icon import GlobalGridParams from icon4py.model.common.grid.vertical import VerticalGrid, VerticalGridConfig -from icon4py.model.common.settings import backend from icon4py.model.common.states.prognostic_state import PrognosticState from icon4py.model.common.test_utils.helpers import ( as_1D_sparse_field, @@ -66,6 +65,7 @@ ) from icon4pytools.common.logger import setup_logger +from icon4pytools.py2fgen.settings import backend, device from icon4pytools.py2fgen.wrappers import common as wrapper_common from icon4pytools.py2fgen.wrappers.wrapper_dimension import ( CellIndexDim, @@ -383,7 +383,7 @@ def solve_nh_run( ndyn_substeps: gtx.float64, idyn_timestep: gtx.int32, ): - logger.info(f"Using Device = {settings.device}") + logger.info(f"Using Device = {device}") prep_adv = dycore_states.PrepAdvection( vn_traj=vn_traj, @@ -484,20 +484,20 @@ def grid_init( num_edges=num_edges, vertical_size=vertical_size, limited_area=limited_area, - on_gpu=True if settings.device == "GPU" else False, - cell_starts=cell_starts, - cell_ends=cell_ends, - vertex_starts=vertex_starts, - vertex_ends=vertex_ends, - edge_starts=edge_starts, - edge_ends=edge_ends, - c2e=c2e, - e2c=e2c, - c2e2c=c2e2c, - e2c2e=e2c2e, - e2v=e2v, - v2e=v2e, - v2c=v2c, - e2c2v=e2c2v, - c2v=c2v, + on_gpu=True if device == "GPU" else False, + cell_starts=cell_starts.ndarray, + cell_ends=cell_ends.ndarray, + vertex_starts=vertex_starts.ndarray, + vertex_ends=vertex_ends.ndarray, + edge_starts=edge_starts.ndarray, + edge_ends=edge_ends.ndarray, + c2e=c2e.ndarray, + e2c=e2c.ndarray, + c2e2c=c2e2c.ndarray, + e2c2e=e2c2e.ndarray, + e2v=e2v.ndarray, + v2e=v2e.ndarray, + v2c=v2c.ndarray, + e2c2v=e2c2v.ndarray, + c2v=c2v.ndarray, ) diff --git a/tools/src/icon4pytools/py2fgen/wrappers/simple.py b/tools/src/icon4pytools/py2fgen/wrappers/simple.py index a474d6f51d..2adf436bbf 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/simple.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/simple.py @@ -15,7 +15,8 @@ from gt4py.next.ffront.decorator import field_operator, program from icon4py.model.common import dimension as dims from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.settings import backend + +from icon4pytools.py2fgen.settings import backend # global profiler object diff --git a/tools/tests/py2fgen/test_codegen.py b/tools/tests/py2fgen/test_codegen.py index 7223fbc64c..d75d24e7a1 100644 --- a/tools/tests/py2fgen/test_codegen.py +++ b/tools/tests/py2fgen/test_codegen.py @@ -249,7 +249,8 @@ def test_python_wrapper(dummy_plugin): import cupy as cp from numpy.typing import NDArray from gt4py.next.iterator.embedded import np_as_located_field -from icon4py.model.common.settings import xp +from icon4pytools.py2fgen.settings import config +xp = config.array_ns from icon4py.model.common import dimension as dims # logger setup diff --git a/tools/tests/py2fgen/utils.py b/tools/tests/py2fgen/utils.py index 823b3cacb0..4ec2828d77 100644 --- a/tools/tests/py2fgen/utils.py +++ b/tools/tests/py2fgen/utils.py @@ -10,9 +10,11 @@ from gt4py.next.embedded.nd_array_field import NdArrayField from icon4py.model.atmosphere.diffusion import diffusion from icon4py.model.atmosphere.dycore import solve_nonhydro as solve_nh -from icon4py.model.common.settings import xp +from icon4pytools.py2fgen.settings import config + +xp = config.array_ns # TODO: the configuration code is replicated across the codebase currently. In future, the configuration should be read from an external file.