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/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_fractional_flux.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_fractional_flux.py index 2ac11aa036..f2098b9ee9 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_fractional_flux.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_fractional_flux.py @@ -25,16 +25,23 @@ def _sum_neighbor_contributions( js: fa.CellKField[ta.wpfloat], p_cc: fa.CellKField[ta.wpfloat], ) -> fa.CellKField[ta.wpfloat]: - p_cc_p0 = where(mask1 & (js == 0.0), p_cc, 0.0) - p_cc_p1 = where(mask1 & (js == 1.0), p_cc(Koff[1]), 0.0) - p_cc_p2 = where(mask1 & (js == 2.0), p_cc(Koff[2]), 0.0) - p_cc_p3 = where(mask1 & (js == 3.0), p_cc(Koff[3]), 0.0) - p_cc_p4 = where(mask1 & (js == 4.0), p_cc(Koff[4]), 0.0) - p_cc_m0 = where(mask2 & (js == 0.0), p_cc(Koff[-1]), 0.0) - p_cc_m1 = where(mask2 & (js == 1.0), p_cc(Koff[-2]), 0.0) - p_cc_m2 = where(mask2 & (js == 2.0), p_cc(Koff[-3]), 0.0) - p_cc_m3 = where(mask2 & (js == 3.0), p_cc(Koff[-4]), 0.0) - p_cc_m4 = where(mask2 & (js == 4.0), p_cc(Koff[-5]), 0.0) + js_eq0 = js == 0.0 + js_eq1 = js == 1.0 + js_eq2 = js == 2.0 + js_eq3 = js == 3.0 + js_eq4 = js == 4.0 + + p_cc_p0 = where(mask1 & js_eq0, p_cc, 0.0) + p_cc_p1 = where(mask1 & js_eq1, p_cc(Koff[1]), 0.0) + p_cc_p2 = where(mask1 & js_eq2, p_cc(Koff[2]), 0.0) + p_cc_p3 = where(mask1 & js_eq3, p_cc(Koff[3]), 0.0) + p_cc_p4 = where(mask1 & js_eq4, p_cc(Koff[4]), 0.0) + p_cc_m0 = where(mask2 & js_eq0, p_cc(Koff[-1]), 0.0) + p_cc_m1 = where(mask2 & js_eq1, p_cc(Koff[-2]), 0.0) + p_cc_m2 = where(mask2 & js_eq2, p_cc(Koff[-3]), 0.0) + p_cc_m3 = where(mask2 & js_eq3, p_cc(Koff[-4]), 0.0) + p_cc_m4 = where(mask2 & js_eq4, p_cc(Koff[-5]), 0.0) + p_cc_jks = ( p_cc_p0 + p_cc_p1 @@ -63,17 +70,21 @@ def _compute_ppm4gpu_fractional_flux( ) -> fa.CellKField[ta.wpfloat]: js = floor(abs(z_cfl)) z_cflfrac = abs(z_cfl) - js + z_cflfrac_nonzero = z_cflfrac != 0.0 z_cfl_pos = z_cfl > 0.0 - z_cfl_neg = not z_cfl_pos + z_cfl_neg = z_cfl < 0.0 wsign = where(z_cfl_pos, 1.0, -1.0) + mask1 = z_cfl_pos & z_cflfrac_nonzero + mask2 = z_cfl_neg & z_cflfrac_nonzero + in_slev_bounds = astype(k, wpfloat) - js >= astype(slev, wpfloat) - p_cc_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cc) - p_cellmass_now_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cellmass_now) - z_delta_q_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, z_delta_q) - z_a1_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, z_a1) + p_cc_jks = _sum_neighbor_contributions(mask1, mask2, js, p_cc) + p_cellmass_now_jks = _sum_neighbor_contributions(mask1, mask2, js, p_cellmass_now) + z_delta_q_jks = _sum_neighbor_contributions(mask1, mask2, js, z_delta_q) + z_a1_jks = _sum_neighbor_contributions(mask1, mask2, js, z_a1) z_q_int = ( p_cc_jks diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_integer_flux.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_integer_flux.py index b6c5b8c7d0..bac2532c25 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_integer_flux.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_integer_flux.py @@ -9,10 +9,8 @@ import gt4py.next as gtx from gt4py.next.ffront.fbuiltins import abs, astype, floor, where -from icon4py.model.atmosphere.advection.stencils.compute_ppm4gpu_fractional_flux import ( - _sum_neighbor_contributions, -) 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.type_alias import wpfloat @@ -20,6 +18,46 @@ # TODO (dastrm): this stencil does not strictly match the fortran code +@gtx.field_operator +def _sum_neighbor_contributions_all( + mask1: fa.CellKField[bool], + mask2: fa.CellKField[bool], + js: fa.CellKField[ta.wpfloat], + p_cc: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], +) -> fa.CellKField[ta.wpfloat]: + js_gt0 = js >= 0.0 + js_gt1 = js >= 1.0 + js_gt2 = js >= 2.0 + js_gt3 = js >= 3.0 + js_gt4 = js >= 4.0 + + prod_p0 = where(mask1 & js_gt0, p_cc * p_cellmass_now, 0.0) + prod_p1 = where(mask1 & js_gt1, p_cc(Koff[1]) * p_cellmass_now(Koff[1]), 0.0) + prod_p2 = where(mask1 & js_gt2, p_cc(Koff[2]) * p_cellmass_now(Koff[2]), 0.0) + prod_p3 = where(mask1 & js_gt3, p_cc(Koff[3]) * p_cellmass_now(Koff[3]), 0.0) + prod_p4 = where(mask1 & js_gt4, p_cc(Koff[4]) * p_cellmass_now(Koff[4]), 0.0) + prod_m0 = where(mask2 & js_gt0, p_cc(Koff[-1]) * p_cellmass_now(Koff[-1]), 0.0) + prod_m1 = where(mask2 & js_gt1, p_cc(Koff[-2]) * p_cellmass_now(Koff[-2]), 0.0) + prod_m2 = where(mask2 & js_gt2, p_cc(Koff[-3]) * p_cellmass_now(Koff[-3]), 0.0) + prod_m3 = where(mask2 & js_gt3, p_cc(Koff[-4]) * p_cellmass_now(Koff[-4]), 0.0) + prod_m4 = where(mask2 & js_gt4, p_cc(Koff[-5]) * p_cellmass_now(Koff[-5]), 0.0) + + prod_jks = ( + prod_p0 + + prod_p1 + + prod_p2 + + prod_p3 + + prod_p4 + + prod_m0 + + prod_m1 + + prod_m2 + + prod_m3 + + prod_m4 + ) + return prod_jks + + @gtx.field_operator def _compute_ppm4gpu_integer_flux( p_cc: fa.CellKField[ta.wpfloat], @@ -33,15 +71,16 @@ def _compute_ppm4gpu_integer_flux( js = floor(abs(z_cfl)) - 1.0 z_cfl_pos = z_cfl > 0.0 - z_cfl_neg = not z_cfl_pos + z_cfl_neg = z_cfl < 0.0 wsign = where(z_cfl_pos, 1.0, -1.0) in_slev_bounds = astype(k, wpfloat) - js >= astype(slev, wpfloat) - p_cc_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cc) - p_cellmass_now_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cellmass_now) + p_cc_cellmass_now_jks = _sum_neighbor_contributions_all( + z_cfl_pos, z_cfl_neg, js, p_cc, p_cellmass_now + ) - z_iflx = wsign * p_cc_jks * p_cellmass_now_jks + z_iflx = wsign * p_cc_cellmass_now_jks p_upflux = p_upflux + where(in_slev_bounds, z_iflx / p_dtime, 0.0) 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 68e4b6bba5..6437de5a2d 100644 --- a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py +++ b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py @@ -5,20 +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.decomposition import definitions from icon4py.model.common.grid import ( - geometry, geometry_attributes as geometry_meta, - icon, 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, @@ -26,7 +22,6 @@ reference_funcs as ref_funcs, serialbox_utils as sb, ) -from icon4py.model.common.utils import gt4py_field_allocation as alloc from .utils import ( compare_dace_orchestration_multiple_steps, @@ -53,25 +48,16 @@ def get_cell_geometry_for_experiment(experiment, backend): def _get_or_initialize(experiment, backend, name): - def _construct_minimal_decomposition_info(grid: icon.IconGrid): - edge_indices = alloc.allocate_indices(dims.EdgeDim, grid) - owner_mask = xp.ones((grid.num_edges,), dtype=bool) - decomposition_info = definitions.DecompositionInfo(klevels=grid.num_levels) - decomposition_info.with_dimension(dims.EdgeDim, edge_indices.ndarray, owner_mask) - return decomposition_info + grid_file = ( + dt_utils.REGIONAL_EXPERIMENT + if experiment == dt_utils.REGIONAL_EXPERIMENT + else dt_utils.R02B04_GLOBAL + ) if not grid_functionality[experiment].get(name): - gm = grid_utils.get_grid_manager_for_experiment(experiment, backend) - grid = gm.grid - decomposition_info = _construct_minimal_decomposition_info(grid) - geometry_ = geometry.GridGeometry( - grid=grid, - decomposition_info=decomposition_info, - backend=backend, - coordinates=gm.coordinates, - extra_fields=gm.geometry, - metadata=geometry_meta.attrs, - ) + geometry_ = grid_utils.get_grid_geometry(backend, experiment, grid_file) + grid = geometry_.grid + cell_params = grid_states.CellParams.from_global_num_cells( cell_center_lat=geometry_.get(geometry_meta.CELL_LAT), cell_center_lon=geometry_.get(geometry_meta.CELL_LON), @@ -174,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 @@ -285,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"] @@ -387,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 @@ -401,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, @@ -414,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) @@ -475,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 @@ -511,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 ###################################################################### @@ -560,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(), @@ -579,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, ) @@ -592,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(), @@ -612,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): @@ -634,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, @@ -647,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) @@ -704,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 a3be7e7d9b..9d9ba6cc3a 100644 --- a/model/common/src/icon4py/model/common/grid/geometry.py +++ b/model/common/src/icon4py/model/common/grid/geometry.py @@ -8,8 +8,10 @@ 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, backend as gtx_backend +from gt4py.next import backend as gtx_backend import icon4py.model.common.grid.geometry_attributes as attrs import icon4py.model.common.math.helpers as math_helpers @@ -26,10 +28,7 @@ horizontal as h_grid, icon, ) -from icon4py.model.common.settings import xp from icon4py.model.common.states import factory, model, utils as state_utils -from icon4py.model.common.states.factory import FieldProvider -from icon4py.model.common.states.model import FieldMetaData InputGeometryFieldType: TypeAlias = Literal[attrs.CELL_AREA, attrs.TANGENT_ORIENTATION] @@ -93,6 +92,7 @@ def __init__( metadata: a dictionary of FieldMetaData for all fields computed in GridGeometry. """ + self._providers = {} self._backend = backend self._allocator = gtx.constructors.zeros.partial(allocator=backend) self._grid = grid @@ -100,7 +100,6 @@ def __init__( self._attrs = metadata self._geometry_type: icon.GeometryType = grid.global_properties.geometry_type self._edge_domain = h_grid.domain(dims.EdgeDim) - self._providers: dict[str, factory.FieldProvider] = {} ( edge_orientation0_lat, @@ -131,11 +130,26 @@ def __init__( input_fields_provider = factory.PrecomputedFieldProvider( { + # TODO (@magdalena) rescaled by grid_length_rescale_factor (mo_grid_tools.f90) + attrs.EDGE_CELL_DISTANCE: extra_fields[gm.GeometryName.EDGE_CELL_DISTANCE], attrs.CELL_AREA: extra_fields[gm.GeometryName.CELL_AREA], + attrs.DUAL_AREA: extra_fields[gm.GeometryName.DUAL_AREA], attrs.TANGENT_ORIENTATION: extra_fields[gm.GeometryName.TANGENT_ORIENTATION], "edge_owner_mask": gtx.as_field( (dims.EdgeDim,), decomposition_info.owner_mask(dims.EdgeDim), dtype=bool ), + attrs.CELL_NORMAL_ORIENTATION: extra_fields[ + gm.GeometryName.CELL_NORMAL_ORIENTATION + ], + attrs.VERTEX_EDGE_ORIENTATION: extra_fields[ + gm.GeometryName.EDGE_ORIENTATION_ON_VERTEX + ], + "vertex_owner_mask": gtx.as_field( + (dims.VertexDim,), decomposition_info.owner_mask(dims.VertexDim) + ), + "cell_owner_mask": gtx.as_field( + (dims.VertexDim,), decomposition_info.owner_mask(dims.CellDim) + ), } ) self.register_provider(input_fields_provider) @@ -445,21 +459,13 @@ def __repr__(self): return f"{self.__class__.__name__} for geometry_type={self._geometry_type._name_} (grid={self._grid.id!r})" @property - def providers(self) -> dict[str, FieldProvider]: - return self._providers - - @property - def metadata(self) -> dict[str, FieldMetaData]: + def metadata(self) -> dict[str, model.FieldMetaData]: return self._attrs @property - def backend(self) -> backend.Backend: + def backend(self) -> gtx_backend.Backend: return self._backend - @property - def grid_provider(self): - return self - @property def grid(self): return self._grid @@ -527,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 @@ -568,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/geometry_attributes.py b/model/common/src/icon4py/model/common/grid/geometry_attributes.py index ca1e81e767..ce9c59833a 100644 --- a/model/common/src/icon4py/model/common/grid/geometry_attributes.py +++ b/model/common/src/icon4py/model/common/grid/geometry_attributes.py @@ -7,6 +7,8 @@ # SPDX-License-Identifier: BSD-3-Clause from typing import Final +import gt4py.next as gtx + from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.states import model @@ -27,7 +29,11 @@ CELL_LAT: Final[str] = "grid_latitude_of_cell_center" CELL_AREA: Final[str] = "cell_area" EDGE_AREA: Final[str] = "edge_area" +DUAL_AREA: Final[str] = "dual_area" +EDGE_CELL_DISTANCE: Final[str] = "edge_midpoint_to_cell_center_distance" TANGENT_ORIENTATION: Final[str] = "edge_orientation" +CELL_NORMAL_ORIENTATION: Final[str] = "orientation_of_normal_to_cell_edges" +VERTEX_EDGE_ORIENTATION: Final[str] = "orientation_of_edges_around_vertex" CORIOLIS_PARAMETER: Final[str] = "coriolis_parameter" @@ -105,6 +111,21 @@ icon_var_name="t_grid_edges%primal_edge_length", dtype=ta.wpfloat, ), + CELL_NORMAL_ORIENTATION: dict( + standard_name=CELL_NORMAL_ORIENTATION, + units="", + dims=(dims.CellDim, dims.C2EDim), + icon_var_name="t_grid_cells%edge_orientation", + dtype=gtx.int32, + ), + EDGE_CELL_DISTANCE: dict( + standard_name=EDGE_CELL_DISTANCE, + long_name="distances between edge midpoint and adjacent triangle midpoints", + units="m", + dims=(dims.EdgeDim, dims.E2CDim), + icon_var_name="t_grid_edges%edge_cell_length", + dtype=ta.wpfloat, + ), DUAL_EDGE_LENGTH: dict( standard_name=DUAL_EDGE_LENGTH, long_name="length of the dual edge", @@ -124,11 +145,27 @@ EDGE_AREA: dict( standard_name=EDGE_AREA, long_name="area of quadrilateral spanned by edge and associated dual edge", - units="m", + units="m2", dims=(dims.EdgeDim,), icon_var_name="t_grid_edges%area_edge", dtype=ta.wpfloat, ), + CELL_AREA: dict( + standard_name=CELL_AREA, + long_name="area of a triangular cell", + units="m2", + dims=(dims.CellDim,), + icon_var_name="t_grid_cells%area", + dtype=ta.wpfloat, + ), + DUAL_AREA: dict( + standard_name=DUAL_AREA, + long_name="area of the dual grid cell (hexagon cell)", + units="m2", + dims=(dims.VertexDim,), + icon_var_name="t_grid_verts%dual_area", + dtype=ta.wpfloat, + ), CORIOLIS_PARAMETER: dict( standard_name=CORIOLIS_PARAMETER, long_name="coriolis parameter at cell edges", @@ -271,6 +308,14 @@ units="1", dims=(dims.EdgeDim,), icon_var_name=f"t_grid_edges%{TANGENT_ORIENTATION}", + dtype=ta.wpfloat, # TODO (@halungge) netcdf: int + ), + VERTEX_EDGE_ORIENTATION: dict( + standard_name=VERTEX_EDGE_ORIENTATION, + long_name="orientation of tangent vector", + units="1", + dims=(dims.VertexDim, dims.V2EDim), + icon_var_name="t_grid_vertex%edge_orientation", dtype=ta.wpfloat, ), } 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 4f40a01316..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__) @@ -192,10 +184,15 @@ class ConnectivityName(FieldName): class GeometryName(FieldName): + # TODO (@halungge) compute from coordinates CELL_AREA = "cell_area" - EDGE_NORMAL_ORIENTATION = "orientation_of_normal" + # TODO (@halungge) compute from coordinates + DUAL_AREA = "dual_area" + CELL_NORMAL_ORIENTATION = "orientation_of_normal" TANGENT_ORIENTATION = "edge_system_orientation" - EDGE_ORIENTATION_ = "edge_orientation" + EDGE_ORIENTATION_ON_VERTEX = "edge_orientation" + # TODO (@halungge) compute from coordinates + EDGE_CELL_DISTANCE = "edge_cell_distance" class CoordinateName(FieldName): @@ -271,15 +268,18 @@ 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}") - data = self.variable(name, indices, dtype=gtx.int32) - return np.transpose(data) if transpose else data + return self.variable(name, indices, transpose=transpose, dtype=gtx.int32) def variable( - self, name: FieldName, indices: np.ndarray = None, dtype: np.dtype = gtx.float64 + self, + name: FieldName, + indices: np.ndarray = None, + transpose=False, + dtype: np.dtype = gtx.float64, ) -> np.ndarray: """Read a field from the grid file. @@ -287,14 +287,16 @@ def variable( Args: name: name of the field to read indices: indices to read + transpose: flag indicateing whether the array needs to be transposed + to match icon4py dimension ordering, defaults to False dtype: datatype of the field """ try: variable = self._dataset.variables[name] - _log.debug(f"reading {name}: {variable}") + _log.debug(f"reading {name}: transposing = {transpose}") data = variable[:] if indices is None else variable[indices] data = np.array(data, dtype=dtype) - return data + return np.transpose(data) if transpose else data except KeyError as err: msg = f"{name} does not exist in dataset" _log.warning(msg) @@ -318,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. @@ -397,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) @@ -456,18 +458,35 @@ def _read_geometry_fields(self, backend: Optional[gtx_backend.Backend]): GeometryName.CELL_AREA.value: gtx.as_field( (dims.CellDim,), self._reader.variable(GeometryName.CELL_AREA), allocator=backend ), + # TODO (@halungge) easily computed from a neighbor_sum V2C over the cell areas? + GeometryName.DUAL_AREA.value: gtx.as_field( + (dims.VertexDim,), self._reader.variable(GeometryName.DUAL_AREA) + ), + GeometryName.EDGE_CELL_DISTANCE.value: gtx.as_field( + (dims.EdgeDim, dims.E2CDim), + self._reader.variable(GeometryName.EDGE_CELL_DISTANCE, transpose=True), + ), + # TODO (@halungge) recompute from coordinates? field in gridfile contains NaN on boundary edges GeometryName.TANGENT_ORIENTATION.value: gtx.as_field( (dims.EdgeDim,), self._reader.variable(GeometryName.TANGENT_ORIENTATION), allocator=backend, ), + GeometryName.CELL_NORMAL_ORIENTATION.value: gtx.as_field( + (dims.CellDim, dims.C2EDim), + self._reader.int_variable(GeometryName.CELL_NORMAL_ORIENTATION, transpose=True), + ), + GeometryName.EDGE_ORIENTATION_ON_VERTEX.value: gtx.as_field( + (dims.VertexDim, dims.V2EDim), + self._reader.int_variable(GeometryName.EDGE_ORIENTATION_ON_VERTEX, transpose=True), + ), } 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], ]: """ " @@ -522,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, @@ -679,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. @@ -718,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. @@ -757,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---- @@ -786,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 02a3bb4f67..46ad1c584d 100644 --- a/model/common/src/icon4py/model/common/grid/icon.py +++ b/model/common/src/icon4py/model/common/grid/icon.py @@ -117,9 +117,20 @@ def __init__(self, id_: uuid.UUID): def __repr__(self): return f"{self.__class__.__name__}: id={self._id}, R{self.global_properties.root}B{self.global_properties.level}" + def __eq__(self, other: "IconGrid"): + """TODO (@halungge) this might not be enough at least for the distributed case: we might additional properties like sizes""" + if isinstance(other, IconGrid): + return self.id == other.id + + else: + return False + @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_attributes.py b/model/common/src/icon4py/model/common/interpolation/interpolation_attributes.py new file mode 100644 index 0000000000..cd8de9139b --- /dev/null +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_attributes.py @@ -0,0 +1,89 @@ +# 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 typing import Final + +from icon4py.model.common import dimension as dims, type_alias as ta +from icon4py.model.common.states import model + + +C_LIN_E: Final[str] = "interpolation_coefficient_from_cell_to_edge" +C_BLN_AVG: Final[str] = "bilinear_cell_average_weight" +GEOFAC_DIV: Final[str] = "geometrical_factor_for_divergence" +GEOFAC_ROT: Final[str] = "geometrical_factor_for_curl" +GEOFAC_N2S: Final[str] = "geometrical_factor_for_nabla_2_scalar" +GEOFAC_GRDIV: Final[str] = "geometrical_factor_for_gradient_of_divergence" +GEOFAC_GRG_X: Final[str] = "geometrical_factor_for_green_gauss_gradient_x" +GEOFAC_GRG_Y: Final[str] = "geometrical_factor_for_green_gauss_gradient_y" + +attrs: dict[str, model.FieldMetaData] = { + C_LIN_E: dict( + standard_name=C_LIN_E, + long_name="interpolation coefficient from cell to edges", + units="", # TODO (@halungge) check or confirm + dims=(dims.EdgeDim, dims.E2CDim), + icon_var_name="c_lin_e", + dtype=ta.wpfloat, + ), + C_BLN_AVG: dict( + standard_name=C_BLN_AVG, + long_name="mass conserving bilinear cell average weight", + units="", # TODO (@halungge) check or confirm + dims=(dims.EdgeDim, dims.E2CDim), + icon_var_name="c_lin_e", + dtype=ta.wpfloat, + ), + GEOFAC_DIV: dict( + standard_name=GEOFAC_DIV, + long_name="geometrical factor for divergence", # TODO (@halungge) find proper description + units="", # TODO (@halungge) check or confirm + dims=(dims.CellDim, dims.C2EDim), + icon_var_name="geofac_div", + dtype=ta.wpfloat, + ), + GEOFAC_ROT: dict( + standard_name=GEOFAC_ROT, + long_name="geometrical factor for curl", + units="", # TODO (@halungge) check or confirm + dims=(dims.VertexDim, dims.V2EDim), + icon_var_name="geofac_rot", + dtype=ta.wpfloat, + ), + GEOFAC_N2S: dict( + standard_name=GEOFAC_N2S, + long_name="geometrical factor nabla-2 scalar", + units="", # TODO (@halungge) check or confirm + dims=(dims.CellDim, dims.C2E2CODim), + icon_var_name="geofac_n2s", + dtype=ta.wpfloat, + ), + GEOFAC_GRDIV: dict( + standard_name=GEOFAC_GRDIV, + long_name="geometrical factor for gradient of divergence", + units="", # TODO (@halungge) check or confirm + dims=(dims.EdgeDim, dims.E2C2EODim), + icon_var_name="geofac_grdiv", + dtype=ta.wpfloat, + ), + GEOFAC_GRG_X: dict( + standard_name=GEOFAC_GRG_X, + long_name="geometrical factor for Green Gauss gradient (first component)", + units="", # TODO (@halungge) check or confirm + dims=(dims.CellDim, dims.C2E2CODim), + icon_var_name="geofac_grg", + dtype=ta.wpfloat, + ), + GEOFAC_GRG_Y: dict( + standard_name=GEOFAC_GRG_Y, + long_name="geometrical factor for Green Gauss gradient (second component)", + units="", # TODO (@halungge) check or confirm + dims=(dims.CellDim, dims.C2E2CODim), + icon_var_name="geofac_grg", + dtype=ta.wpfloat, + ), +} diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py new file mode 100644 index 0000000000..62f4f6e85f --- /dev/null +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py @@ -0,0 +1,201 @@ +# 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 +import functools + +import gt4py.next as gtx +from gt4py.next import backend as gtx_backend + +from icon4py.model.common import dimension as dims +from icon4py.model.common.decomposition import definitions +from icon4py.model.common.grid import ( + geometry, + geometry_attributes as geometry_attrs, + horizontal as h_grid, + icon, +) +from icon4py.model.common.interpolation import ( + interpolation_attributes as attrs, + interpolation_fields, +) +from icon4py.model.common.states import factory, model +from icon4py.model.common.utils import gt4py_field_allocation as alloc + + +cell_domain = h_grid.domain(dims.CellDim) +edge_domain = h_grid.domain(dims.EdgeDim) + + +class InterpolationFieldsFactory(factory.FieldSource, factory.GridProvider): + def __init__( + self, + grid: icon.IconGrid, + decomposition_info: definitions.DecompositionInfo, + geometry_source: geometry.GridGeometry, + backend: gtx_backend.Backend, + metadata: dict[str, model.FieldMetaData], + ): + self._backend = backend + self._xp = alloc.import_array_ns(backend) + self._allocator = gtx.constructors.zeros.partial(allocator=backend) + self._grid = grid + self._decomposition_info = decomposition_info + self._attrs = metadata + self._providers: dict[str, factory.FieldProvider] = {} + self._geometry = geometry_source + # TODO @halungge: Dummy config dict - to be replaced by real configuration + self._config = {"divavg_cntrwgt": 0.5} + self._register_computed_fields() + + def __repr__(self): + return f"{self.__class__.__name__} on (grid={self._grid!r}) providing fields f{self.metadata.keys()}" + + @property + def _sources(self) -> factory.FieldSource: + return factory.CompositeSource(self, (self._geometry,)) + + def _register_computed_fields(self): + geofac_div = factory.FieldOperatorProvider( + # needs to be computed on fieldview-embedded backend + func=interpolation_fields.compute_geofac_div.with_backend(None), + domain=(dims.CellDim, dims.C2EDim), + fields={attrs.GEOFAC_DIV: attrs.GEOFAC_DIV}, + deps={ + "primal_edge_length": geometry_attrs.EDGE_LENGTH, + "edge_orientation": geometry_attrs.CELL_NORMAL_ORIENTATION, + "area": geometry_attrs.CELL_AREA, + }, + ) + self.register_provider(geofac_div) + geofac_rot = factory.FieldOperatorProvider( + # needs to be computed on fieldview-embedded backend + func=interpolation_fields.compute_geofac_rot.with_backend(None), + domain=(dims.VertexDim, dims.V2EDim), + fields={attrs.GEOFAC_ROT: attrs.GEOFAC_ROT}, + deps={ + "dual_edge_length": geometry_attrs.DUAL_EDGE_LENGTH, + "edge_orientation": geometry_attrs.VERTEX_EDGE_ORIENTATION, + "dual_area": geometry_attrs.DUAL_AREA, + "owner_mask": "vertex_owner_mask", + }, + ) + self.register_provider(geofac_rot) + + geofac_n2s = factory.NumpyFieldsProvider( + func=functools.partial(interpolation_fields.compute_geofac_n2s, array_ns=self._xp), + fields=(attrs.GEOFAC_N2S,), + domain=(dims.CellDim, dims.C2E2CODim), + deps={ + "dual_edge_length": geometry_attrs.DUAL_EDGE_LENGTH, + "geofac_div": attrs.GEOFAC_DIV, + }, + connectivities={"c2e": dims.C2EDim, "e2c": dims.E2CDim, "c2e2c": dims.C2E2CDim}, + params={ + "horizontal_start": self._grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ) + }, + ) + self.register_provider(geofac_n2s) + + geofac_grdiv = factory.NumpyFieldsProvider( + func=functools.partial(interpolation_fields.compute_geofac_grdiv, array_ns=self._xp), + fields=(attrs.GEOFAC_GRDIV,), + domain=(dims.EdgeDim, dims.E2C2EODim), + deps={ + "geofac_div": attrs.GEOFAC_DIV, + "inv_dual_edge_length": f"inverse_of_{geometry_attrs.DUAL_EDGE_LENGTH}", + "owner_mask": "edge_owner_mask", + }, + connectivities={"c2e": dims.C2EDim, "e2c": dims.E2CDim, "e2c2e": dims.E2C2EDim}, + params={ + "horizontal_start": self._grid.start_index( + edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ) + }, + ) + + self.register_provider(geofac_grdiv) + + cell_average_weight = factory.NumpyFieldsProvider( + func=functools.partial( + interpolation_fields.compute_mass_conserving_bilinear_cell_average_weight, + array_ns=self._xp, + ), + fields=(attrs.C_BLN_AVG,), + domain=(dims.CellDim, dims.C2E2CODim), + deps={ + "lat": geometry_attrs.CELL_LAT, + "lon": geometry_attrs.CELL_LON, + "cell_areas": geometry_attrs.CELL_AREA, + "cell_owner_mask": "cell_owner_mask", + }, + connectivities={"c2e2c0": dims.C2E2CODim}, + params={ + "horizontal_start": self.grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ), + "horizontal_start_level_3": self.grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) + ), + "divavg_cntrwgt": self._config["divavg_cntrwgt"], + }, + ) + self.register_provider(cell_average_weight) + + c_lin_e = factory.NumpyFieldsProvider( + func=functools.partial(interpolation_fields.compute_c_lin_e, array_ns=self._xp), + fields=(attrs.C_LIN_E,), + domain=(dims.EdgeDim, dims.E2CDim), + deps={ + "edge_cell_length": geometry_attrs.EDGE_CELL_DISTANCE, + "inv_dual_edge_length": f"inverse_of_{geometry_attrs.DUAL_EDGE_LENGTH}", + "edge_owner_mask": "edge_owner_mask", + }, + params={ + "horizontal_start": self._grid.start_index( + edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ) + }, + ) + self.register_provider(c_lin_e) + geofac_grg = factory.NumpyFieldsProvider( + func=functools.partial(interpolation_fields.compute_geofac_grg, array_ns=self._xp), + fields=(attrs.GEOFAC_GRG_X, attrs.GEOFAC_GRG_Y), + domain=(dims.CellDim, dims.C2E2CODim), + deps={ + "primal_normal_cell_x": geometry_attrs.EDGE_NORMAL_CELL_U, + "primal_normal_cell_y": geometry_attrs.EDGE_NORMAL_CELL_V, + "owner_mask": "cell_owner_mask", + "geofac_div": attrs.GEOFAC_DIV, + "c_lin_e": attrs.C_LIN_E, + }, + connectivities={"c2e": dims.C2EDim, "e2c": dims.E2CDim, "c2e2c": dims.C2E2CDim}, + params={ + "horizontal_start": self.grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ) + }, + ) + + self.register_provider(geofac_grg) + + @property + def metadata(self) -> dict[str, model.FieldMetaData]: + return self._attrs + + @property + def backend(self) -> gtx_backend.Backend: + return self._backend + + @property + def grid(self): + return self._grid + + @property + def vertical_grid(self): + return None 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 497c275b33..8244bc5880 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py @@ -5,6 +5,10 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +import functools +import math +from types import ModuleType + import gt4py.next as gtx import numpy as np from gt4py.next import where @@ -15,31 +19,33 @@ 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 field_alloc def compute_c_lin_e( - edge_cell_length: np.ndarray, - inv_dual_edge_length: np.ndarray, - owner_mask: np.ndarray, + edge_cell_length: field_alloc.NDArray, + inv_dual_edge_length: field_alloc.NDArray, + edge_owner_mask: field_alloc.NDArray, horizontal_start: np.int32, -) -> np.ndarray: + array_ns: ModuleType = np, +) -> field_alloc.NDArray: """ Compute E2C average inverse distance. Args: - edge_cell_length: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], ta.wpfloat] - inv_dual_edge_length: inverse dual edge length, numpy array representing a gtx.Field[gtx.Dims[EdgeDim], ta.wpfloat] - owner_mask: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim], bool]boolean field, True for all edges owned by this compute node - horizontal_start: start index of the 2nd boundary line: c_lin_e is not calculated for the first boundary layer - + edge_cell_length: ndarray, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], ta.wpfloat] + inv_dual_edge_length: ndarray, inverse dual edge length, numpy array representing a gtx.Field[gtx.Dims[EdgeDim], ta.wpfloat] + edge_owner_mask: ndarray, representing a gtx.Field[gtx.Dims[EdgeDim], bool]boolean field, True for all edges owned by this compute node + horizontal_start: start index from the field is computed: c_lin_e is not calculated for the first boundary layer + array_ns: ModuleType to use for the computation, numpy or cupy, defaults to cupy Returns: c_lin_e: numpy array, representing gtx.Field[gtx.Dims[EdgeDim, E2CDim], ta.wpfloat] """ c_lin_e_ = edge_cell_length[:, 1] * inv_dual_edge_length - c_lin_e = np.transpose(np.vstack((c_lin_e_, (1.0 - c_lin_e_)))) + c_lin_e = array_ns.transpose(array_ns.vstack((c_lin_e_, (1.0 - c_lin_e_)))) c_lin_e[0:horizontal_start, :] = 0.0 - mask = np.transpose(np.tile(owner_mask, (2, 1))) - return np.where(mask, c_lin_e, 0.0) + mask = array_ns.transpose(array_ns.tile(edge_owner_mask, (2, 1))) + return array_ns.where(mask, c_lin_e, 0.0) @gtx.field_operator @@ -85,134 +91,148 @@ def compute_geofac_rot( def compute_geofac_n2s( - dual_edge_length: np.ndarray, - geofac_div: np.ndarray, - c2e: np.ndarray, - e2c: np.ndarray, - c2e2c: np.ndarray, + 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, -) -> np.ndarray: + array_ns: ModuleType = np, +) -> field_alloc.NDArray: """ Compute geometric factor for nabla2-scalar. Args: - dual_edge_length: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim], ta.wpfloat] - geofac_div: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], ta.wpfloat] - c2e: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], gtx.int32] - e2c: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], gtx.int32] - c2e2c: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2E2CDim], gtx.int32] - horizontal_start: + dual_edge_length: ndarray, representing a gtx.Field[gtx.Dims[EdgeDim], ta.wpfloat] + geofac_div: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], ta.wpfloat] + c2e: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], gtx.int32] + e2c: ndarray, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], gtx.int32] + c2e2c: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2E2CDim], gtx.int32] + horizontal_start: start index from where the field is computed + array_ns: python module, numpy or cpu defaults to numpy Returns: geometric factor for nabla2-scalar, Field[CellDim, C2E2CODim] """ - llb = horizontal_start - geofac_n2s = np.zeros([c2e.shape[0], 4]) - index = np.transpose( - np.vstack( + num_cells = c2e.shape[0] + geofac_n2s = array_ns.zeros([num_cells, 4]) + index = array_ns.transpose( + array_ns.vstack( ( - np.arange(c2e.shape[0]), - np.arange(c2e.shape[0]), - np.arange(c2e.shape[0]), + array_ns.arange(num_cells), + array_ns.arange(num_cells), + array_ns.arange(num_cells), ) ) ) mask = e2c[c2e, 0] == index - geofac_n2s[llb:, 0] = geofac_n2s[llb:, 0] - np.sum( - mask[llb:] * (geofac_div / dual_edge_length[c2e])[llb:], axis=1 + geofac_n2s[horizontal_start:, 0] = geofac_n2s[horizontal_start:, 0] - array_ns.sum( + mask[horizontal_start:] * (geofac_div / dual_edge_length[c2e])[horizontal_start:], axis=1 ) mask = e2c[c2e, 1] == index - geofac_n2s[llb:, 0] = geofac_n2s[llb:, 0] + np.sum( - mask[llb:] * (geofac_div / dual_edge_length[c2e])[llb:], axis=1 + geofac_n2s[horizontal_start:, 0] = geofac_n2s[horizontal_start:, 0] + array_ns.sum( + mask[horizontal_start:] * (geofac_div / dual_edge_length[c2e])[horizontal_start:], axis=1 ) mask = e2c[c2e, 0] == c2e2c - geofac_n2s[llb:, 1:] = ( - geofac_n2s[llb:, 1:] - mask[llb:, :] * (geofac_div / dual_edge_length[c2e])[llb:, :] + geofac_n2s[horizontal_start:, 1:] = ( + geofac_n2s[horizontal_start:, 1:] + - mask[horizontal_start:, :] * (geofac_div / dual_edge_length[c2e])[horizontal_start:, :] ) mask = e2c[c2e, 1] == c2e2c - geofac_n2s[llb:, 1:] = ( - geofac_n2s[llb:, 1:] + mask[llb:, :] * (geofac_div / dual_edge_length[c2e])[llb:, :] + geofac_n2s[horizontal_start:, 1:] = ( + geofac_n2s[horizontal_start:, 1:] + + mask[horizontal_start:, :] * (geofac_div / dual_edge_length[c2e])[horizontal_start:, :] ) return geofac_n2s -def compute_primal_normal_ec( - primal_normal_cell_x: np.ndarray, - primal_normal_cell_y: np.ndarray, - owner_mask: np.ndarray, - c2e: np.ndarray, - e2c: np.ndarray, +def _compute_primal_normal_ec( + 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, -) -> np.ndarray: + array_ns: ModuleType = np, +) -> field_alloc.NDArray: """ Compute primal_normal_ec. Args: - primal_normal_cell_x: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], gtx.int32] - primal_normal_cell_y: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], gtx.int32] - owner_mask: numpy array, representing a gtx.Field[gtx.Dims[CellDim], bool] - c2e: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], gtx.int32] - e2c: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], gtx.int32] - horizontal_start: - + primal_normal_cell_x: ndarray, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], gtx.int32] + primal_normal_cell_y: ndarray, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], gtx.int32] + owner_mask: ndarray, representing a gtx.Field[gtx.Dims[CellDim], bool] + c2e: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], gtx.int32] + e2c: ndarray, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], gtx.int32] + horizontal_start: start index to compute from + array_ns: module - the array interface implementation to compute on, defaults to numpy Returns: primal_normal_ec: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim, 2], ta.wpfloat] """ - llb = horizontal_start + num_cells = c2e.shape[0] primal_normal_ec = np.zeros([c2e.shape[0], c2e.shape[1], 2]) index = np.transpose( np.vstack( ( - np.arange(c2e.shape[0]), - np.arange(c2e.shape[0]), - np.arange(c2e.shape[0]), + array_ns.arange(num_cells), + array_ns.arange(num_cells), + array_ns.arange(num_cells), ) ) ) + owned = np.vstack((owner_mask, owner_mask, owner_mask)).T for i in range(2): mask = e2c[c2e, i] == index - primal_normal_ec[llb:, :, 0] = primal_normal_ec[llb:, :, 0] + np.where( - owner_mask, mask[llb:, :] * primal_normal_cell_x[c2e[llb:], i], 0.0 + primal_normal_ec[horizontal_start:, :, 0] = primal_normal_ec[ + horizontal_start:, :, 0 + ] + array_ns.where( + owned[horizontal_start:, :], + mask[horizontal_start:, :] * primal_normal_cell_x[c2e[horizontal_start:], i], + 0.0, ) - primal_normal_ec[llb:, :, 1] = primal_normal_ec[llb:, :, 1] + np.where( - owner_mask, mask[llb:, :] * primal_normal_cell_y[c2e[llb:], i], 0.0 + primal_normal_ec[horizontal_start:, :, 1] = primal_normal_ec[ + horizontal_start:, :, 1 + ] + array_ns.where( + owned[horizontal_start:, :], + mask[horizontal_start:, :] * primal_normal_cell_y[c2e[horizontal_start:], i], + 0.0, ) return primal_normal_ec -def compute_geofac_grg( - primal_normal_ec: np.ndarray, - geofac_div: np.ndarray, - c_lin_e: np.ndarray, - c2e: np.ndarray, - e2c: np.ndarray, - c2e2c: np.ndarray, +def _compute_geofac_grg( + 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, -) -> np.ndarray: + array_ns: ModuleType = np, +) -> tuple[field_alloc.NDArray, field_alloc.NDArray]: """ Compute geometrical factor for Green-Gauss gradient. Args: - primal_normal_ec: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim, 2], ta.wpfloat] - geofac_div: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], ta.wpfloat] - c_lin_e: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], ta.wpfloat] - c2e: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], gtx.int32] - e2c: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], gtx.int32] - c2e2c: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2E2CDim], gtx.int32] - horizontal_start: - + primal_normal_ec: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2EDim, 2], ta.wpfloat] + geofac_div: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], ta.wpfloat] + c_lin_e: ndarray, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], ta.wpfloat] + c2e: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], gtx.int32] + e2c: ndarray, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], gtx.int32] + c2e2c: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2E2CDim], gtx.int32] + horizontal_start: start index from where the computation is done + array_ns: module - the array interface implementation to compute on, defaults to numpy Returns: - geofac_grg: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim + 1, 2], ta.wpfloat] + geofac_grg: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2EDim + 1, 2], ta.wpfloat] """ - llb = horizontal_start num_cells = c2e.shape[0] - geofac_grg = np.zeros([num_cells, c2e.shape[1] + 1, primal_normal_ec.shape[2]]) - index = np.transpose( - np.vstack( + geofac_grg = array_ns.zeros([num_cells, c2e.shape[1] + 1, primal_normal_ec.shape[2]]) + index = array_ns.transpose( + array_ns.vstack( ( - np.arange(num_cells), - np.arange(num_cells), - np.arange(num_cells), + array_ns.arange(num_cells), + array_ns.arange(num_cells), + array_ns.arange(num_cells), ) ) ) @@ -220,84 +240,117 @@ def compute_geofac_grg( mask = e2c[c2e, k] == index for i in range(primal_normal_ec.shape[2]): for j in range(c2e.shape[1]): - geofac_grg[llb:, 0, i] = ( - geofac_grg[llb:, 0, i] - + mask[llb:, j] - * (primal_normal_ec[:, :, i] * geofac_div * c_lin_e[c2e, k])[llb:, j] + geofac_grg[horizontal_start:, 0, i] = ( + geofac_grg[horizontal_start:, 0, i] + + mask[horizontal_start:, j] + * (primal_normal_ec[:, :, i] * geofac_div * c_lin_e[c2e, k])[ + horizontal_start:, j + ] ) for k in range(e2c.shape[1]): mask = e2c[c2e, k] == c2e2c for i in range(primal_normal_ec.shape[2]): for j in range(c2e.shape[1]): - geofac_grg[llb:, 1 + j, i] = ( - geofac_grg[llb:, 1 + j, i] - + mask[llb:, j] - * (primal_normal_ec[:, :, i] * geofac_div * c_lin_e[c2e, k])[llb:, j] + geofac_grg[horizontal_start:, 1 + j, i] = ( + geofac_grg[horizontal_start:, 1 + j, i] + + mask[horizontal_start:, j] + * (primal_normal_ec[:, :, i] * geofac_div * c_lin_e[c2e, k])[ + horizontal_start:, j + ] ) - return geofac_grg + return geofac_grg[:, :, 0], geofac_grg[:, :, 1] + + +def compute_geofac_grg( + 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[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 + ) + return functools.partial(_compute_geofac_grg, array_ns=array_ns)( + primal_normal_ec, geofac_div, c_lin_e, c2e, e2c, c2e2c, horizontal_start + ) def compute_geofac_grdiv( - geofac_div: np.ndarray, - inv_dual_edge_length: np.ndarray, - owner_mask: np.ndarray, - c2e: np.ndarray, - e2c: np.ndarray, - e2c2e: np.ndarray, + 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, -) -> np.ndarray: + array_ns: ModuleType = np, +) -> field_alloc.NDArray: """ Compute geometrical factor for gradient of divergence (triangles only). Args: - geofac_div: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], ta.wpfloat] - inv_dual_edge_length: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim], ta.wpfloat] - owner_mask: numpy array, representing a gtx.Field[gtx.Dims[CellDim], bool] - c2e: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], gtx.int32] - e2c: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], gtx.int32] - e2c2e: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2C2EDim], gtx.int32] + geofac_div: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], ta.wpfloat] + inv_dual_edge_length: ndarray, representing a gtx.Field[gtx.Dims[EdgeDim], ta.wpfloat] + owner_mask: ndarray, representing a gtx.Field[gtx.Dims[CellDim], bool] + c2e: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], gtx.int32] + e2c: ndarray, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], gtx.int32] + e2c2e: ndarray, representing a gtx.Field[gtx.Dims[EdgeDim, E2C2EDim], gtx.int32] horizontal_start: + array_ns: module either used or array computations defaults to numpy Returns: - geofac_grdiv: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2C2EODim], ta.wpfloat] + geofac_grdiv: ndarray, representing a gtx.Field[gtx.Dims[EdgeDim, E2C2EODim], ta.wpfloat] """ - llb = horizontal_start num_edges = e2c.shape[0] - geofac_grdiv = np.zeros([num_edges, 1 + 2 * e2c.shape[1]]) - index = np.arange(llb, num_edges) + geofac_grdiv = array_ns.zeros([num_edges, 1 + 2 * e2c.shape[1]]) + index = array_ns.arange(horizontal_start, num_edges) for j in range(c2e.shape[1]): - mask = np.where(c2e[e2c[llb:, 1], j] == index, owner_mask[llb:], False) - geofac_grdiv[llb:, 0] = np.where(mask, geofac_div[e2c[llb:, 1], j], geofac_grdiv[llb:, 0]) + mask = array_ns.where( + c2e[e2c[horizontal_start:, 1], j] == index, owner_mask[horizontal_start:], False + ) + geofac_grdiv[horizontal_start:, 0] = array_ns.where( + mask, geofac_div[e2c[horizontal_start:, 1], j], geofac_grdiv[horizontal_start:, 0] + ) for j in range(c2e.shape[1]): - mask = np.where(c2e[e2c[llb:, 0], j] == index, owner_mask[llb:], False) - geofac_grdiv[llb:, 0] = np.where( + mask = array_ns.where( + c2e[e2c[horizontal_start:, 0], j] == index, owner_mask[horizontal_start:], False + ) + geofac_grdiv[horizontal_start:, 0] = array_ns.where( mask, - (geofac_grdiv[llb:, 0] - geofac_div[e2c[llb:, 0], j]) * inv_dual_edge_length[llb:], - geofac_grdiv[llb:, 0], + (geofac_grdiv[horizontal_start:, 0] - geofac_div[e2c[horizontal_start:, 0], j]) + * inv_dual_edge_length[horizontal_start:], + geofac_grdiv[horizontal_start:, 0], ) for j in range(e2c.shape[1]): for k in range(c2e.shape[1]): - mask = c2e[e2c[llb:, 0], k] == e2c2e[llb:, j] - geofac_grdiv[llb:, e2c.shape[1] - 1 + j] = np.where( + mask = c2e[e2c[horizontal_start:, 0], k] == e2c2e[horizontal_start:, j] + geofac_grdiv[horizontal_start:, e2c.shape[1] - 1 + j] = array_ns.where( mask, - -geofac_div[e2c[llb:, 0], k] * inv_dual_edge_length[llb:], - geofac_grdiv[llb:, e2c.shape[1] - 1 + j], + -geofac_div[e2c[horizontal_start:, 0], k] * inv_dual_edge_length[horizontal_start:], + geofac_grdiv[horizontal_start:, e2c.shape[1] - 1 + j], ) - mask = c2e[e2c[llb:, 1], k] == e2c2e[llb:, e2c.shape[1] + j] - geofac_grdiv[llb:, 2 * e2c.shape[1] - 1 + j] = np.where( + mask = c2e[e2c[horizontal_start:, 1], k] == e2c2e[horizontal_start:, e2c.shape[1] + j] + geofac_grdiv[horizontal_start:, 2 * e2c.shape[1] - 1 + j] = array_ns.where( mask, - geofac_div[e2c[llb:, 1], k] * inv_dual_edge_length[llb:], - geofac_grdiv[llb:, 2 * e2c.shape[1] - 1 + j], + geofac_div[e2c[horizontal_start:, 1], k] * inv_dual_edge_length[horizontal_start:], + geofac_grdiv[horizontal_start:, 2 * e2c.shape[1] - 1 + j], ) return geofac_grdiv def rotate_latlon( - lat: np.ndarray, - lon: np.ndarray, - pollat: np.ndarray, - pollon: np.ndarray, -) -> tuple[np.ndarray, np.ndarray]: + lat: field_alloc.NDArray, + lon: field_alloc.NDArray, + pollat: field_alloc.NDArray, + pollon: field_alloc.NDArray, + array_ns: ModuleType = np, +) -> tuple[field_alloc.NDArray, field_alloc.NDArray]: """ (Compute rotation of lattitude and longitude.) @@ -309,29 +362,35 @@ def rotate_latlon( lon: scalar or numpy array pollat: scalar or numpy array pollon: scalar or numpy array + array_ns array namespace to be used, defaults to numpy Returns: rotlat: rotlon: """ - rotlat = np.arcsin( - np.sin(lat) * np.sin(pollat) + np.cos(lat) * np.cos(pollat) * np.cos(lon - pollon) + rotlat = array_ns.arcsin( + array_ns.sin(lat) * array_ns.sin(pollat) + + array_ns.cos(lat) * array_ns.cos(pollat) * array_ns.cos(lon - pollon) ) - rotlon = np.arctan2( - np.cos(lat) * np.sin(lon - pollon), - (np.cos(lat) * np.sin(pollat) * np.cos(lon - pollon) - np.sin(lat) * np.cos(pollat)), + rotlon = array_ns.arctan2( + array_ns.cos(lat) * array_ns.sin(lon - pollon), + ( + array_ns.cos(lat) * array_ns.sin(pollat) * array_ns.cos(lon - pollon) + - array_ns.sin(lat) * array_ns.cos(pollat) + ), ) return (rotlat, rotlon) -def weighting_factors( - ytemp: np.ndarray, - xtemp: np.ndarray, - yloc: np.ndarray, - xloc: np.ndarray, +def _weighting_factors( + ytemp: field_alloc.NDArray, + xtemp: field_alloc.NDArray, + yloc: field_alloc.NDArray, + xloc: field_alloc.NDArray, wgt_loc: ta.wpfloat, -) -> np.ndarray: + array_ns: ModuleType = np, +) -> field_alloc.NDArray: """ Compute weighting factors. The weighting factors are based on the requirement that sum(w(i)*x(i)) = 0 @@ -350,33 +409,36 @@ def weighting_factors( yloc: \\ numpy array of size [[flexible], ta.wpfloat] xloc: // wgt_loc: + array_ns: array namespace to be used defaults to numpy Returns: wgt: numpy array of size [[3, flexible], ta.wpfloat] """ - pollat = np.where(yloc >= 0.0, yloc - np.pi * 0.5, yloc + np.pi * 0.5) + rotate = functools.partial(rotate_latlon, array_ns=array_ns) + + pollat = array_ns.where(yloc >= 0.0, yloc - math.pi * 0.5, yloc + math.pi * 0.5) pollon = xloc - (yloc, xloc) = rotate_latlon(yloc, xloc, pollat, pollon) - x = np.zeros([ytemp.shape[0], ytemp.shape[1]]) - y = np.zeros([ytemp.shape[0], ytemp.shape[1]]) - wgt = np.zeros([ytemp.shape[0], ytemp.shape[1]]) + (yloc, xloc) = rotate(yloc, xloc, pollat, pollon) + x = array_ns.zeros([ytemp.shape[0], ytemp.shape[1]]) + y = array_ns.zeros([ytemp.shape[0], ytemp.shape[1]]) + wgt = array_ns.zeros([ytemp.shape[0], ytemp.shape[1]]) for i in range(ytemp.shape[0]): - (ytemp[i], xtemp[i]) = rotate_latlon(ytemp[i], xtemp[i], pollat, pollon) + (ytemp[i], xtemp[i]) = rotate(ytemp[i], xtemp[i], pollat, pollon) y[i] = ytemp[i] - yloc x[i] = xtemp[i] - xloc # This is needed when the date line is crossed - x[i] = np.where(x[i] > 3.5, x[i] - np.pi * 2, x[i]) - x[i] = np.where(x[i] < -3.5, x[i] + np.pi * 2, x[i]) + x[i] = array_ns.where(x[i] > 3.5, x[i] - math.pi * 2, x[i]) + x[i] = array_ns.where(x[i] < -3.5, x[i] + math.pi * 2, x[i]) - mask = np.logical_and(abs(x[1] - x[0]) > 1.0e-11, abs(y[2] - y[0]) > 1.0e-11) + mask = array_ns.logical_and(abs(x[1] - x[0]) > 1.0e-11, abs(y[2] - y[0]) > 1.0e-11) wgt_1_no_mask = ( 1.0 / ((y[1] - y[0]) - (x[1] - x[0]) * (y[2] - y[0]) / (x[2] - x[0])) * (1.0 - wgt_loc) * (-y[0] + x[0] * (y[2] - y[0]) / (x[2] - x[0])) ) - wgt[2] = np.where( + wgt[2] = array_ns.where( mask, 1.0 / ((y[2] - y[0]) - (x[2] - x[0]) * (y[1] - y[0]) / (x[1] - x[0])) @@ -384,7 +446,7 @@ def weighting_factors( * (-y[0] + x[0] * (y[1] - y[0]) / (x[1] - x[0])), (-(1.0 - wgt_loc) * x[0] - wgt_1_no_mask * (x[1] - x[0])) / (x[2] - x[0]), ) - wgt[1] = np.where( + wgt[1] = array_ns.where( mask, (-(1.0 - wgt_loc) * x[0] - wgt[2] * (x[2] - x[0])) / (x[1] - x[0]), wgt_1_no_mask, @@ -394,12 +456,13 @@ def weighting_factors( def _compute_c_bln_avg( - c2e2c: np.ndarray, - lat: np.ndarray, - lon: np.ndarray, + c2e2c: field_alloc.NDArray, + lat: field_alloc.NDArray, + lon: field_alloc.NDArray, divavg_cntrwgt: ta.wpfloat, horizontal_start: np.int32, -) -> np.ndarray: + array_ns: ModuleType = np, +) -> field_alloc.NDArray: """ Compute bilinear cell average weight. @@ -415,21 +478,22 @@ def _compute_c_bln_avg( c_bln_avg: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], ta.wpfloat] """ num_cells = c2e2c.shape[0] - ytemp = np.zeros([c2e2c.shape[1], num_cells - horizontal_start]) - xtemp = np.zeros([c2e2c.shape[1], num_cells - horizontal_start]) + ytemp = array_ns.zeros([c2e2c.shape[1], num_cells - horizontal_start]) + xtemp = array_ns.zeros([c2e2c.shape[1], num_cells - horizontal_start]) for i in range(ytemp.shape[0]): ytemp[i] = lat[c2e2c[horizontal_start:, i]] xtemp[i] = lon[c2e2c[horizontal_start:, i]] - wgt = weighting_factors( + wgt = _weighting_factors( ytemp, xtemp, lat[horizontal_start:], lon[horizontal_start:], divavg_cntrwgt, + array_ns=array_ns, ) - c_bln_avg = np.zeros((c2e2c.shape[0], c2e2c.shape[1] + 1)) + c_bln_avg = array_ns.zeros((c2e2c.shape[0], c2e2c.shape[1] + 1)) c_bln_avg[horizontal_start:, 0] = divavg_cntrwgt c_bln_avg[horizontal_start:, 1] = wgt[0] c_bln_avg[horizontal_start:, 2] = wgt[1] @@ -438,14 +502,15 @@ def _compute_c_bln_avg( def _force_mass_conservation_to_c_bln_avg( - c2e2c0: np.ndarray, - c_bln_avg: np.ndarray, - cell_areas: np.ndarray, - cell_owner_mask: np.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: np.int32, + horizontal_start: gtx.int32, + array_ns: ModuleType = np, niter: int = 1000, -) -> np.ndarray: +) -> field_alloc.NDArray: """ Iteratively enforce mass conservation to the input field c_bln_avg. @@ -469,7 +534,9 @@ def _force_mass_conservation_to_c_bln_avg( """ - def _compute_local_weights(c_bln_avg, cell_areas, c2e2c0, inverse_neighbor_idx) -> np.ndarray: + def _compute_local_weights( + c_bln_avg, cell_areas, c2e2c0, inverse_neighbor_idx + ) -> field_alloc.NDArray: """ Compute the total weight which each local point contributes to the sum. @@ -480,26 +547,28 @@ def _compute_local_weights(c_bln_avg, cell_areas, c2e2c0, inverse_neighbor_idx) Returns: ndarray of CellDim, containing the sum of weigh contributions for each local cell index """ - weights = np.sum(c_bln_avg[c2e2c0, inverse_neighbor_idx] * cell_areas[c2e2c0], axis=1) + weights = array_ns.sum(c_bln_avg[c2e2c0, inverse_neighbor_idx] * cell_areas[c2e2c0], axis=1) return weights def _compute_residual_to_mass_conservation( - owner_mask: np.ndarray, local_weight: np.ndarray, cell_area: np.ndarray - ) -> np.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" assert horizontal_size == cell_area.shape[0], "Fields do not have the same shape" - residual = np.where(owner_mask, local_weight / cell_area - 1.0, 0.0) + residual = array_ns.where(owner_mask, local_weight / cell_area - 1.0, 0.0) return residual def _apply_correction( - c_bln_avg: np.ndarray, - residual: np.ndarray, - c2e2c0: np.ndarray, + c_bln_avg: field_alloc.NDArray, + residual: field_alloc.NDArray, + c2e2c0: field_alloc.NDArray, divavg_cntrwgt: float, horizontal_start: gtx.int32, - ) -> np.ndarray: + ) -> 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 @@ -507,35 +576,39 @@ def _apply_correction( c_bln_avg[horizontal_start:, :] = ( c_bln_avg[horizontal_start:, :] - relax_coeff * residual[c2e2c0][horizontal_start:, :] ) - local_weight = np.sum(c_bln_avg, axis=1) - 1.0 + local_weight = array_ns.sum(c_bln_avg, axis=1) - 1.0 c_bln_avg[horizontal_start:, :] = c_bln_avg[horizontal_start:, :] - ( 0.25 * local_weight[horizontal_start:, np.newaxis] ) # avoid runaway condition: - c_bln_avg[horizontal_start:, 0] = np.maximum(c_bln_avg[horizontal_start:, 0], minwgt_loc) - c_bln_avg[horizontal_start:, 0] = np.minimum(c_bln_avg[horizontal_start:, 0], maxwgt_loc) + c_bln_avg[horizontal_start:, 0] = array_ns.maximum( + c_bln_avg[horizontal_start:, 0], minwgt_loc + ) + c_bln_avg[horizontal_start:, 0] = array_ns.minimum( + c_bln_avg[horizontal_start:, 0], maxwgt_loc + ) return c_bln_avg def _enforce_mass_conservation( - c_bln_avg: np.ndarray, - residual: np.ndarray, - owner_mask: np.ndarray, + c_bln_avg: field_alloc.NDArray, + residual: field_alloc.NDArray, + owner_mask: field_alloc.NDArray, horizontal_start: gtx.int32, - ) -> np.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] = np.where( + c_bln_avg[horizontal_start:, 0] = array_ns.where( owner_mask[horizontal_start:], c_bln_avg[horizontal_start:, 0] - residual[horizontal_start:], c_bln_avg[horizontal_start:, 0], ) return c_bln_avg - local_summed_weights = np.zeros(c_bln_avg.shape[0]) - residual = np.zeros(c_bln_avg.shape[0]) - inverse_neighbor_idx = create_inverse_neighbor_index(c2e2c0) + local_summed_weights = array_ns.zeros(c_bln_avg.shape[0]) + residual = array_ns.zeros(c_bln_avg.shape[0]) + inverse_neighbor_idx = create_inverse_neighbor_index(c2e2c0, array_ns=array_ns) for iteration in range(niter): local_summed_weights[horizontal_start:] = _compute_local_weights( @@ -546,7 +619,7 @@ def _enforce_mass_conservation( cell_owner_mask, local_summed_weights, cell_areas )[horizontal_start:] - max_ = np.max(residual) + max_ = array_ns.max(residual) if iteration >= (niter - 1) or max_ < 1e-9: print(f"number of iterations: {iteration} - max residual={max_}") c_bln_avg = _enforce_mass_conservation( @@ -566,32 +639,42 @@ def _enforce_mass_conservation( def compute_mass_conserving_bilinear_cell_average_weight( - c2e2c0: np.ndarray, - lat: np.ndarray, - lon: np.ndarray, - cell_areas: np.ndarray, - cell_owner_mask: np.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: np.int32, - horizontal_start_level_3, -) -> np.ndarray: - c_bln_avg = _compute_c_bln_avg(c2e2c0[:, 1:], lat, lon, divavg_cntrwgt, horizontal_start) + horizontal_start: gtx.int32, + horizontal_start_level_3: gtx.int32, + array_ns: ModuleType = np, +) -> field_alloc.NDArray: + c_bln_avg = _compute_c_bln_avg( + c2e2c0[:, 1:], lat, lon, divavg_cntrwgt, horizontal_start, array_ns + ) return _force_mass_conservation_to_c_bln_avg( - c2e2c0, c_bln_avg, cell_areas, cell_owner_mask, divavg_cntrwgt, horizontal_start_level_3 + c2e2c0, + c_bln_avg, + cell_areas, + cell_owner_mask, + divavg_cntrwgt, + horizontal_start_level_3, + array_ns, ) -def create_inverse_neighbor_index(c2e2c0): - inv_neighbor_idx = -1 * np.ones(c2e2c0.shape, dtype=int) +def create_inverse_neighbor_index(c2e2c0, array_ns: ModuleType = np): + inv_neighbor_idx = -1 * array_ns.ones(c2e2c0.shape, dtype=int) for jc in range(c2e2c0.shape[0]): for i in range(c2e2c0.shape[1]): if c2e2c0[jc, i] >= 0: - inv_neighbor_idx[jc, i] = np.argwhere(c2e2c0[c2e2c0[jc, i], :] == jc)[0, 0] + inv_neighbor_idx[jc, i] = array_ns.argwhere(c2e2c0[c2e2c0[jc, i], :] == jc)[0, 0] return inv_neighbor_idx +# TODO (@halungge) this can be simplified using only def compute_e_flx_avg( c_bln_avg: np.ndarray, geofac_div: np.ndarray, @@ -735,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, - horizontal_start_vertex: ta.wpfloat, -) -> 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, +) -> field_alloc.NDArray: """ Compute cells_aw_verts. @@ -762,7 +845,7 @@ def compute_cells_aw_verts( aw_verts: numpy array, representing a gtx.Field[gtx.Dims[VertexDim, 6], ta.wpfloat] """ cells_aw_verts = np.zeros(v2e.shape) - for jv in range(horizontal_start_vertex, cells_aw_verts.shape[0]): + for jv in range(horizontal_start, cells_aw_verts.shape[0]): cells_aw_verts[jv, :] = 0.0 for je in range(v2e.shape[1]): # INVALID_INDEX @@ -796,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. @@ -829,7 +912,7 @@ def compute_e_bln_c_s( ytemp[i] = edges_lat[c2e[llb:, i]] xtemp[i] = edges_lon[c2e[llb:, i]] - wgt = weighting_factors( + wgt = _weighting_factors( ytemp, xtemp, yloc, 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 d7b62ed3fb..b759e67709 100644 --- a/model/common/src/icon4py/model/common/states/factory.py +++ b/model/common/src/icon4py/model/common/states/factory.py @@ -7,47 +7,47 @@ # SPDX-License-Identifier: BSD-3-Clause """ -Provide a FieldFactory that can serve as a simple in memory database for Fields. +Provides Protocols and default implementations for Fields factories, which can be used to compute static +fields and manage their dependencies -Once setup, the factory can be queried for fields using a string name for the field. Three query modes are available: +- `FieldSource`: allows to query for a field, by a `.get(field_name, retrieval_type)` method: + +Three `RetrievalMode` s are available: _ `FIELD`: return the buffer containing the computed values as a GT4Py `Field` -- `METADATA`: return metadata such as units, CF standard_name or similar, dimensions... +- `METADATA`: return metadata (`FieldMetaData`) such as units, CF standard_name or similar, dimensions... - `DATA_ARRAY`: combination of the two above in the form of `xarray.dataarray` The factory can be used to "store" already computed fields or register functions and call arguments -and only compute the fields lazily upon request. In order to do so the user registers the fields computation with factory. +and only compute the fields lazily upon request. In order to do so the user registers the fields +computation with factory by setting up a `FieldProvider` It should be possible to setup the factory and computations and the factory independent of concrete runtime parameters that define the computation, passing those only once they are defined at runtime, for example --- -factory = Factory(metadata) -foo_provider = FieldProvider("foo", func = f1, dependencies = []) +factory = Factory(metadata, ...) +foo_provider = FieldProvider("foo", func = f1, dependencies, fields) bar_provider = FieldProvider("bar", func = f2, dependencies = ["foo"]) factory.register_provider(foo_provider) factory.register_provider(bar_provider) (...) ---- -def main(backend, grid) -factory.with_backend(backend).with_grid(grid) - val = factory.get("foo", RetrievalType.DATA_ARRAY) -TODO (halungge): except for domain parameters and other fields managed by the same factory we currently lack the ability to specify - other input sources in the factory for lazy evaluation. - factory.with_sources({"geometry": x}, where x:FieldSourceN - -TODO: for the numpy functions we might have to work on the func interfaces to make them a bit more uniform. +TODO: @halungge: allow to read configuration data """ +import collections import enum +import functools import inspect +from functools import cached_property from typing import ( Any, Callable, Mapping, + MutableMapping, Optional, Protocol, Sequence, @@ -69,10 +69,8 @@ def main(backend, grid) 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.states.model import FieldMetaData -from icon4py.model.common.states.utils import FieldType, to_data_array +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc DomainType = TypeVar("DomainType", h_grid.Domain, v_grid.Domain) @@ -130,62 +128,106 @@ class RetrievalType(enum.Enum): METADATA = 2 -class FieldSource(Protocol): +class FieldSource(GridProvider, Protocol): """ Protocol for object that can be queried for fields and field metadata Provides a default implementation of the get method. """ - @property - def metadata(self) -> dict[str, FieldMetaData]: - ... + _providers: MutableMapping[str, FieldProvider] = {} # noqa: RUF012 instance variable @property - def providers(self) -> dict[str, FieldProvider]: - ... + def _sources(self) -> "FieldSource": + return self @property - def backend(self) -> backend.Backend: + def metadata(self) -> MutableMapping[str, model.FieldMetaData]: + """Returns metadata for the fields that this field source provides.""" ... + # TODO @halungge: this is the target Backend: not necessarily the one that the field is computed and + # there are fields which need to be computed on a specific backend, which can be different from the + # general run backend @property - def grid_provider(self) -> GridProvider: + def backend(self) -> backend.Backend: ... def get( self, field_name: str, type_: RetrievalType = RetrievalType.FIELD - ) -> Union[FieldType, xa.DataArray, model.FieldMetaData]: - if field_name not in self.providers: + ) -> Union[state_utils.FieldType, xa.DataArray, model.FieldMetaData]: + """ + Get a field or its metadata from the factory. + + Fields are computed upon first call to `get`. + Args: + field_name: + type_: RetrievalType, determines whether only the field (databuffer) or Metadata or both will be returned + + Returns: + gt4py field containing allocated using this factories backend, a fields metadata or a + dataarray containing both. + + """ + if field_name not in self._providers: raise ValueError(f"Field '{field_name}' not provided by the source '{self.__class__}'") match type_: case RetrievalType.METADATA: return self.metadata[field_name] case RetrievalType.FIELD | RetrievalType.DATA_ARRAY: - provider = self.providers[field_name] + provider = self._providers[field_name] if field_name not in provider.fields: raise ValueError( f"Field {field_name} not provided by f{provider.func.__name__}." ) - buffer = provider(field_name, self, self.backend, self.grid_provider) + buffer = provider(field_name, self._sources, self.backend, self) return ( buffer if type_ == RetrievalType.FIELD - else to_data_array(buffer, self.metadata[field_name]) + else state_utils.to_data_array(buffer, self.metadata[field_name]) ) case _: raise ValueError(f"Invalid retrieval type {type_}") + def _provided_by_source(self, name): + return name in self._sources._providers or name in self._sources.metadata.keys() + def register_provider(self, provider: FieldProvider): + # dependencies must be provider by this field source or registered in sources for dependency in provider.dependencies: - if dependency not in self.providers.keys(): + if not (dependency in self._providers.keys() or self._provided_by_source(dependency)): raise ValueError( - f"Dependency '{dependency}' not found in registered providers of source {self.__class__}" + f"Missing dependency: '{dependency}' in registered of sources {self.__class__}" ) for field in provider.fields: - self.providers[field] = provider + self._providers[field] = provider + + +class CompositeSource(FieldSource): + def __init__(self, me: FieldSource, others: tuple[FieldSource, ...]): + self._backend = me.backend + self._grid = me.grid + self._vertical_grid = me.vertical_grid + self._metadata = collections.ChainMap(me.metadata, *(s.metadata for s in others)) + self._providers = collections.ChainMap(me._providers, *(s._providers for s in others)) + + @cached_property + def metadata(self) -> MutableMapping[str, model.FieldMetaData]: + return self._metadata + + @property + def backend(self) -> backend.Backend: + return self._backend + + @property + def vertical_grid(self) -> Optional[v_grid.VerticalGrid]: + return self._vertical_grid + + @property + def grid(self) -> Optional[icon_grid.IconGrid]: + return self._grid class PrecomputedFieldProvider(FieldProvider): @@ -213,6 +255,130 @@ def func(self) -> Callable: return lambda: self.fields +class FieldOperatorProvider(FieldProvider): + """Provider that calls a GT4Py Fieldoperator. + + # TODO (@halungge) for now to be used only on FieldView Embedded GT4Py backend. + - restrictions: + - (if only called on FieldView-Embedded, this is not a necessary restriction) + calls field operators without domain args, so it can only be used for full field computations + - plus: + - can write sparse/local fields + """ + + def __init__( + self, + func: gtx_decorator.FieldOperator, + domain: tuple[gtx.Dimension, ...], + fields: dict[str, str], # keyword arg to (field_operator, field_name) + deps: dict[str, str], # keyword arg to (field_operator, field_name) need: src + params: Optional[ + dict[str, state_utils.ScalarType] + ] = None, # keyword arg to (field_operator, field_name) + ): + self._func = func + self._dims = domain + self._dependencies = deps + self._output = fields + self._params = {} if params is None else params + self._fields: dict[str, Optional[gtx.Field | state_utils.ScalarType]] = { + name: None for name in fields.values() + } + + @property + def dependencies(self) -> Sequence[str]: + return list(self._dependencies.values()) + + @property + def fields(self) -> Mapping[str, state_utils.FieldType]: + return self._fields + + @property + def func(self) -> Callable: + return self._func + + def __call__( + self, + field_name: str, + field_src: Optional[FieldSource], + backend: Optional[gtx_backend.Backend], + grid: GridProvider, + ) -> state_utils.FieldType: + if any([f is None for f in self.fields.values()]): + self._compute(field_src, grid) + return self.fields[field_name] + + def _compute(self, factory, grid_provider): + # allocate output buffer + compute_backend = self._func.backend + try: + metadata = {k: factory.get(k, RetrievalType.METADATA) for k, v in self._output.items()} + dtype = metadata["dtype"] + except (ValueError, KeyError): + dtype = ta.wpfloat + self._fields = self._allocate(compute_backend, grid_provider, dtype=dtype) + # call field operator + # construct dependencies + deps = {k: factory.get(v) for k, v in self._dependencies.items()} + + out_fields = self._unravel_output_fields() + + self._func( + **deps, out=out_fields, offset_provider=self._get_offset_providers(grid_provider.grid) + ) + # transfer to target backend, the fields might have been computed on a compute backend + for f in self._fields.values(): + gtx.as_field(f.domain, f.ndarray, allocator=factory.backend) + + def _unravel_output_fields(self): + out_fields = tuple(self._fields.values()) + if len(out_fields) == 1: + out_fields = out_fields[0] + return out_fields + + # TODO: do we need that here? + def _get_offset_providers(self, grid: icon_grid.IconGrid) -> dict[str, gtx.FieldOffset]: + offset_providers = {} + for dim in self._dims: + if dim.kind == gtx.DimensionKind.HORIZONTAL: + horizontal_offsets = { + k: v + for k, v in grid.offset_providers.items() + if isinstance(v, gtx.NeighborTableOffsetProvider) + and v.origin_axis.kind == gtx.DimensionKind.HORIZONTAL + } + offset_providers.update(horizontal_offsets) + if dim.kind == gtx.DimensionKind.VERTICAL: + vertical_offsets = { + k: v + for k, v in grid.offset_providers.items() + if isinstance(v, gtx.Dimension) and v.kind == gtx.DimensionKind.VERTICAL + } + offset_providers.update(vertical_offsets) + return offset_providers + + def _allocate( + self, + backend: gtx_backend.Backend, + grid: GridProvider, + dtype: state_utils.ScalarType = ta.wpfloat, + ) -> dict[str, state_utils.FieldType]: + def _map_size(dim: gtx.Dimension, grid: GridProvider) -> int: + if dim.kind == gtx.DimensionKind.VERTICAL: + size = grid.vertical_grid.num_levels + return size + 1 if dims == dims.KHalfDim else size + return grid.grid.size[dim] + + def _map_dim(dim: gtx.Dimension) -> gtx.Dimension: + if dim == dims.KHalfDim: + return dims.KDim + return dim + + allocate = gtx.constructors.zeros.partial(allocator=backend) + field_domain = {_map_dim(dim): (0, _map_size(dim, grid)) for dim in self._dims} + return {k: allocate(field_domain, dtype=dtype) for k in self._fields.keys()} + + class ProgramFieldProvider(FieldProvider): """ Computes a field defined by a GT4Py Program. @@ -227,7 +393,7 @@ class ProgramFieldProvider(FieldProvider): the out arguments used in the program and the value the name the field is registered under and declared in the metadata. deps: dict[str, str], input fields used for computing this stencil: - the key is the variable name used in the program and the value the name + the key is the variable name used in the `gtx.program` and the value the name of the field it depends on. params: scalar parameters used in the program """ @@ -242,6 +408,7 @@ def __init__( ): self._func = func self._compute_domain = domain + self._dims = domain.keys() self._dependencies = deps self._output = fields self._params = params if params is not None else {} @@ -249,13 +416,10 @@ def __init__( name: None for name in fields.values() } - def _unallocated(self) -> bool: - return not all(self._fields.values()) - def _allocate( self, backend: gtx_backend.Backend, - grid: base_grid.BaseGrid, + grid: base_grid.BaseGrid, # TODO @halungge: change to vertical grid dtype: state_utils.ScalarType = ta.wpfloat, ) -> dict[str, state_utils.FieldType]: def _map_size(dim: gtx.Dimension, grid: base_grid.BaseGrid) -> int: @@ -269,9 +433,7 @@ def _map_dim(dim: gtx.Dimension) -> gtx.Dimension: return dim allocate = gtx.constructors.zeros.partial(allocator=backend) - field_domain = { - _map_dim(dim): (0, _map_size(dim, grid)) for dim in self._compute_domain.keys() - } + field_domain = {_map_dim(dim): (0, _map_size(dim, grid)) for dim in self._dims} return {k: allocate(field_domain, dtype=dtype) for k in self._fields.keys()} # TODO (@halungge) this can be simplified when completely disentangling vertical and horizontal grid. @@ -369,11 +531,6 @@ class NumpyFieldsProvider(FieldProvider): """ Computes a field defined by a numpy function. - TODO (halungge): - need to specify a parameter source to be able to postpone evaluation: paramters are mostly - configuration values - - need to able to access fields from several sources. - - Args: func: numpy function that computes the fields domain: the compute domain used for the stencil computation @@ -388,18 +545,17 @@ class NumpyFieldsProvider(FieldProvider): def __init__( self, func: Callable, - domain: dict[gtx.Dimension : tuple[DomainType, DomainType]], + domain: Sequence[gtx.Dimension], fields: Sequence[str], deps: dict[str, str], connectivities: Optional[dict[str, gtx.Dimension]] = None, params: Optional[dict[str, state_utils.ScalarType]] = None, ): self._func = func - self._compute_domain = domain - self._dims = domain.keys() + self._dims = domain self._fields: dict[str, Optional[state_utils.FieldType]] = {name: None for name in fields} self._dependencies = deps - self.connectivities = connectivities if connectivities is not None else {} + self._connectivities = connectivities if connectivities is not None else {} self._params = params if params is not None else {} def __call__( @@ -426,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) @@ -437,19 +593,22 @@ def _validate_dependencies(self): parameters = func_signature.parameters for dep_key in self._dependencies.keys(): parameter_definition = parameters.get(dep_key) - assert parameter_definition.annotation == xp.ndarray, ( - f"Dependency {dep_key} in function {self._func.__name__}: does not exist or has " - f"wrong type ('expected xp.ndarray') in {func_signature}." + checked = _check_union(parameter_definition, union=field_alloc.NDArray) + assert checked, ( + f"Dependency '{dep_key}' in function '{_func_name(self._func)}': does not exist or has " + f"wrong type ('expected ndarray') but was '{parameter_definition}'." ) for param_key, param_value in self._params.items(): parameter_definition = parameters.get(param_key) - checked = _check( + checked = _check_union_and_type( parameter_definition, param_value, union=state_utils.IntegerType - ) or _check(parameter_definition, param_value, union=state_utils.FloatType) + ) or _check_union_and_type( + parameter_definition, param_value, union=state_utils.FloatType + ) assert checked, ( - f"Parameter {param_key} in function {self._func.__name__} does not " - f"exist or has the wrong type: {type(param_value)}." + f"Parameter '{param_key}' in function '{_func_name(self._func)}' does not " + f"exist or has the wrong type: '{type(param_value)}'." ) @property @@ -465,14 +624,35 @@ def fields(self) -> Mapping[str, state_utils.FieldType]: return self._fields -def _check( +def _check_union_and_type( parameter_definition: inspect.Parameter, value: Union[state_utils.ScalarType, gtx.Field], union: Union, ) -> bool: + _check_union(parameter_definition, union) and type(value) in get_args(union) members = get_args(union) return ( parameter_definition is not None and parameter_definition.annotation in members and type(value) in members ) + + +def _check_union( + parameter_definition: inspect.Parameter, + union: Union, +) -> bool: + members = get_args(union) + # fix for unions with only one member, which implicitly are not Union but fallback to the type + # fix for unions with only one member, which implicitly are not Union but fallback to the type + if not members: + members = (union,) + annotation = parameter_definition.annotation + return parameter_definition is not None and (annotation == union or annotation in members) + + +def _func_name(callable_: Callable[..., Any]) -> str: + if isinstance(callable_, functools.partial): + return callable_.func.__name__ + else: + return callable_.__name__ diff --git a/model/common/src/icon4py/model/common/states/utils.py b/model/common/src/icon4py/model/common/states/utils.py index 736d0abb7f..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) @@ -22,9 +22,10 @@ T = TypeVar("T", ta.wpfloat, ta.vpfloat, float, bool, gtx.int32, gtx.int64) -FieldType: TypeAlias = Union[gtx.Field[Sequence[gtx.Dims[DimT]], T], xp.ndarray] +GTXFieldType: TypeAlias = gtx.Field[Sequence[gtx.Dims[DimT]], T] +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/grid_utils.py b/model/common/src/icon4py/model/common/test_utils/grid_utils.py index 4491c3e09e..3595d487d1 100644 --- a/model/common/src/icon4py/model/common/test_utils/grid_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/grid_utils.py @@ -134,6 +134,7 @@ def get_grid_geometry( on_gpu = alloc.is_cupy_device(backend) xp = alloc.array_ns(on_gpu) num_levels = get_num_levels(experiment) + register_name = experiment.join(backend.name) def construct_decomposition_info(grid: icon.IconGrid) -> definitions.DecompositionInfo: def _add_dimension(dim: gtx.Dimension): @@ -157,8 +158,8 @@ def construct_grid_geometry(grid_file: str): ) return geometry_source - if not grid_geometries.get(grid_file): - grid_geometries[grid_file] = construct_grid_geometry( + if not grid_geometries.get(register_name): + grid_geometries[register_name] = construct_grid_geometry( str(resolve_full_grid_file_name(grid_file)) ) - return grid_geometries[grid_file] + return grid_geometries[register_name] 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 638b91a8f0..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,7 +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/__init__.py b/model/common/tests/__init__.py new file mode 100644 index 0000000000..80b673df7e --- /dev/null +++ b/model/common/tests/__init__.py @@ -0,0 +1,8 @@ +# 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 + diff --git a/model/common/tests/conftest.py b/model/common/tests/conftest.py index da4c5f728b..658856c2e2 100644 --- a/model/common/tests/conftest.py +++ b/model/common/tests/conftest.py @@ -11,7 +11,11 @@ import pytest -from icon4py.model.common.test_utils.helpers import backend, grid # noqa: F401 # fixtures +from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package + decomposition_info, + experiment, +) +from icon4py.model.common.test_utils.helpers import backend # noqa: F401 # fixtures @pytest.fixture diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py index 19bef3ac04..8413486204 100644 --- a/model/common/tests/grid_tests/test_grid_manager.py +++ b/model/common/tests/grid_tests/test_grid_manager.py @@ -6,11 +6,11 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from __future__ import annotations import logging import typing +import gt4py.next as gtx import numpy as np import pytest from gt4py.next import backend as gtx_backend @@ -18,6 +18,7 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import ( grid_manager as gm, + horizontal as h_grid, refinement as refin, vertical as v_grid, ) @@ -86,7 +87,7 @@ def test_grid_file_dimension(global_grid_file): "grid_file, experiment", [ (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), - (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) def test_grid_file_vertex_cell_edge_dimensions(grid_savepoint, grid_file): @@ -142,7 +143,7 @@ def test_grid_file_index_fields(global_grid_file, caplog, icon_grid): "grid_file, experiment", [ (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), - (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) def test_grid_manager_eval_v2e(caplog, grid_savepoint, experiment, grid_file, backend): @@ -165,7 +166,7 @@ def test_grid_manager_eval_v2e(caplog, grid_savepoint, experiment, grid_file, ba "grid_file, experiment", [ (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), - (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) @pytest.mark.parametrize("dim", [dims.CellDim, dims.EdgeDim, dims.VertexDim]) @@ -185,7 +186,7 @@ def test_grid_manager_refin_ctrl(grid_savepoint, grid_file, experiment, dim, bac "grid_file, experiment", [ (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), - (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) def test_grid_manager_eval_v2c(caplog, grid_savepoint, experiment, grid_file, backend): @@ -237,7 +238,7 @@ def reset_invalid_index(index_array: np.ndarray): "grid_file, experiment", [ (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), - (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) def test_grid_manager_eval_e2v(caplog, grid_savepoint, grid_file, experiment, backend): @@ -256,7 +257,7 @@ def has_invalid_index(ar: np.ndarray): return np.any(invalid_index(ar)) -def invalid_index(ar): +def invalid_index(ar: np.ndarray): return np.where(ar == gm.GridFile.INVALID_INDEX) @@ -288,7 +289,7 @@ def assert_invalid_indices(e2c_table: np.ndarray, grid_file: str): "grid_file, experiment", [ (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), - (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) def test_grid_manager_eval_e2c(caplog, grid_savepoint, grid_file, experiment, backend): @@ -309,7 +310,7 @@ def test_grid_manager_eval_e2c(caplog, grid_savepoint, grid_file, experiment, ba "grid_file, experiment", [ (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), - (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) def test_grid_manager_eval_c2e(caplog, grid_savepoint, grid_file, experiment, backend): @@ -370,7 +371,7 @@ def test_grid_manager_eval_c2e2cO(caplog, grid_savepoint, grid_file, experiment, "grid_file, experiment", [ (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), - (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) def test_grid_manager_eval_e2c2e(caplog, grid_savepoint, grid_file, experiment, backend): @@ -382,18 +383,17 @@ def test_grid_manager_eval_e2c2e(caplog, grid_savepoint, grid_file, experiment, assert_invalid_indices(serialized_e2c2e, grid_file) e2c2e_table = grid.get_offset_provider("E2C2E").table - e2c2e0_table = grid.get_offset_provider("E2C2EO").table + e2c2eO_table = grid.get_offset_provider("E2C2EO").table assert_invalid_indices(e2c2e_table, grid_file) # ICON calculates diamond edges only from rl_start = 2 (lateral_boundary(dims.EdgeDim) + 1 for # boundaries all values are INVALID even though the half diamond exists (see mo_model_domimp_setup.f90 ll 163ff.) - assert_unless_invalid(e2c2e_table, serialized_e2c2e) - assert_unless_invalid(e2c2e0_table, serialized_e2c2eO) - - -def assert_unless_invalid(table, serialized_ref): - invalid_positions = invalid_index(table) - np.allclose(table[invalid_positions], serialized_ref[invalid_positions]) + start_index = grid.start_index( + h_grid.domain(dims.EdgeDim)(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) + ) + # e2c2e in ICON (quad_idx) has a different neighbor ordering than the e2c2e constructed in grid_manager.py + assert_up_to_order(e2c2e_table, serialized_e2c2e, start_index) + assert_up_to_order(e2c2eO_table, serialized_e2c2eO, start_index) @pytest.mark.datatest @@ -402,17 +402,21 @@ def assert_unless_invalid(table, serialized_ref): "grid_file, experiment", [ (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), - (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) def test_grid_manager_eval_e2c2v(caplog, grid_savepoint, grid_file, backend): caplog.set_level(logging.DEBUG) grid = _run_grid_manager(grid_file, backend).grid + serialized_ref = grid_savepoint.e2c2v() # the "far" (adjacent to edge normal ) is not always there, because ICON only calculates those starting from # (lateral_boundary(dims.EdgeDim) + 1) to end(dims.EdgeDim) (see mo_intp_coeffs.f90) and only for owned cells - serialized_ref = grid_savepoint.e2c2v() table = grid.get_offset_provider("E2C2V").table - assert_unless_invalid(table, serialized_ref) + start_index = grid.start_index( + h_grid.domain(dims.EdgeDim)(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ) + # e2c2e in ICON (quad_idx) has a different neighbor ordering than the e2c2e constructed in grid_manager.py + assert_up_to_order(table, serialized_ref, start_index) assert np.allclose(table[:, :2], grid.get_offset_provider("E2V").table) @@ -422,7 +426,7 @@ def test_grid_manager_eval_e2c2v(caplog, grid_savepoint, grid_file, backend): "grid_file, experiment", [ (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), - (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) def test_grid_manager_eval_c2v(caplog, grid_savepoint, grid_file, backend): @@ -442,14 +446,18 @@ def test_grid_manager_eval_c2v(caplog, grid_savepoint, grid_file, backend): ) @pytest.mark.with_netcdf def test_grid_manager_grid_size(dim, size, backend): - grid = _run_grid_manager(utils.R02B04_GLOBAL, backend=backend).grid + grid = _run_grid_manager(dt_utils.R02B04_GLOBAL, backend=backend).grid assert size == grid.size[dim] -def assert_up_to_order(table, diamond_table): - assert table.shape == diamond_table.shape - for n in range(table.shape[0]): - assert np.all(np.in1d(table[n, :], diamond_table[n, :])) +def assert_up_to_order(table: np.ndarray, reference_table: np.ndarray, start_index: gtx.int = 0): + assert table.shape == reference_table.shape, "arrays need to have the same shape" + reduced_table = table[start_index:, :] + reduced_reference = reference_table[start_index:, :] + for n in range(reduced_table.shape[0]): + assert np.all( + np.in1d(reduced_table[n, :], reduced_reference[n, :]) + ), f"values in row {n+start_index} are not equal: {reduced_table[n, :]} vs ref= {reduced_reference[n, :]}." @pytest.mark.with_netcdf @@ -477,7 +485,7 @@ def test_gt4py_transform_offset_by_1_where_valid(size): @pytest.mark.parametrize( "grid_file, global_num_cells", [ - (utils.R02B04_GLOBAL, R02B04_GLOBAL_NUM_CELLS), + (dt_utils.R02B04_GLOBAL, R02B04_GLOBAL_NUM_CELLS), (dt_utils.REGIONAL_EXPERIMENT, MCH_CH_RO4B09_GLOBAL_NUM_CELLS), ], ) @@ -489,7 +497,7 @@ def test_grid_manager_grid_level_and_root(grid_file, global_num_cells, backend): @pytest.mark.with_netcdf @pytest.mark.parametrize( "grid_file, experiment", - [(utils.R02B04_GLOBAL, dt_utils.JABW_EXPERIMENT)], + [(dt_utils.R02B04_GLOBAL, dt_utils.JABW_EXPERIMENT)], ) def test_grid_manager_eval_c2e2c2e(caplog, grid_savepoint, grid_file, backend): caplog.set_level(logging.DEBUG) @@ -507,7 +515,7 @@ def test_grid_manager_eval_c2e2c2e(caplog, grid_savepoint, grid_file, backend): @pytest.mark.parametrize( "grid_file, experiment", [ - (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), ], ) @@ -546,8 +554,8 @@ def test_grid_manager_start_end_index(caplog, grid_file, experiment, dim, icon_g ) def test_read_geometry_fields(grid_savepoint, grid_file, backend): manager = _run_grid_manager(grid_file, backend=backend) - cell_area = manager.geometry[GeometryName.CELL_AREA.value] - tangent_orientation = manager.geometry[GeometryName.TANGENT_ORIENTATION.value] + cell_area = manager.geometry[gm.GeometryName.CELL_AREA.value] + tangent_orientation = manager.geometry[gm.GeometryName.TANGENT_ORIENTATION.value] assert helpers.dallclose(cell_area.asnumpy(), grid_savepoint.cell_areas().asnumpy()) assert helpers.dallclose( @@ -580,10 +588,79 @@ def test_coordinates(grid_savepoint, grid_file, experiment, dim, backend): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_tangent_orientation(experiment, grid_file, grid_savepoint, backend): +def test_tangent_orientation(grid_file, grid_savepoint, backend): expected = grid_savepoint.tangent_orientation() manager = _run_grid_manager(grid_file, backend=backend) geometry_fields = manager.geometry assert helpers.dallclose( - geometry_fields[GeometryName.TANGENT_ORIENTATION].asnumpy(), expected.asnumpy() + geometry_fields[gm.GeometryName.TANGENT_ORIENTATION].asnumpy(), expected.asnumpy() + ) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +def test_cell_normal_orientation(grid_file, grid_savepoint, backend): + expected = grid_savepoint.edge_orientation() + manager = _run_grid_manager(grid_file, backend=backend) + geometry_fields = manager.geometry + assert helpers.dallclose( + geometry_fields[GeometryName.CELL_NORMAL_ORIENTATION].ndarray, expected.ndarray + ) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +def test_edge_orientation_on_vertex(grid_file, grid_savepoint, backend): + expected = grid_savepoint.vertex_edge_orientation() + manager = _run_grid_manager(grid_file, backend=backend) + geometry_fields = manager.geometry + assert helpers.dallclose( + geometry_fields[GeometryName.EDGE_ORIENTATION_ON_VERTEX].ndarray, expected.ndarray + ) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +def test_dual_area(grid_file, grid_savepoint, backend): + expected = grid_savepoint.vertex_dual_area() + manager = _run_grid_manager(grid_file, backend=backend) + geometry_fields = manager.geometry + assert helpers.dallclose(geometry_fields[GeometryName.DUAL_AREA].ndarray, expected.ndarray) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +def test_edge_cell_distance(grid_file, grid_savepoint, backend): + expected = grid_savepoint.edge_cell_length() + manager = _run_grid_manager(grid_file, backend=backend) + geometry_fields = manager.geometry + + assert helpers.dallclose( + geometry_fields[GeometryName.EDGE_CELL_DISTANCE].asnumpy(), + expected.asnumpy(), + equal_nan=True, ) diff --git a/model/common/tests/grid_tests/utils.py b/model/common/tests/grid_tests/utils.py index 3b1463dc36..788bc9da50 100644 --- a/model/common/tests/grid_tests/utils.py +++ b/model/common/tests/grid_tests/utils.py @@ -9,17 +9,13 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import horizontal as h_grid -from icon4py.model.common.test_utils.datatest_utils import ( - GRIDS_PATH, - R02B04_GLOBAL, - REGIONAL_EXPERIMENT, -) +from icon4py.model.common.test_utils import datatest_utils as dt_utils -r04b09_dsl_grid_path = GRIDS_PATH.joinpath(REGIONAL_EXPERIMENT) +r04b09_dsl_grid_path = dt_utils.GRIDS_PATH.joinpath(dt_utils.REGIONAL_EXPERIMENT) r04b09_dsl_data_file = r04b09_dsl_grid_path.joinpath("mch_ch_r04b09_dsl_grids_v1.tar.gz").name -r02b04_global_grid_path = GRIDS_PATH.joinpath(R02B04_GLOBAL) +r02b04_global_grid_path = dt_utils.GRIDS_PATH.joinpath(dt_utils.R02B04_GLOBAL) r02b04_global_data_file = r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_R.tar.gz").name diff --git a/model/common/tests/interpolation_tests/__init__.py b/model/common/tests/interpolation_tests/__init__.py new file mode 100644 index 0000000000..80b673df7e --- /dev/null +++ b/model/common/tests/interpolation_tests/__init__.py @@ -0,0 +1,8 @@ +# 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 + diff --git a/model/common/tests/interpolation_tests/test_interpolation_factory.py b/model/common/tests/interpolation_tests/test_interpolation_factory.py new file mode 100644 index 0000000000..1077da297b --- /dev/null +++ b/model/common/tests/interpolation_tests/test_interpolation_factory.py @@ -0,0 +1,226 @@ +# 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 +import numpy as np +import pytest + +import icon4py.model.common.states.factory as factory +from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import horizontal as h_grid +from icon4py.model.common.interpolation import ( + interpolation_attributes as attrs, + interpolation_factory, +) +from icon4py.model.common.test_utils import ( + datatest_utils as dt_utils, + grid_utils as gridtest_utils, + helpers as test_helpers, +) + + +V2E_SIZE = 6 + +C2E_SIZE = 3 +E2C_SIZE = 2 + + +interpolation_factories = {} + +vertex_domain = h_grid.domain(dims.VertexDim) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_raises_error_on_unknown_field(grid_file, experiment, backend, decomposition_info): + geometry = gridtest_utils.get_grid_geometry(backend, experiment, grid_file) + interpolation_source = interpolation_factory.InterpolationFieldsFactory( + grid=geometry.grid, + decomposition_info=decomposition_info, + geometry_source=geometry, + backend=backend, + metadata=attrs.attrs, + ) + with pytest.raises(ValueError) as error: + interpolation_source.get("foo", factory.RetrievalType.METADATA) + assert "unknown field" in error.value + + +@pytest.mark.parametrize( + "grid_file, experiment, rtol", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT, 5e-9), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT, 1e-11), + ], +) +@pytest.mark.datatest +def test_get_c_lin_e(interpolation_savepoint, grid_file, experiment, backend, rtol): + field_ref = interpolation_savepoint.c_lin_e() + factory = get_interpolation_factory(backend, experiment, grid_file) + grid = factory.grid + field = factory.get(attrs.C_LIN_E) + assert field.shape == (grid.num_edges, E2C_SIZE) + assert test_helpers.dallclose(field.asnumpy(), field_ref.asnumpy(), rtol=rtol) + + +def get_interpolation_factory( + backend, experiment, grid_file +) -> interpolation_factory.InterpolationFieldsFactory: + registry_key = experiment.join(backend.name) + factory = interpolation_factories.get(registry_key) + if not factory: + geometry = gridtest_utils.get_grid_geometry(backend, experiment, grid_file) + + factory = interpolation_factory.InterpolationFieldsFactory( + grid=geometry.grid, + decomposition_info=geometry._decomposition_info, + geometry_source=geometry, + backend=backend, + metadata=attrs.attrs, + ) + interpolation_factories[registry_key] = factory + return factory + + +@pytest.mark.parametrize( + "grid_file, experiment, rtol", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT, 1e-9), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT, 1e-12), + ], +) +@pytest.mark.datatest +def test_get_geofac_div(interpolation_savepoint, grid_file, experiment, backend, rtol): + field_ref = interpolation_savepoint.geofac_div() + factory = get_interpolation_factory(backend, experiment, grid_file) + grid = factory.grid + field = factory.get(attrs.GEOFAC_DIV) + assert field.shape == (grid.num_cells, C2E_SIZE) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=rtol) + + +## FIXME: does not validate +# -> connectivity order between reference from serialbox and computed value is different +@pytest.mark.parametrize( + "grid_file, experiment, rtol", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT, 5e-9), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT, 1e-11), + ], +) +@pytest.mark.datatest +def test_get_geofac_grdiv(interpolation_savepoint, grid_file, experiment, backend, rtol): + field_ref = interpolation_savepoint.geofac_grdiv() + factory = get_interpolation_factory(backend, experiment, grid_file) + grid = factory.grid + field = factory.get(attrs.GEOFAC_GRDIV) + assert field.shape == (grid.num_edges, 5) + # FIXME: e2c2e constructed from grid file has different ordering than the serialized one + assert_reordered(field.asnumpy(), field_ref.asnumpy(), rtol) + + +def assert_reordered(val: np.ndarray, ref: np.ndarray, rtol): + assert val.shape == ref.shape, f"arrays do not have the same shape: {val.shape} vs {ref.shape}" + s_val = np.argsort(val) + s_ref = np.argsort(ref) + for i in range(val.shape[0]): + assert test_helpers.dallclose( + val[i, s_val[i, :]], ref[i, s_ref[i, :]], rtol=rtol + ), f"assertion failed for row {i}" + + +@pytest.mark.parametrize( + "grid_file, experiment, rtol", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT, 5e-9), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT, 1e-11), + ], +) +@pytest.mark.datatest +def test_get_geofac_rot(interpolation_savepoint, grid_file, experiment, backend, rtol): + field_ref = interpolation_savepoint.geofac_rot() + factory = get_interpolation_factory(backend, experiment, grid_file) + grid = factory.grid + field = factory.get(attrs.GEOFAC_ROT) + horizontal_start = grid.start_index(vertex_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)) + assert field.shape == (grid.num_vertices, V2E_SIZE) + assert test_helpers.dallclose( + field_ref.asnumpy()[horizontal_start:, :], field.asnumpy()[horizontal_start:, :], rtol=rtol + ) + + +@pytest.mark.parametrize( + "grid_file, experiment, rtol", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT, 5e-9), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT, 1e-11), + ], +) +@pytest.mark.datatest +def test_get_geofac_n2s(interpolation_savepoint, grid_file, experiment, backend, rtol): + field_ref = interpolation_savepoint.geofac_n2s() + factory = get_interpolation_factory(backend, experiment, grid_file) + grid = factory.grid + field = factory.get(attrs.GEOFAC_N2S) + assert field.shape == (grid.num_cells, 4) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=rtol) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_get_geofac_grg(interpolation_savepoint, grid_file, experiment, backend): + field_ref = interpolation_savepoint.geofac_grg() + factory = get_interpolation_factory(backend, experiment, grid_file) + grid = factory.grid + field_x = factory.get(attrs.GEOFAC_GRG_X) + assert field_x.shape == (grid.num_cells, 4) + field_y = factory.get(attrs.GEOFAC_GRG_Y) + assert field_y.shape == (grid.num_cells, 4) + # TODO (@halungge) tolerances are high, especially in the 0th (central) component, check stencil + # this passes due to the atol which is too large for the values + assert test_helpers.dallclose( + field_ref[0].asnumpy(), + field_x.asnumpy(), + rtol=1e-7, + atol=1e-6, + ) + assert test_helpers.dallclose( + field_ref[1].asnumpy(), + field_y.asnumpy(), + rtol=1e-7, + atol=1e-6, + ) + + +@pytest.mark.parametrize( + "grid_file, experiment, rtol", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT, 5e-9), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT, 1e-11), + ], +) +@pytest.mark.datatest +def test_get_mass_conserving_cell_average_weight( + interpolation_savepoint, grid_file, experiment, backend, rtol +): + field_ref = interpolation_savepoint.c_bln_avg() + factory = get_interpolation_factory(backend, experiment, grid_file) + grid = factory.grid + field = factory.get(attrs.C_BLN_AVG) + + assert field.shape == (grid.num_cells, 4) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=rtol) diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields.py b/model/common/tests/interpolation_tests/test_interpolation_fields.py index 24ef28264a..d7aa99ceb0 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_fields.py +++ b/model/common/tests/interpolation_tests/test_interpolation_fields.py @@ -5,6 +5,7 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +import functools import numpy as np import pytest @@ -25,7 +26,6 @@ compute_geofac_rot, compute_mass_conserving_bilinear_cell_average_weight, compute_pos_on_tplane_e_x_y, - compute_primal_normal_ec, ) from icon4py.model.common.test_utils import datatest_utils as dt_utils from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package @@ -35,6 +35,7 @@ processor_props, ranked_data_path, ) +from icon4py.model.common.utils import gt4py_field_allocation as alloc cell_domain = h_grid.domain(dims.CellDim) @@ -44,45 +45,61 @@ @pytest.mark.datatest @pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) -def test_compute_c_lin_e(grid_savepoint, interpolation_savepoint, icon_grid): # fixture +def test_compute_c_lin_e(grid_savepoint, interpolation_savepoint, icon_grid, backend): # fixture + xp = alloc.import_array_ns(backend) + func = functools.partial(compute_c_lin_e, array_ns=xp) inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() edge_cell_length = grid_savepoint.edge_cell_length() - owner_mask = grid_savepoint.e_owner_mask() + edge_owner_mask = grid_savepoint.e_owner_mask() c_lin_e_ref = interpolation_savepoint.c_lin_e() + horizontal_start = icon_grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)) + c_lin_e = compute_c_lin_e( - edge_cell_length.asnumpy(), - inv_dual_edge_length.asnumpy(), - owner_mask.asnumpy(), + edge_cell_length.ndarray, + inv_dual_edge_length.ndarray, + edge_owner_mask.ndarray, horizontal_start, + xp, ) + assert test_helpers.dallclose(alloc.as_numpy(c_lin_e), c_lin_e_ref.asnumpy()) - assert test_helpers.dallclose(c_lin_e, c_lin_e_ref.asnumpy()) + c_lin_e_partial = func( + edge_cell_length.ndarray, + inv_dual_edge_length.ndarray, + edge_owner_mask.ndarray, + horizontal_start, + ) + assert test_helpers.dallclose(alloc.as_numpy(c_lin_e_partial), c_lin_e_ref.asnumpy()) @pytest.mark.datatest @pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) -def test_compute_geofac_div(grid_savepoint, interpolation_savepoint, icon_grid): +def test_compute_geofac_div(grid_savepoint, interpolation_savepoint, icon_grid, backend): + if backend is not None: + pytest.xfail("writes a sparse fields: only runs in field view embedded") mesh = icon_grid primal_edge_length = grid_savepoint.primal_edge_length() edge_orientation = grid_savepoint.edge_orientation() area = grid_savepoint.cell_areas() geofac_div_ref = interpolation_savepoint.geofac_div() geofac_div = test_helpers.zero_field(mesh, dims.CellDim, dims.C2EDim) - compute_geofac_div( - primal_edge_length, - edge_orientation, - area, - out=geofac_div, + compute_geofac_div.with_backend(backend)( + primal_edge_length=primal_edge_length, + edge_orientation=edge_orientation, + area=area, + out=(geofac_div), offset_provider={"C2E": mesh.get_offset_provider("C2E")}, ) - assert test_helpers.dallclose(geofac_div.asnumpy(), geofac_div_ref.asnumpy()) @pytest.mark.datatest @pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) -def test_compute_geofac_rot(grid_savepoint, interpolation_savepoint, icon_grid): +def test_compute_geofac_rot(grid_savepoint, interpolation_savepoint, icon_grid, backend): + if backend is not None: + pytest.xfail("writes a sparse fields: only runs in field view embedded") + mesh = icon_grid dual_edge_length = grid_savepoint.dual_edge_length() edge_orientation = grid_savepoint.vertex_edge_orientation() @@ -106,7 +123,8 @@ def test_compute_geofac_rot(grid_savepoint, interpolation_savepoint, icon_grid): @pytest.mark.datatest @pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) -def test_compute_geofac_n2s(grid_savepoint, interpolation_savepoint, icon_grid): +def test_compute_geofac_n2s(grid_savepoint, interpolation_savepoint, icon_grid, backend): + xp = alloc.import_array_ns(backend) dual_edge_length = grid_savepoint.dual_edge_length() geofac_div = interpolation_savepoint.geofac_div() geofac_n2s_ref = interpolation_savepoint.geofac_n2s() @@ -114,15 +132,15 @@ def test_compute_geofac_n2s(grid_savepoint, interpolation_savepoint, icon_grid): e2c = icon_grid.connectivities[dims.E2CDim] c2e2c = icon_grid.connectivities[dims.C2E2CDim] horizontal_start = icon_grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)) - geofac_n2s = compute_geofac_n2s( - dual_edge_length.asnumpy(), - geofac_div.asnumpy(), + geofac_n2s = functools.partial(compute_geofac_n2s, array_ns=xp)( + dual_edge_length.ndarray, + geofac_div.ndarray, c2e, e2c, c2e2c, horizontal_start, ) - assert test_helpers.dallclose(geofac_n2s, geofac_n2s_ref.asnumpy()) + assert test_helpers.dallclose(alloc.as_numpy(geofac_n2s), geofac_n2s_ref.asnumpy()) @pytest.mark.datatest @@ -138,16 +156,11 @@ def test_compute_geofac_grg(grid_savepoint, interpolation_savepoint, icon_grid): e2c = icon_grid.connectivities[dims.E2CDim] c2e2c = icon_grid.connectivities[dims.C2E2CDim] horizontal_start = icon_grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)) - primal_normal_ec = compute_primal_normal_ec( + + geofac_grg_0, geofac_grg_1 = compute_geofac_grg( primal_normal_cell_x, primal_normal_cell_y, - owner_mask, - c2e, - e2c, - horizontal_start, - ) - geofac_grg = compute_geofac_grg( - primal_normal_ec, + owner_mask.asnumpy(), geofac_div.asnumpy(), c_lin_e.asnumpy(), c2e, @@ -156,10 +169,10 @@ def test_compute_geofac_grg(grid_savepoint, interpolation_savepoint, icon_grid): horizontal_start, ) assert test_helpers.dallclose( - geofac_grg[:, :, 0], geofac_grg_ref[0].asnumpy(), atol=1e-6, rtol=1e-7 + alloc.as_numpy(geofac_grg_0), geofac_grg_ref[0].asnumpy(), atol=1e-6, rtol=1e-7 ) assert test_helpers.dallclose( - geofac_grg[:, :, 1], geofac_grg_ref[1].asnumpy(), atol=1e-6, rtol=1e-7 + alloc.as_numpy(geofac_grg_1), geofac_grg_ref[1].asnumpy(), atol=1e-6, rtol=1e-7 ) @@ -220,7 +233,7 @@ def test_compute_c_bln_avg(grid_savepoint, interpolation_savepoint, icon_grid, a @pytest.mark.datatest @pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) -def test_compute_e_flx_avg(grid_savepoint, interpolation_savepoint, icon_grid): +def test_compute_e_flx_avg(grid_savepoint, interpolation_savepoint, icon_grid, backend): e_flx_avg_ref = interpolation_savepoint.e_flx_avg().asnumpy() c_bln_avg = interpolation_savepoint.c_bln_avg().asnumpy() geofac_div = interpolation_savepoint.geofac_div().asnumpy() @@ -277,7 +290,7 @@ def test_compute_cells_aw_verts( e2v=e2v, v2c=v2c, e2c=e2c, - horizontal_start_vertex=horizontal_start_vertex, + horizontal_start=horizontal_start_vertex, ) assert test_helpers.dallclose(cells_aw_verts, cells_aw_verts_ref, atol=1e-3) 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/common/tests/states_test/test_factory.py b/model/common/tests/states_test/test_factory.py new file mode 100644 index 0000000000..08d266a03a --- /dev/null +++ b/model/common/tests/states_test/test_factory.py @@ -0,0 +1,237 @@ +# 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 typing import Optional + +import gt4py.next as gtx +import pytest + +from icon4py.model.common import dimension as dims, utils as common_utils +from icon4py.model.common.grid import horizontal as h_grid, icon, vertical as v_grid +from icon4py.model.common.math import helpers as math_helpers +from icon4py.model.common.metrics import metric_fields as metrics +from icon4py.model.common.states import factory, model, utils as state_utils +from icon4py.model.common.test_utils import helpers as test_helpers + + +cell_domain = h_grid.domain(dims.CellDim) +k_domain = v_grid.domain(dims.KDim) + + +class SimpleFieldSource(factory.FieldSource): + def __init__( + self, + data_: dict[str, tuple[state_utils.FieldType, model.FieldMetaData]], + backend, + grid: icon.IconGrid, + vertical_grid: v_grid.VerticalGrid = None, + ): + self._providers = {} + self._backend = backend + self._grid = grid + self._vertical_grid = vertical_grid + self._metadata = {} + self._initial_data = data_ + + for key, value in data_.items(): + self.register_provider(factory.PrecomputedFieldProvider({key: value[0]})) + self._metadata[key] = value[1] + + def _register_initial_fields(self): + for key, value in self._initial_data.items(): + self.register_provider(factory.PrecomputedFieldProvider({key: value[0]})) + self._metadata[key] = value[1] + + def reset(self): + self._providers = {} + self._metadata = {} + self._register_initial_fields() + + @property + def metadata(self): + return self._metadata + + @property + def _sources(self) -> factory.FieldSource: + return self + + @property + def grid(self): + return self._grid + + @property + def vertical_grid(self) -> Optional[v_grid.VerticalGrid]: + return self._vertical_grid + + @property + def backend(self): + return self._backend + + +@pytest.fixture(scope="function") +def cell_coordinate_source(grid_savepoint, backend): + on_gpu = common_utils.gt4py_field_allocation.is_cupy_device(backend) + grid = grid_savepoint.construct_icon_grid(on_gpu) + lat = grid_savepoint.lat(dims.CellDim) + lon = grid_savepoint.lon(dims.CellDim) + data = { + "lat": (lat, {"standard_name": "lat", "units": ""}), + "lon": (lon, {"standard_name": "lon", "units": ""}), + } + + coordinate_source = SimpleFieldSource(data_=data, backend=backend, grid=grid) + yield coordinate_source + coordinate_source.reset() + + +@pytest.fixture(scope="function") +def height_coordinate_source(metrics_savepoint, grid_savepoint, backend): + on_gpu = common_utils.gt4py_field_allocation.is_cupy_device(backend) + grid = grid_savepoint.construct_icon_grid(on_gpu) + z_ifc = metrics_savepoint.z_ifc() + vct_a = grid_savepoint.vct_a() + vct_b = grid_savepoint.vct_b() + data = {"height_coordinate": (z_ifc, {"standard_name": "height_coordinate", "units": ""})} + vertical_grid = v_grid.VerticalGrid(v_grid.VerticalGridConfig(num_levels=10), vct_a, vct_b) + field_source = SimpleFieldSource( + data_=data, backend=backend, grid=grid, vertical_grid=vertical_grid + ) + yield field_source + field_source.reset() + + +@pytest.mark.datatest +def test_field_operator_provider(cell_coordinate_source): + field_op = math_helpers.geographical_to_cartesian_on_cells.with_backend(None) + domain = {dims.CellDim: (cell_domain(h_grid.Zone.LOCAL), cell_domain(h_grid.Zone.LOCAL))} + deps = {"lat": "lat", "lon": "lon"} + fields = {"x": "x", "y": "y", "z": "z"} + + provider = factory.FieldOperatorProvider(field_op, domain, fields, deps) + provider("x", cell_coordinate_source, cell_coordinate_source.backend, cell_coordinate_source) + x = provider.fields["x"] + assert isinstance(x, gtx.Field) + assert dims.CellDim in x.domain.dims + + +@pytest.mark.datatest +def test_program_provider(height_coordinate_source): + program = metrics.compute_z_mc + domain = { + dims.CellDim: (cell_domain(h_grid.Zone.LOCAL), cell_domain(h_grid.Zone.LOCAL)), + dims.KDim: (k_domain(v_grid.Zone.TOP), k_domain(v_grid.Zone.BOTTOM)), + } + deps = { + "z_ifc": "height_coordinate", + } + fields = {"z_mc": "output_f"} + provider = factory.ProgramFieldProvider(program, domain, fields, deps) + provider( + "output_f", + height_coordinate_source, + height_coordinate_source.backend, + height_coordinate_source, + ) + x = provider.fields["output_f"] + assert isinstance(x, gtx.Field) + assert dims.CellDim in x.domain.dims + + +def test_field_source_raise_error_on_register(cell_coordinate_source): + program = metrics.compute_z_mc + domain = { + dims.CellDim: (cell_domain(h_grid.Zone.LOCAL), cell_domain(h_grid.Zone.LOCAL)), + dims.KDim: (k_domain(v_grid.Zone.TOP), k_domain(v_grid.Zone.BOTTOM)), + } + deps = { + "z_ifc": "height_coordinate", + } + fields = {"z_mc": "output_f"} + provider = factory.ProgramFieldProvider(func=program, domain=domain, fields=fields, deps=deps) + with pytest.raises(ValueError) as err: + cell_coordinate_source.register_provider(provider) + assert "not provided by source " in err.value + + +def test_composite_field_source_contains_all_metadata( + cell_coordinate_source, height_coordinate_source +): + backend = cell_coordinate_source.backend + grid = cell_coordinate_source.grid + foo = test_helpers.random_field(grid, dims.CellDim, dims.KDim) + bar = test_helpers.random_field(grid, dims.EdgeDim, dims.KDim) + data = { + "foo": (foo, {"standard_name": "foo", "units": ""}), + "bar": (bar, {"standard_name": "bar", "units": ""}), + } + + test_source = SimpleFieldSource(data_=data, grid=grid, backend=backend) + composite = factory.CompositeSource( + test_source, (cell_coordinate_source, height_coordinate_source) + ) + + assert composite.backend == test_source.backend + assert composite.grid.id == test_source.grid.id + assert test_source.metadata.items() <= composite.metadata.items() + assert height_coordinate_source.metadata.items() <= composite.metadata.items() + assert cell_coordinate_source.metadata.items() <= composite.metadata.items() + + +def test_composite_field_source_get_all_fields(cell_coordinate_source, height_coordinate_source): + backend = cell_coordinate_source.backend + grid = cell_coordinate_source.grid + foo = test_helpers.random_field(grid, dims.CellDim, dims.KDim) + bar = test_helpers.random_field(grid, dims.EdgeDim, dims.KDim) + data = { + "foo": (foo, {"standard_name": "foo", "units": ""}), + "bar": (bar, {"standard_name": "bar", "units": ""}), + } + + test_source = SimpleFieldSource(data_=data, grid=grid, backend=backend) + composite = factory.CompositeSource( + test_source, (cell_coordinate_source, height_coordinate_source) + ) + foo = composite.get("foo") + assert isinstance(foo, gtx.Field) + assert {dims.CellDim, dims.KDim}.issubset(foo.domain.dims) + + bar = composite.get("bar") + assert len(bar.domain.dims) == 2 + assert isinstance(bar, gtx.Field) + assert {dims.EdgeDim, dims.KDim}.issubset(bar.domain.dims) + + lon = composite.get("lon") + assert isinstance(lon, gtx.Field) + assert dims.CellDim in lon.domain.dims + assert len(lon.domain.dims) == 1 + + lat = composite.get("height_coordinate") + assert isinstance(lat, gtx.Field) + assert dims.KDim in lat.domain.dims + assert len(lat.domain.dims) == 2 + + +def test_composite_field_source_raises_upon_get_unknown_field( + cell_coordinate_source, height_coordinate_source +): + backend = cell_coordinate_source.backend + grid = cell_coordinate_source.grid + foo = test_helpers.random_field(grid, dims.CellDim, dims.KDim) + bar = test_helpers.random_field(grid, dims.EdgeDim, dims.KDim) + data = { + "foo": (foo, {"standard_name": "foo", "units": ""}), + "bar": (bar, {"standard_name": "bar", "units": ""}), + } + + test_source = SimpleFieldSource(data_=data, grid=grid, backend=backend) + composite = factory.CompositeSource( + test_source, (cell_coordinate_source, height_coordinate_source) + ) + with pytest.raises(ValueError) as err: + composite.get("alice") + assert "not provided by source " in err.value 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/tools/src/icon4pytools/py2fgen/plugin.py b/tools/src/icon4pytools/py2fgen/plugin.py index 95c8a8a1fd..95db0f778b 100644 --- a/tools/src/icon4pytools/py2fgen/plugin.py +++ b/tools/src/icon4pytools/py2fgen/plugin.py @@ -8,7 +8,6 @@ import logging import math -import typing from pathlib import Path from typing import Tuple @@ -16,10 +15,11 @@ import gt4py.next as gtx import numpy as np from cffi import FFI -from icon4py.model.common.settings import xp from numpy.typing import NDArray from icon4pytools.common.logger import setup_logger +from icon4pytools.py2fgen.settings import xp + ffi = FFI() # needed for unpack and unpack_gpu functions @@ -183,7 +183,7 @@ def unpack_and_cache_pointer( unpacked = xp.ones((1,) * len(sizes), dtype=dtype, order="F") else: unpacked = unpack_gpu(pointer, *sizes) if backend == "GPU" else unpack(pointer, *sizes) - + if is_bool: unpacked = int_array_to_bool_array(unpacked) diff --git a/model/common/src/icon4py/model/common/config.py b/tools/src/icon4pytools/py2fgen/settings.py similarity index 88% rename from model/common/src/icon4py/model/common/config.py rename to tools/src/icon4pytools/py2fgen/settings.py index fba138bf14..ffc64028a7 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,11 @@ 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 +xp = config.array_ns diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index bbf93f82a7..ab6b46bb81 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -12,19 +12,21 @@ 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, BUILTIN_TO_ISO_C_TYPE, BUILTIN_TO_NUMPY_TYPE, ) +from icon4pytools.py2fgen import settings 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", @@ -102,7 +104,7 @@ class PythonWrapper(CffiPlugin): is_gt4py_program_present: bool = datamodels.field(init=False) def __post_init__(self, *args: Any, **kwargs: Any) -> None: - self.gt4py_backend = GT4PyBackend[self.backend].value + self.gt4py_backend = settings.GT4PyBackend[self.backend].value self.is_gt4py_program_present = any(func.is_gt4py_program for func in self.functions) self.uninitialised_arrays = get_uninitialised_arrays(self.limited_area) 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.