diff --git a/ci/distributed.yml b/ci/distributed.yml index b4f7e22c08..63dd82d334 100644 --- a/ci/distributed.yml +++ b/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"]' diff --git a/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_benchmark_diffusion.py b/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_benchmark_diffusion.py index 52bf93c95a..40b25be3f7 100644 --- a/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_benchmark_diffusion.py +++ b/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( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py index ba08105e05..984d2c02c2 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py +++ b/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, 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 f93599d4d3..29cc3cb661 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 @@ -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 diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py index 2395cda5d9..9f8e62eb7b 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py +++ b/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), diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index 9ed4926167..011800753a 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/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 diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_velocity_advection.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_velocity_advection.py index d9529e32b7..cacfda444c 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_velocity_advection.py +++ b/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() diff --git a/model/atmosphere/dycore/tests/dycore/mpi_tests/test_parallel_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/mpi_tests/test_parallel_solve_nonhydro.py index ea57fde4c1..da5ffc3504 100644 --- a/model/atmosphere/dycore/tests/dycore/mpi_tests/test_parallel_solve_nonhydro.py +++ b/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))}" diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py index 94849c4b3b..0ddbe7cf4e 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py @@ -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) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py index 8b38daf130..3c32009b11 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py @@ -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 diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_averaged_vn_and_fluxes_and_prepare_tracer_advection.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_averaged_vn_and_fluxes_and_prepare_tracer_advection.py index d07d1e25bf..cf503c38ca 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_averaged_vn_and_fluxes_and_prepare_tracer_advection.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_averaged_vn_and_fluxes_and_prepare_tracer_advection.py @@ -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 diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_diagnostics_from_normal_wind.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_diagnostics_from_normal_wind.py index addb22cc65..c47e4e5332 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_diagnostics_from_normal_wind.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_diagnostics_from_normal_wind.py @@ -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 diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py index b6ec8475fd..8c66e8ab86 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py @@ -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 diff --git a/model/atmosphere/dycore/tests/dycore/utils.py b/model/atmosphere/dycore/tests/dycore/utils.py index a7eb6568bd..efe09d1b39 100644 --- a/model/atmosphere/dycore/tests/dycore/utils.py +++ b/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(), diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py index 0551cf977c..0149930441 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py +++ b/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.", ) diff --git a/model/common/pyproject.toml b/model/common/pyproject.toml index 914543da24..98b04233fa 100644 --- a/model/common/pyproject.toml +++ b/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', diff --git a/model/common/src/icon4py/model/common/grid/geometry.py b/model/common/src/icon4py/model/common/grid/geometry.py index 2563491963..fdbe124370 100644 --- a/model/common/src/icon4py/model/common/grid/geometry.py +++ b/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") diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py index 8c196d0d9e..cfc192fe3a 100644 --- a/model/common/src/icon4py/model/common/grid/vertical.py +++ b/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) ), ] diff --git a/model/common/src/icon4py/model/common/io/io.py b/model/common/src/icon4py/model/common/io/io.py index 634bf7f11b..1685dc9919 100644 --- a/model/common/src/icon4py/model/common/io/io.py +++ b/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" 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 f0ad660e27..18b71093e1 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 @@ -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 diff --git a/model/common/src/icon4py/model/common/metrics/compute_weight_factors.py b/model/common/src/icon4py/model/common/metrics/compute_weight_factors.py index d97d9cc13b..e1dee3d4f0 100644 --- a/model/common/src/icon4py/model/common/metrics/compute_weight_factors.py +++ b/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:] diff --git a/model/common/src/icon4py/model/common/metrics/metrics_attributes.py b/model/common/src/icon4py/model/common/metrics/metrics_attributes.py index d1416c4157..331f0acdae 100644 --- a/model/common/src/icon4py/model/common/metrics/metrics_attributes.py +++ b/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( diff --git a/model/common/src/icon4py/model/common/metrics/metrics_factory.py b/model/common/src/icon4py/model/common/metrics/metrics_factory.py index 9e47493b0b..1cc16324a3 100644 --- a/model/common/src/icon4py/model/common/metrics/metrics_factory.py +++ b/model/common/src/icon4py/model/common/metrics/metrics_factory.py @@ -715,7 +715,7 @@ def _register_computed_fields(self) -> None: # noqa: PLR0915 [too-many-statemen vertical_domain(v_grid.Zone.BOTTOM), ), }, - fields={"pg_exdist_dsl": attrs.PG_EDGEDIST_DSL}, + fields={"pg_exdist_dsl": attrs.PG_EXDIST_DSL}, do_exchange=False, ) self.register_provider(pressure_gradient_fields) @@ -811,7 +811,12 @@ def _register_computed_fields(self) -> None: # noqa: PLR0915 [too-many-statemen compute_wgtfacq_c = factory.NumpyDataProvider( func=functools.partial(weight_factors.compute_wgtfacq_c_dsl, array_ns=self._xp), - domain=(dims.CellDim, dims.KDim), + domain=gtx.domain( + { + dims.CellDim: (0, self._grid.num_cells), + dims.KDim: (self._grid.num_levels - 3, self._grid.num_levels), + } + ), fields=(attrs.WGTFACQ_C,), deps={"z_ifc": attrs.CELL_HEIGHT_ON_HALF_LEVEL}, params={"nlev": self._grid.num_levels}, @@ -831,7 +836,12 @@ def _register_computed_fields(self) -> None: # noqa: PLR0915 [too-many-statemen "wgtfacq_c_dsl": attrs.WGTFACQ_C, }, connectivities={"e2c": dims.E2CDim}, - domain=(dims.EdgeDim, dims.KDim), + domain=gtx.domain( + { + dims.EdgeDim: (0, self._grid.num_edges), + dims.KDim: (self._grid.num_levels - 3, self._grid.num_levels), + } + ), fields=(attrs.WGTFACQ_E,), params={"n_edges": self._grid.num_edges, "nlev": self._grid.num_levels}, ) @@ -912,7 +922,7 @@ def _register_computed_fields(self) -> None: # noqa: PLR0915 [too-many-statemen }, connectivities={"c2e2c": dims.C2E2CDim}, domain=(dims.CellDim, dims.KDim), - fields=(attrs.ZD_DIFFCOEF_DSL,), + fields=(attrs.ZD_DIFFCOEF,), params={ "thslp_zdiffu": self._config["thslp_zdiffu"], "thhgtd_zdiffu": self._config["thhgtd_zdiffu"], @@ -940,8 +950,8 @@ def _register_computed_fields(self) -> None: # noqa: PLR0915 [too-many-statemen connectivities={"c2e2c": dims.C2E2CDim}, domain=(dims.CellDim, dims.C2E2CDim, dims.KDim), fields=( - attrs.ZD_INTCOEF_DSL, - attrs.ZD_VERTOFFSET_DSL, + attrs.ZD_INTCOEF, + attrs.ZD_VERTOFFSET, ), params={ "thslp_zdiffu": self._config["thslp_zdiffu"], diff --git a/model/common/src/icon4py/model/common/utils/data_allocation.py b/model/common/src/icon4py/model/common/utils/data_allocation.py index 011d1f72b9..6f3ee8d6a4 100644 --- a/model/common/src/icon4py/model/common/utils/data_allocation.py +++ b/model/common/src/icon4py/model/common/utils/data_allocation.py @@ -12,10 +12,11 @@ from types import ModuleType from typing import TYPE_CHECKING, Any, TypeAlias, TypeGuard, TypeVar +import array_api_compat +import gt4py.next as gtx import gt4py.next.typing as gtx_typing import numpy as np import numpy.typing as npt -from gt4py import next as gtx from icon4py.model.common import type_alias as ta from icon4py.model.common.utils import device_utils @@ -23,6 +24,7 @@ if TYPE_CHECKING: from icon4py.model.common.grid import base as grid_base + from icon4py.model.common.states import utils as state_utils try: @@ -196,3 +198,39 @@ def index_field( xp = import_array_ns(allocator) shapex = _shape(grid, dim, extend=extend)[0] return gtx.as_field((dim,), xp.arange(shapex, dtype=dtype), allocator=allocator) + + +def array_namespace(array: NDArray) -> ModuleType: + """ + Returns the array namespace for a given array. + """ + return array_api_compat.array_namespace(array) + + +def list2field( + domain: gtx.Domain, + values: NDArray, + indices: tuple[NDArray, ...], + default_value: state_utils.ScalarType, + allocator: gtx_typing.Allocator, +) -> gtx.Field: + if len(domain) != len(indices): + raise RuntimeError("The number of indices must match the shape of the domain.") + assert all(index.shape == indices[0].shape for index in indices if not isinstance(index, slice)) + xp = array_namespace(values) + arr = xp.full(domain.shape, fill_value=default_value, dtype=values.dtype) + arr[indices] = values + return gtx.as_field(domain, arr, allocator=allocator) + + +def kflip_wgtfacq( + arr: NDArray, + domain: gtx.Domain, + allocator: gtx_typing.Allocator, +) -> gtx.Field: + return gtx.as_field(domain, arr[:, ::-1], allocator=allocator) # type: ignore [arg-type] # type "ndarray[Any, dtype[Any] | Any"; expected "NDArrayObject" + + +def adjust_fortran_indices(inp: NDArray) -> NDArray: + """For some Fortran arrays we need to subtract 1 to be compatible with Python indexing.""" + return inp - 1 diff --git a/model/common/tests/common/decomposition/mpi_tests/test_mpi_decomposition.py b/model/common/tests/common/decomposition/mpi_tests/test_mpi_decomposition.py index 0f0d4c59f1..61d8fee2d0 100644 --- a/model/common/tests/common/decomposition/mpi_tests/test_mpi_decomposition.py +++ b/model/common/tests/common/decomposition/mpi_tests/test_mpi_decomposition.py @@ -267,7 +267,7 @@ def test_exchange_on_dummy_data( result = input_field.asnumpy() print(f"rank={processor_props.rank} - num of halo points ={halo_points.shape}") print( - f" rank={processor_props.rank} - exchanged points: {np.sum(result != number)/grid.num_levels}" + f" rank={processor_props.rank} - exchanged points: {np.sum(result != number) / grid.num_levels}" ) print(f"rank={processor_props.rank} - halo points: {halo_points}") changed_points = np.argwhere(result[:, 2] != number) diff --git a/model/common/tests/common/grid/unit_tests/test_grid_manager.py b/model/common/tests/common/grid/unit_tests/test_grid_manager.py index 41cfb9b238..f1ac4e6071 100644 --- a/model/common/tests/common/grid/unit_tests/test_grid_manager.py +++ b/model/common/tests/common/grid/unit_tests/test_grid_manager.py @@ -334,7 +334,7 @@ def assert_up_to_order( for n in range(reduced_table.shape[0]): assert np.all( np.isin(reduced_table[n, :], reduced_reference[n, :]) - ), f"values in row {n+start_index} are not equal: {reduced_table[n, :]} vs ref= {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 diff --git a/model/common/tests/common/metrics/mpi_tests/test_parallel_metrics.py b/model/common/tests/common/metrics/mpi_tests/test_parallel_metrics.py index 1cab116d5a..b75e217d08 100644 --- a/model/common/tests/common/metrics/mpi_tests/test_parallel_metrics.py +++ b/model/common/tests/common/metrics/mpi_tests/test_parallel_metrics.py @@ -101,7 +101,7 @@ def test_distributed_metrics_attrs( (attrs.DDXT_Z_FULL, "ddxt_z_full"), (attrs.EXNER_W_IMPLICIT_WEIGHT_PARAMETER, "vwind_impl_wgt"), (attrs.EXNER_W_EXPLICIT_WEIGHT_PARAMETER, "vwind_expl_wgt"), - (attrs.PG_EDGEDIST_DSL, "pg_exdist"), + (attrs.PG_EXDIST_DSL, "pg_exdist_dsl"), (attrs.MASK_PROG_HALO_C, "mask_prog_halo_c"), (attrs.HORIZONTAL_MASK_FOR_3D_DIVDAMP, "hmask_dd3d"), (attrs.WGTFAC_C, "wgtfac_c"), @@ -135,9 +135,9 @@ def test_distributed_metrics_attrs_no_halo( @pytest.mark.parametrize( "attrs_name, metrics_name", [ - (attrs.ZD_DIFFCOEF_DSL, "zd_diffcoef"), - (attrs.ZD_INTCOEF_DSL, "zd_intcoef"), - (attrs.ZD_VERTOFFSET_DSL, "zd_vertoffset"), + (attrs.ZD_DIFFCOEF, "zd_diffcoef"), + (attrs.ZD_INTCOEF, "zd_intcoef"), + (attrs.ZD_VERTOFFSET, "zd_vertoffset"), ], ) def test_distributed_metrics_attrs_no_halo_regional( @@ -184,7 +184,7 @@ def test_distributed_metrics_wgtfacq_e( factory = metrics_factory_from_savepoint field = factory.get(attrs.WGTFACQ_E).asnumpy() - field_ref = metrics_savepoint.wgtfacq_e_dsl(field.shape[1]).asnumpy() + field_ref = metrics_savepoint.wgtfacq_e_dsl().asnumpy() assert test_utils.dallclose(field, field_ref) diff --git a/model/common/tests/common/metrics/unit_tests/test_compute_diffusion_metrics.py b/model/common/tests/common/metrics/unit_tests/test_compute_diffusion_metrics.py index c2821c5e71..c176edc7d6 100644 --- a/model/common/tests/common/metrics/unit_tests/test_compute_diffusion_metrics.py +++ b/model/common/tests/common/metrics/unit_tests/test_compute_diffusion_metrics.py @@ -111,7 +111,7 @@ def test_compute_diffusion_mask_and_coeff( offset_provider={"C2E2C": icon_grid.get_connectivity("C2E2C")}, ) - zd_diffcoef_dsl = compute_diffusion_mask_and_coef( + zd_diffcoef = compute_diffusion_mask_and_coef( c2e2c=c2e2c, z_mc=z_mc.asnumpy(), max_nbhgt=max_nbhgt.asnumpy(), @@ -124,7 +124,7 @@ def test_compute_diffusion_mask_and_coeff( nlev=nlev, ) assert test_utils.dallclose( - zd_diffcoef_dsl, metrics_savepoint.zd_diffcoef().asnumpy(), rtol=1.0e-11 + zd_diffcoef, metrics_savepoint.zd_diffcoef().asnumpy(), rtol=1.0e-11 ) @@ -194,7 +194,7 @@ def test_compute_diffusion_intcoef_and_vertoffset( offset_provider={"C2E2C": icon_grid.get_connectivity("C2E2C")}, ) - zd_intcoef_dsl, zd_vertoffset_dsl = compute_diffusion_intcoef_and_vertoffset( + zd_intcoef, zd_vertoffset = compute_diffusion_intcoef_and_vertoffset( c2e2c=c2e2c, z_mc=z_mc.asnumpy(), max_nbhgt=max_nbhgt.asnumpy(), @@ -206,5 +206,5 @@ def test_compute_diffusion_intcoef_and_vertoffset( cell_nudging=cell_nudging, nlev=nlev, ) - assert test_utils.dallclose(zd_vertoffset_dsl, metrics_savepoint.zd_vertoffset().asnumpy()) - assert test_utils.dallclose(zd_intcoef_dsl, metrics_savepoint.zd_intcoef().asnumpy()) + assert test_utils.dallclose(zd_vertoffset, metrics_savepoint.zd_vertoffset().asnumpy()) + assert test_utils.dallclose(zd_intcoef, metrics_savepoint.zd_intcoef().asnumpy()) diff --git a/model/common/tests/common/metrics/unit_tests/test_compute_weight_factors.py b/model/common/tests/common/metrics/unit_tests/test_compute_weight_factors.py index 2df409e01c..746307ae84 100644 --- a/model/common/tests/common/metrics/unit_tests/test_compute_weight_factors.py +++ b/model/common/tests/common/metrics/unit_tests/test_compute_weight_factors.py @@ -79,14 +79,14 @@ def test_compute_wgtfacq_e_dsl( icon_grid: base_grid.Grid, backend: gtx_typing.Backend | None, ) -> None: - wgtfacq_e_dsl_ref = metrics_savepoint.wgtfacq_e_dsl(icon_grid.num_levels + 1) - wgtfacq_c_dsl = metrics_savepoint.wgtfacq_c_dsl() + wgtfacq_e_dsl_ref = metrics_savepoint.wgtfacq_e_dsl() + wgtfacq_c_dsl_ref = metrics_savepoint.wgtfacq_c_dsl() xp = data_alloc.import_array_ns(backend) - wgtfacq_e_dsl_full = weight_factors.compute_wgtfacq_e_dsl( + wgtfacq_e_dsl = weight_factors.compute_wgtfacq_e_dsl( e2c=icon_grid.get_connectivity("E2C").ndarray, z_ifc=metrics_savepoint.z_ifc().ndarray, - wgtfacq_c_dsl=wgtfacq_c_dsl.ndarray, + wgtfacq_c_dsl=wgtfacq_c_dsl_ref.ndarray, c_lin_e=interpolation_savepoint.c_lin_e().ndarray, n_edges=icon_grid.num_edges, nlev=icon_grid.num_levels, @@ -94,9 +94,7 @@ def test_compute_wgtfacq_e_dsl( array_ns=xp, ) - assert test_utils.dallclose( - data_alloc.as_numpy(wgtfacq_e_dsl_full), wgtfacq_e_dsl_ref.asnumpy() - ) + assert test_utils.dallclose(data_alloc.as_numpy(wgtfacq_e_dsl), wgtfacq_e_dsl_ref.asnumpy()) @pytest.mark.datatest @@ -105,12 +103,12 @@ def test_compute_wgtfacq_c_dsl( metrics_savepoint: sb.MetricSavepoint, backend: gtx_typing.Backend | None, ) -> None: - wgtfacq_c_dsl = metrics_savepoint.wgtfacq_c_dsl() + wgtfacq_c_dsl_ref = metrics_savepoint.wgtfacq_c_dsl() xp = data_alloc.import_array_ns(backend) - wgtfacq_c_dsl_ndarray = weight_factors.compute_wgtfacq_c_dsl( + wgtfacq_c_dsl = weight_factors.compute_wgtfacq_c_dsl( z_ifc=metrics_savepoint.z_ifc().ndarray, nlev=icon_grid.num_levels, array_ns=xp, ) - assert test_utils.dallclose(data_alloc.as_numpy(wgtfacq_c_dsl_ndarray), wgtfacq_c_dsl.asnumpy()) + assert test_utils.dallclose(data_alloc.as_numpy(wgtfacq_c_dsl), wgtfacq_c_dsl_ref.asnumpy()) diff --git a/model/common/tests/common/metrics/unit_tests/test_metric_fields.py b/model/common/tests/common/metrics/unit_tests/test_metric_fields.py index 22836428b0..9caf3fa81c 100644 --- a/model/common/tests/common/metrics/unit_tests/test_metric_fields.py +++ b/model/common/tests/common/metrics/unit_tests/test_metric_fields.py @@ -406,7 +406,7 @@ def test_compute_pressure_gradient_downward_extrapolation_mask_distance( grid_savepoint: sb.IconGridSavepoint, backend: gtx_typing.Backend, ) -> None: - pg_exdist_ref = metrics_savepoint.pg_exdist() + pg_exdist_ref = metrics_savepoint.pg_exdist_dsl() nlev = icon_grid.num_levels z_mc = metrics_savepoint.z_mc() diff --git a/model/common/tests/common/metrics/unit_tests/test_metrics_factory.py b/model/common/tests/common/metrics/unit_tests/test_metrics_factory.py index 4688f56b54..516f7d2d5a 100644 --- a/model/common/tests/common/metrics/unit_tests/test_metrics_factory.py +++ b/model/common/tests/common/metrics/unit_tests/test_metrics_factory.py @@ -430,14 +430,14 @@ def test_factory_pressure_gradient_fields( experiment: definitions.Experiment, backend: gtx_typing.Backend | None, ) -> None: - field_1_ref = metrics_savepoint.pg_exdist() + field_1_ref = metrics_savepoint.pg_exdist_dsl() factory = _get_metrics_factory( backend=backend, experiment=experiment, grid_savepoint=grid_savepoint, topography_savepoint=topography_savepoint, ) - field_1 = factory.get(attrs.PG_EDGEDIST_DSL) + field_1 = factory.get(attrs.PG_EXDIST_DSL) assert test_helpers.dallclose(field_1_ref.asnumpy(), field_1.asnumpy(), atol=1.0e-5) @@ -545,8 +545,11 @@ def test_factory_wgtfacq_e( topography_savepoint=topography_savepoint, ) field = factory.get(attrs.WGTFACQ_E) - field_ref = metrics_savepoint.wgtfacq_e_dsl(field.shape[1]) - assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=1e-9) + field_ref = metrics_savepoint.wgtfacq_e_dsl() + # TODO: upgrade the dallclose such that it verifies the domain ranges. + # This field is defined on k (nlev-3, nlev) an converting to numpy + # doesn't know if it's there or (whatever-3, whatever) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy()) @pytest.mark.level("integration") @@ -607,7 +610,7 @@ def test_factory_compute_diffusion_mask_and_coef( grid_savepoint=grid_savepoint, topography_savepoint=topography_savepoint, ) - field = factory.get(attrs.ZD_DIFFCOEF_DSL) + field = factory.get(attrs.ZD_DIFFCOEF) assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), atol=1.0e-10) @@ -630,7 +633,7 @@ def test_factory_compute_diffusion_intcoeff_and_vertoffset( grid_savepoint=grid_savepoint, topography_savepoint=topography_savepoint, ) - field_1 = factory.get(attrs.ZD_INTCOEF_DSL) - field_2 = factory.get(attrs.ZD_VERTOFFSET_DSL) + field_1 = factory.get(attrs.ZD_INTCOEF) + field_2 = factory.get(attrs.ZD_VERTOFFSET) assert test_helpers.dallclose(field_ref_1.asnumpy(), field_1.asnumpy(), atol=1.0e-8) assert test_helpers.dallclose(field_ref_2.asnumpy(), field_2.asnumpy()) diff --git a/model/driver/src/icon4py/model/driver/initialization_utils.py b/model/driver/src/icon4py/model/driver/initialization_utils.py index 6812886ff8..a78be8f6c3 100644 --- a/model/driver/src/icon4py/model/driver/initialization_utils.py +++ b/model/driver/src/icon4py/model/driver/initialization_utils.py @@ -507,11 +507,11 @@ def read_static_fields( 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(), diff --git a/model/driver/tests/driver/integration_tests/test_icon4py.py b/model/driver/tests/driver/integration_tests/test_icon4py.py index cecc1d70ff..9191b7d75a 100644 --- a/model/driver/tests/driver/integration_tests/test_icon4py.py +++ b/model/driver/tests/driver/integration_tests/test_icon4py.py @@ -225,11 +225,11 @@ def test_run_timeloop_single_step( 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(icon_grid.num_levels), + 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(), diff --git a/model/standalone_driver/src/icon4py/model/standalone_driver/driver_states.py b/model/standalone_driver/src/icon4py/model/standalone_driver/driver_states.py index 473f176d74..7d5ef5892a 100644 --- a/model/standalone_driver/src/icon4py/model/standalone_driver/driver_states.py +++ b/model/standalone_driver/src/icon4py/model/standalone_driver/driver_states.py @@ -179,10 +179,5 @@ def show_timer_report( ) else: log.info( - f"|{timer_name:^30}|" - f"{'not started':^23}|" - f"{'':^23}|" - f"{'':^23}|" - f"{'':^23}|" - f"{'':^23}|" + f"|{timer_name:^30}|{'not started':^23}|{'':^23}|{'':^23}|{'':^23}|{'':^23}|" ) diff --git a/model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py b/model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py index 553ecba494..0d7dec40a6 100644 --- a/model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py +++ b/model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py @@ -223,9 +223,9 @@ def initialize_granules( diffusion_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), ) log.info("creating solve nonhydro interpolation state") @@ -293,7 +293,7 @@ def initialize_granules( 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), diff --git a/model/testing/src/icon4py/model/testing/definitions.py b/model/testing/src/icon4py/model/testing/definitions.py index a6a011b9dd..b2f820c298 100644 --- a/model/testing/src/icon4py/model/testing/definitions.py +++ b/model/testing/src/icon4py/model/testing/definitions.py @@ -194,7 +194,7 @@ class Experiment: description: str grid: GridDescription num_levels: int - version: int = 0 + version: int = 2 class Experiments: diff --git a/model/testing/src/icon4py/model/testing/filters.py b/model/testing/src/icon4py/model/testing/filters.py index 8886e1a2c2..0821945056 100644 --- a/model/testing/src/icon4py/model/testing/filters.py +++ b/model/testing/src/icon4py/model/testing/filters.py @@ -38,8 +38,8 @@ class ItemFilter(NamedTuple): action=functools.partial(pytest.skip, "currently only runs on CPU"), ), pytest.mark.embedded_only.name: ItemFilter( - condition=lambda item: not test_utils.is_embedded( - test_utils.get_backend_fixture_value(item) + condition=lambda item: ( + not test_utils.is_embedded(test_utils.get_backend_fixture_value(item)) ), action=functools.partial(pytest.skip, "stencil runs only on embedded backend"), ), @@ -70,8 +70,10 @@ class ItemFilter(NamedTuple): action=functools.partial(pytest.skip, "GTFN compilation is too slow for this test."), ), pytest.mark.skip_value_error.name: ItemFilter( - condition=lambda item: (grid := test_utils.get_fixture_value("grid", item)).limited_area - or grid.geometry_type == base.GeometryType.ICOSAHEDRON, + condition=lambda item: ( + (grid := test_utils.get_fixture_value("grid", item)).limited_area + or grid.geometry_type == base.GeometryType.ICOSAHEDRON + ), action=functools.partial( pytest.skip, "Stencil does not support domain containing skip values. Consider shrinking domain.", diff --git a/model/testing/src/icon4py/model/testing/serialbox.py b/model/testing/src/icon4py/model/testing/serialbox.py index 5f7ea52d54..c25988cfc8 100644 --- a/model/testing/src/icon4py/model/testing/serialbox.py +++ b/model/testing/src/icon4py/model/testing/serialbox.py @@ -17,7 +17,7 @@ import icon4py.model.common.decomposition.definitions as decomposition import icon4py.model.common.field_type_aliases as fa import icon4py.model.common.grid.states as grid_states -from icon4py.model.common import dimension as dims, type_alias +from icon4py.model.common import dimension as dims, model_backends, type_alias from icon4py.model.common.grid import base, horizontal as h_grid, icon, utils as grid_utils from icon4py.model.common.states import prognostic_state from icon4py.model.common.utils import data_allocation as data_alloc @@ -663,10 +663,6 @@ def geofac_grg(self): (dims.CellDim, dims.C2E2CODim), grg[:num_cells, :, 1], allocator=self.backend ) - @IconSavepoint.optionally_registered() - def zd_intcoef(self): - return self._get_field("vcoef", dims.CellDim, dims.C2E2CDim, dims.KDim) - def geofac_n2s(self): return self._get_field("geofac_n2s", dims.CellDim, dims.C2E2CODim) @@ -748,8 +744,42 @@ def ddqz_z_full(self): def mask_prog_halo_c(self): return self._get_field("mask_prog_halo_c", dims.CellDim, dtype=bool) + @IconSavepoint.optionally_registered() + def pg_edgeidx(self): + return np.squeeze(self.serializer.read("pg_edgeidx", self.savepoint)) + + @IconSavepoint.optionally_registered() + def pg_vertidx(self): + return np.squeeze(self.serializer.read("pg_vertidx", self.savepoint)) + + @IconSavepoint.optionally_registered() def pg_exdist(self): - return self._get_field("pg_exdist_dsl", dims.EdgeDim, dims.KDim) + return np.squeeze(self.serializer.read("pg_exdist", self.savepoint)) + + def pg_exdist_dsl(self): + pg_edgeidx = self.pg_edgeidx() + pg_vertidx = self.pg_vertidx() + pg_exdist = self.pg_exdist() + domain = self.rho_ref_me().domain + default_value = gtx.float64(0.0) + if (pg_edgeidx is None) or (pg_vertidx is None) or (pg_exdist is None): + # if any of the fields is missing, return a zero field with the correct shape + return gtx.as_field( + domain, + self.xp.full(domain.shape, fill_value=default_value, dtype=gtx.float64), + allocator=model_backends.get_allocator(self.backend), + ) + else: + return data_alloc.list2field( + domain=domain, + values=pg_exdist, + indices=( + data_alloc.adjust_fortran_indices(pg_edgeidx), + data_alloc.adjust_fortran_indices(pg_vertidx), + ), + default_value=default_value, + allocator=model_backends.get_allocator(self.backend), + ) def rayleigh_w(self): return self._get_field("rayleigh_w", dims.KDim) @@ -781,8 +811,27 @@ def vwind_expl_wgt(self): def vwind_impl_wgt(self): return self._get_field("vwind_impl_wgt", dims.CellDim) + def wgtfacq_c(self): + return self._get_field("wgtfacq_c", dims.CellDim, dims.KDim) + def wgtfacq_c_dsl(self): - return self._get_field("wgtfacq_c_dsl", dims.CellDim, dims.KDim) + ar = self.wgtfacq_c().ndarray + k = ar.shape[1] + wgtfac_c = self.wgtfac_c() + cell_range = wgtfac_c.domain[dims.CellDim].unit_range + nlev = wgtfac_c.domain[dims.KDim].unit_range.stop - 1 + k_range = (nlev - k, nlev) + cell_kflip_domain = gtx.domain( + { + dims.CellDim: cell_range, + dims.KDim: k_range, + } + ) + return data_alloc.kflip_wgtfacq( + arr=ar, + domain=cell_kflip_domain, + allocator=model_backends.get_allocator(self.backend), + ) def zdiff_gradp(self): return self._get_field("zdiff_gradp_dsl", dims.EdgeDim, dims.E2CDim, dims.KDim) @@ -822,42 +871,106 @@ def wgtfac_c(self): def wgtfac_e(self): return self._get_field("wgtfac_e", dims.EdgeDim, dims.KDim) - def wgtfacq_e_dsl(self, k_level): - ar = np.squeeze(self.serializer.read("wgtfacq_e", self.savepoint)) - k = k_level - 3 - 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) - def zd_diffcoef(self): - return self._get_field("zd_diffcoef", dims.CellDim, dims.KDim) - - @IconSavepoint.optionally_registered(dims.CellDim, dims.C2E2CDim, dims.KDim) - def zd_intcoef(self): - return self._read_and_reorder_sparse_field("vcoef") + def wgtfacq_e(self): + return self._get_field("wgtfacq_e", dims.EdgeDim, dims.KDim) + + def wgtfacq_e_dsl(self): + ar = self.wgtfacq_e().ndarray + k = ar.shape[1] + wgtfac_e = self.wgtfac_e() + edge_range = wgtfac_e.domain[dims.EdgeDim].unit_range + nlev = wgtfac_e.domain[dims.KDim].unit_range.stop - 1 + k_range = (nlev - k, nlev) + edge_kflip_domain = gtx.domain( + { + dims.EdgeDim: edge_range, + dims.KDim: k_range, + } + ) + return data_alloc.kflip_wgtfacq( + arr=ar, + domain=edge_kflip_domain, + allocator=model_backends.get_allocator(self.backend), + ) 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 = 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 = np.moveaxis(ser_input, 1, -1) + @IconSavepoint.optionally_registered() + def zd_cellidx(self): + return np.squeeze(self.serializer.read("zd_cellidx", self.savepoint)) - return gtx.as_field( - (dims.CellDim, dims.C2E2CDim, dims.KDim), ser_input, allocator=self.backend - ) + @IconSavepoint.optionally_registered() + def zd_vertidx(self): + # this is the k list (with fortran 1-based indexing) for the central point of the C2E2C stencil + return np.squeeze(self.serializer.read("zd_vertidx", self.savepoint))[0, :] @IconSavepoint.optionally_registered(dims.CellDim, dims.C2E2CDim, dims.KDim, dtype=gtx.int32) def zd_vertoffset(self): - return self._read_and_reorder_sparse_field("zd_vertoffset") + zd_cellidx = self.zd_cellidx() + zd_vertidx = self.zd_vertidx() + # these are the three k offsets for the C2E2C neighbors + zd_vertoffset = ( + np.squeeze(self.serializer.read("zd_vertidx", self.savepoint))[1:, :] - zd_vertidx + ) + cell_c2e2c_k_domain = gtx.domain( + { + dims.CellDim: self.theta_ref_mc().domain[dims.CellDim].unit_range, + dims.C2E2CDim: 3, + dims.KDim: self.theta_ref_mc().domain[dims.KDim].unit_range, + } + ) + return data_alloc.list2field( + domain=cell_c2e2c_k_domain, + values=zd_vertoffset.T, + indices=( + data_alloc.adjust_fortran_indices(zd_cellidx), + slice(None), + data_alloc.adjust_fortran_indices(zd_vertidx), + ), + default_value=gtx.int32(0), + allocator=model_backends.get_allocator(self.backend), + ) - def zd_vertidx(self): - return np.squeeze(self.serializer.read("zd_vertidx", self.savepoint)) + @IconSavepoint.optionally_registered(dims.CellDim, dims.C2E2CDim, dims.KDim) + def zd_intcoef(self): + zd_cellidx = self.zd_cellidx() + zd_vertidx = self.zd_vertidx() + zd_intcoef = np.squeeze(self.serializer.read("zd_intcoef", self.savepoint)) + cell_c2e2c_k_domain = gtx.domain( + { + dims.CellDim: self.theta_ref_mc().domain[dims.CellDim].unit_range, + dims.C2E2CDim: 3, + dims.KDim: self.theta_ref_mc().domain[dims.KDim].unit_range, + } + ) + return data_alloc.list2field( + domain=cell_c2e2c_k_domain, + values=zd_intcoef.T, + indices=( + data_alloc.adjust_fortran_indices(zd_cellidx), + slice(None), + data_alloc.adjust_fortran_indices(zd_vertidx), + ), + default_value=gtx.float64(0.0), + allocator=model_backends.get_allocator(self.backend), + ) - def zd_indlist(self): - return np.squeeze(self.serializer.read("zd_indlist", self.savepoint)) + @IconSavepoint.optionally_registered(dims.CellDim, dims.KDim) + def zd_diffcoef(self): + zd_cellidx = self.zd_cellidx() + zd_vertidx = self.zd_vertidx() + zd_diffcoef = np.squeeze(self.serializer.read("zd_diffcoef", self.savepoint)) + return data_alloc.list2field( + domain=self.geopot().domain, + values=zd_diffcoef, + indices=( + data_alloc.adjust_fortran_indices(zd_cellidx), + data_alloc.adjust_fortran_indices(zd_vertidx), + ), + default_value=gtx.float64(0.0), + allocator=model_backends.get_allocator(self.backend), + ) class AdvectionInitSavepoint(IconSavepoint): diff --git a/scripts/compare_icon_icon4py.py b/scripts/compare_icon_icon4py.py index 60e2549e12..58c996b26d 100644 --- a/scripts/compare_icon_icon4py.py +++ b/scripts/compare_icon_icon4py.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 diff --git a/tach.toml b/tach.toml index a9c93be531..d92122081b 100644 --- a/tach.toml +++ b/tach.toml @@ -27,6 +27,7 @@ exclude = [ "uxarray", "cftime", "viztracer", + "clang_format", ] rename = ["serialbox:serialbox4py"] diff --git a/tools/pyproject.toml b/tools/pyproject.toml index c634fe2646..ec006979f9 100644 --- a/tools/pyproject.toml +++ b/tools/pyproject.toml @@ -29,6 +29,7 @@ dependencies = [ 'icon4py-atmosphere-dycore>=0.0.6', 'icon4py-common>=0.0.6', # external dependencies + 'clang-format>=22.1.0', 'cffi>=1.5', 'gt4py==1.1.6', "numpy>=1.23.3", diff --git a/tools/src/icon4py/tools/py2fgen/_codegen.py b/tools/src/icon4py/tools/py2fgen/_codegen.py index 39b7e7e1b4..ef6df41c36 100644 --- a/tools/src/icon4py/tools/py2fgen/_codegen.py +++ b/tools/src/icon4py/tools/py2fgen/_codegen.py @@ -256,7 +256,7 @@ def visit_Parameter(self, name: str, param: _definitions.ParamDescriptor, **kwar def _render_parameter_declaration(name: str, attributes: Sequence[str | None]) -> str: - return f"{','.join(attribute for attribute in attributes if attribute is not None)} :: {name}" + return f"{','.join(attribute for attribute in attributes if attribute is not None)} :: {name}" def _size_arg_name(name: str, i: int) -> str: diff --git a/tools/src/icon4py/tools/py2fgen/wrappers/common.py b/tools/src/icon4py/tools/py2fgen/wrappers/common.py index d40d857ba8..09db91b72a 100644 --- a/tools/src/icon4py/tools/py2fgen/wrappers/common.py +++ b/tools/src/icon4py/tools/py2fgen/wrappers/common.py @@ -11,18 +11,21 @@ import logging from collections.abc import Callable, Iterable from types import ModuleType -from typing import TYPE_CHECKING, Final, TypeAlias +from typing import TYPE_CHECKING, Annotated, Final, TypeAlias import gt4py.next as gtx import gt4py.next.typing as gtx_typing import numpy as np from gt4py import eve from gt4py._core import definitions as gt4py_definitions +from gt4py.next.type_system import type_specifications as ts from icon4py.model.common import dimension as dims, model_backends from icon4py.model.common.decomposition import definitions, mpi_decomposition from icon4py.model.common.grid import base, horizontal as h_grid, icon from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.common.utils.data_allocation import adjust_fortran_indices +from icon4py.tools import py2fgen if TYPE_CHECKING: @@ -44,6 +47,76 @@ log = logging.getLogger(__name__) +Int32Array1D: TypeAlias = Annotated[ + data_alloc.NDArray, + py2fgen.ArrayParamDescriptor( + rank=1, + dtype=ts.ScalarKind.INT32, + memory_space=py2fgen.MemorySpace.MAYBE_DEVICE, + is_optional=False, + ), +] + +OptionalInt32Array1D: TypeAlias = Annotated[ + data_alloc.NDArray | None, + py2fgen.ArrayParamDescriptor( + rank=1, + dtype=ts.ScalarKind.INT32, + memory_space=py2fgen.MemorySpace.MAYBE_DEVICE, + is_optional=True, + ), +] + +OptionalInt32Array2D: TypeAlias = Annotated[ + data_alloc.NDArray | None, + py2fgen.ArrayParamDescriptor( + rank=2, + dtype=ts.ScalarKind.INT32, + memory_space=py2fgen.MemorySpace.MAYBE_DEVICE, + is_optional=True, + ), +] + +Float64Array1D: TypeAlias = Annotated[ + data_alloc.NDArray, + py2fgen.ArrayParamDescriptor( + rank=1, + dtype=ts.ScalarKind.FLOAT64, + memory_space=py2fgen.MemorySpace.MAYBE_DEVICE, + is_optional=False, + ), +] + +Float64Array2D: TypeAlias = Annotated[ + data_alloc.NDArray, + py2fgen.ArrayParamDescriptor( + rank=2, + dtype=ts.ScalarKind.FLOAT64, + memory_space=py2fgen.MemorySpace.MAYBE_DEVICE, + is_optional=False, + ), +] + +OptionalFloat64Array1D: TypeAlias = Annotated[ + data_alloc.NDArray | None, + py2fgen.ArrayParamDescriptor( + rank=1, + dtype=ts.ScalarKind.FLOAT64, + memory_space=py2fgen.MemorySpace.MAYBE_DEVICE, + is_optional=True, + ), +] + +OptionalFloat64Array2D: TypeAlias = Annotated[ + data_alloc.NDArray, + py2fgen.ArrayParamDescriptor( + rank=2, + dtype=ts.ScalarKind.FLOAT64, + memory_space=py2fgen.MemorySpace.MAYBE_DEVICE, + is_optional=True, + ), +] + class BackendIntEnum(eve.IntEnum): DEFAULT = 0 @@ -80,11 +153,6 @@ def impl(_name: str, domain: gtx.Domain, dtype: gt4py_definitions.DType) -> gtx. return impl -def adjust_fortran_indices(inp: NDArray) -> NDArray: - """For some Fortran arrays we need to subtract 1 to be compatible with Python indexing.""" - return inp - 1 - - def shrink_to_dimension( sizes: dict[gtx.Dimension, int], tables: dict[gtx.FieldOffset, NDArray] ) -> dict[gtx.FieldOffset, NDArray]: diff --git a/tools/src/icon4py/tools/py2fgen/wrappers/diffusion_wrapper.py b/tools/src/icon4py/tools/py2fgen/wrappers/diffusion_wrapper.py index 2702c0b5f1..936f14604a 100644 --- a/tools/src/icon4py/tools/py2fgen/wrappers/diffusion_wrapper.py +++ b/tools/src/icon4py/tools/py2fgen/wrappers/diffusion_wrapper.py @@ -22,6 +22,7 @@ import gt4py.next as gtx import numpy as np +import icon4py.model.common.utils.data_allocation as data_alloc from icon4py.model.atmosphere.diffusion.diffusion import ( Diffusion, DiffusionConfig, @@ -64,9 +65,10 @@ def diffusion_init( nudgecoeff_e: fa.EdgeField[wpfloat], rbf_coeff_1: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], gtx.float64], rbf_coeff_2: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], gtx.float64], - zd_diffcoef: fa.CellKField[wpfloat] | None, - zd_vertoffset: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CDim, dims.KDim], gtx.int32] | None, - zd_intcoef: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CDim, dims.KDim], gtx.float64] | None, + zd_cellidx: wrapper_common.OptionalInt32Array2D, + zd_vertidx: wrapper_common.OptionalInt32Array2D, + zd_intcoef: wrapper_common.OptionalFloat64Array2D, + zd_diffcoef: wrapper_common.OptionalFloat64Array1D, ndyn_substeps: gtx.int32, diffusion_type: gtx.int32, hdiff_w: bool, @@ -122,20 +124,67 @@ def diffusion_init( diffusion_params = DiffusionParams(config) nlev = wgtfac_c.domain[dims.KDim].unit_range.stop - 1 # wgtfac_c has nlevp1 levels - cell_k_domain = {dims.CellDim: wgtfac_c.domain[dims.CellDim].unit_range, dims.KDim: nlev} + cell_k_domain = gtx.domain( + {dims.CellDim: wgtfac_c.domain[dims.CellDim].unit_range, dims.KDim: nlev} + ) c2e2c_size = geofac_grg_x.domain[dims.C2E2CODim].unit_range.stop - 1 - cell_c2e2c_k_domain = { - dims.CellDim: wgtfac_c.domain[dims.CellDim].unit_range, - dims.C2E2CDim: c2e2c_size, - dims.KDim: nlev, - } + cell_c2e2c_k_domain = gtx.domain( + { + dims.CellDim: wgtfac_c.domain[dims.CellDim].unit_range, + dims.C2E2CDim: c2e2c_size, + dims.KDim: nlev, + } + ) xp = wgtfac_c.array_ns - if zd_diffcoef is None: + + if zd_cellidx is None: + # then l_zdiffu_t = .false. and these are all not initialized zd_diffcoef = gtx.zeros(cell_k_domain, dtype=theta_ref_mc.dtype) - if zd_intcoef is None: zd_intcoef = gtx.zeros(cell_c2e2c_k_domain, dtype=wgtfac_c.dtype) - if zd_vertoffset is None: zd_vertoffset = gtx.zeros(cell_c2e2c_k_domain, dtype=xp.int32) + else: + # transform lists to fields + # + # only the first row is needed, the others are for C2E2C neighbors, but slicing in fortran causes issues + zd_cellidx = zd_cellidx[0, :] + # these are the three k offsets for the C2E2C neighbors + zd_vertoffset = zd_vertidx[1:, :] - zd_vertidx[0, :] + # this is the k list (with fortran 1-based indexing) for the central point of the C2E2C stencil + zd_vertidx = zd_vertidx[0, :] + + zd_diffcoef = data_alloc.list2field( + domain=cell_k_domain, + values=zd_diffcoef, + indices=( + data_alloc.adjust_fortran_indices(zd_cellidx), + data_alloc.adjust_fortran_indices(zd_vertidx), + ), + default_value=gtx.float64(0.0), + allocator=model_backends.get_allocator(actual_backend), + ) + zd_intcoef = data_alloc.list2field( + domain=cell_c2e2c_k_domain, + values=zd_intcoef.T, + indices=( + data_alloc.adjust_fortran_indices(zd_cellidx), + slice(None), + data_alloc.adjust_fortran_indices(zd_vertidx), + ), + default_value=gtx.float64(0.0), + allocator=model_backends.get_allocator(actual_backend), + ) + zd_vertoffset = data_alloc.list2field( + domain=cell_c2e2c_k_domain, + values=zd_vertoffset.T, + indices=( + data_alloc.adjust_fortran_indices(zd_cellidx), + slice(None), + data_alloc.adjust_fortran_indices(zd_vertidx), + ), + default_value=gtx.int32(0), + allocator=model_backends.get_allocator(actual_backend), + ) + # Metric state metric_state = DiffusionMetricState( theta_ref_mc=theta_ref_mc, diff --git a/tools/src/icon4py/tools/py2fgen/wrappers/dycore_wrapper.py b/tools/src/icon4py/tools/py2fgen/wrappers/dycore_wrapper.py index 8b8dec2492..4fceb153a5 100644 --- a/tools/src/icon4py/tools/py2fgen/wrappers/dycore_wrapper.py +++ b/tools/src/icon4py/tools/py2fgen/wrappers/dycore_wrapper.py @@ -85,7 +85,9 @@ def solve_nh_init( ddxn_z_full: gtx.Field[gtx.Dims[dims.EdgeDim, dims.KDim], gtx.float64], zdiff_gradp: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim, dims.KDim], gtx.float64], vertoffset_gradp: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim, dims.KDim], gtx.int32], - pg_exdist: gtx.Field[gtx.Dims[dims.EdgeDim, dims.KDim], gtx.float64], + pg_edgeidx: wrapper_common.OptionalInt32Array1D, + pg_vertidx: wrapper_common.OptionalInt32Array1D, + pg_exdist: wrapper_common.OptionalFloat64Array1D, ddqz_z_full_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.KDim], gtx.float64], ddxt_z_full: gtx.Field[gtx.Dims[dims.EdgeDim, dims.KDim], gtx.float64], wgtfac_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.KDim], gtx.float64], @@ -134,6 +136,28 @@ def solve_nh_init( backend_name = actual_backend.name if hasattr(actual_backend, "name") else actual_backend logger.info(f"Using Backend {backend_name} with on_gpu={on_gpu}") + xp = rho_ref_me.array_ns + domain = rho_ref_me.domain + default_value = gtx.float64(0.0) + if (pg_edgeidx is None) or (pg_vertidx is None) or (pg_exdist is None): + # if any of the fields is missing, return a zero field with the correct shape + pg_exdist_dsl = gtx.as_field( + domain, + xp.full(domain.shape, fill_value=default_value, dtype=gtx.float64), + allocator=model_backends.get_allocator(actual_backend), + ) + else: + pg_exdist_dsl = data_alloc.list2field( + domain=domain, + values=pg_exdist, + indices=( + data_alloc.adjust_fortran_indices(pg_edgeidx), + data_alloc.adjust_fortran_indices(pg_vertidx), + ), + default_value=default_value, + allocator=model_backends.get_allocator(actual_backend), + ) + config = solve_nonhydro.NonHydrostaticConfig( itime_scheme=itime_scheme, iadv_rhotheta=iadv_rhotheta, @@ -181,6 +205,32 @@ def solve_nh_init( nudgecoeff_e=nudgecoeff_e, ) + nlev = wgtfac_c.domain[dims.KDim].unit_range.stop - 1 + k = wgtfacq_c.ndarray.shape[1] + cell_kflip_domain = gtx.domain( + { + dims.CellDim: wgtfac_c.domain[dims.CellDim].unit_range, + dims.KDim: (nlev - k, nlev), + } + ) + k = wgtfacq_e.ndarray.shape[1] + edge_kflip_domain = gtx.domain( + { + dims.EdgeDim: wgtfac_e.domain[dims.EdgeDim].unit_range, + dims.KDim: (nlev - k, nlev), + } + ) + wgtfacq_c = data_alloc.kflip_wgtfacq( + arr=wgtfacq_c.ndarray, + domain=cell_kflip_domain, + allocator=model_backends.get_allocator(actual_backend), + ) + wgtfacq_e = data_alloc.kflip_wgtfacq( + arr=wgtfacq_e.ndarray, + domain=edge_kflip_domain, + allocator=model_backends.get_allocator(actual_backend), + ) + metric_state_nonhydro = dycore_states.MetricStateNonHydro( mask_prog_halo_c=mask_prog_halo_c, rayleigh_w=rayleigh_w, @@ -203,7 +253,7 @@ def solve_nh_init( zdiff_gradp=zdiff_gradp, vertoffset_gradp=vertoffset_gradp, nflat_gradp=gtx.int32(nflat_gradp - 1), # Fortran vs Python indexing - pg_exdist=pg_exdist, + pg_exdist=pg_exdist_dsl, ddqz_z_full_e=ddqz_z_full_e, ddxt_z_full=ddxt_z_full, wgtfac_e=wgtfac_e, diff --git a/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.f90 b/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.f90 index 4bbc897946..537ce6e91b 100644 --- a/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.f90 +++ b/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.f90 @@ -132,17 +132,17 @@ function diffusion_init_wrapper(theta_ref_mc, & rbf_coeff_2, & rbf_coeff_2_size_0, & rbf_coeff_2_size_1, & - zd_diffcoef, & - zd_diffcoef_size_0, & - zd_diffcoef_size_1, & - zd_vertoffset, & - zd_vertoffset_size_0, & - zd_vertoffset_size_1, & - zd_vertoffset_size_2, & + zd_cellidx, & + zd_cellidx_size_0, & + zd_cellidx_size_1, & + zd_vertidx, & + zd_vertidx_size_0, & + zd_vertidx_size_1, & zd_intcoef, & zd_intcoef_size_0, & zd_intcoef_size_1, & - zd_intcoef_size_2, & + zd_diffcoef, & + zd_diffcoef_size_0, & ndyn_substeps, & diffusion_type, & hdiff_w, & @@ -223,19 +223,17 @@ function diffusion_init_wrapper(theta_ref_mc, & integer(c_int), value :: rbf_coeff_2_size_1 - type(c_ptr), value, target :: zd_diffcoef - - integer(c_int), value :: zd_diffcoef_size_0 + type(c_ptr), value, target :: zd_cellidx - integer(c_int), value :: zd_diffcoef_size_1 + integer(c_int), value :: zd_cellidx_size_0 - type(c_ptr), value, target :: zd_vertoffset + integer(c_int), value :: zd_cellidx_size_1 - integer(c_int), value :: zd_vertoffset_size_0 + type(c_ptr), value, target :: zd_vertidx - integer(c_int), value :: zd_vertoffset_size_1 + integer(c_int), value :: zd_vertidx_size_0 - integer(c_int), value :: zd_vertoffset_size_2 + integer(c_int), value :: zd_vertidx_size_1 type(c_ptr), value, target :: zd_intcoef @@ -243,7 +241,9 @@ function diffusion_init_wrapper(theta_ref_mc, & integer(c_int), value :: zd_intcoef_size_1 - integer(c_int), value :: zd_intcoef_size_2 + type(c_ptr), value, target :: zd_diffcoef + + integer(c_int), value :: zd_diffcoef_size_0 integer(c_int), value, target :: ndyn_substeps @@ -488,9 +488,10 @@ subroutine diffusion_init(theta_ref_mc, & nudgecoeff_e, & rbf_coeff_1, & rbf_coeff_2, & - zd_diffcoef, & - zd_vertoffset, & + zd_cellidx, & + zd_vertidx, & zd_intcoef, & + zd_diffcoef, & ndyn_substeps, & diffusion_type, & hdiff_w, & @@ -532,11 +533,13 @@ subroutine diffusion_init(theta_ref_mc, & real(c_double), dimension(:, :), target :: rbf_coeff_2 - real(c_double), dimension(:, :), pointer :: zd_diffcoef + integer(c_int), dimension(:, :), pointer :: zd_cellidx + + integer(c_int), dimension(:, :), pointer :: zd_vertidx - integer(c_int), dimension(:, :, :), pointer :: zd_vertoffset + real(c_double), dimension(:, :), pointer :: zd_intcoef - real(c_double), dimension(:, :, :), pointer :: zd_intcoef + real(c_double), dimension(:), pointer :: zd_diffcoef integer(c_int), value, target :: ndyn_substeps @@ -614,37 +617,39 @@ subroutine diffusion_init(theta_ref_mc, & integer(c_int) :: rbf_coeff_2_size_1 - integer(c_int) :: zd_diffcoef_size_0 - - integer(c_int) :: zd_diffcoef_size_1 + integer(c_int) :: zd_cellidx_size_0 - integer(c_int) :: zd_vertoffset_size_0 + integer(c_int) :: zd_cellidx_size_1 - integer(c_int) :: zd_vertoffset_size_1 + integer(c_int) :: zd_vertidx_size_0 - integer(c_int) :: zd_vertoffset_size_2 + integer(c_int) :: zd_vertidx_size_1 integer(c_int) :: zd_intcoef_size_0 integer(c_int) :: zd_intcoef_size_1 - integer(c_int) :: zd_intcoef_size_2 + integer(c_int) :: zd_diffcoef_size_0 integer(c_int) :: rc ! Stores the return code ! ptrs - type(c_ptr) :: zd_diffcoef_ptr + type(c_ptr) :: zd_cellidx_ptr - type(c_ptr) :: zd_vertoffset_ptr + type(c_ptr) :: zd_vertidx_ptr type(c_ptr) :: zd_intcoef_ptr - zd_diffcoef_ptr = c_null_ptr + type(c_ptr) :: zd_diffcoef_ptr - zd_vertoffset_ptr = c_null_ptr + zd_cellidx_ptr = c_null_ptr + + zd_vertidx_ptr = c_null_ptr zd_intcoef_ptr = c_null_ptr + zd_diffcoef_ptr = c_null_ptr + !$acc host_data use_device(theta_ref_mc) !$acc host_data use_device(wgtfac_c) !$acc host_data use_device(e_bln_c_s) @@ -655,9 +660,10 @@ subroutine diffusion_init(theta_ref_mc, & !$acc host_data use_device(nudgecoeff_e) !$acc host_data use_device(rbf_coeff_1) !$acc host_data use_device(rbf_coeff_2) - !$acc host_data use_device(zd_diffcoef) if(associated(zd_diffcoef)) - !$acc host_data use_device(zd_vertoffset) if(associated(zd_vertoffset)) + !$acc host_data use_device(zd_cellidx) if(associated(zd_cellidx)) + !$acc host_data use_device(zd_vertidx) if(associated(zd_vertidx)) !$acc host_data use_device(zd_intcoef) if(associated(zd_intcoef)) + !$acc host_data use_device(zd_diffcoef) if(associated(zd_diffcoef)) #ifdef _OPENACC on_gpu = .True. @@ -694,24 +700,27 @@ subroutine diffusion_init(theta_ref_mc, & rbf_coeff_2_size_0 = SIZE(rbf_coeff_2, 1) rbf_coeff_2_size_1 = SIZE(rbf_coeff_2, 2) - if (associated(zd_diffcoef)) then - zd_diffcoef_ptr = c_loc(zd_diffcoef) - zd_diffcoef_size_0 = SIZE(zd_diffcoef, 1) - zd_diffcoef_size_1 = SIZE(zd_diffcoef, 2) + if (associated(zd_cellidx)) then + zd_cellidx_ptr = c_loc(zd_cellidx) + zd_cellidx_size_0 = SIZE(zd_cellidx, 1) + zd_cellidx_size_1 = SIZE(zd_cellidx, 2) end if - if (associated(zd_vertoffset)) then - zd_vertoffset_ptr = c_loc(zd_vertoffset) - zd_vertoffset_size_0 = SIZE(zd_vertoffset, 1) - zd_vertoffset_size_1 = SIZE(zd_vertoffset, 2) - zd_vertoffset_size_2 = SIZE(zd_vertoffset, 3) + if (associated(zd_vertidx)) then + zd_vertidx_ptr = c_loc(zd_vertidx) + zd_vertidx_size_0 = SIZE(zd_vertidx, 1) + zd_vertidx_size_1 = SIZE(zd_vertidx, 2) end if if (associated(zd_intcoef)) then zd_intcoef_ptr = c_loc(zd_intcoef) zd_intcoef_size_0 = SIZE(zd_intcoef, 1) zd_intcoef_size_1 = SIZE(zd_intcoef, 2) - zd_intcoef_size_2 = SIZE(zd_intcoef, 3) + end if + + if (associated(zd_diffcoef)) then + zd_diffcoef_ptr = c_loc(zd_diffcoef) + zd_diffcoef_size_0 = SIZE(zd_diffcoef, 1) end if rc = diffusion_init_wrapper(theta_ref_mc=c_loc(theta_ref_mc), & @@ -743,17 +752,17 @@ subroutine diffusion_init(theta_ref_mc, & rbf_coeff_2=c_loc(rbf_coeff_2), & rbf_coeff_2_size_0=rbf_coeff_2_size_0, & rbf_coeff_2_size_1=rbf_coeff_2_size_1, & - zd_diffcoef=zd_diffcoef_ptr, & - zd_diffcoef_size_0=zd_diffcoef_size_0, & - zd_diffcoef_size_1=zd_diffcoef_size_1, & - zd_vertoffset=zd_vertoffset_ptr, & - zd_vertoffset_size_0=zd_vertoffset_size_0, & - zd_vertoffset_size_1=zd_vertoffset_size_1, & - zd_vertoffset_size_2=zd_vertoffset_size_2, & + zd_cellidx=zd_cellidx_ptr, & + zd_cellidx_size_0=zd_cellidx_size_0, & + zd_cellidx_size_1=zd_cellidx_size_1, & + zd_vertidx=zd_vertidx_ptr, & + zd_vertidx_size_0=zd_vertidx_size_0, & + zd_vertidx_size_1=zd_vertidx_size_1, & zd_intcoef=zd_intcoef_ptr, & zd_intcoef_size_0=zd_intcoef_size_0, & zd_intcoef_size_1=zd_intcoef_size_1, & - zd_intcoef_size_2=zd_intcoef_size_2, & + zd_diffcoef=zd_diffcoef_ptr, & + zd_diffcoef_size_0=zd_diffcoef_size_0, & ndyn_substeps=ndyn_substeps, & diffusion_type=diffusion_type, & hdiff_w=hdiff_w, & @@ -786,6 +795,7 @@ subroutine diffusion_init(theta_ref_mc, & !$acc end host_data !$acc end host_data !$acc end host_data + !$acc end host_data end subroutine diffusion_init end module \ No newline at end of file diff --git a/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.h b/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.h index 69010bdf4b..5c35ce2466 100644 --- a/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.h +++ b/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.h @@ -16,13 +16,13 @@ extern int diffusion_init_wrapper( double *geofac_n2s, int geofac_n2s_size_0, int geofac_n2s_size_1, double *nudgecoeff_e, int nudgecoeff_e_size_0, double *rbf_coeff_1, int rbf_coeff_1_size_0, int rbf_coeff_1_size_1, double *rbf_coeff_2, - int rbf_coeff_2_size_0, int rbf_coeff_2_size_1, double *zd_diffcoef, - int zd_diffcoef_size_0, int zd_diffcoef_size_1, int *zd_vertoffset, - int zd_vertoffset_size_0, int zd_vertoffset_size_1, - int zd_vertoffset_size_2, double *zd_intcoef, int zd_intcoef_size_0, - int zd_intcoef_size_1, int zd_intcoef_size_2, int ndyn_substeps, - int diffusion_type, int hdiff_w, int hdiff_vn, int hdiff_smag_w, - int zdiffu_t, int type_t_diffu, int type_vn_diffu, double hdiff_efdt_ratio, + int rbf_coeff_2_size_0, int rbf_coeff_2_size_1, int *zd_cellidx, + int zd_cellidx_size_0, int zd_cellidx_size_1, int *zd_vertidx, + int zd_vertidx_size_0, int zd_vertidx_size_1, double *zd_intcoef, + int zd_intcoef_size_0, int zd_intcoef_size_1, double *zd_diffcoef, + int zd_diffcoef_size_0, int ndyn_substeps, int diffusion_type, int hdiff_w, + int hdiff_vn, int hdiff_smag_w, int zdiffu_t, int type_t_diffu, + int type_vn_diffu, double hdiff_efdt_ratio, double smagorinski_scaling_factor, int hdiff_temp, double thslp_zdiffu, double thhgtd_zdiffu, double denom_diffu_v, double nudge_max_coeff, int itype_sher, int ltkeshs, int backend, int on_gpu); \ No newline at end of file diff --git a/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.py b/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.py index 17419df15f..0de4b5bddd 100644 --- a/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.py +++ b/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.py @@ -371,17 +371,17 @@ def diffusion_init_wrapper( rbf_coeff_2, rbf_coeff_2_size_0, rbf_coeff_2_size_1, - zd_diffcoef, - zd_diffcoef_size_0, - zd_diffcoef_size_1, - zd_vertoffset, - zd_vertoffset_size_0, - zd_vertoffset_size_1, - zd_vertoffset_size_2, + zd_cellidx, + zd_cellidx_size_0, + zd_cellidx_size_1, + zd_vertidx, + zd_vertidx_size_0, + zd_vertidx_size_1, zd_intcoef, zd_intcoef_size_0, zd_intcoef_size_1, - zd_intcoef_size_2, + zd_diffcoef, + zd_diffcoef_size_0, ndyn_substeps, diffusion_type, hdiff_w, @@ -505,22 +505,21 @@ def diffusion_init_wrapper( False, ) - zd_diffcoef = ( - zd_diffcoef, + zd_cellidx = ( + zd_cellidx, ( - zd_diffcoef_size_0, - zd_diffcoef_size_1, + zd_cellidx_size_0, + zd_cellidx_size_1, ), on_gpu, True, ) - zd_vertoffset = ( - zd_vertoffset, + zd_vertidx = ( + zd_vertidx, ( - zd_vertoffset_size_0, - zd_vertoffset_size_1, - zd_vertoffset_size_2, + zd_vertidx_size_0, + zd_vertidx_size_1, ), on_gpu, True, @@ -531,12 +530,13 @@ def diffusion_init_wrapper( ( zd_intcoef_size_0, zd_intcoef_size_1, - zd_intcoef_size_2, ), on_gpu, True, ) + zd_diffcoef = (zd_diffcoef, (zd_diffcoef_size_0,), on_gpu, True) + if __debug__: if runtime_config.PROFILING: allocate_end_time = _runtime.perf_counter() @@ -564,9 +564,10 @@ def diffusion_init_wrapper( nudgecoeff_e=nudgecoeff_e, rbf_coeff_1=rbf_coeff_1, rbf_coeff_2=rbf_coeff_2, - zd_diffcoef=zd_diffcoef, - zd_vertoffset=zd_vertoffset, + zd_cellidx=zd_cellidx, + zd_vertidx=zd_vertidx, zd_intcoef=zd_intcoef, + zd_diffcoef=zd_diffcoef, ndyn_substeps=ndyn_substeps, diffusion_type=diffusion_type, hdiff_w=hdiff_w, @@ -763,34 +764,34 @@ def diffusion_init_wrapper( ) logger.debug(msg) - zd_diffcoef_arr = ( - _conversion.as_array(ffi, zd_diffcoef, _definitions.FLOAT64) - if zd_diffcoef is not None + zd_cellidx_arr = ( + _conversion.as_array(ffi, zd_cellidx, _definitions.INT32) + if zd_cellidx is not None else None ) - msg = "shape of zd_diffcoef after computation = %s" % str( - zd_diffcoef_arr.shape if zd_diffcoef is not None else "None" + msg = "shape of zd_cellidx after computation = %s" % str( + zd_cellidx_arr.shape if zd_cellidx is not None else "None" ) logger.debug(msg) msg = ( - "zd_diffcoef after computation: %s" % str(zd_diffcoef_arr) - if zd_diffcoef is not None + "zd_cellidx after computation: %s" % str(zd_cellidx_arr) + if zd_cellidx is not None else "None" ) logger.debug(msg) - zd_vertoffset_arr = ( - _conversion.as_array(ffi, zd_vertoffset, _definitions.INT32) - if zd_vertoffset is not None + zd_vertidx_arr = ( + _conversion.as_array(ffi, zd_vertidx, _definitions.INT32) + if zd_vertidx is not None else None ) - msg = "shape of zd_vertoffset after computation = %s" % str( - zd_vertoffset_arr.shape if zd_vertoffset is not None else "None" + msg = "shape of zd_vertidx after computation = %s" % str( + zd_vertidx_arr.shape if zd_vertidx is not None else "None" ) logger.debug(msg) msg = ( - "zd_vertoffset after computation: %s" % str(zd_vertoffset_arr) - if zd_vertoffset is not None + "zd_vertidx after computation: %s" % str(zd_vertidx_arr) + if zd_vertidx is not None else "None" ) logger.debug(msg) @@ -811,6 +812,22 @@ def diffusion_init_wrapper( ) logger.debug(msg) + zd_diffcoef_arr = ( + _conversion.as_array(ffi, zd_diffcoef, _definitions.FLOAT64) + if zd_diffcoef is not None + else None + ) + msg = "shape of zd_diffcoef after computation = %s" % str( + zd_diffcoef_arr.shape if zd_diffcoef is not None else "None" + ) + logger.debug(msg) + msg = ( + "zd_diffcoef after computation: %s" % str(zd_diffcoef_arr) + if zd_diffcoef is not None + else "None" + ) + logger.debug(msg) + if __debug__: logger.info("Python execution of diffusion_init completed.") diff --git a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.f90 b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.f90 index 795dfa5035..6b29e00d95 100644 --- a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.f90 +++ b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.f90 @@ -461,9 +461,12 @@ function solve_nh_init_wrapper(c_lin_e, & vertoffset_gradp_size_0, & vertoffset_gradp_size_1, & vertoffset_gradp_size_2, & + pg_edgeidx, & + pg_edgeidx_size_0, & + pg_vertidx, & + pg_vertidx_size_0, & pg_exdist, & pg_exdist_size_0, & - pg_exdist_size_1, & ddqz_z_full_e, & ddqz_z_full_e_size_0, & ddqz_z_full_e_size_1, & @@ -735,12 +738,18 @@ function solve_nh_init_wrapper(c_lin_e, & integer(c_int), value :: vertoffset_gradp_size_2 + type(c_ptr), value, target :: pg_edgeidx + + integer(c_int), value :: pg_edgeidx_size_0 + + type(c_ptr), value, target :: pg_vertidx + + integer(c_int), value :: pg_vertidx_size_0 + type(c_ptr), value, target :: pg_exdist integer(c_int), value :: pg_exdist_size_0 - integer(c_int), value :: pg_exdist_size_1 - type(c_ptr), value, target :: ddqz_z_full_e integer(c_int), value :: ddqz_z_full_e_size_0 @@ -1492,6 +1501,8 @@ subroutine solve_nh_init(c_lin_e, & ddxn_z_full, & zdiff_gradp, & vertoffset_gradp, & + pg_edgeidx, & + pg_vertidx, & pg_exdist, & ddqz_z_full_e, & ddxt_z_full, & @@ -1605,7 +1616,11 @@ subroutine solve_nh_init(c_lin_e, & integer(c_int), dimension(:, :, :), target :: vertoffset_gradp - real(c_double), dimension(:, :), target :: pg_exdist + integer(c_int), dimension(:), pointer :: pg_edgeidx + + integer(c_int), dimension(:), pointer :: pg_vertidx + + real(c_double), dimension(:), pointer :: pg_exdist real(c_double), dimension(:, :), target :: ddqz_z_full_e @@ -1823,9 +1838,11 @@ subroutine solve_nh_init(c_lin_e, & integer(c_int) :: vertoffset_gradp_size_2 - integer(c_int) :: pg_exdist_size_0 + integer(c_int) :: pg_edgeidx_size_0 + + integer(c_int) :: pg_vertidx_size_0 - integer(c_int) :: pg_exdist_size_1 + integer(c_int) :: pg_exdist_size_0 integer(c_int) :: ddqz_z_full_e_size_0 @@ -1866,6 +1883,18 @@ subroutine solve_nh_init(c_lin_e, & integer(c_int) :: rc ! Stores the return code ! ptrs + type(c_ptr) :: pg_edgeidx_ptr + + type(c_ptr) :: pg_vertidx_ptr + + type(c_ptr) :: pg_exdist_ptr + + pg_edgeidx_ptr = c_null_ptr + + pg_vertidx_ptr = c_null_ptr + + pg_exdist_ptr = c_null_ptr + !$acc host_data use_device(c_lin_e) !$acc host_data use_device(c_intp) !$acc host_data use_device(e_flx_avg) @@ -1902,7 +1931,6 @@ subroutine solve_nh_init(c_lin_e, & !$acc host_data use_device(ddxn_z_full) !$acc host_data use_device(zdiff_gradp) !$acc host_data use_device(vertoffset_gradp) - !$acc host_data use_device(pg_exdist) !$acc host_data use_device(ddqz_z_full_e) !$acc host_data use_device(ddxt_z_full) !$acc host_data use_device(wgtfac_e) @@ -1914,6 +1942,9 @@ subroutine solve_nh_init(c_lin_e, & !$acc host_data use_device(coeff2_dwdz) !$acc host_data use_device(coeff_gradekin) !$acc host_data use_device(c_owner_mask) + !$acc host_data use_device(pg_edgeidx) if(associated(pg_edgeidx)) + !$acc host_data use_device(pg_vertidx) if(associated(pg_vertidx)) + !$acc host_data use_device(pg_exdist) if(associated(pg_exdist)) #ifdef _OPENACC on_gpu = .True. @@ -2027,9 +2058,6 @@ subroutine solve_nh_init(c_lin_e, & vertoffset_gradp_size_1 = SIZE(vertoffset_gradp, 2) vertoffset_gradp_size_2 = SIZE(vertoffset_gradp, 3) - pg_exdist_size_0 = SIZE(pg_exdist, 1) - pg_exdist_size_1 = SIZE(pg_exdist, 2) - ddqz_z_full_e_size_0 = SIZE(ddqz_z_full_e, 1) ddqz_z_full_e_size_1 = SIZE(ddqz_z_full_e, 2) @@ -2059,6 +2087,21 @@ subroutine solve_nh_init(c_lin_e, & c_owner_mask_size_0 = SIZE(c_owner_mask, 1) + if (associated(pg_edgeidx)) then + pg_edgeidx_ptr = c_loc(pg_edgeidx) + pg_edgeidx_size_0 = SIZE(pg_edgeidx, 1) + end if + + if (associated(pg_vertidx)) then + pg_vertidx_ptr = c_loc(pg_vertidx) + pg_vertidx_size_0 = SIZE(pg_vertidx, 1) + end if + + if (associated(pg_exdist)) then + pg_exdist_ptr = c_loc(pg_exdist) + pg_exdist_size_0 = SIZE(pg_exdist, 1) + end if + rc = solve_nh_init_wrapper(c_lin_e=c_loc(c_lin_e), & c_lin_e_size_0=c_lin_e_size_0, & c_lin_e_size_1=c_lin_e_size_1, & @@ -2165,9 +2208,12 @@ subroutine solve_nh_init(c_lin_e, & vertoffset_gradp_size_0=vertoffset_gradp_size_0, & vertoffset_gradp_size_1=vertoffset_gradp_size_1, & vertoffset_gradp_size_2=vertoffset_gradp_size_2, & - pg_exdist=c_loc(pg_exdist), & + pg_edgeidx=pg_edgeidx_ptr, & + pg_edgeidx_size_0=pg_edgeidx_size_0, & + pg_vertidx=pg_vertidx_ptr, & + pg_vertidx_size_0=pg_vertidx_size_0, & + pg_exdist=pg_exdist_ptr, & pg_exdist_size_0=pg_exdist_size_0, & - pg_exdist_size_1=pg_exdist_size_1, & ddqz_z_full_e=c_loc(ddqz_z_full_e), & ddqz_z_full_e_size_0=ddqz_z_full_e_size_0, & ddqz_z_full_e_size_1=ddqz_z_full_e_size_1, & @@ -2272,6 +2318,8 @@ subroutine solve_nh_init(c_lin_e, & !$acc end host_data !$acc end host_data !$acc end host_data + !$acc end host_data + !$acc end host_data end subroutine solve_nh_init end module \ No newline at end of file diff --git a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.h b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.h index 40ecd7ef05..660a2fc279 100644 --- a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.h +++ b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.h @@ -75,8 +75,9 @@ extern int solve_nh_init_wrapper( double *ddxn_z_full, int ddxn_z_full_size_0, int ddxn_z_full_size_1, double *zdiff_gradp, int zdiff_gradp_size_0, int zdiff_gradp_size_1, int zdiff_gradp_size_2, int *vertoffset_gradp, int vertoffset_gradp_size_0, - int vertoffset_gradp_size_1, int vertoffset_gradp_size_2, double *pg_exdist, - int pg_exdist_size_0, int pg_exdist_size_1, double *ddqz_z_full_e, + int vertoffset_gradp_size_1, int vertoffset_gradp_size_2, int *pg_edgeidx, + int pg_edgeidx_size_0, int *pg_vertidx, int pg_vertidx_size_0, + double *pg_exdist, int pg_exdist_size_0, double *ddqz_z_full_e, int ddqz_z_full_e_size_0, int ddqz_z_full_e_size_1, double *ddxt_z_full, int ddxt_z_full_size_0, int ddxt_z_full_size_1, double *wgtfac_e, int wgtfac_e_size_0, int wgtfac_e_size_1, double *wgtfacq_e, diff --git a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.py b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.py index a936b6dbc3..15ead38e9a 100644 --- a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.py +++ b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.py @@ -1276,9 +1276,12 @@ def solve_nh_init_wrapper( vertoffset_gradp_size_0, vertoffset_gradp_size_1, vertoffset_gradp_size_2, + pg_edgeidx, + pg_edgeidx_size_0, + pg_vertidx, + pg_vertidx_size_0, pg_exdist, pg_exdist_size_0, - pg_exdist_size_1, ddqz_z_full_e, ddqz_z_full_e_size_0, ddqz_z_full_e_size_1, @@ -1677,15 +1680,11 @@ def solve_nh_init_wrapper( False, ) - pg_exdist = ( - pg_exdist, - ( - pg_exdist_size_0, - pg_exdist_size_1, - ), - on_gpu, - False, - ) + pg_edgeidx = (pg_edgeidx, (pg_edgeidx_size_0,), on_gpu, True) + + pg_vertidx = (pg_vertidx, (pg_vertidx_size_0,), on_gpu, True) + + pg_exdist = (pg_exdist, (pg_exdist_size_0,), on_gpu, True) ddqz_z_full_e = ( ddqz_z_full_e, @@ -1818,6 +1817,8 @@ def solve_nh_init_wrapper( ddxn_z_full=ddxn_z_full, zdiff_gradp=zdiff_gradp, vertoffset_gradp=vertoffset_gradp, + pg_edgeidx=pg_edgeidx, + pg_vertidx=pg_vertidx, pg_exdist=pg_exdist, ddqz_z_full_e=ddqz_z_full_e, ddxt_z_full=ddxt_z_full, @@ -2450,6 +2451,38 @@ def solve_nh_init_wrapper( ) logger.debug(msg) + pg_edgeidx_arr = ( + _conversion.as_array(ffi, pg_edgeidx, _definitions.INT32) + if pg_edgeidx is not None + else None + ) + msg = "shape of pg_edgeidx after computation = %s" % str( + pg_edgeidx_arr.shape if pg_edgeidx is not None else "None" + ) + logger.debug(msg) + msg = ( + "pg_edgeidx after computation: %s" % str(pg_edgeidx_arr) + if pg_edgeidx is not None + else "None" + ) + logger.debug(msg) + + pg_vertidx_arr = ( + _conversion.as_array(ffi, pg_vertidx, _definitions.INT32) + if pg_vertidx is not None + else None + ) + msg = "shape of pg_vertidx after computation = %s" % str( + pg_vertidx_arr.shape if pg_vertidx is not None else "None" + ) + logger.debug(msg) + msg = ( + "pg_vertidx after computation: %s" % str(pg_vertidx_arr) + if pg_vertidx is not None + else "None" + ) + logger.debug(msg) + pg_exdist_arr = ( _conversion.as_array(ffi, pg_exdist, _definitions.FLOAT64) if pg_exdist is not None diff --git a/tools/tests/tools/py2fgen/wrappers/test_diffusion_wrapper.py b/tools/tests/tools/py2fgen/wrappers/test_diffusion_wrapper.py index 6a4cad0c63..41a63f4fee 100644 --- a/tools/tests/tools/py2fgen/wrappers/test_diffusion_wrapper.py +++ b/tools/tests/tools/py2fgen/wrappers/test_diffusion_wrapper.py @@ -71,22 +71,23 @@ def test_diffusion_wrapper_granule_inputs( # --- Extract Metric State Parameters --- theta_ref_mc = test_utils.array_to_array_info(metrics_savepoint.theta_ref_mc().ndarray) wgtfac_c = test_utils.array_to_array_info(metrics_savepoint.wgtfac_c().ndarray) - zd_diffcoef = test_utils.array_to_array_info(metrics_savepoint.zd_diffcoef().ndarray) - # TODO(): special handling, determine if this is necessary for Fortran arrays too - zd_vertoffset = np.squeeze( - metrics_savepoint.serializer.read("zd_vertoffset", metrics_savepoint.savepoint) + # The wrapper expects [cellidx, c2e2c_ids] and then extracts `zd_cellidx[0,:]` because it only needs the cellidxs + # (this is because slicing causes issue in the bindings, but not for serialization) + zd_cellidx = test_utils.array_to_array_info( + np.squeeze(metrics_savepoint.serializer.read("zd_cellidx", metrics_savepoint.savepoint))[ + np.newaxis, : + ] ) - zd_vertoffset = metrics_savepoint._reduce_to_dim_size( - zd_vertoffset, (dims.CellDim, dims.C2E2CDim, dims.KDim) + zd_vertidx = test_utils.array_to_array_info( + np.squeeze(metrics_savepoint.serializer.read("zd_vertidx", metrics_savepoint.savepoint)) ) - zd_vertoffset = test_utils.array_to_array_info(zd_vertoffset) - - zd_intcoef = np.squeeze(metrics_savepoint.serializer.read("vcoef", metrics_savepoint.savepoint)) - zd_intcoef = metrics_savepoint._reduce_to_dim_size( - zd_intcoef, (dims.CellDim, dims.C2E2CDim, dims.KDim) + zd_intcoef = test_utils.array_to_array_info( + np.squeeze(metrics_savepoint.serializer.read("zd_intcoef", metrics_savepoint.savepoint)) + ) + zd_diffcoef = test_utils.array_to_array_info( + np.squeeze(metrics_savepoint.serializer.read("zd_diffcoef", metrics_savepoint.savepoint)) ) - zd_intcoef = test_utils.array_to_array_info(zd_intcoef) # --- Extract Interpolation State Parameters --- e_bln_c_s = test_utils.array_to_array_info(interpolation_savepoint.e_bln_c_s().ndarray) @@ -160,9 +161,10 @@ def test_diffusion_wrapper_granule_inputs( nudgecoeff_e=nudgecoeff_e, rbf_coeff_1=rbf_coeff_1, rbf_coeff_2=rbf_coeff_2, - zd_diffcoef=zd_diffcoef, - zd_vertoffset=zd_vertoffset, + zd_cellidx=zd_cellidx, + zd_vertidx=zd_vertidx, zd_intcoef=zd_intcoef, + zd_diffcoef=zd_diffcoef, ndyn_substeps=ndyn_substeps, diffusion_type=diffusion_type, hdiff_w=hdiff_w, @@ -300,22 +302,23 @@ def test_diffusion_wrapper_single_step( # Metric state parameters theta_ref_mc = test_utils.array_to_array_info(metrics_savepoint.theta_ref_mc().ndarray) wgtfac_c = test_utils.array_to_array_info(metrics_savepoint.wgtfac_c().ndarray) - zd_diffcoef = test_utils.array_to_array_info(metrics_savepoint.zd_diffcoef().ndarray) - # TODO(): special handling, determine if this is necessary for Fortran arrays too - zd_vertoffset = np.squeeze( - metrics_savepoint.serializer.read("zd_vertoffset", metrics_savepoint.savepoint) + # The wrapper expects [cellidx, c2e2c_ids] and then extracts `zd_cellidx[0,:]` because it only needs the cellidxs + # (this is because slicing causes issue in the bindings, but not for serialization) + zd_cellidx = test_utils.array_to_array_info( + np.squeeze(metrics_savepoint.serializer.read("zd_cellidx", metrics_savepoint.savepoint))[ + np.newaxis, : + ] ) - zd_vertoffset = metrics_savepoint._reduce_to_dim_size( - zd_vertoffset, (dims.CellDim, dims.C2E2CDim, dims.KDim) + zd_vertidx = test_utils.array_to_array_info( + np.squeeze(metrics_savepoint.serializer.read("zd_vertidx", metrics_savepoint.savepoint)) ) - zd_vertoffset = test_utils.array_to_array_info(zd_vertoffset) - - zd_intcoef = np.squeeze(metrics_savepoint.serializer.read("vcoef", metrics_savepoint.savepoint)) - zd_intcoef = metrics_savepoint._reduce_to_dim_size( - zd_intcoef, (dims.CellDim, dims.C2E2CDim, dims.KDim) + zd_intcoef = test_utils.array_to_array_info( + np.squeeze(metrics_savepoint.serializer.read("zd_intcoef", metrics_savepoint.savepoint)) + ) + zd_diffcoef = test_utils.array_to_array_info( + np.squeeze(metrics_savepoint.serializer.read("zd_diffcoef", metrics_savepoint.savepoint)) ) - zd_intcoef = test_utils.array_to_array_info(zd_intcoef) # Interpolation state parameters e_bln_c_s = test_utils.array_to_array_info(interpolation_savepoint.e_bln_c_s().ndarray) @@ -357,9 +360,10 @@ def test_diffusion_wrapper_single_step( nudgecoeff_e=nudgecoeff_e, rbf_coeff_1=rbf_coeff_1, rbf_coeff_2=rbf_coeff_2, - zd_diffcoef=zd_diffcoef, - zd_vertoffset=zd_vertoffset, + zd_cellidx=zd_cellidx, + zd_vertidx=zd_vertidx, zd_intcoef=zd_intcoef, + zd_diffcoef=zd_diffcoef, ndyn_substeps=ndyn_substeps, diffusion_type=diffusion_type, hdiff_w=hdiff_w, diff --git a/tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py b/tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py index a1a8170e3e..726a86e7bd 100644 --- a/tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py +++ b/tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py @@ -64,7 +64,6 @@ def solve_nh_init( divdamp_z4 = 80000.0 # vertical grid params - vertical_size = grid_savepoint.num(dims.KDim) nflat_gradp = gtx.int32( grid_savepoint.nflat_gradp() + 1 ) # undo the -1 to go back to Fortran value @@ -75,7 +74,7 @@ def solve_nh_init( exner_exfac = test_utils.array_to_array_info(metrics_savepoint.exner_exfac().ndarray) exner_ref_mc = test_utils.array_to_array_info(metrics_savepoint.exner_ref_mc().ndarray) wgtfac_c = test_utils.array_to_array_info(metrics_savepoint.wgtfac_c().ndarray) - wgtfacq_c = test_utils.array_to_array_info(metrics_savepoint.wgtfacq_c_dsl().ndarray) + wgtfacq_c = test_utils.array_to_array_info(metrics_savepoint.wgtfacq_c().ndarray) inv_ddqz_z_full = test_utils.array_to_array_info(metrics_savepoint.inv_ddqz_z_full().ndarray) rho_ref_mc = test_utils.array_to_array_info(metrics_savepoint.rho_ref_mc().ndarray) theta_ref_mc = test_utils.array_to_array_info(metrics_savepoint.theta_ref_mc().ndarray) @@ -101,13 +100,13 @@ def solve_nh_init( ) vertoffset_gradp = test_utils.array_to_array_info(vertoffset_gradp_field.ndarray) - pg_exdist = test_utils.array_to_array_info(metrics_savepoint.pg_exdist().ndarray) + pg_edgeidx = test_utils.array_to_array_info(metrics_savepoint.pg_edgeidx()) + pg_vertidx = test_utils.array_to_array_info(metrics_savepoint.pg_vertidx()) + pg_exdist = test_utils.array_to_array_info(metrics_savepoint.pg_exdist()) ddqz_z_full_e = test_utils.array_to_array_info(metrics_savepoint.ddqz_z_full_e().ndarray) ddxt_z_full = test_utils.array_to_array_info(metrics_savepoint.ddxt_z_full().ndarray) wgtfac_e = test_utils.array_to_array_info(metrics_savepoint.wgtfac_e().ndarray) - wgtfacq_e = test_utils.array_to_array_info( - metrics_savepoint.wgtfacq_e_dsl(vertical_size).ndarray - ) + wgtfacq_e = test_utils.array_to_array_info(metrics_savepoint.wgtfacq_e().ndarray) vwind_impl_wgt = test_utils.array_to_array_info(metrics_savepoint.vwind_impl_wgt().ndarray) hmask_dd3d = test_utils.array_to_array_info(metrics_savepoint.hmask_dd3d().ndarray) scalfac_dd3d = test_utils.array_to_array_info(metrics_savepoint.scalfac_dd3d().ndarray) @@ -189,6 +188,8 @@ def solve_nh_init( ddxn_z_full=ddxn_z_full, zdiff_gradp=zdiff_gradp, vertoffset_gradp=vertoffset_gradp, + pg_edgeidx=pg_edgeidx, + pg_vertidx=pg_vertidx, pg_exdist=pg_exdist, ddqz_z_full_e=ddqz_z_full_e, ddxt_z_full=ddxt_z_full, @@ -299,7 +300,6 @@ def test_dycore_wrapper_granule_inputs( divdamp_z4 = 80000.0 # vertical grid params - vertical_size = grid_savepoint.num(dims.KDim) nflat_gradp = gtx.int32( grid_savepoint.nflat_gradp() + 1 ) # undo the -1 to go back to Fortran value @@ -314,7 +314,7 @@ def test_dycore_wrapper_granule_inputs( exner_exfac = test_utils.array_to_array_info(metrics_savepoint.exner_exfac().ndarray) exner_ref_mc = test_utils.array_to_array_info(metrics_savepoint.exner_ref_mc().ndarray) wgtfac_c = test_utils.array_to_array_info(metrics_savepoint.wgtfac_c().ndarray) - wgtfacq_c = test_utils.array_to_array_info(metrics_savepoint.wgtfacq_c_dsl().ndarray) + wgtfacq_c = test_utils.array_to_array_info(metrics_savepoint.wgtfacq_c().ndarray) inv_ddqz_z_full = test_utils.array_to_array_info(metrics_savepoint.inv_ddqz_z_full().ndarray) rho_ref_mc = test_utils.array_to_array_info(metrics_savepoint.rho_ref_mc().ndarray) theta_ref_mc = test_utils.array_to_array_info(metrics_savepoint.theta_ref_mc().ndarray) @@ -340,13 +340,13 @@ def test_dycore_wrapper_granule_inputs( ) vertoffset_gradp = test_utils.array_to_array_info(vertoffset_gradp_field.ndarray) - pg_exdist = test_utils.array_to_array_info(metrics_savepoint.pg_exdist().ndarray) + pg_edgeidx = test_utils.array_to_array_info(metrics_savepoint.pg_edgeidx()) + pg_vertidx = test_utils.array_to_array_info(metrics_savepoint.pg_vertidx()) + pg_exdist = test_utils.array_to_array_info(metrics_savepoint.pg_exdist()) ddqz_z_full_e = test_utils.array_to_array_info(metrics_savepoint.ddqz_z_full_e().ndarray) ddxt_z_full = test_utils.array_to_array_info(metrics_savepoint.ddxt_z_full().ndarray) wgtfac_e = test_utils.array_to_array_info(metrics_savepoint.wgtfac_e().ndarray) - wgtfacq_e = test_utils.array_to_array_info( - metrics_savepoint.wgtfacq_e_dsl(vertical_size).ndarray - ) + wgtfacq_e = test_utils.array_to_array_info(metrics_savepoint.wgtfacq_e().ndarray) vwind_impl_wgt = test_utils.array_to_array_info(metrics_savepoint.vwind_impl_wgt().ndarray) hmask_dd3d = test_utils.array_to_array_info(metrics_savepoint.hmask_dd3d().ndarray) scalfac_dd3d = test_utils.array_to_array_info(metrics_savepoint.scalfac_dd3d().ndarray) @@ -489,11 +489,11 @@ def test_dycore_wrapper_granule_inputs( zdiff_gradp=metrics_savepoint.zdiff_gradp(), nflat_gradp=grid_savepoint.nflat_gradp(), vertoffset_gradp=metrics_savepoint.vertoffset_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(vertical_size), + 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(), @@ -609,6 +609,8 @@ def test_dycore_wrapper_granule_inputs( ddxn_z_full=ddxn_z_full, zdiff_gradp=zdiff_gradp, vertoffset_gradp=vertoffset_gradp, + pg_edgeidx=pg_edgeidx, + pg_vertidx=pg_vertidx, pg_exdist=pg_exdist, ddqz_z_full_e=ddqz_z_full_e, ddxt_z_full=ddxt_z_full, diff --git a/uv.lock b/uv.lock index 734eb4eea4..94ac4e7607 100644 --- a/uv.lock +++ b/uv.lock @@ -563,6 +563,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bf/9b/08c0432272d77b04803958a4598a51e2a4b51c06640af8b8f0f908c18bf2/charset_normalizer-3.4.0-py3-none-any.whl", hash = "sha256:fe9f97feb71aa9896b81973a7bbada8c49501dc73e58a10fcef6663af95e5079", size = 49446, upload-time = "2024-10-09T07:40:19.383Z" }, ] +[[package]] +name = "clang-format" +version = "22.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e5/84/8ab81082e59115e00a162365b50495bc608b49bbaffed479c44b8abd8dde/clang_format-22.1.0.tar.gz", hash = "sha256:41c535251ebc6f1ff824ca3d92f36ec7f784eb7c355b3cd6fc3963f3d72ebca8", size = 11508, upload-time = "2026-02-24T22:12:09.578Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d7/c2/db042a9f7a9497fba214935551d8fef098368147be39d9c23161bcb24eb0/clang_format-22.1.0-py2.py3-none-macosx_10_9_x86_64.whl", hash = "sha256:b11aaeb35e19948b290335f65b56e7e8d057137ac1094156584de4eec6b92c54", size = 1492625, upload-time = "2026-02-24T22:11:38.605Z" }, + { url = "https://files.pythonhosted.org/packages/35/9a/0a8634a452d6bb5609854c6cc4fc2ad417ba76c7d1dd5ca6bec66a464f8c/clang_format-22.1.0-py2.py3-none-macosx_11_0_arm64.whl", hash = "sha256:2dbbfc5f275121f067c76459de16120fb70b483927d7c93320eb93238b9b6ae4", size = 1478843, upload-time = "2026-02-24T22:11:40.525Z" }, + { url = "https://files.pythonhosted.org/packages/17/6e/722918070bf3d6ce3afdc4972a74bda4c4645f16f8b2d8246ff0818fe9d1/clang_format-22.1.0-py2.py3-none-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:97d12f9785f7b909ef737da3de95471d1915462908004d041869f41129de467d", size = 1757448, upload-time = "2026-02-24T22:11:42.314Z" }, + { url = "https://files.pythonhosted.org/packages/1f/e0/68257bff84d0fe8d855d8f255b1bc192cd21db702b551bc002100a46a211/clang_format-22.1.0-py2.py3-none-manylinux_2_26_i686.manylinux_2_28_i686.whl", hash = "sha256:4695288f19a2374e22c09eb615d2874b697dd7fc11c0efe4eae12cfe739083c1", size = 1888405, upload-time = "2026-02-24T22:11:44.242Z" }, + { url = "https://files.pythonhosted.org/packages/5e/29/393907895dbaffe6cbf527890e04d51e80c750b359e1d97cc8c2dfa815d6/clang_format-22.1.0-py2.py3-none-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:91ddd90a5745320e16ab06f6dd279bfe3b304dbe603980278a1035bcb2d9667d", size = 2070418, upload-time = "2026-02-24T22:11:46.229Z" }, + { url = "https://files.pythonhosted.org/packages/fe/d4/908b17c926eb3a90109dfc4f5195b438f9d2e4c1da3de0b955eb93c14e43/clang_format-22.1.0-py2.py3-none-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:50b85b2ef55474676465ea3460e63437a9dd0aa2b1e3f9e485daef6508dee086", size = 2101315, upload-time = "2026-02-24T22:11:48.287Z" }, + { url = "https://files.pythonhosted.org/packages/ba/f8/550085c7e0b64178bf409234afd2a789195b320e8a16d4ba060cdac4a6d0/clang_format-22.1.0-py2.py3-none-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ae07624f13f404b4405a2a20bfaffa33ada5602dc75dea6f9c9ab5ecfe6e07b3", size = 1840809, upload-time = "2026-02-24T22:11:50.176Z" }, + { url = "https://files.pythonhosted.org/packages/9d/be/a3e8ba3365c21135673566bb03b7d78f1fb95a4a2511ddca2c8794239c3b/clang_format-22.1.0-py2.py3-none-manylinux_2_31_armv7l.whl", hash = "sha256:5a7f76065a1e80e0397d64d2def12e3c7d5854a13c1108283e65ad7dc5673653", size = 1684794, upload-time = "2026-02-24T22:11:51.615Z" }, + { url = "https://files.pythonhosted.org/packages/55/4c/b3682ada8332959546adc92f8222b0a9fa560e0fa4d591db25fa7a06bb2e/clang_format-22.1.0-py2.py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:f136fee157901366c092ee9e629aa51b337e3c5c45fa7c15522fbf3aee3c8c53", size = 2735158, upload-time = "2026-02-24T22:11:53.406Z" }, + { url = "https://files.pythonhosted.org/packages/7b/a5/14f60ee2a0f88bc59b381472c271eba12877a65dbb4ad5bc425f9a0cb3df/clang_format-22.1.0-py2.py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:44aa67e948a6f03d62f5b4076041033f9fdbdb7b3a48df3906d7cceb73a4a61e", size = 2514611, upload-time = "2026-02-24T22:11:55.464Z" }, + { url = "https://files.pythonhosted.org/packages/c7/97/2bad239dd072a37f3fc71e1bc4950fafffc3f3fac6e8e7ff1bb83e510c0a/clang_format-22.1.0-py2.py3-none-musllinux_1_2_i686.whl", hash = "sha256:e6f4700aeb76b2057de7905bdef988e52a51270d66dc05daadb5c17680395437", size = 2985306, upload-time = "2026-02-24T22:11:57.409Z" }, + { url = "https://files.pythonhosted.org/packages/1a/d2/d6dcd1e0400fa1b8a7d784f6382291f055c57264daee1502549942f27d50/clang_format-22.1.0-py2.py3-none-musllinux_1_2_ppc64le.whl", hash = "sha256:8ab962542105ee4e3aabee3f8ef21241f2bb922b4353d5defb76df0363e6eb92", size = 3119154, upload-time = "2026-02-24T22:11:59.121Z" }, + { url = "https://files.pythonhosted.org/packages/be/c3/65de01a192d7237c36e5b73e0ff4a5f81e4bd7991669a38955fc50f93bbd/clang_format-22.1.0-py2.py3-none-musllinux_1_2_s390x.whl", hash = "sha256:05950eeb29d7181075b9000b86c2fc787cf706a23671cc9ee6fbc28e3ac6b3e2", size = 3217108, upload-time = "2026-02-24T22:12:00.863Z" }, + { url = "https://files.pythonhosted.org/packages/c1/8d/8d41438ef6fda40058cd9b85cbbd4fabd3bb1e67f5324b2a234bdf70e72e/clang_format-22.1.0-py2.py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:a5d7cc77bf7454b8c9f308e7205adf734c99a0dcb5114de76dedebc53c4a84eb", size = 2848307, upload-time = "2026-02-24T22:12:02.793Z" }, + { url = "https://files.pythonhosted.org/packages/19/e0/74ca2c38f2203a729e8f4ff2f3ffeec84216b294f65ce55f7da416d238cf/clang_format-22.1.0-py2.py3-none-win32.whl", hash = "sha256:beeae9d7fc3490cce6ffac0eb5643d508991678cf5b8217ca9a90a16b83a5a1d", size = 1299281, upload-time = "2026-02-24T22:12:04.446Z" }, + { url = "https://files.pythonhosted.org/packages/60/f2/4391a14b872b7c52b68c09bb34f8e22aa068ea93dbbdd1be55962ac58bd2/clang_format-22.1.0-py2.py3-none-win_amd64.whl", hash = "sha256:962c7b1c9362ac028e664ba8b1c4f0d71eab13ceade80f19f84ed8a66d126db4", size = 1460209, upload-time = "2026-02-24T22:12:06.431Z" }, + { url = "https://files.pythonhosted.org/packages/dd/c3/4573ef27c004b30fabb23a82c79eb09d76ac7d6b94422cc767483fa42c72/clang_format-22.1.0-py2.py3-none-win_arm64.whl", hash = "sha256:be5e43f2263ab6d6f9d76acaab39976b210bbdcf658623386560d5aa8f8deeda", size = 1344909, upload-time = "2026-02-24T22:12:08.321Z" }, +] + [[package]] name = "click" version = "8.1.7" @@ -1840,6 +1865,7 @@ name = "icon4py-common" version = "0.0.6" source = { editable = "model/common" } dependencies = [ + { name = "array-api-compat" }, { name = "dace" }, { name = "gt4py" }, { name = "numpy" }, @@ -1888,6 +1914,7 @@ io = [ [package.metadata] requires-dist = [ + { name = "array-api-compat", specifier = ">=1.13.0" }, { name = "cartopy", marker = "extra == 'io'", specifier = ">=0.22.0" }, { name = "cftime", marker = "extra == 'io'", specifier = ">=1.6.3" }, { name = "cupy-cuda11x", marker = "extra == 'cuda11'", specifier = ">=13.0" }, @@ -2006,6 +2033,7 @@ version = "0.0.6" source = { editable = "tools" } dependencies = [ { name = "cffi" }, + { name = "clang-format" }, { name = "click" }, { name = "configargparse" }, { name = "fprettify" }, @@ -2034,6 +2062,7 @@ profiling = [ [package.metadata] requires-dist = [ { name = "cffi", specifier = ">=1.5" }, + { name = "clang-format", specifier = ">=22.1.0" }, { name = "click", specifier = ">=8.1.7" }, { name = "configargparse", specifier = ">=1.7.1" }, { name = "cupy-cuda11x", marker = "extra == 'cuda11'", specifier = ">=13.0" },