Cleanup _dsl fields

ci/distributed.yml

@@ -64,7 +64,7 @@ build_distributed_cpu:
     SLURM_NTASKS: 4
     ICON4PY_TEST_DATA_PATH: "/icon4py/testdata"
     ICON4PY_ENABLE_GRID_DOWNLOAD: false
-    ICON4PY_ENABLE_TESTDATA_DOWNLOAD: false
+    ICON4PY_ENABLE_TESTDATA_DOWNLOAD: true
     PYTEST_ADDOPTS: "--durations=0"
     CSCS_ADDITIONAL_MOUNTS: '["/capstor/store/cscs/userlab/cwci02/icon4py/ci/testdata:$ICON4PY_TEST_DATA_PATH"]'
 

model/atmosphere/diffusion/tests/diffusion/integration_tests/test_benchmark_diffusion.py

@@ -135,9 +135,9 @@ def test_diffusion_benchmark(
     metric_state = diffusion_states.DiffusionMetricState(
         theta_ref_mc=metrics_field_source.get(metrics_attributes.THETA_REF_MC),
         wgtfac_c=metrics_field_source.get(metrics_attributes.WGTFAC_C),
-        zd_intcoef=metrics_field_source.get(metrics_attributes.ZD_INTCOEF_DSL),
-        zd_vertoffset=metrics_field_source.get(metrics_attributes.ZD_VERTOFFSET_DSL),
-        zd_diffcoef=metrics_field_source.get(metrics_attributes.ZD_DIFFCOEF_DSL),
+        zd_intcoef=metrics_field_source.get(metrics_attributes.ZD_INTCOEF),
+        zd_vertoffset=metrics_field_source.get(metrics_attributes.ZD_VERTOFFSET),
+        zd_diffcoef=metrics_field_source.get(metrics_attributes.ZD_DIFFCOEF),
     )
     # initialization of the diagnostic and prognostic state
     diagnostic_state = diffusion_states.DiffusionDiagnosticState(

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py

@@ -5,7 +5,7 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
-# ruff: noqa: ERA001, B008
+# ruff: noqa: B008
 
 import dataclasses
 import logging
@@ -1071,7 +1071,6 @@ def time_step(
             exner_new=prognostic_states.next.exner,
         )
 
-    # flake8: noqa: C901
     def run_predictor_step(
         self,
         diagnostic_state_nh: dycore_states.DiagnosticStateNonHydro,

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py

@@ -5,7 +5,6 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
-# ruff: noqa: ERA001
 
 from __future__ import annotations
 

model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py

@@ -175,7 +175,7 @@ def solve_nonhydro(
         zdiff_gradp=metrics_field_source.get(metrics_attributes.ZDIFF_GRADP),
         vertoffset_gradp=metrics_field_source.get(metrics_attributes.VERTOFFSET_GRADP),
         nflat_gradp=metrics_field_source.get(metrics_attributes.NFLAT_GRADP),
-        pg_exdist=metrics_field_source.get(metrics_attributes.PG_EDGEDIST_DSL),
+        pg_exdist=metrics_field_source.get(metrics_attributes.PG_EXDIST_DSL),
         ddqz_z_full_e=metrics_field_source.get(metrics_attributes.DDQZ_Z_FULL_E),
         ddxt_z_full=metrics_field_source.get(metrics_attributes.DDXT_Z_FULL),
         wgtfac_e=metrics_field_source.get(metrics_attributes.WGTFAC_E),

model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py

@@ -1513,7 +1513,7 @@ def test_compute_rho_theta_pgrad_and_update_vn(
         c_lin_e=interpolation_savepoint.c_lin_e(),
         ikoffset=metrics_savepoint.vertoffset_gradp(),
         zdiff_gradp=metrics_savepoint.zdiff_gradp(),
-        pg_exdist=metrics_savepoint.pg_exdist(),
+        pg_exdist=metrics_savepoint.pg_exdist_dsl(),
         inv_dual_edge_length=grid_savepoint.inv_dual_edge_length(),
         dtime=savepoint_nonhydro_init.get_metadata("dtime").get("dtime"),
         iau_wgt_dyn=iau_wgt_dyn,
@@ -1768,7 +1768,7 @@ def test_compute_horizontal_velocity_quantities_and_fluxes(
     ddxn_z_full = metrics_savepoint.ddxn_z_full()
     ddxt_z_full = metrics_savepoint.ddxt_z_full()
     wgtfac_e = metrics_savepoint.wgtfac_e()
-    wgtfacq_e = metrics_savepoint.wgtfacq_e_dsl(icon_grid.num_levels)
+    wgtfacq_e = metrics_savepoint.wgtfacq_e_dsl()
     rbf_vec_coeff_e = interpolation_savepoint.rbf_vec_coeff_e()
     geofac_grdiv = interpolation_savepoint.geofac_grdiv()
     nflatlev = vertical_params.nflatlev

model/atmosphere/dycore/tests/dycore/integration_tests/test_velocity_advection.py

@@ -494,7 +494,7 @@ def test_compute_diagnostics_from_normal_wind(
     ddxn_z_full = metrics_savepoint.ddxn_z_full()
     ddxt_z_full = metrics_savepoint.ddxt_z_full()
     contravariant_correction_at_edges_on_model_levels = savepoint_velocity_init.z_w_concorr_me()
-    wgtfacq_e = metrics_savepoint.wgtfacq_e_dsl(icon_grid.num_levels)
+    wgtfacq_e = metrics_savepoint.wgtfacq_e_dsl()
     nflatlev = grid_savepoint.nflatlev()
     c_intp = interpolation_savepoint.c_intp()
     inv_dual_edge_length = grid_savepoint.inv_dual_edge_length()

model/atmosphere/dycore/tests/dycore/mpi_tests/test_parallel_solve_nonhydro.py

@@ -35,7 +35,7 @@
             2,
             "2021-06-20T12:00:10.000",
             1,
-        )
+        ),
     ],
 )
 @pytest.mark.mpi
@@ -83,7 +83,7 @@ def test_run_solve_nonhydro_single_step(
         f"rank={processor_props.rank}/{processor_props.comm_size}: number of halo cells {np.count_nonzero(np.invert(owned_cells))}"
     )
     print(
-        f"rank={processor_props.rank}/{processor_props.comm_size}: number of halo edges {np.count_nonzero(np.invert( decomposition_info.owner_mask(dims.EdgeDim)))}"
+        f"rank={processor_props.rank}/{processor_props.comm_size}: number of halo edges {np.count_nonzero(np.invert(decomposition_info.owner_mask(dims.EdgeDim)))}"
     )
     print(
         f"rank={processor_props.rank}/{processor_props.comm_size}: number of halo cells {np.count_nonzero(np.invert(owned_cells))}"

model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py

@@ -202,7 +202,9 @@ def reference(
                 ),  # For benchmarking against MCH experiments
             ]
         ],
-        ids=lambda param: f"divdamp_order[{param['divdamp_order']}]__is_iau_active[{param['is_iau_active']}]__second_order_divdamp_factor[{param['second_order_divdamp_factor']}]",
+        ids=lambda param: (
+            f"divdamp_order[{param['divdamp_order']}]__is_iau_active[{param['is_iau_active']}]__second_order_divdamp_factor[{param['second_order_divdamp_factor']}]"
+        ),
     )
     def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict:
         current_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim)

model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py

@@ -626,7 +626,9 @@ def reference(
         params=[
             {"skip_compute_predictor_vertical_advection": value} for value in [True, False]
         ],  # True for benchmarking, False for testing
-        ids=lambda param: f"skip_compute_predictor_vertical_advection[{param['skip_compute_predictor_vertical_advection']}]",
+        ids=lambda param: (
+            f"skip_compute_predictor_vertical_advection[{param['skip_compute_predictor_vertical_advection']}]"
+        ),
     )
     def input_data(
         self, grid: base.Grid, request: pytest.FixtureRequest

model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_averaged_vn_and_fluxes_and_prepare_tracer_advection.py

@@ -162,7 +162,9 @@ def reference(
                 (True, False),
             ]
         ],
-        ids=lambda p: f"prepare_advection[{p['prepare_advection']}]__at_first_substep[{p['at_first_substep']}]",
+        ids=lambda p: (
+            f"prepare_advection[{p['prepare_advection']}]__at_first_substep[{p['at_first_substep']}]"
+        ),
     )
     def input_data(
         self, request: pytest.FixtureRequest, grid: base.Grid

model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_diagnostics_from_normal_wind.py

@@ -294,7 +294,9 @@ def reference(
         params=[
             {"skip_compute_predictor_vertical_advection": value} for value in [True, False]
         ],  # True for benchmarking, False for testing
-        ids=lambda param: f"skip_compute_predictor_vertical_advection[{param['skip_compute_predictor_vertical_advection']}]",
+        ids=lambda param: (
+            f"skip_compute_predictor_vertical_advection[{param['skip_compute_predictor_vertical_advection']}]"
+        ),
     )
     def input_data(
         self, grid: base.Grid, request: pytest.FixtureRequest

model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py

@@ -405,7 +405,9 @@ def reference(
                 (False, False, 32),  # For benchmarking against MCH experiments
             ]
         ],
-        ids=lambda param: f"at_first_substep[{param['at_first_substep']}]__is_iau_active[{param['is_iau_active']}]__divdamp_type[{param['divdamp_type']}]",
+        ids=lambda param: (
+            f"at_first_substep[{param['at_first_substep']}]__is_iau_active[{param['is_iau_active']}]__divdamp_type[{param['divdamp_type']}]"
+        ),
     )
     def input_data(
         self, request: pytest.FixtureRequest, grid: base.Grid

model/atmosphere/dycore/tests/dycore/utils.py

@@ -67,11 +67,11 @@ def construct_metric_state(
         zdiff_gradp=metrics_savepoint.zdiff_gradp(),
         vertoffset_gradp=metrics_savepoint.vertoffset_gradp(),
         nflat_gradp=grid_savepoint.nflat_gradp(),
-        pg_exdist=metrics_savepoint.pg_exdist(),
+        pg_exdist=metrics_savepoint.pg_exdist_dsl(),
         ddqz_z_full_e=metrics_savepoint.ddqz_z_full_e(),
         ddxt_z_full=metrics_savepoint.ddxt_z_full(),
         wgtfac_e=metrics_savepoint.wgtfac_e(),
-        wgtfacq_e=metrics_savepoint.wgtfacq_e_dsl(grid_savepoint.num(dims.KDim)),
+        wgtfacq_e=metrics_savepoint.wgtfacq_e_dsl(),
         exner_w_implicit_weight_parameter=metrics_savepoint.vwind_impl_wgt(),
         horizontal_mask_for_3d_divdamp=metrics_savepoint.hmask_dd3d(),
         scaling_factor_for_3d_divdamp=metrics_savepoint.scalfac_dd3d(),

model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py

@@ -43,7 +43,7 @@ def get_args():
         "--masking",
         dest="enable_masking",
         choices=[True, False],
-        type=lambda x: (str(x).lower() == "true"),
+        type=lambda x: str(x).lower() == "true",
         default=True,
         help="Enable compatibility with reference implementation.",
     )

model/common/pyproject.toml

@@ -23,6 +23,7 @@ classifiers = [
   'Topic :: Scientific/Engineering :: Physics'
 ]
 dependencies = [
+  'array-api-compat>=1.13.0',
   'gt4py==1.1.6',
   'packaging>=20.0',
   'typing-extensions>=4.11.0',

model/common/src/icon4py/model/common/grid/geometry.py

@@ -48,7 +48,12 @@ class GridGeometry(factory.FieldSource):
 
     Examples:
         >>> geometry = GridGeometry(
-        ...     grid, decomposition_info, backend, coordinates, extra_fields, geometry_attributes.attrs
+        ...     grid,
+        ...     decomposition_info,
+        ...     backend,
+        ...     coordinates,
+        ...     extra_fields,
+        ...     geometry_attributes.attrs,
         ... )
         GridGeometry for geometry_type=SPHERE grid=f2e06839-694a-cca1-a3d5-028e0ff326e0 : R9B4
         >>> geometry.get("edge_length")

model/common/src/icon4py/model/common/grid/vertical.py

@@ -183,7 +183,7 @@ def __str__(self) -> str:
         array_value = [
             f"   0   {vct_a_array[0]:12.3f}             ",
             *(
-                f"{k+1:4d}   {vct_a_array[k+1]:12.3f} {dvct[k]:12.3f}"
+                f"{k + 1:4d}   {vct_a_array[k + 1]:12.3f} {dvct[k]:12.3f}"
                 for k in range(vct_a_array.shape[0] - 1)
             ),
         ]

model/common/src/icon4py/model/common/io/io.py

@@ -346,5 +346,5 @@ def close(self) -> None:
 
 
 def generate_name(fname: str, counter: int) -> str:
-    stem = fname.split(".")[0]
+    stem = fname.split(".", maxsplit=1)[0]
     return f"{stem}_{counter:0>4}.nc"

model/common/src/icon4py/model/common/metrics/compute_diffusion_metrics.py

@@ -120,7 +120,7 @@ def compute_diffusion_mask_and_coef(
     array_ns: ModuleType = np,
 ) -> tuple[data_alloc.NDArray, data_alloc.NDArray]:
     n_cells = c2e2c.shape[0]
-    zd_diffcoef_dsl = array_ns.zeros(shape=(n_cells, nlev))
+    zd_diffcoef = array_ns.zeros(shape=(n_cells, nlev))
     k_start, k_end, _ = _compute_k_start_end(
         z_mc=z_mc,
         max_nbhgt=max_nbhgt,
@@ -140,7 +140,7 @@ def compute_diffusion_mask_and_coef(
         if kend > kstart:
             k_range = range(kstart, kend)
 
-            zd_diffcoef_dsl_var = array_ns.maximum(
+            zd_diffcoef_var = array_ns.maximum(
                 0.0,
                 array_ns.maximum(
                     array_ns.sqrt(array_ns.maximum(0.0, maxslp_avg[jc, k_range] - thslp_zdiffu))
@@ -151,9 +151,9 @@ def compute_diffusion_mask_and_coef(
                     ),
                 ),
             )
-            zd_diffcoef_dsl[jc, k_range] = array_ns.minimum(0.002, zd_diffcoef_dsl_var)
+            zd_diffcoef[jc, k_range] = array_ns.minimum(0.002, zd_diffcoef_var)
 
-    return zd_diffcoef_dsl
+    return zd_diffcoef
 
 
 def compute_diffusion_intcoef_and_vertoffset(
@@ -175,8 +175,8 @@ def compute_diffusion_intcoef_and_vertoffset(
     nbidx = array_ns.ones(shape=(n_cells, n_c2e2c, nlev), dtype=int)
 
     z_vintcoeff = array_ns.zeros(shape=(n_cells, n_c2e2c, nlev))
-    zd_vertoffset_dsl = array_ns.zeros(shape=(n_cells, n_c2e2c, nlev), dtype=gtx.int32)
-    zd_intcoef_dsl = array_ns.zeros(shape=(n_cells, n_c2e2c, nlev))
+    zd_vertoffset = array_ns.zeros(shape=(n_cells, n_c2e2c, nlev), dtype=gtx.int32)
+    zd_intcoef = array_ns.zeros(shape=(n_cells, n_c2e2c, nlev))
 
     k_start, k_end, _ = _compute_k_start_end(
         z_mc=z_mc,
@@ -201,10 +201,10 @@ def compute_diffusion_intcoef_and_vertoffset(
                 k_range, z_mc, z_mc_off, z_vintcoeff, jc, nlev
             )
 
-            zd_intcoef_dsl[jc, :, k_range] = z_vintcoeff[jc, :, k_range]
+            zd_intcoef[jc, :, k_range] = z_vintcoeff[jc, :, k_range]
 
-            zd_vertoffset_dsl[jc, :, k_range] = (
+            zd_vertoffset[jc, :, k_range] = (
                 nbidx[jc, :, k_range] - array_ns.tile(array_ns.array(k_range), (3, 1)).T
             )
 
-    return zd_intcoef_dsl, zd_vertoffset_dsl
+    return zd_intcoef, zd_vertoffset

model/common/src/icon4py/model/common/metrics/compute_weight_factors.py

@@ -115,7 +115,7 @@ def compute_wgtfacq_c_dsl(
     wgtfacq_c_dsl[:, nlev - 2] = wgtfacq_c[:, 1]
     wgtfacq_c_dsl[:, nlev - 3] = wgtfacq_c[:, 2]
 
-    return wgtfacq_c_dsl
+    return wgtfacq_c_dsl[:, -3:]
 
 
 def compute_wgtfacq_e_dsl(
@@ -127,7 +127,7 @@ def compute_wgtfacq_e_dsl(
     nlev: int,
     exchange: Callable[[data_alloc.NDArray], None],
     array_ns: ModuleType = np,
-):
+) -> data_alloc.NDArray:
     """
     Compute weighting factor for quadratic interpolation to surface.
 
@@ -145,8 +145,8 @@ def compute_wgtfacq_e_dsl(
     z_aux_c = array_ns.zeros((z_ifc.shape[0], 6))
     z1, z2, z3 = _compute_z1_z2_z3(z_ifc, nlev, nlev - 1, nlev - 2, nlev - 3)
     z_aux_c[:, 2] = z1 * z2 / (z2 - z3) / (z1 - z3)
-    z_aux_c[:, 1] = (z1 - wgtfacq_c_dsl[:, nlev - 3] * (z1 - z3)) / (z1 - z2)
-    z_aux_c[:, 0] = 1.0 - (wgtfacq_c_dsl[:, nlev - 2] + wgtfacq_c_dsl[:, nlev - 3])
+    z_aux_c[:, 1] = (z1 - wgtfacq_c_dsl[:, 0] * (z1 - z3)) / (z1 - z2)
+    z_aux_c[:, 0] = 1.0 - (wgtfacq_c_dsl[:, 1] + wgtfacq_c_dsl[:, 0])
 
     z1, z2, z3 = _compute_z1_z2_z3(z_ifc, 0, 1, 2, 3)
     z_aux_c[:, 5] = z1 * z2 / (z2 - z3) / (z1 - z3)
@@ -161,4 +161,4 @@ def compute_wgtfacq_e_dsl(
     wgtfacq_e_dsl[:, nlev - 1] = z_aux_e[:, 1]
     wgtfacq_e_dsl[:, nlev - 2] = z_aux_e[:, 2]
 
-    return wgtfacq_e_dsl
+    return wgtfacq_e_dsl[:, -3:]

model/common/src/icon4py/model/common/metrics/metrics_attributes.py

@@ -49,7 +49,7 @@
 WGTFAC_E: Final[str] = "wgtfac_e"
 FLAT_IDX_MAX: Final[str] = "flat_idx_max"
 NFLAT_GRADP: Final[str] = "nflat_gradp"
-PG_EDGEDIST_DSL: Final[str] = "distance_for_pressure_gradient_extrapolation"
+PG_EXDIST_DSL: Final[str] = "distance_for_pressure_gradient_extrapolation"
 MASK_PROG_HALO_C: Final[str] = "mask_prog_halo_c"
 HORIZONTAL_MASK_FOR_3D_DIVDAMP: Final[str] = "horizontal_mask_for_3d_divdamp"
 ZDIFF_GRADP: Final[str] = "zdiff_gradp"
@@ -62,9 +62,9 @@
 MAXSLP_AVG: Final[str] = "maxslp_avg"
 MAXHGTD_AVG: Final[str] = "maxhgtd_avg"
 MAX_NBHGT: Final[str] = "max_nbhgt"
-ZD_DIFFCOEF_DSL: Final[str] = "zd_diffcoef_dsl"
-ZD_INTCOEF_DSL: Final[str] = "zd_intcoef_dsl"
-ZD_VERTOFFSET_DSL: Final[str] = "zd_vertoffset_dsl"
+ZD_DIFFCOEF: Final[str] = "zd_diffcoef"
+ZD_INTCOEF: Final[str] = "zd_intcoef"
+ZD_VERTOFFSET: Final[str] = "zd_vertoffset"
 CELL_HEIGHT_ON_HALF_LEVEL: Final[str] = "vertical_coordinates_on_half_levels"
 DEEPATMO_DIVH: Final[str] = "deepatmo_divh"
 DEEPATMO_DIVZL: Final[str] = "deepatmo_divzL"
@@ -303,8 +303,8 @@
         icon_var_name="flat_idx_max",
         dtype=ta.wpfloat,
     ),
-    PG_EDGEDIST_DSL: dict(
-        standard_name=PG_EDGEDIST_DSL,
+    PG_EXDIST_DSL: dict(
+        standard_name=PG_EXDIST_DSL,
         long_name="extrapolation distance for pressure gradient downward extrapolation",
         units="",
         dims=(dims.EdgeDim, dims.KDim),
@@ -407,28 +407,28 @@
         icon_var_name="max_nbhgt",
         dtype=ta.wpfloat,
     ),
-    ZD_DIFFCOEF_DSL: dict(
-        standard_name=ZD_DIFFCOEF_DSL,
-        long_name="zd_diffcoef_dsl",
+    ZD_DIFFCOEF: dict(
+        standard_name=ZD_DIFFCOEF,
+        long_name="zd_diffcoef",
         units="",
         dims=(dims.CellDim, dims.KDim),
         icon_var_name="zd_diffcoef",
         dtype=ta.wpfloat,
     ),
-    ZD_INTCOEF_DSL: dict(
-        standard_name=ZD_INTCOEF_DSL,
-        long_name="zd_intcoef_dsl",
+    ZD_INTCOEF: dict(
+        standard_name=ZD_INTCOEF,
+        long_name="zd_intcoef",
         units="",
         dims=(dims.CellDim, dims.C2E2CDim, dims.KDim),
         icon_var_name="zd_intcoef",
         dtype=ta.wpfloat,
     ),
-    ZD_VERTOFFSET_DSL: dict(
-        standard_name=ZD_VERTOFFSET_DSL,
-        long_name="zd_vertoffset_dsl",
+    ZD_VERTOFFSET: dict(
+        standard_name=ZD_VERTOFFSET,
+        long_name="zd_vertoffset",
         units="",
         dims=(dims.CellDim, dims.C2E2CDim, dims.KDim),
-        icon_var_name="zd_vertoffset_dsl",
+        icon_var_name="zd_vertoffset",
         dtype=gtx.int32,
     ),
     CELL_HEIGHT_ON_HALF_LEVEL: dict(