diff --git a/ci/benchmark_bencher.yml b/ci/benchmark_bencher.yml index 120c144277..e63f3aca02 100644 --- a/ci/benchmark_bencher.yml +++ b/ci/benchmark_bencher.yml @@ -17,7 +17,7 @@ include: stage: benchmark script: - !reference [.bencher_setup_env, setup] - - nox -s __bencher_feature_branch_CI-3.10 -- --backend=$BACKEND --grid=$GRID + - nox -s __bencher_feature_branch_CI-3.10 -- --backend=$BACKEND --grid=$GRID -k "compile_time_domain" parallel: matrix: - BACKEND: [dace_cpu, dace_gpu, gtfn_cpu, gtfn_gpu] diff --git a/ci/default.yml b/ci/default.yml index cba7314809..db7d29b3ac 100644 --- a/ci/default.yml +++ b/ci/default.yml @@ -62,3 +62,47 @@ test_tools_datatests_aarch64: # extends: [.test_model_datatests, .test_template_x86_64] test_model_datatests_aarch64: extends: [.test_model_datatests, .test_template_aarch64] + + +.test_single_precision: + stage: test + allow_failure: true + script: + - nox -s "__test_file-3.10" -- --testfile=model/atmosphere/dycore/tests/dycore/integration_tests/test_velocity_advection.py --backend=$BACKEND -k "test_compute_advection_in_horizontal_momentum_equation or test_compute_advection_in_vertical_momentum_equation" + variables: + FLOAT_PRECISION: single + SLURM_TIMELIMIT: '00:20:00' + parallel: + matrix: + - BACKEND: [embedded, dace_gpu, gtfn_cpu, gtfn_gpu] + +test_single_precision_aarch64: + extends: [.test_single_precision, .test_template_aarch64] + +# TODO(pstark): remove one version of the single precision tests +.test_single_precision_v2: + stage: test + script: + - nox -s "test_model-3.10(datatest, $COMPONENT)" -- --single-precision --backend=$BACKEND --level=$LEVEL + rules: + - if: $BACKEND == 'dace_gpu' && $COMPONENT != 'dycore' + when: never # run only in daily CI, to save compute resources + - if: $COMPONENT == 'common' && $LEVEL == 'integration' + variables: + NUM_PROCESSES: 1 + SLURM_TIMELIMIT: '00:45:00' + - if: $BACKEND == 'dace_gpu' + variables: + NUM_PROCESSES: 8 + SLURM_TIMELIMIT: '01:00:00' + - if: $BACKEND == 'embedded' + variables: + SLURM_TIMELIMIT: '00:15:00' + - when: on_success + variables: + SLURM_TIMELIMIT: '00:30:00' + parallel: + matrix: + - COMPONENT: [dycore] + BACKEND: [embedded, dace_gpu, gtfn_cpu, gtfn_gpu] + LEVEL: [integration] diff --git a/ci/docker/base.Dockerfile b/ci/docker/base.Dockerfile index 099ab1d06c..1b476e8bf3 100644 --- a/ci/docker/base.Dockerfile +++ b/ci/docker/base.Dockerfile @@ -37,6 +37,10 @@ ENV PATH="/root/.cargo/bin:${PATH}" RUN rustc --version && which rustc && cargo --version && which cargo # Install Bencher for performance monitoring +# Update the following comment to trigger a rebuild to update the CLI: +# last update: 2025-09-05 +# This is necessary because the cloud version and the CLI version have to match +# but obviously, version changes do not register in the Dockerfile hash. RUN curl --proto '=https' --tlsv1.2 -sSfL https://bencher.dev/download/install-cli.sh | sh RUN bencher --version && which bencher diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py index fae19d65bf..5de341da42 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py @@ -127,7 +127,7 @@ def apply_flux_limiter( p_mflx_tracer_h=p_mflx_tracer_h, r_m=self._r_m, p_dtime=dtime, - dbl_eps=constants.DBL_EPS, + wp_eps=constants.WP_EPS, horizontal_start=self._start_cell_lateral_boundary_level_2, # originally i_rlstart_c = get_startrow_c(startrow_e=5) = 2 horizontal_end=self._end_cell_local, vertical_start=0, diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py index d62d759e5a..672544fec1 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py @@ -785,7 +785,7 @@ def _compute_numerical_flux( k=self._k_field, slevp1_ti=self._slevp1_ti, nlev=self._nlev, - dbl_eps=constants.DBL_EPS, + wp_eps=constants.WP_EPS, p_dtime=dtime, horizontal_start=horizontal_start, horizontal_end=horizontal_end, diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_monotone_horizontal_multiplicative_flux_factors.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_monotone_horizontal_multiplicative_flux_factors.py index 34e1ea5e75..ff1e02937b 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_monotone_horizontal_multiplicative_flux_factors.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_monotone_horizontal_multiplicative_flux_factors.py @@ -40,10 +40,10 @@ def _compute_monotone_horizontal_multiplicative_flux_factors_p_m( z_tracer_new_low: fa.CellKField[ta.wpfloat], z_max: fa.CellKField[ta.vpfloat], z_min: fa.CellKField[ta.vpfloat], - dbl_eps: ta.wpfloat, + wp_eps: ta.wpfloat, ) -> tuple[fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat]]: - r_p = (astype(z_max, wpfloat) - z_tracer_new_low) / (astype(z_mflx_anti_in, wpfloat) + dbl_eps) - r_m = (z_tracer_new_low - astype(z_min, wpfloat)) / (astype(z_mflx_anti_out, wpfloat) + dbl_eps) + r_p = (astype(z_max, wpfloat) - z_tracer_new_low) / (astype(z_mflx_anti_in, wpfloat) + wp_eps) + r_m = (z_tracer_new_low - astype(z_min, wpfloat)) / (astype(z_mflx_anti_out, wpfloat) + wp_eps) return r_p, r_m @@ -57,7 +57,7 @@ def _compute_monotone_horizontal_multiplicative_flux_factors( z_tracer_new_low: fa.CellKField[ta.wpfloat], beta_fct: ta.wpfloat, r_beta_fct: ta.wpfloat, - dbl_eps: ta.wpfloat, + wp_eps: ta.wpfloat, ) -> tuple[fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat]]: z_max, z_min = _compute_monotone_horizontal_multiplicative_flux_factors_min_max( z_tracer_max, z_tracer_min, beta_fct, r_beta_fct @@ -69,7 +69,7 @@ def _compute_monotone_horizontal_multiplicative_flux_factors( z_tracer_new_low, z_max, z_min, - dbl_eps, + wp_eps, ) return r_p, r_m @@ -85,7 +85,7 @@ def compute_monotone_horizontal_multiplicative_flux_factors( r_m: fa.CellKField[ta.wpfloat], beta_fct: ta.wpfloat, r_beta_fct: ta.wpfloat, - dbl_eps: ta.wpfloat, + wp_eps: ta.wpfloat, horizontal_start: gtx.int32, horizontal_end: gtx.int32, vertical_start: gtx.int32, @@ -99,7 +99,7 @@ def compute_monotone_horizontal_multiplicative_flux_factors( z_tracer_new_low, beta_fct, r_beta_fct, - dbl_eps, + wp_eps, out=(r_p, r_m), domain={ dims.CellDim: (horizontal_start, horizontal_end), diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_positive_definite_horizontal_multiplicative_flux_factor.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_positive_definite_horizontal_multiplicative_flux_factor.py index 670aa88fe7..e73d58a717 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_positive_definite_horizontal_multiplicative_flux_factor.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_positive_definite_horizontal_multiplicative_flux_factor.py @@ -20,10 +20,10 @@ def _compute_positive_definite_horizontal_multiplicative_flux_factor( p_rhodz_now: fa.CellKField[ta.wpfloat], p_mflx_tracer_h: fa.EdgeKField[ta.wpfloat], p_dtime: ta.wpfloat, - dbl_eps: ta.wpfloat, + wp_eps: ta.wpfloat, ) -> fa.CellKField[ta.wpfloat]: p_m = neighbor_sum(maximum(0.0, p_mflx_tracer_h(C2E) * geofac_div * p_dtime), axis=C2EDim) - r_m = minimum(1.0, (p_cc * p_rhodz_now) / (p_m + dbl_eps)) + r_m = minimum(1.0, (p_cc * p_rhodz_now) / (p_m + wp_eps)) return r_m @@ -35,7 +35,7 @@ def compute_positive_definite_horizontal_multiplicative_flux_factor( p_mflx_tracer_h: fa.EdgeKField[ta.wpfloat], r_m: fa.CellKField[ta.wpfloat], p_dtime: ta.wpfloat, - dbl_eps: ta.wpfloat, + wp_eps: ta.wpfloat, horizontal_start: gtx.int32, horizontal_end: gtx.int32, vertical_start: gtx.int32, @@ -47,7 +47,7 @@ def compute_positive_definite_horizontal_multiplicative_flux_factor( p_rhodz_now, p_mflx_tracer_h, p_dtime, - dbl_eps, + wp_eps, out=r_m, domain={ dims.CellDim: (horizontal_start, horizontal_end), diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_courant_number.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_courant_number.py index ce681f0d2c..033be9abbf 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_courant_number.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_courant_number.py @@ -9,8 +9,9 @@ import gt4py.next as gtx from gt4py.next import abs, where # noqa: A004 -from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import Koff +from icon4py.model.common.type_alias import wpfloat # TODO(dastrm): this stencil has no test @@ -19,13 +20,13 @@ @gtx.field_operator def _compute_courant_number_below( - p_cellmass_now: fa.CellKField[ta.wpfloat], - z_mass: fa.CellKField[ta.wpfloat], - z_cfl: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[wpfloat], + z_mass: fa.CellKField[wpfloat], + z_cfl: fa.CellKField[wpfloat], k: fa.KField[gtx.int32], nlev: gtx.int32, - dbl_eps: ta.wpfloat, -) -> fa.CellKField[ta.wpfloat]: + wp_eps: wpfloat, +) -> fa.CellKField[wpfloat]: z_mass_pos = z_mass > 0.0 in_bounds_p0 = k <= nlev - 1 @@ -63,20 +64,20 @@ def _compute_courant_number_below( p_cellmass_now_jks = where(mass_gt_cellmass_p3, p_cellmass_now(Koff[4]), p_cellmass_now_jks) z_cflfrac = where(z_mass_pos, z_mass / p_cellmass_now_jks, 0.0) - z_cfl = z_cfl + where(z_cflfrac < 1.0, z_cflfrac, 1.0 - dbl_eps) + z_cfl = z_cfl + where(z_cflfrac < 1.0, z_cflfrac, 1.0 - wp_eps) return z_cfl @gtx.field_operator def _compute_courant_number_above( - p_cellmass_now: fa.CellKField[ta.wpfloat], - z_mass: fa.CellKField[ta.wpfloat], - z_cfl: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[wpfloat], + z_mass: fa.CellKField[wpfloat], + z_cfl: fa.CellKField[wpfloat], k: fa.KField[gtx.int32], slevp1_ti: gtx.int32, - dbl_eps: ta.wpfloat, -) -> fa.CellKField[ta.wpfloat]: + wp_eps: wpfloat, +) -> fa.CellKField[wpfloat]: z_mass_neg = z_mass <= 0.0 in_bounds_m0 = k >= slevp1_ti + 1 @@ -116,26 +117,26 @@ def _compute_courant_number_above( z_cfl = z_cfl - where(mass_gt_cellmass_m3, 1.0, 0.0) z_cflfrac = where(z_mass_neg, z_mass / p_cellmass_now_jks, 0.0) - z_cfl = z_cfl + where(abs(z_cflfrac) < 1.0, z_cflfrac, dbl_eps - 1.0) + z_cfl = z_cfl + where(abs(z_cflfrac) < 1.0, z_cflfrac, wp_eps - 1.0) return z_cfl @gtx.field_operator def _compute_ppm4gpu_courant_number( - p_mflx_contra_v: fa.CellKField[ta.wpfloat], - p_cellmass_now: fa.CellKField[ta.wpfloat], - z_cfl: fa.CellKField[ta.wpfloat], + p_mflx_contra_v: fa.CellKField[wpfloat], + p_cellmass_now: fa.CellKField[wpfloat], + z_cfl: fa.CellKField[wpfloat], k: fa.KField[gtx.int32], slevp1_ti: gtx.int32, nlev: gtx.int32, - dbl_eps: ta.wpfloat, - p_dtime: ta.wpfloat, -) -> fa.CellKField[ta.wpfloat]: + wp_eps: wpfloat, + p_dtime: wpfloat, +) -> fa.CellKField[wpfloat]: z_mass = p_dtime * p_mflx_contra_v - cfl_below = _compute_courant_number_below(p_cellmass_now, z_mass, z_cfl, k, nlev, dbl_eps) - cfl_above = _compute_courant_number_above(p_cellmass_now, z_mass, z_cfl, k, slevp1_ti, dbl_eps) + cfl_below = _compute_courant_number_below(p_cellmass_now, z_mass, z_cfl, k, nlev, wp_eps) + cfl_above = _compute_courant_number_above(p_cellmass_now, z_mass, z_cfl, k, slevp1_ti, wp_eps) z_cfl = cfl_below + cfl_above @@ -144,14 +145,14 @@ def _compute_ppm4gpu_courant_number( @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_ppm4gpu_courant_number( - p_mflx_contra_v: fa.CellKField[ta.wpfloat], - p_cellmass_now: fa.CellKField[ta.wpfloat], - z_cfl: fa.CellKField[ta.wpfloat], + p_mflx_contra_v: fa.CellKField[wpfloat], + p_cellmass_now: fa.CellKField[wpfloat], + z_cfl: fa.CellKField[wpfloat], k: fa.KField[gtx.int32], slevp1_ti: gtx.int32, nlev: gtx.int32, - dbl_eps: ta.wpfloat, - p_dtime: ta.wpfloat, + wp_eps: wpfloat, + p_dtime: wpfloat, horizontal_start: gtx.int32, horizontal_end: gtx.int32, vertical_start: gtx.int32, @@ -164,7 +165,7 @@ def compute_ppm4gpu_courant_number( k, slevp1_ti, nlev, - dbl_eps, + wp_eps, p_dtime, out=z_cfl, domain={ diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/prepare_numerical_quadrature_for_cubic_reconstruction.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/prepare_numerical_quadrature_for_cubic_reconstruction.py index 566ded8fca..6c45a32d8b 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/prepare_numerical_quadrature_for_cubic_reconstruction.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/prepare_numerical_quadrature_for_cubic_reconstruction.py @@ -51,7 +51,7 @@ def _prepare_numerical_quadrature_for_cubic_reconstruction( wgt_zeta_2: ta.wpfloat, wgt_eta_1: ta.wpfloat, wgt_eta_2: ta.wpfloat, - dbl_eps: ta.wpfloat, + wp_eps: ta.wpfloat, eps: ta.wpfloat, ) -> tuple[ fa.EdgeKField[ta.vpfloat], @@ -97,7 +97,7 @@ def _prepare_numerical_quadrature_for_cubic_reconstruction( p_coords_dreg_v_3_y_wp = astype(p_coords_dreg_v_3_y, wpfloat) p_coords_dreg_v_4_y_wp = astype(p_coords_dreg_v_4_y, wpfloat) - wgt_t_detjac_1 = dbl_eps + z_wgt_1 * ( + wgt_t_detjac_1 = wp_eps + z_wgt_1 * ( ( z_eta_1_1 * (p_coords_dreg_v_2_x_wp - p_coords_dreg_v_1_x_wp) + z_eta_1_2 * (p_coords_dreg_v_3_x_wp - p_coords_dreg_v_4_x_wp) @@ -115,7 +115,7 @@ def _prepare_numerical_quadrature_for_cubic_reconstruction( - z_eta_1_4 * (p_coords_dreg_v_2_x_wp - p_coords_dreg_v_3_x_wp) ) ) - wgt_t_detjac_2 = dbl_eps + z_wgt_2 * ( + wgt_t_detjac_2 = wp_eps + z_wgt_2 * ( ( z_eta_2_1 * (p_coords_dreg_v_2_x_wp - p_coords_dreg_v_1_x_wp) + z_eta_2_2 * (p_coords_dreg_v_3_x_wp - p_coords_dreg_v_4_x_wp) @@ -133,7 +133,7 @@ def _prepare_numerical_quadrature_for_cubic_reconstruction( - z_eta_2_4 * (p_coords_dreg_v_2_x_wp - p_coords_dreg_v_3_x_wp) ) ) - wgt_t_detjac_3 = dbl_eps + z_wgt_3 * ( + wgt_t_detjac_3 = wp_eps + z_wgt_3 * ( ( z_eta_3_1 * (p_coords_dreg_v_2_x_wp - p_coords_dreg_v_1_x_wp) + z_eta_3_2 * (p_coords_dreg_v_3_x_wp - p_coords_dreg_v_4_x_wp) @@ -151,7 +151,7 @@ def _prepare_numerical_quadrature_for_cubic_reconstruction( - z_eta_3_4 * (p_coords_dreg_v_2_x_wp - p_coords_dreg_v_3_x_wp) ) ) - wgt_t_detjac_4 = dbl_eps + z_wgt_4 * ( + wgt_t_detjac_4 = wp_eps + z_wgt_4 * ( ( z_eta_4_1 * (p_coords_dreg_v_2_x_wp - p_coords_dreg_v_1_x_wp) + z_eta_4_2 * (p_coords_dreg_v_3_x_wp - p_coords_dreg_v_4_x_wp) @@ -342,7 +342,7 @@ def prepare_numerical_quadrature_for_cubic_reconstruction( wgt_zeta_2: ta.wpfloat, wgt_eta_1: ta.wpfloat, wgt_eta_2: ta.wpfloat, - dbl_eps: ta.wpfloat, + wp_eps: ta.wpfloat, eps: ta.wpfloat, horizontal_start: gtx.int32, horizontal_end: gtx.int32, @@ -386,7 +386,7 @@ def prepare_numerical_quadrature_for_cubic_reconstruction( wgt_zeta_2, wgt_eta_1, wgt_eta_2, - dbl_eps, + wp_eps, eps, out=( p_quad_vector_sum_1, diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/prepare_numerical_quadrature_list_for_cubic_reconstruction.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/prepare_numerical_quadrature_list_for_cubic_reconstruction.py index 747d5cc1f7..c0017ef2cf 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/prepare_numerical_quadrature_list_for_cubic_reconstruction.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/prepare_numerical_quadrature_list_for_cubic_reconstruction.py @@ -53,7 +53,7 @@ def _prepare_numerical_quadrature_list_for_cubic_reconstruction( wgt_zeta_2: ta.wpfloat, wgt_eta_1: ta.wpfloat, wgt_eta_2: ta.wpfloat, - dbl_eps: ta.wpfloat, + wp_eps: ta.wpfloat, eps: ta.wpfloat, ) -> tuple[ fa.EdgeKField[ta.vpfloat], @@ -110,7 +110,7 @@ def _prepare_numerical_quadrature_list_for_cubic_reconstruction( wgt_t_detjac_1 = where( famask_bool, - dbl_eps + wp_eps + z_wgt_1 * ( ( @@ -134,7 +134,7 @@ def _prepare_numerical_quadrature_list_for_cubic_reconstruction( ) wgt_t_detjac_2 = where( famask_bool, - dbl_eps + wp_eps + z_wgt_2 * ( ( @@ -158,7 +158,7 @@ def _prepare_numerical_quadrature_list_for_cubic_reconstruction( ) wgt_t_detjac_3 = where( famask_bool, - dbl_eps + wp_eps + z_wgt_3 * ( ( @@ -182,7 +182,7 @@ def _prepare_numerical_quadrature_list_for_cubic_reconstruction( ) wgt_t_detjac_4 = where( famask_bool, - dbl_eps + wp_eps + z_wgt_4 * ( ( @@ -378,7 +378,7 @@ def prepare_numerical_quadrature_list_for_cubic_reconstruction( wgt_zeta_2: ta.wpfloat, wgt_eta_1: ta.wpfloat, wgt_eta_2: ta.wpfloat, - dbl_eps: ta.wpfloat, + wp_eps: ta.wpfloat, eps: ta.wpfloat, horizontal_start: gtx.int32, horizontal_end: gtx.int32, @@ -424,7 +424,7 @@ def prepare_numerical_quadrature_list_for_cubic_reconstruction( wgt_zeta_2, wgt_eta_1, wgt_eta_2, - dbl_eps, + wp_eps, eps, out=( p_quad_vector_sum_1, diff --git a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py index b0654ee0d2..3e1fded686 100644 --- a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py +++ b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py @@ -151,7 +151,7 @@ def test_advection_run_single_step( prep_adv = construct_prep_adv(advection_init_savepoint) p_tracer_now = advection_init_savepoint.tracer(ntracer) p_tracer_new = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, allocator=backend) - dtime = advection_init_savepoint.get_metadata("dtime").get("dtime") + dtime = advection_init_savepoint.dtime() log_serialized(diagnostic_state, prep_adv, p_tracer_now, dtime) diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_positive_definite_horizontal_multiplicative_flux_factor.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_positive_definite_horizontal_multiplicative_flux_factor.py index c6ee1ad78e..89867e1a44 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_positive_definite_horizontal_multiplicative_flux_factor.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_compute_positive_definite_horizontal_multiplicative_flux_factor.py @@ -30,7 +30,7 @@ def reference( p_rhodz_now: np.ndarray, p_mflx_tracer_h: np.ndarray, p_dtime, - dbl_eps, + wp_eps, **kwargs, ) -> dict: c2e = connectivities[dims.C2EDim] @@ -49,7 +49,7 @@ def reference( ) p_m = p_m_0 + p_m_1 + p_m_2 - r_m = np.minimum(1.0, p_cc * p_rhodz_now / (p_m + dbl_eps)) + r_m = np.minimum(1.0, p_cc * p_rhodz_now / (p_m + wp_eps)) return dict(r_m=r_m) @@ -61,14 +61,14 @@ def input_data(self, grid) -> dict: p_mflx_tracer_h = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) r_m = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) p_dtime = np.float64(5) - dbl_eps = np.float64(1e-9) + wp_eps = np.float64(1e-9) return dict( geofac_div=geofac_div, p_cc=p_cc, p_rhodz_now=p_rhodz_now, p_mflx_tracer_h=p_mflx_tracer_h, p_dtime=p_dtime, - dbl_eps=dbl_eps, + wp_eps=wp_eps, r_m=r_m, horizontal_start=0, horizontal_end=gtx.int32(grid.num_cells), diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_numerical_quadrature_for_cubic_reconstruction.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_numerical_quadrature_for_cubic_reconstruction.py index efe12e6b23..30bd4110df 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_numerical_quadrature_for_cubic_reconstruction.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_numerical_quadrature_for_cubic_reconstruction.py @@ -42,7 +42,7 @@ def _compute_wgt_t_detjac( wgt_zeta_2, wgt_eta_1, wgt_eta_2, - dbl_eps, + wp_eps, p_coords_dreg_v_1_x, p_coords_dreg_v_2_x, p_coords_dreg_v_3_x, @@ -92,7 +92,7 @@ def _compute_wgt_t_detjac( 1.0 + zeta_4, ) - wgt_t_detjac_1 = dbl_eps + z_wgt_1 * ( + wgt_t_detjac_1 = wp_eps + z_wgt_1 * ( ( z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) @@ -110,7 +110,7 @@ def _compute_wgt_t_detjac( - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) ) ) - wgt_t_detjac_2 = dbl_eps + z_wgt_2 * ( + wgt_t_detjac_2 = wp_eps + z_wgt_2 * ( ( z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) @@ -128,7 +128,7 @@ def _compute_wgt_t_detjac( - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) ) ) - wgt_t_detjac_3 = dbl_eps + z_wgt_3 * ( + wgt_t_detjac_3 = wp_eps + z_wgt_3 * ( ( z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) @@ -146,7 +146,7 @@ def _compute_wgt_t_detjac( - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) ) ) - wgt_t_detjac_4 = dbl_eps + z_wgt_4 * ( + wgt_t_detjac_4 = wp_eps + z_wgt_4 * ( ( z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) @@ -375,7 +375,7 @@ def reference( wgt_zeta_2: float, wgt_eta_1: float, wgt_eta_2: float, - dbl_eps: float, + wp_eps: float, eps: float, **kwargs: Any, ) -> dict: @@ -384,7 +384,7 @@ def reference( wgt_zeta_2, wgt_eta_1, wgt_eta_2, - dbl_eps, + wp_eps, p_coords_dreg_v_1_x, p_coords_dreg_v_2_x, p_coords_dreg_v_3_x, @@ -534,7 +534,7 @@ def input_data(self, grid: base.Grid) -> dict: wgt_zeta_2 = 0.003 wgt_eta_1 = 0.002 wgt_eta_2 = 0.007 - dbl_eps = np.float64(0.1) + wp_eps = np.float64(0.1) eps = 0.1 return dict( p_coords_dreg_v_1_x=p_coords_dreg_v_1_x, @@ -584,7 +584,7 @@ def input_data(self, grid: base.Grid) -> dict: wgt_zeta_2=wgt_zeta_2, wgt_eta_1=wgt_eta_1, wgt_eta_2=wgt_eta_2, - dbl_eps=dbl_eps, + wp_eps=wp_eps, eps=eps, horizontal_start=0, horizontal_end=gtx.int32(grid.num_edges), diff --git a/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_numerical_quadrature_list_for_cubic_reconstruction.py b/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_numerical_quadrature_list_for_cubic_reconstruction.py index 66b376c37a..2a9aa8826b 100644 --- a/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_numerical_quadrature_list_for_cubic_reconstruction.py +++ b/model/atmosphere/advection/tests/advection/stencil_tests/test_prepare_numerical_quadrature_list_for_cubic_reconstruction.py @@ -59,7 +59,7 @@ def _compute_wgt_t_detjac( zeta_3, zeta_4, famask_int, - dbl_eps, + wp_eps, ): z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 @@ -104,7 +104,7 @@ def _compute_wgt_t_detjac( wgt_t_detjac_1 = np.where( famask_bool, - dbl_eps + wp_eps + z_wgt_1 * ( ( @@ -129,7 +129,7 @@ def _compute_wgt_t_detjac( wgt_t_detjac_2 = np.where( famask_bool, - dbl_eps + wp_eps + z_wgt_2 * ( ( @@ -154,7 +154,7 @@ def _compute_wgt_t_detjac( wgt_t_detjac_3 = np.where( famask_bool, - dbl_eps + wp_eps + z_wgt_3 * ( ( @@ -178,7 +178,7 @@ def _compute_wgt_t_detjac( ) wgt_t_detjac_4 = np.where( famask_bool, - dbl_eps + wp_eps + z_wgt_4 * ( ( @@ -414,7 +414,7 @@ def reference( wgt_zeta_2: float, wgt_eta_1: float, wgt_eta_2: float, - dbl_eps: float, + wp_eps: float, **kwargs: Any, ) -> dict: wgt_t_detjac_1, wgt_t_detjac_2, wgt_t_detjac_3, wgt_t_detjac_4 = cls._compute_wgt_t_detjac( @@ -439,7 +439,7 @@ def reference( zeta_3, zeta_4, famask_int, - dbl_eps, + wp_eps, ) ( @@ -570,7 +570,7 @@ def input_data(self, grid: base.Grid) -> dict: wgt_zeta_2 = 0.003 wgt_eta_1 = 0.002 wgt_eta_2 = 0.007 - dbl_eps = np.float64(0.1) + wp_eps = np.float64(0.1) eps = 0.1 return dict( famask_int=famask_int, @@ -622,7 +622,7 @@ def input_data(self, grid: base.Grid) -> dict: wgt_zeta_2=wgt_zeta_2, wgt_eta_1=wgt_eta_1, wgt_eta_2=wgt_eta_2, - dbl_eps=dbl_eps, + wp_eps=wp_eps, eps=eps, horizontal_start=0, horizontal_end=gtx.int32(grid.num_edges), diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py index 2f4fba2226..43ba44ffcd 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -10,8 +10,6 @@ import enum import functools import logging -import math -import sys from typing import Final import gt4py.next as gtx @@ -56,6 +54,7 @@ ) from icon4py.model.common.model_options import setup_program from icon4py.model.common.orchestration import decorator as dace_orchestration +from icon4py.model.common.type_alias import vpfloat, wpfloat from icon4py.model.common.utils import data_allocation as data_alloc @@ -65,6 +64,15 @@ Supports only diffusion_type (=hdiff_order) 5 from the diffusion namelist. """ + +# TODO(pstark): switch back to gtx.sqrt once gt4py can return gtx.float32 in gtx.sqrt +def sqrt(value): + if isinstance(value, float): + return value**0.5 + else: + return value ** (wpfloat(0.5)) + + log = logging.getLogger(__name__) @@ -121,18 +129,18 @@ def __init__( type_vn_diffu: int = 1, smag_3d: bool = False, type_t_diffu: int = 2, - hdiff_efdt_ratio: float = 36.0, - hdiff_w_efdt_ratio: float = 15.0, - smagorinski_scaling_factor: float = 0.015, + hdiff_efdt_ratio: wpfloat = 36.0, + hdiff_w_efdt_ratio: wpfloat = 15.0, + smagorinski_scaling_factor: wpfloat = 0.015, n_substeps: int = 5, zdiffu_t: bool = True, - thslp_zdiffu: float = 0.025, - thhgtd_zdiffu: float = 200.0, - velocity_boundary_diffusion_denom: float = 200.0, - temperature_boundary_diffusion_denom: float = 135.0, - _nudge_max_coeff: float | None = None, # default is set in __init__ - max_nudging_coefficient: float | None = None, # default is set in __init__ - nudging_decay_rate: float = 2.0, + thslp_zdiffu: wpfloat = 0.025, + thhgtd_zdiffu: wpfloat = 200.0, + velocity_boundary_diffusion_denom: wpfloat = 200.0, # denom_diffu_v + temperature_boundary_diffusion_denom: wpfloat = 135.0, # denom_diffu_t + _nudge_max_coeff: wpfloat | None = None, # default is set in __init__ + max_nudging_coefficient: wpfloat | None = None, # default is set in __init__ + nudging_decay_rate: wpfloat = 2.0, shear_type: TurbulenceShearForcingType = TurbulenceShearForcingType.VERTICAL_OF_HORIZONTAL_WIND, ltkeshs: bool = True, ): @@ -167,24 +175,24 @@ def __init__( #: Ratio of e-folding time to (2*)time step #: Called 'hdiff_efdt_ratio' in mo_diffusion_nml.f90 - self.hdiff_efdt_ratio: float = hdiff_efdt_ratio + self.hdiff_efdt_ratio: wpfloat = wpfloat(hdiff_efdt_ratio) #: Ratio of e-folding time to time step for w diffusion (NH only) #: Called 'hdiff_w_efdt_ratio' in mo_diffusion_nml.f90. - self.hdiff_w_efdt_ratio: float = hdiff_w_efdt_ratio + self.hdiff_w_efdt_ratio: wpfloat = wpfloat(hdiff_w_efdt_ratio) #: Scaling factor for Smagorinsky diffusion at height hdiff_smag_z and below #: Called 'hdiff_smag_fac' in mo_diffusion_nml.f90 - self.smagorinski_scaling_factor: float = smagorinski_scaling_factor + self.smagorinski_scaling_factor: wpfloat = wpfloat(smagorinski_scaling_factor) #: If True, apply truly horizontal temperature diffusion over steep slopes #: Called 'l_zdiffu_t' in mo_nonhydrostatic_nml.f90 self.apply_zdiffusion_t: bool = zdiffu_t #:slope threshold (temperature diffusion): is used to build up an index list for application of truly horizontal diffusion in mo_vertical_grid.f90 - self.thslp_zdiffu = thslp_zdiffu + self.thslp_zdiffu: wpfloat = wpfloat(thslp_zdiffu) #: threshold [m] for height difference between adjacent grid points, defaults to 200m (temperature diffusion) - self.thhgtd_zdiffu = thhgtd_zdiffu + self.thhgtd_zdiffu: wpfloat = wpfloat(thhgtd_zdiffu) # from other namelists: # from parent namelist mo_nonhydrostatic_nml @@ -197,13 +205,15 @@ def __init__( #: Denominator for temperature boundary diffusion #: Called 'denom_diffu_t' in mo_gridref_nml.f90 - self.temperature_boundary_diffusion_denominator: float = ( + self.temperature_boundary_diffusion_denominator: wpfloat = wpfloat( temperature_boundary_diffusion_denom ) #: Denominator for velocity boundary diffusion #: Called 'denom_diffu_v' in mo_gridref_nml.f90 - self.velocity_boundary_diffusion_denominator: float = velocity_boundary_diffusion_denom + self.velocity_boundary_diffusion_denominator: wpfloat = wpfloat( + velocity_boundary_diffusion_denom + ) # parameters from namelist: mo_interpol_nml.f90 @@ -221,19 +231,19 @@ def __init__( "Cannot set both '_max_nudging_coefficient' and 'scaled_max_nudging_coefficient'." ) elif max_nudging_coefficient is not None: - self.max_nudging_coefficient: float = max_nudging_coefficient + self.max_nudging_coefficient: wpfloat = wpfloat(max_nudging_coefficient) elif _nudge_max_coeff is not None: - self.max_nudging_coefficient: float = ( - constants.DEFAULT_DYNAMICS_TO_PHYSICS_TIMESTEP_RATIO * _nudge_max_coeff + self.max_nudging_coefficient: wpfloat = ( + constants.DEFAULT_DYNAMICS_TO_PHYSICS_TIMESTEP_RATIO * wpfloat(_nudge_max_coeff) ) else: # default value in ICON - self.max_nudging_coefficient: float = ( - constants.DEFAULT_DYNAMICS_TO_PHYSICS_TIMESTEP_RATIO * 0.02 + self.max_nudging_coefficient: wpfloat = ( + constants.DEFAULT_DYNAMICS_TO_PHYSICS_TIMESTEP_RATIO * wpfloat(0.02) ) #: Exponential decay rate (in units of cell rows) of the lateral boundary nudging coefficients #: Called 'nudge_efold_width' in mo_interpol_nml.f90 - self.nudge_efold_width: float = nudging_decay_rate + self.nudge_efold_width: wpfloat = wpfloat(nudging_decay_rate) #: Type of shear forcing used in turbulence #: Called 'itype_shear' in mo_turbdiff_nml.f90 @@ -270,7 +280,7 @@ def _validate(self): @functools.cached_property def substep_as_float(self): - return float(self.ndyn_substeps) + return wpfloat(self.ndyn_substeps) @dataclasses.dataclass(frozen=True) @@ -278,25 +288,33 @@ class DiffusionParams: """Calculates derived quantities depending on the diffusion config.""" config: dataclasses.InitVar[DiffusionConfig] - K2: Final[float] = dataclasses.field(init=False) - K4: Final[float] = dataclasses.field(init=False) - K6: Final[float] = dataclasses.field(init=False) - K4W: Final[float] = dataclasses.field(init=False) - smagorinski_factor: Final[float] = dataclasses.field(init=False) - smagorinski_height: Final[float] = dataclasses.field(init=False) + K2: Final[wpfloat] = dataclasses.field(init=False) + K4: Final[wpfloat] = dataclasses.field(init=False) + K6: Final[wpfloat] = dataclasses.field(init=False) + K4W: Final[wpfloat] = dataclasses.field(init=False) + smagorinski_factor: Final[wpfloat] = dataclasses.field(init=False) + smagorinski_height: Final[wpfloat] = dataclasses.field(init=False) def __post_init__(self, config): object.__setattr__( self, "K2", - (1.0 / (config.hdiff_efdt_ratio * 8.0) if config.hdiff_efdt_ratio > 0.0 else 0.0), - ) - object.__setattr__(self, "K4", self.K2 / 8.0) - object.__setattr__(self, "K6", self.K2 / 64.0) + ( + wpfloat(1.0) / (config.hdiff_efdt_ratio * wpfloat(8.0)) + if config.hdiff_efdt_ratio > wpfloat(0.0) + else wpfloat(0.0) + ), + ) + object.__setattr__(self, "K4", self.K2 / wpfloat(8.0)) + object.__setattr__(self, "K6", self.K2 / wpfloat(64.0)) object.__setattr__( self, "K4W", - (1.0 / (config.hdiff_w_efdt_ratio * 36.0) if config.hdiff_w_efdt_ratio > 0 else 0.0), + ( + wpfloat(1.0) / (config.hdiff_w_efdt_ratio * wpfloat(36.0)) + if config.hdiff_w_efdt_ratio > wpfloat(0.0) + else wpfloat(0.0) + ), ) ( @@ -326,7 +344,7 @@ def _determine_smagorinski_factor(self, config: DiffusionConfig): smagorinski_factor = ( config.smagorinski_scaling_factor if config.smagorinski_scaling_factor - else 0.15, + else wpfloat(0.15), ) smagorinski_height = None case _: @@ -343,12 +361,12 @@ def diffusion_type_5_smagorinski_factor(config: DiffusionConfig): The calculation and magic numbers are taken from mo_diffusion_nml.f90 """ - magic_sqrt = math.sqrt(1600.0 * (1600 + 50000.0)) + magic_sqrt = sqrt(1600.0 * (1600.0 + 50000.0)) magic_fac2_value = 2e-6 * (1600.0 + 25000.0 + magic_sqrt) magic_z2 = 1600.0 + 50000.0 + magic_sqrt factor = (config.smagorinski_scaling_factor, magic_fac2_value, 0.0, 1.0) heights = (32500.0, magic_z2, 50000.0, 90000.0) - return factor, heights + return gtx.astype((factor, heights), wpfloat) class Diffusion: @@ -386,23 +404,30 @@ def __init__( self.halo_exchange_wait = decomposition.create_halo_exchange_wait( self._exchange ) # wait on a communication handle - self.rd_o_cvd: float = constants.GAS_CONSTANT_DRY_AIR / ( - constants.CPD - constants.GAS_CONSTANT_DRY_AIR + self.rd_o_cvd: vpfloat = gtx.astype( + constants.GAS_CONSTANT_DRY_AIR / (constants.CPD - constants.GAS_CONSTANT_DRY_AIR), + vpfloat, ) #: threshold temperature deviation from neighboring grid points hat activates extra diffusion against runaway cooling - self.thresh_tdiff: float = -5.0 + self.thresh_tdiff: wpfloat = wpfloat(-5.0) self._horizontal_start_index_w_diffusion: gtx.int32 = gtx.int32(0) - self.nudgezone_diff: float = 0.04 / ( - config.max_nudging_coefficient + sys.float_info.epsilon + self.nudgezone_diff: vpfloat = gtx.astype( + wpfloat(0.04) / (config.max_nudging_coefficient + constants.WP_EPS), vpfloat ) - self.bdy_diff: float = 0.015 / (config.max_nudging_coefficient + sys.float_info.epsilon) - self.fac_bdydiff_v: float = ( - math.sqrt(config.substep_as_float) / config.velocity_boundary_diffusion_denominator + self.bdy_diff: wpfloat = wpfloat(0.015) / ( + config.max_nudging_coefficient + constants.WP_EPS + ) + self.fac_bdydiff_v: wpfloat = ( + sqrt(config.substep_as_float) / config.velocity_boundary_diffusion_denominator ) - self.smag_offset: float = 0.25 * params.K4 * config.substep_as_float - self.diff_multfac_w: float = min(1.0 / 48.0, params.K4W * config.substep_as_float) + self.smag_offset: vpfloat = gtx.astype( + wpfloat(0.25) * params.K4 * config.substep_as_float, vpfloat + ) + self.diff_multfac_w: wpfloat = gtx.astype( + min(wpfloat(1.0) / wpfloat(48.0), params.K4W * config.substep_as_float), wpfloat + ) self._determine_horizontal_domains() self.mo_intp_rbf_rbf_vec_interpol_vertex = setup_program( @@ -511,7 +536,7 @@ def __init__( constant_args={ "theta_ref_mc": self._metric_state.theta_ref_mc, "thresh_tdiff": self.thresh_tdiff, - "smallest_vpfloat": constants.DBL_EPS, + "smallest_vpfloat": constants.VP_EPS, }, horizontal_sizes={ "horizontal_start": self._edge_start_nudging, @@ -574,6 +599,7 @@ def __init__( self.enh_smag_fac, offset_provider={"Koff": dims.KDim}, ) + setup_program( backend=backend, program=diffusion_utils.init_nabla2_factor_in_upper_damping_zone, @@ -585,10 +611,10 @@ def __init__( "vertical_start": 1, "vertical_end": gtx.int32(self._vertical_grid.end_index_of_damping_layer + 1), "end_index_of_damping_layer": self._vertical_grid.end_index_of_damping_layer, - "heights_1": self._vertical_grid.interface_physical_height.ndarray[1].item(), - "heights_nrd_shift": self._vertical_grid.interface_physical_height.ndarray[ + "heights_1": self._vertical_grid.interface_physical_height[1].as_scalar(), + "heights_nrd_shift": self._vertical_grid.interface_physical_height[ self._vertical_grid.end_index_of_damping_layer + 1 - ].item(), + ].as_scalar(), }, )(diff_multfac_n2w=self.diff_multfac_n2w) @@ -673,7 +699,7 @@ def initial_run( self, diagnostic_state: diffusion_states.DiffusionDiagnosticState, prognostic_state: prognostics.PrognosticState, - dtime: float, + dtime: wpfloat, ): """ Calculate initial diffusion step. @@ -696,7 +722,7 @@ def initial_run( smag_limit, ) self._do_diffusion_step( - diagnostic_state, prognostic_state, dtime, diff_multfac_vn, smag_limit, 0.0 + diagnostic_state, prognostic_state, dtime, diff_multfac_vn, smag_limit, wpfloat(0.0) ) self._sync_cell_fields(prognostic_state) @@ -704,7 +730,7 @@ def run( self, diagnostic_state: diffusion_states.DiffusionDiagnosticState, prognostic_state: prognostics.PrognosticState, - dtime: float, + dtime: wpfloat, ): """ Do one diffusion step within regular time loop. @@ -742,10 +768,10 @@ def _do_diffusion_step( self, diagnostic_state: diffusion_states.DiffusionDiagnosticState, prognostic_state: prognostics.PrognosticState, - dtime: float, - diff_multfac_vn: fa.KField[float], - smag_limit: fa.KField[float], - smag_offset: float, + dtime: wpfloat, + diff_multfac_vn: fa.KField[wpfloat], + smag_limit: fa.KField[wpfloat], + smag_offset: wpfloat, ): """ Run a diffusion step. diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py index 166cec8e76..977013c5b9 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py @@ -13,54 +13,55 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import KDim from icon4py.model.common.math.smagorinsky import _en_smag_fac_for_zero_nshift +from icon4py.model.common.type_alias import wpfloat @gtx.field_operator -def _identity_c_k(field: fa.CellKField[float]) -> fa.CellKField[float]: +def _identity_c_k(field: fa.CellKField[wpfloat]) -> fa.CellKField[wpfloat]: return field @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def copy_field(old_f: fa.CellKField[float], new_f: fa.CellKField[float]): +def copy_field(old_f: fa.CellKField[wpfloat], new_f: fa.CellKField[wpfloat]): _identity_c_k(old_f, out=new_f) @gtx.field_operator -def _identity_e_k(field: fa.EdgeKField[float]) -> fa.EdgeKField[float]: +def _identity_e_k(field: fa.EdgeKField[wpfloat]) -> fa.EdgeKField[wpfloat]: return field @gtx.field_operator -def _scale_k(field: fa.KField[float], factor: float) -> fa.KField[float]: +def _scale_k(field: fa.KField[wpfloat], factor: wpfloat) -> fa.KField[wpfloat]: return field * factor @gtx.program -def scale_k(field: fa.KField[float], factor: float, scaled_field: fa.KField[float]): +def scale_k(field: fa.KField[wpfloat], factor: wpfloat, scaled_field: fa.KField[wpfloat]): _scale_k(field, factor, out=scaled_field) @gtx.field_operator -def _setup_smag_limit(diff_multfac_vn: fa.KField[float]) -> fa.KField[float]: - return 0.125 - 4.0 * diff_multfac_vn +def _setup_smag_limit(diff_multfac_vn: fa.KField[wpfloat]) -> fa.KField[wpfloat]: + return wpfloat("0.125") - wpfloat("4.0") * diff_multfac_vn @gtx.field_operator -def _setup_runtime_diff_multfac_vn(k4: float, dyn_substeps: float) -> fa.KField[float]: - con = 1.0 / 128.0 - dyn = k4 * dyn_substeps / 3.0 +def _setup_runtime_diff_multfac_vn(k4: wpfloat, dyn_substeps: wpfloat) -> fa.KField[wpfloat]: + con = wpfloat("1.0") / wpfloat("128.0") + dyn = k4 * dyn_substeps / wpfloat("3.0") return broadcast(minimum(con, dyn), (KDim,)) @gtx.field_operator -def _setup_initial_diff_multfac_vn(k4: float, hdiff_efdt_ratio: float) -> fa.KField[float]: - return broadcast(k4 / 3.0 * hdiff_efdt_ratio, (KDim,)) +def _setup_initial_diff_multfac_vn(k4: wpfloat, hdiff_efdt_ratio: wpfloat) -> fa.KField[wpfloat]: + return broadcast(k4 / wpfloat(wpfloat("3.0")) * hdiff_efdt_ratio, (KDim,)) @gtx.field_operator def _setup_fields_for_initial_step( - k4: float, hdiff_efdt_ratio: float -) -> tuple[fa.KField[float], fa.KField[float]]: + k4: wpfloat, hdiff_efdt_ratio: wpfloat +) -> tuple[fa.KField[wpfloat], fa.KField[wpfloat]]: diff_multfac_vn = _setup_initial_diff_multfac_vn(k4, hdiff_efdt_ratio) smag_limit = _setup_smag_limit(diff_multfac_vn) return diff_multfac_vn, smag_limit @@ -68,28 +69,28 @@ def _setup_fields_for_initial_step( @gtx.program def setup_fields_for_initial_step( - k4: float, - hdiff_efdt_ratio: float, - diff_multfac_vn: fa.KField[float], - smag_limit: fa.KField[float], + k4: wpfloat, + hdiff_efdt_ratio: wpfloat, + diff_multfac_vn: fa.KField[wpfloat], + smag_limit: fa.KField[wpfloat], ): _setup_fields_for_initial_step(k4, hdiff_efdt_ratio, out=(diff_multfac_vn, smag_limit)) @gtx.field_operator def _init_diffusion_local_fields_for_regular_timestemp( - k4: float, - dyn_substeps: float, - hdiff_smag_fac: float, - hdiff_smag_fac2: float, - hdiff_smag_fac3: float, - hdiff_smag_fac4: float, - hdiff_smag_z: float, - hdiff_smag_z2: float, - hdiff_smag_z3: float, - hdiff_smag_z4: float, - vect_a: fa.KField[float], -) -> tuple[fa.KField[float], fa.KField[float], fa.KField[float]]: + k4: wpfloat, + dyn_substeps: wpfloat, + hdiff_smag_fac: wpfloat, + hdiff_smag_fac2: wpfloat, + hdiff_smag_fac3: wpfloat, + hdiff_smag_fac4: wpfloat, + hdiff_smag_z: wpfloat, + hdiff_smag_z2: wpfloat, + hdiff_smag_z3: wpfloat, + hdiff_smag_z4: wpfloat, + vect_a: fa.KField[wpfloat], +) -> tuple[fa.KField[wpfloat], fa.KField[wpfloat], fa.KField[wpfloat]]: diff_multfac_vn = _setup_runtime_diff_multfac_vn(k4, dyn_substeps) smag_limit = _setup_smag_limit(diff_multfac_vn) enh_smag_fac = _en_smag_fac_for_zero_nshift( @@ -112,20 +113,20 @@ def _init_diffusion_local_fields_for_regular_timestemp( @gtx.program def init_diffusion_local_fields_for_regular_timestep( - k4: float, - dyn_substeps: float, - hdiff_smag_fac: float, - hdiff_smag_fac2: float, - hdiff_smag_fac3: float, - hdiff_smag_fac4: float, - hdiff_smag_z: float, - hdiff_smag_z2: float, - hdiff_smag_z3: float, - hdiff_smag_z4: float, - vect_a: fa.KField[float], - diff_multfac_vn: fa.KField[float], - smag_limit: fa.KField[float], - enh_smag_fac: fa.KField[float], + k4: wpfloat, + dyn_substeps: wpfloat, + hdiff_smag_fac: wpfloat, + hdiff_smag_fac2: wpfloat, + hdiff_smag_fac3: wpfloat, + hdiff_smag_fac4: wpfloat, + hdiff_smag_z: wpfloat, + hdiff_smag_z2: wpfloat, + hdiff_smag_z3: wpfloat, + hdiff_smag_z4: wpfloat, + vect_a: fa.KField[wpfloat], + diff_multfac_vn: fa.KField[wpfloat], + smag_limit: fa.KField[wpfloat], + enh_smag_fac: fa.KField[wpfloat], ): _init_diffusion_local_fields_for_regular_timestemp( k4, @@ -149,31 +150,33 @@ def init_diffusion_local_fields_for_regular_timestep( @gtx.field_operator def _init_nabla2_factor_in_upper_damping_zone( - physical_heights: fa.KField[float], + physical_heights: fa.KField[wpfloat], end_index_of_damping_layer: gtx.int32, nshift: gtx.int32, - heights_nrd_shift: float, - heights_1: float, -) -> fa.KField[float]: + heights_nrd_shift: wpfloat, + heights_1: wpfloat, +) -> fa.KField[wpfloat]: height_sliced = concat_where( ((1 + nshift) <= dims.KDim) & (dims.KDim < (nshift + end_index_of_damping_layer + 1)), physical_heights, - 0.0, + wpfloat("0.0"), ) diff_multfac_n2w = ( - 1.0 / 12.0 * ((height_sliced - heights_nrd_shift) / (heights_1 - heights_nrd_shift)) ** 4 + wpfloat("1.0") + / wpfloat("12.0") + * ((height_sliced - heights_nrd_shift) / (heights_1 - heights_nrd_shift)) ** wpfloat("4") ) return diff_multfac_n2w @gtx.program def init_nabla2_factor_in_upper_damping_zone( - physical_heights: fa.KField[float], - diff_multfac_n2w: fa.KField[float], + physical_heights: fa.KField[wpfloat], + diff_multfac_n2w: fa.KField[wpfloat], end_index_of_damping_layer: gtx.int32, nshift: gtx.int32, - heights_nrd_shift: float, - heights_1: float, + heights_nrd_shift: wpfloat, + heights_1: wpfloat, vertical_start: gtx.int32, vertical_end: gtx.int32, ): diff --git a/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_diffusion.py b/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_diffusion.py index 630c560ff5..f635a9f108 100644 --- a/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_diffusion.py +++ b/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_diffusion.py @@ -258,7 +258,7 @@ def test_diffusion_init( def _verify_init_values_against_savepoint( savepoint: sb.IconDiffusionInitSavepoint, diffusion_granule: diffusion.Diffusion, backend ): - dtime = savepoint.get_metadata("dtime")["dtime"] + dtime = savepoint.dtime() assert savepoint.nudgezone_diff() == diffusion_granule.nudgezone_diff assert savepoint.bdy_diff() == diffusion_granule.bdy_diff @@ -394,7 +394,7 @@ def test_run_diffusion_single_step( cell_geometry = get_cell_geometry_for_experiment(experiment, backend) edge_geometry = get_edge_geometry_for_experiment(experiment, backend) - dtime = savepoint_diffusion_init.get_metadata("dtime").get("dtime") + dtime = savepoint_diffusion_init.dtime() diagnostic_state = diffusion_states.DiffusionDiagnosticState( hdef_ic=savepoint_diffusion_init.hdef_ic(), @@ -478,7 +478,7 @@ def test_run_diffusion_multiple_steps( ###################################################################### # Diffusion initialization ###################################################################### - dtime = savepoint_diffusion_init.get_metadata("dtime").get("dtime") + dtime = savepoint_diffusion_init.dtime() edge_geometry: grid_states.EdgeParams = grid_savepoint.construct_edge_geometry() cell_geometry: grid_states.CellParams = grid_savepoint.construct_cell_geometry() @@ -597,7 +597,7 @@ def test_run_diffusion_initial_step( grid = get_grid_for_experiment(experiment, backend) cell_geometry = get_cell_geometry_for_experiment(experiment, backend) edge_geometry = get_edge_geometry_for_experiment(experiment, backend) - dtime = savepoint_diffusion_init.get_metadata("dtime").get("dtime") + dtime = savepoint_diffusion_init.dtime() vertical_config = v_grid.VerticalGridConfig( grid.num_levels, diff --git a/model/atmosphere/diffusion/tests/diffusion/mpi_tests/test_parallel_diffusion.py b/model/atmosphere/diffusion/tests/diffusion/mpi_tests/test_parallel_diffusion.py index c2971203c4..3104573df8 100644 --- a/model/atmosphere/diffusion/tests/diffusion/mpi_tests/test_parallel_diffusion.py +++ b/model/atmosphere/diffusion/tests/diffusion/mpi_tests/test_parallel_diffusion.py @@ -78,7 +78,7 @@ def test_parallel_diffusion( f"rank={processor_props.rank}/{processor_props.comm_size}: using local grid with {icon_grid.num_cells} Cells, {icon_grid.num_edges} Edges, {icon_grid.num_vertices} Vertices" ) config = definitions.construct_diffusion_config(experiment, ndyn_substeps=ndyn_substeps) - dtime = savepoint_diffusion_init.get_metadata("dtime").get("dtime") + dtime = savepoint_diffusion_init.dtime() print( f"rank={processor_props.rank}/{processor_props.comm_size}: setup: using {processor_props.comm_name} with {processor_props.comm_size} nodes" ) @@ -215,7 +215,7 @@ def test_parallel_diffusion_multiple_steps( ) config = definitions.construct_diffusion_config(experiment, ndyn_substeps=ndyn_substeps) diffusion_params = diffusion_.DiffusionParams(config) - dtime = savepoint_diffusion_init.get_metadata("dtime").get("dtime") + dtime = savepoint_diffusion_init.dtime() print( f"rank={processor_props.rank}/{processor_props.comm_size}: setup: using {processor_props.comm_name} with {processor_props.comm_size} nodes" ) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py index 946f1132a5..30a6926887 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py @@ -14,8 +14,9 @@ from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_vn import apply_diffusion_to_vn from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid +from icon4py.model.common.type_alias import vpfloat, wpfloat from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest from .test_apply_nabla2_and_nabla4_global_to_vn import apply_nabla2_and_nabla4_global_to_vn_numpy from .test_apply_nabla2_and_nabla4_to_vn import apply_nabla2_and_nabla4_to_vn_numpy @@ -25,10 +26,29 @@ from .test_calculate_nabla4 import calculate_nabla4_numpy +@pytest.mark.single_precision_ready @pytest.mark.uses_concat_where class TestApplyDiffusionToVn(StencilTest): PROGRAM = apply_diffusion_to_vn OUTPUTS = ("vn",) + STATIC_PARAMS = { + StandardStaticVariants.NONE: (), + StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "start_2nd_nudge_line_idx_e", + "vertical_start", + "vertical_end", + "limited_area", + # "nudgezone_diff", #TODO(pstark): Why not these two? + # "fac_bdydiff_v", + ), + StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + "limited_area", + ), + } @staticmethod def reference( @@ -100,25 +120,29 @@ def reference( @pytest.fixture def input_data(self, grid: base.Grid) -> dict: - u_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim) - v_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim) + u_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim, dtype=vpfloat) + v_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim, dtype=vpfloat) - primal_normal_vert_v1 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2C2VDim) - primal_normal_vert_v2 = data_alloc.random_field(grid, dims.EdgeDim, dims.E2C2VDim) + primal_normal_vert_v1 = data_alloc.random_field( + grid, dims.EdgeDim, dims.E2C2VDim, dtype=wpfloat + ) + primal_normal_vert_v2 = data_alloc.random_field( + grid, dims.EdgeDim, dims.E2C2VDim, dtype=wpfloat + ) - inv_vert_vert_length = data_alloc.random_field(grid, dims.EdgeDim) - inv_primal_edge_length = data_alloc.random_field(grid, dims.EdgeDim) + inv_vert_vert_length = data_alloc.random_field(grid, dims.EdgeDim, dtype=wpfloat) + inv_primal_edge_length = data_alloc.random_field(grid, dims.EdgeDim, dtype=wpfloat) - area_edge = data_alloc.random_field(grid, dims.EdgeDim) - kh_smag_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - z_nabla2_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - diff_multfac_vn = data_alloc.random_field(grid, dims.KDim) - vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - nudgecoeff_e = data_alloc.random_field(grid, dims.EdgeDim) + area_edge = data_alloc.random_field(grid, dims.EdgeDim, dtype=wpfloat) + kh_smag_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) + z_nabla2_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + diff_multfac_vn = data_alloc.random_field(grid, dims.KDim, dtype=wpfloat) + vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + nudgecoeff_e = data_alloc.random_field(grid, dims.EdgeDim, dtype=wpfloat) limited_area = grid.limited_area if hasattr(grid, "limited_area") else True - fac_bdydiff_v = 5.0 - nudgezone_diff = 9.0 + fac_bdydiff_v = wpfloat(5.0) + nudgezone_diff = vpfloat(9.0) edge_domain = h_grid.domain(dims.EdgeDim) start_2nd_nudge_line_idx_e = grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) @@ -147,3 +171,8 @@ def input_data(self, grid: base.Grid) -> dict: vertical_start=0, vertical_end=grid.num_levels, ) + + +@pytest.mark.continuous_benchmarking +class TestApplyDiffusionToVnContinuousBenchmarking(TestApplyDiffusionToVn): + pass diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py index 45d054a47f..b1a90e54d5 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py @@ -8,14 +8,16 @@ import gt4py.next as gtx import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence import ( apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import base +from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing import definitions +from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest from .test_apply_nabla2_to_w import apply_nabla2_to_w_numpy from .test_apply_nabla2_to_w_in_upper_damping_layer import ( @@ -31,6 +33,25 @@ class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence(StencilTest): PROGRAM = apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence OUTPUTS = ("w", "dwdx", "dwdy") + STATIC_PARAMS = { + StandardStaticVariants.NONE: (), + StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "halo_idx", + "interior_idx", + "vertical_start", + "vertical_end", + "nrdmax", + "type_shear", + ), + StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + "nrdmax", + "type_shear", + ), + } @staticmethod def reference( @@ -48,12 +69,21 @@ def reference( nrdmax, interior_idx, halo_idx, + horizontal_start, + horizontal_end, + vertical_start, + vertical_end, **kwargs, ) -> dict: k = np.arange(w_old.shape[1]) cell = np.arange(w_old.shape[0]) reshaped_k = k[np.newaxis, :] reshaped_cell = cell[:, np.newaxis] + out_w, out_dwdx, out_dwdy = ( + np.zeros_like(w_old), + dwdx.copy(), + dwdy.copy(), + ) # create output arrays to update only the necessary slices if type_shear == 2: dwdx, dwdy = np.where( reshaped_k > 0, @@ -81,15 +111,32 @@ def reference( apply_nabla2_to_w_in_upper_damping_layer_numpy(w, diff_multfac_n2w, area, z_nabla2_c), w, ) - return dict(w=w, dwdx=dwdx, dwdy=dwdy) + subset = (slice(horizontal_start, horizontal_end), slice(vertical_start, vertical_end)) + out_w[subset] = w[subset] + out_dwdx[subset] = dwdx[subset] + out_dwdy[subset] = dwdy[subset] + return dict(w=out_w, dwdx=out_dwdx, dwdy=out_dwdy) @pytest.fixture def input_data(self, grid: base.Grid) -> dict: nrdmax = 13 - interior_idx = 1 - halo_idx = 5 + cell_domain = h_grid.domain(dims.CellDim) + interior_idx = grid.start_index(cell_domain(h_grid.Zone.INTERIOR)) # 0 for simple grid + halo_idx = grid.end_index( + cell_domain(h_grid.Zone.LOCAL) + ) # same as horizontal_end for simple grid type_shear = 2 + def _get_start_index_for_w_diffusion() -> int32: + return ( + grid.start_index(cell_domain(h_grid.Zone.NUDGING)) + if grid.limited_area + else grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4)) + ) + + horizontal_start = _get_start_index_for_w_diffusion() + horizontal_end = grid.end_index(cell_domain(h_grid.Zone.HALO)) + geofac_grg_x = random_field(grid, dims.CellDim, dims.C2E2CODim) geofac_grg_y = random_field(grid, dims.CellDim, dims.C2E2CODim) diff_multfac_n2w = random_field(grid, dims.KDim) @@ -99,8 +146,8 @@ def input_data(self, grid: base.Grid) -> dict: diff_multfac_w = 5.0 w = zero_field(grid, dims.CellDim, dims.KDim) - dwdx = zero_field(grid, dims.CellDim, dims.KDim) - dwdy = zero_field(grid, dims.CellDim, dims.KDim) + dwdx = random_field(grid, dims.CellDim, dims.KDim) + dwdy = random_field(grid, dims.CellDim, dims.KDim) return dict( area=area, @@ -117,8 +164,15 @@ def input_data(self, grid: base.Grid) -> dict: w=w, dwdx=dwdx, dwdy=dwdy, - horizontal_start=0, - horizontal_end=grid.num_cells, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, vertical_start=0, vertical_end=grid.num_levels, ) + + +@pytest.mark.continuous_benchmarking +class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulenceContinuousBenchmarking( + TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence +): + pass diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py index 0c4629f7af..fa6659390b 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py @@ -14,7 +14,7 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field +from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing.stencil_tests import StencilTest @@ -53,11 +53,11 @@ def reference( @pytest.fixture def input_data(self, grid): - w = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - geofac_grg_x = random_field(grid, dims.CellDim, dims.C2E2CODim, dtype=wpfloat) - geofac_grg_y = random_field(grid, dims.CellDim, dims.C2E2CODim, dtype=wpfloat) - dwdx = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - dwdy = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + w = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) + geofac_grg_x = data_alloc.random_field(grid, dims.CellDim, dims.C2E2CODim, dtype=wpfloat) + geofac_grg_y = data_alloc.random_field(grid, dims.CellDim, dims.C2E2CODim, dtype=wpfloat) + dwdx = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + dwdy = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) return dict( w=w, diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index 83d006abb2..5d7a01bf01 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -14,14 +14,27 @@ calculate_nabla2_and_smag_coefficients_for_vn, ) from icon4py.model.common import dimension as dims, type_alias as ta +from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import stencil_tests +from icon4py.model.testing import definitions, stencil_tests -@pytest.mark.skip_value_error class TestCalculateNabla2AndSmagCoefficientsForVn(stencil_tests.StencilTest): PROGRAM = calculate_nabla2_and_smag_coefficients_for_vn OUTPUTS = ("kh_smag_e", "kh_smag_ec", "z_nabla2_e") + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( @@ -144,7 +157,7 @@ def reference( return dict(kh_smag_e=kh_smag_e, kh_smag_ec=kh_smag_ec, z_nabla2_e=z_nabla2_e) @pytest.fixture - def input_data(self, grid): + def input_data(self, grid: base.Grid) -> dict: u_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim, dtype=ta.vpfloat) v_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim, dtype=ta.vpfloat) smag_offset = ta.vpfloat("9.0") @@ -194,3 +207,19 @@ def input_data(self, grid): vertical_start=0, vertical_end=grid.num_levels, ) + + +@pytest.mark.continuous_benchmarking +class TestCalculateNabla2AndSmagCoefficientsForVnContinuousBenchmarking( + TestCalculateNabla2AndSmagCoefficientsForVn +): + @pytest.fixture + def input_data(self, grid: base.Grid) -> dict: + base_data = TestCalculateNabla2AndSmagCoefficientsForVn.input_data.__wrapped__(self, grid) + edge_domain = h_grid.domain(dims.EdgeDim) + horizontal_start = grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5)) + horizontal_end = grid.end_index(edge_domain(h_grid.Zone.HALO_LEVEL_2)) + assert horizontal_start < horizontal_end + base_data["horizontal_start"] = 0 + base_data["horizontal_end"] = grid.num_edges + return base_data diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py index 305cf3f319..064c59e23f 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py @@ -55,12 +55,11 @@ def calculate_nabla4_numpy( return z_nabla4_e2 -@pytest.mark.continuous_benchmarking class TestCalculateNabla4(StencilTest): PROGRAM = calculate_nabla4 OUTPUTS = ("z_nabla4_e2",) STATIC_PARAMS = { - StandardStaticVariants.NONE: None, + StandardStaticVariants.NONE: (), StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", @@ -129,3 +128,8 @@ def input_data(self, grid) -> dict: vertical_start=0, vertical_end=gtx.int32(grid.num_levels), ) + + +@pytest.mark.continuous_benchmarking +class TestCalculateNabla4ContinuousBenchmarking(TestCalculateNabla4): + pass diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index afb7a4ff3e..93f47096d6 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -58,7 +58,6 @@ def truly_horizontal_diffusion_nabla_of_theta_over_steep_points_numpy( @pytest.mark.uses_as_offset -@pytest.mark.skip_value_error class TestTrulyHorizontalDiffusionNablaOfThetaOverSteepPoints(StencilTest): PROGRAM = truly_horizontal_diffusion_nabla_of_theta_over_steep_points OUTPUTS = ("z_temp",) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_utils.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_utils.py index 9a8ac84a84..34fa40b7d7 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_utils.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_utils.py @@ -10,52 +10,57 @@ from icon4py.model.common import field_type_aliases as fa from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.type_alias import wpfloat +from icon4py.model.common.type_alias import vpfloat, wpfloat @gtx.field_operator -def _scale_k(field: fa.KField[float], factor: float) -> fa.KField[float]: +def _scale_k(field: fa.KField[wpfloat], factor: wpfloat) -> fa.KField[wpfloat]: return field * factor @gtx.program -def scale_k(field: fa.KField[float], factor: float, scaled_field: fa.KField[float]): +def scale_k(field: fa.KField[wpfloat], factor: wpfloat, scaled_field: fa.KField[wpfloat]): _scale_k(field, factor, out=scaled_field) @gtx.field_operator -def _broadcast_zero_to_three_edge_kdim_fields_wp() -> ( +def _broadcast_zero_to_three_edge_kdim_fields_2wp1vp() -> ( tuple[ fa.EdgeKField[wpfloat], fa.EdgeKField[wpfloat], - fa.EdgeKField[wpfloat], + fa.EdgeKField[vpfloat], ] ): return ( broadcast(wpfloat("0.0"), (EdgeDim, KDim)), broadcast(wpfloat("0.0"), (EdgeDim, KDim)), - broadcast(wpfloat("0.0"), (EdgeDim, KDim)), + broadcast(vpfloat("0.0"), (EdgeDim, KDim)), ) @gtx.field_operator def _calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary( - fourth_order_divdamp_scaling_coeff: fa.KField[float], - max_nudging_coefficient: float, - dbl_eps: float, -) -> fa.KField[float]: - return 0.75 / (max_nudging_coefficient + dbl_eps) * abs(fourth_order_divdamp_scaling_coeff) + fourth_order_divdamp_scaling_coeff: fa.KField[wpfloat], + max_nudging_coefficient: wpfloat, + wp_eps: wpfloat, +) -> fa.KField[wpfloat]: + return ( + wpfloat(0.75) / (max_nudging_coefficient + wp_eps) * abs(fourth_order_divdamp_scaling_coeff) + ) @gtx.field_operator def _calculate_fourth_order_divdamp_scaling_coeff( - interpolated_fourth_order_divdamp_factor: fa.KField[float], + interpolated_fourth_order_divdamp_factor: fa.KField[wpfloat], divdamp_order: gtx.int32, - mean_cell_area: float, - second_order_divdamp_factor: float, -) -> fa.KField[float]: + mean_cell_area: wpfloat, + second_order_divdamp_factor: wpfloat, +) -> fa.KField[wpfloat]: interpolated_fourth_order_divdamp_factor = ( - maximum(0.0, interpolated_fourth_order_divdamp_factor - 0.25 * second_order_divdamp_factor) + maximum( + wpfloat(0.0), + interpolated_fourth_order_divdamp_factor - wpfloat(0.25) * second_order_divdamp_factor, + ) if divdamp_order == 24 else interpolated_fourth_order_divdamp_factor ) @@ -64,13 +69,13 @@ def _calculate_fourth_order_divdamp_scaling_coeff( @gtx.field_operator def _calculate_divdamp_fields( - interpolated_fourth_order_divdamp_factor: fa.KField[float], + interpolated_fourth_order_divdamp_factor: fa.KField[wpfloat], divdamp_order: gtx.int32, - mean_cell_area: float, - second_order_divdamp_factor: float, - max_nudging_coefficient: float, - dbl_eps: float, -) -> tuple[fa.KField[float], fa.KField[float]]: + mean_cell_area: wpfloat, + second_order_divdamp_factor: wpfloat, + max_nudging_coefficient: wpfloat, + wp_eps: wpfloat, +) -> tuple[fa.KField[wpfloat], fa.KField[wpfloat]]: fourth_order_divdamp_scaling_coeff = _calculate_fourth_order_divdamp_scaling_coeff( interpolated_fourth_order_divdamp_factor, divdamp_order, @@ -79,7 +84,7 @@ def _calculate_divdamp_fields( ) reduced_fourth_order_divdamp_coeff_at_nest_boundary = ( _calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary( - fourth_order_divdamp_scaling_coeff, max_nudging_coefficient, dbl_eps + fourth_order_divdamp_scaling_coeff, max_nudging_coefficient, wp_eps ) ) return (fourth_order_divdamp_scaling_coeff, reduced_fourth_order_divdamp_coeff_at_nest_boundary) @@ -87,14 +92,14 @@ def _calculate_divdamp_fields( @gtx.program def calculate_divdamp_fields( - interpolated_fourth_order_divdamp_factor: fa.KField[float], - fourth_order_divdamp_scaling_coeff: fa.KField[float], - reduced_fourth_order_divdamp_coeff_at_nest_boundary: fa.KField[float], + interpolated_fourth_order_divdamp_factor: fa.KField[wpfloat], + fourth_order_divdamp_scaling_coeff: fa.KField[wpfloat], + reduced_fourth_order_divdamp_coeff_at_nest_boundary: fa.KField[wpfloat], divdamp_order: gtx.int32, - mean_cell_area: float, - second_order_divdamp_factor: float, - max_nudging_coefficient: float, - dbl_eps: float, + mean_cell_area: wpfloat, + second_order_divdamp_factor: wpfloat, + max_nudging_coefficient: wpfloat, + wp_eps: wpfloat, ): _calculate_divdamp_fields( interpolated_fourth_order_divdamp_factor, @@ -102,7 +107,7 @@ def calculate_divdamp_fields( mean_cell_area, second_order_divdamp_factor, max_nudging_coefficient, - dbl_eps, + wp_eps, out=( fourth_order_divdamp_scaling_coeff, reduced_fourth_order_divdamp_coeff_at_nest_boundary, @@ -112,13 +117,13 @@ def calculate_divdamp_fields( @gtx.field_operator def _compute_rayleigh_damping_factor( - rayleigh_w: fa.KField[float], dtime: float -) -> fa.KField[float]: - return 1.0 / (1.0 + dtime * rayleigh_w) + rayleigh_w: fa.KField[wpfloat], dtime: wpfloat +) -> fa.KField[wpfloat]: + return wpfloat(1.0) / (wpfloat(1.0) + dtime * rayleigh_w) @gtx.program def compute_rayleigh_damping_factor( - rayleigh_w: fa.KField[float], dtime: float, rayleigh_damping_factor: fa.KField[float] + rayleigh_w: fa.KField[wpfloat], dtime: wpfloat, rayleigh_damping_factor: fa.KField[wpfloat] ): _compute_rayleigh_damping_factor(rayleigh_w, dtime, out=rayleigh_damping_factor) 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 ce765d35de..5bc25a17e2 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,6 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -# ruff: noqa: ERA001, B008 import dataclasses import logging @@ -51,7 +50,6 @@ dimension as dims, field_type_aliases as fa, model_backends, - type_alias as ta, ) from icon4py.model.common.decomposition import definitions as decomposition from icon4py.model.common.grid import ( @@ -63,6 +61,7 @@ from icon4py.model.common.math import smagorinsky from icon4py.model.common.model_options import setup_program from icon4py.model.common.states import prognostic_state as prognostics +from icon4py.model.common.type_alias import vpfloat, wpfloat from icon4py.model.common.utils import data_allocation as data_alloc @@ -80,31 +79,31 @@ class IntermediateFields: contain state that is built up over the predictor and corrector part in a timestep. """ - horizontal_pressure_gradient: fa.EdgeKField[ta.vpfloat] + horizontal_pressure_gradient: fa.EdgeKField[vpfloat] """ Declared as z_gradh_exner in ICON. """ - rho_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat] + rho_at_edges_on_model_levels: fa.EdgeKField[wpfloat] """ Declared as z_rho_e in ICON. """ - theta_v_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat] + theta_v_at_edges_on_model_levels: fa.EdgeKField[wpfloat] """ Declared as z_theta_v_e in ICON. """ - horizontal_kinetic_energy_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat] + horizontal_kinetic_energy_at_edges_on_model_levels: fa.EdgeKField[vpfloat] """ Declared as z_kin_hor_e in ICON. """ - tangential_wind_on_half_levels: fa.EdgeKField[ta.vpfloat] + tangential_wind_on_half_levels: fa.EdgeKField[vpfloat] """ Declared as z_vt_ie in ICON. Tangential wind at edge on k-half levels. NOTE THAT IT ONLY HAS nlev LEVELS because it is only used for computing horizontal advection of w and thus level nlevp1 is not needed because w[nlevp1-1] is diagnostic. """ - horizontal_gradient_of_normal_wind_divergence: fa.EdgeKField[ta.vpfloat] + horizontal_gradient_of_normal_wind_divergence: fa.EdgeKField[vpfloat] """ Declared as z_graddiv_vn in ICON. """ - dwdz_at_cells_on_model_levels: fa.CellKField[ta.vpfloat] + dwdz_at_cells_on_model_levels: fa.CellKField[vpfloat] """ Declared as z_dwdz_dd in ICON. """ @@ -155,26 +154,26 @@ def __init__( iadv_rhotheta: dycore_states.RhoThetaAdvectionType = dycore_states.RhoThetaAdvectionType.MIURA, igradp_method: dycore_states.HorizontalPressureDiscretizationType = dycore_states.HorizontalPressureDiscretizationType.TAYLOR_HYDRO, rayleigh_type: constants.RayleighType = constants.RayleighType.KLEMP, - rayleigh_coeff: float = 0.05, + rayleigh_coeff: wpfloat | float = 0.05, divdamp_order: dycore_states.DivergenceDampingOrder = dycore_states.DivergenceDampingOrder.COMBINED, # the ICON default is 4, is_iau_active: bool = False, - iau_wgt_dyn: float = 0.0, + iau_wgt_dyn: wpfloat | float = 0.0, divdamp_type: dycore_states.DivergenceDampingType = dycore_states.DivergenceDampingType.THREE_DIMENSIONAL, - divdamp_trans_start: float = 12500.0, - divdamp_trans_end: float = 17500.0, + divdamp_trans_start: wpfloat | float = 12500.0, + divdamp_trans_end: wpfloat | float = 17500.0, l_vert_nested: bool = False, - rhotheta_offctr: float = -0.1, - veladv_offctr: float = 0.25, - _nudge_max_coeff: float | None = None, # default is set in __init__ - max_nudging_coefficient: float | None = None, # default is set in __init__ - fourth_order_divdamp_factor: float = 0.0025, - fourth_order_divdamp_factor2: float = 0.004, - fourth_order_divdamp_factor3: float = 0.004, - fourth_order_divdamp_factor4: float = 0.004, - fourth_order_divdamp_z: float = 32500.0, - fourth_order_divdamp_z2: float = 40000.0, - fourth_order_divdamp_z3: float = 60000.0, - fourth_order_divdamp_z4: float = 80000.0, + rhotheta_offctr: wpfloat | float = -0.1, + veladv_offctr: wpfloat | float = 0.25, + _nudge_max_coeff: wpfloat | float | None = None, # default is set in __init__ + max_nudging_coefficient: wpfloat | float | None = None, # default is set in __init__ + fourth_order_divdamp_factor: wpfloat | float = 0.0025, + fourth_order_divdamp_factor2: wpfloat | float = 0.004, + fourth_order_divdamp_factor3: wpfloat | float = 0.004, + fourth_order_divdamp_factor4: wpfloat | float = 0.004, + fourth_order_divdamp_z: wpfloat | float = 32500.0, + fourth_order_divdamp_z2: wpfloat | float = 40000.0, + fourth_order_divdamp_z3: wpfloat | float = 60000.0, + fourth_order_divdamp_z4: wpfloat | float = 80000.0, ): # parameters from namelist diffusion_nml self.itime_scheme: int = itime_scheme @@ -188,7 +187,7 @@ def __init__( #: type of Rayleigh damping self.rayleigh_type: constants.RayleighType = rayleigh_type # used for calculation of rayleigh_w, rayleigh_vn in mo_vertical_grid.f90 - self.rayleigh_coeff: float = rayleigh_coeff + self.rayleigh_coeff: wpfloat = wpfloat(rayleigh_coeff) #: order of divergence damping self.divdamp_order: dycore_states.DivergenceDampingOrder = divdamp_order @@ -196,57 +195,57 @@ def __init__( #: type of divergence damping self.divdamp_type: dycore_states.DivergenceDampingType = divdamp_type #: Lower and upper bound of transition zone between 2D and 3D divergence damping in case of divdamp_type = 32 [m] - self.divdamp_trans_start: float = divdamp_trans_start - self.divdamp_trans_end: float = divdamp_trans_end + self.divdamp_trans_start: wpfloat = wpfloat(divdamp_trans_start) + self.divdamp_trans_end: wpfloat = wpfloat(divdamp_trans_end) #: off-centering for density and potential temperature at interface levels. #: Specifying a negative value here reduces the amount of vertical #: wind off-centering needed for stability of sound waves. - self.rhotheta_offctr: float = rhotheta_offctr + self.rhotheta_offctr: wpfloat = wpfloat(rhotheta_offctr) #: off-centering of velocity advection in corrector step - self.veladv_offctr: float = veladv_offctr + self.veladv_offctr: wpfloat = wpfloat(veladv_offctr) #: scaling factor for divergence damping - self.fourth_order_divdamp_factor: float = fourth_order_divdamp_factor + self.fourth_order_divdamp_factor: wpfloat = wpfloat(fourth_order_divdamp_factor) """ Declared as divdamp_fac in ICON. It is a scaling factor for fourth order divergence damping between heights of fourth_order_divdamp_z and fourth_order_divdamp_z2. """ - self.fourth_order_divdamp_factor2: float = fourth_order_divdamp_factor2 + self.fourth_order_divdamp_factor2: wpfloat = wpfloat(fourth_order_divdamp_factor2) """ Declared as divdamp_fac2 in ICON. It is a scaling factor for fourth order divergence damping between heights of fourth_order_divdamp_z and fourth_order_divdamp_z2. Divergence damping factor reaches fourth_order_divdamp_factor2 at fourth_order_divdamp_z2. """ - self.fourth_order_divdamp_factor3: float = fourth_order_divdamp_factor3 + self.fourth_order_divdamp_factor3: wpfloat = wpfloat(fourth_order_divdamp_factor3) """ Declared as divdamp_fac3 in ICON. It is a scaling factor to determine the quadratic vertical profile of fourth order divergence damping factor between heights of fourth_order_divdamp_z2 and fourth_order_divdamp_z4. """ - self.fourth_order_divdamp_factor4: float = fourth_order_divdamp_factor4 + self.fourth_order_divdamp_factor4: wpfloat = wpfloat(fourth_order_divdamp_factor4) """ Declared as divdamp_fac4 in ICON. It is a scaling factor to determine the quadratic vertical profile of fourth order divergence damping factor between heights of fourth_order_divdamp_z2 and fourth_order_divdamp_z4. Divergence damping factor reaches fourth_order_divdamp_factor4 at fourth_order_divdamp_z4. """ - self.fourth_order_divdamp_z: float = fourth_order_divdamp_z + self.fourth_order_divdamp_z: wpfloat = wpfloat(fourth_order_divdamp_z) """ Declared as divdamp_z in ICON. The upper limit in height where divergence damping factor is a constant. """ - self.fourth_order_divdamp_z2: float = fourth_order_divdamp_z2 + self.fourth_order_divdamp_z2: wpfloat = wpfloat(fourth_order_divdamp_z2) """ Declared as divdamp_z2 in ICON. The upper limit in height above fourth_order_divdamp_z where divergence damping factor decreases as a linear function of height. """ - self.fourth_order_divdamp_z3: float = fourth_order_divdamp_z3 + self.fourth_order_divdamp_z3: wpfloat = wpfloat(fourth_order_divdamp_z3) """ Declared as divdamp_z3 in ICON. Am intermediate height between fourth_order_divdamp_z2 and fourth_order_divdamp_z4 where divergence damping factor decreases quadratically with height. """ - self.fourth_order_divdamp_z4: float = fourth_order_divdamp_z4 + self.fourth_order_divdamp_z4: wpfloat = wpfloat(fourth_order_divdamp_z4) """ Declared as divdamp_z4 in ICON. The upper limit in height where divergence damping factor decreases quadratically with height. @@ -270,14 +269,14 @@ def __init__( "Cannot set both '_max_nudging_coefficient' and 'scaled_max_nudging_coefficient'." ) elif max_nudging_coefficient is not None: - self.max_nudging_coefficient: float = max_nudging_coefficient + self.max_nudging_coefficient: wpfloat = wpfloat(max_nudging_coefficient) elif _nudge_max_coeff is not None: - self.max_nudging_coefficient: float = ( - constants.DEFAULT_DYNAMICS_TO_PHYSICS_TIMESTEP_RATIO * _nudge_max_coeff + self.max_nudging_coefficient: wpfloat = ( + constants.DEFAULT_DYNAMICS_TO_PHYSICS_TIMESTEP_RATIO * wpfloat(_nudge_max_coeff) ) else: # default value in ICON - self.max_nudging_coefficient: float = ( - constants.DEFAULT_DYNAMICS_TO_PHYSICS_TIMESTEP_RATIO * 0.02 + self.max_nudging_coefficient: wpfloat = ( + constants.DEFAULT_DYNAMICS_TO_PHYSICS_TIMESTEP_RATIO * wpfloat(0.02) ) #: from mo_run_nml.f90 @@ -288,7 +287,7 @@ def __init__( #: whether IAU is active at current time self.is_iau_active: bool = is_iau_active #: IAU weight for dynamics fields - self.iau_wgt_dyn: float = iau_wgt_dyn + self.iau_wgt_dyn: wpfloat = wpfloat(iau_wgt_dyn) self._validate() @@ -323,11 +322,15 @@ def __init__(self, config: NonHydrostaticConfig): #: Weighting coefficients for velocity advection if tendency averaging is used #: The off-centering specified here turned out to be beneficial to numerical #: stability in extreme situations - self.advection_explicit_weight_parameter: Final[float] = 0.5 - config.veladv_offctr + self.advection_explicit_weight_parameter: Final[wpfloat] = ( + wpfloat(0.5) - config.veladv_offctr + ) """ Declared as wgt_nnow_vel in ICON. """ - self.advection_implicit_weight_parameter: Final[float] = 0.5 + config.veladv_offctr + self.advection_implicit_weight_parameter: Final[wpfloat] = ( + wpfloat(0.5) + config.veladv_offctr + ) """ Declared as wgt_nnew_vel in ICON. """ @@ -335,12 +338,14 @@ def __init__(self, config: NonHydrostaticConfig): #: Weighting coefficients for rho and theta at interface levels in the corrector step #: This empirically determined weighting minimizes the vertical wind off-centering #: needed for numerical stability of vertical sound wave propagation - self.rhotheta_implicit_weight_parameter: Final[float] = 0.5 + config.rhotheta_offctr + self.rhotheta_implicit_weight_parameter: Final[wpfloat] = ( + wpfloat(0.5) + config.rhotheta_offctr + ) """ Declared as wgt_nnew_rth in ICON. """ - self.rhotheta_explicit_weight_parameter: Final[float] = ( - 1.0 - self.rhotheta_implicit_weight_parameter + self.rhotheta_explicit_weight_parameter: Final[wpfloat] = ( + wpfloat(1.0) - self.rhotheta_implicit_weight_parameter ) """ Declared as wgt_nnow_rth in ICON. @@ -363,9 +368,9 @@ def __init__( | model_backends.DeviceType | model_backends.BackendDescriptor | None, - exchange: decomposition.ExchangeRuntime = decomposition.SingleNodeExchange(), + exchange: decomposition.ExchangeRuntime | None = None, ): - self._exchange = exchange + self._exchange = exchange if (exchange is not None) else decomposition.SingleNodeExchange() self._grid = grid self._config = config @@ -415,9 +420,7 @@ def __init__( self._update_theta_v = setup_program( backend=backend, program=update_theta_v, - constant_args={ - "mask_prog_halo_c": self._metric_state_nonhydro.mask_prog_halo_c, - }, + constant_args={"mask_prog_halo_c": self._metric_state_nonhydro.mask_prog_halo_c}, horizontal_sizes={ "horizontal_start": self._start_cell_halo, "horizontal_end": self._end_cell_end, @@ -587,9 +590,7 @@ def __init__( "rayleigh_type": self._config.rayleigh_type, "divdamp_type": self._config.divdamp_type, }, - variants={ - "at_first_substep": [False, True], - }, + variants={"at_first_substep": [False, True]}, horizontal_sizes={ "start_cell_index_nudging": self._start_cell_nudging, "end_cell_index_local": self._end_cell_local, @@ -642,9 +643,7 @@ def __init__( self._compute_dwdz_for_divergence_damping = setup_program( backend=backend, program=compute_dwdz_for_divergence_damping, - constant_args={ - "inv_ddqz_z_full": self._metric_state_nonhydro.inv_ddqz_z_full, - }, + constant_args={"inv_ddqz_z_full": self._metric_state_nonhydro.inv_ddqz_z_full}, horizontal_sizes={ "horizontal_start": self._start_cell_lateral_boundary, "horizontal_end": self._end_cell_lateral_boundary_level_4, @@ -689,17 +688,15 @@ def __init__( program=dycore_utils.calculate_divdamp_fields, constant_args={ "divdamp_order": gtx.int32(self._config.divdamp_order), - "mean_cell_area": self._grid.global_properties.mean_cell_area, + "mean_cell_area": wpfloat(self._grid.global_properties.mean_cell_area), "max_nudging_coefficient": self._config.max_nudging_coefficient, - "dbl_eps": constants.DBL_EPS, + "wp_eps": constants.WP_EPS, }, ) self._compute_rayleigh_damping_factor = setup_program( backend=backend, program=dycore_utils.compute_rayleigh_damping_factor, - constant_args={ - "rayleigh_w": self._metric_state_nonhydro.rayleigh_w, - }, + constant_args={"rayleigh_w": self._metric_state_nonhydro.rayleigh_w}, ) self._compute_perturbed_quantities_and_interpolation = setup_program( @@ -814,7 +811,7 @@ def __init__( self._allocate_local_fields(model_backends.get_allocator(backend)) self._en_smag_fac_for_zero_nshift( - enh_smag_fac=self.interpolated_fourth_order_divdamp_factor, + enh_smag_fac=self.interpolated_fourth_order_divdamp_factor ) self.p_test_run = False @@ -824,7 +821,7 @@ def _allocate_local_fields(self, allocator: gtx_allocators.FieldBufferAllocation self._grid, dims.CellDim, dims.KDim, - dtype=ta.vpfloat, + dtype=vpfloat, extend={dims.KDim: 1}, allocator=allocator, ) @@ -835,7 +832,7 @@ def _allocate_local_fields(self, allocator: gtx_allocators.FieldBufferAllocation self._grid, dims.CellDim, dims.KDim, - dtype=ta.vpfloat, + dtype=vpfloat, extend={dims.KDim: 1}, allocator=allocator, ) @@ -844,7 +841,7 @@ def _allocate_local_fields(self, allocator: gtx_allocators.FieldBufferAllocation """ self.ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat, allocator=allocator + self._grid, dims.CellDim, dims.KDim, dtype=vpfloat, allocator=allocator ) ) """ @@ -854,7 +851,7 @@ def _allocate_local_fields(self, allocator: gtx_allocators.FieldBufferAllocation self._grid, dims.CellDim, dims.KDim, - dtype=ta.vpfloat, + dtype=vpfloat, extend={dims.KDim: 1}, allocator=allocator, ) @@ -863,7 +860,7 @@ def _allocate_local_fields(self, allocator: gtx_allocators.FieldBufferAllocation Declared as z_theta_v_pr_ic in ICON. """ self.pressure_buoyancy_acceleration_at_cells_on_half_levels = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat, allocator=allocator + self._grid, dims.CellDim, dims.KDim, dtype=vpfloat, allocator=allocator ) """ Declared as z_th_ddz_exner_c in ICON. theta' dpi0/dz + theta (1 - eta_impl) dpi'/dz. @@ -872,78 +869,78 @@ def _allocate_local_fields(self, allocator: gtx_allocators.FieldBufferAllocation term for updating w, and w at model top/bottom is diagnosed. """ self.perturbed_rho_at_cells_on_model_levels = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat, allocator=allocator + self._grid, dims.CellDim, dims.KDim, dtype=vpfloat, allocator=allocator ) """ Declared as z_rth_pr_1 in ICON. """ self.perturbed_theta_v_at_cells_on_model_levels = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat, allocator=allocator + self._grid, dims.CellDim, dims.KDim, dtype=vpfloat, allocator=allocator ) """ Declared as z_rth_pr_2 in ICON. """ self.d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat, allocator=allocator + self._grid, dims.CellDim, dims.KDim, dtype=vpfloat, allocator=allocator ) ) """ Declared as z_dexner_dz_c_2 in ICON. """ self.z_vn_avg = data_alloc.zero_field( - self._grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat, allocator=allocator + self._grid, dims.EdgeDim, dims.KDim, dtype=wpfloat, allocator=allocator ) self.theta_v_flux_at_edges_on_model_levels = data_alloc.zero_field( - self._grid, dims.EdgeDim, dims.KDim, dtype=ta.wpfloat, allocator=allocator + self._grid, dims.EdgeDim, dims.KDim, dtype=wpfloat, allocator=allocator ) """ Declared as z_theta_v_fl_e in ICON. """ self.z_rho_v = data_alloc.zero_field( - self._grid, dims.VertexDim, dims.KDim, dtype=ta.wpfloat, allocator=allocator + self._grid, dims.VertexDim, dims.KDim, dtype=wpfloat, allocator=allocator ) self.z_theta_v_v = data_alloc.zero_field( - self._grid, dims.VertexDim, dims.KDim, dtype=ta.wpfloat, allocator=allocator + self._grid, dims.VertexDim, dims.KDim, dtype=wpfloat, allocator=allocator ) self.k_field = data_alloc.index_field( self._grid, dims.KDim, extend={dims.KDim: 1}, allocator=allocator ) self._contravariant_correction_at_edges_on_model_levels = data_alloc.zero_field( - self._grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat, allocator=allocator + self._grid, dims.EdgeDim, dims.KDim, dtype=vpfloat, allocator=allocator ) """ Declared as z_w_concorr_me in ICON. vn dz/dn + vt dz/dt, z is topography height """ self.hydrostatic_correction_on_lowest_level = data_alloc.zero_field( - self._grid, dims.EdgeDim, dtype=ta.vpfloat, allocator=allocator + self._grid, dims.EdgeDim, dtype=vpfloat, allocator=allocator ) self.hydrostatic_correction = data_alloc.zero_field( - self._grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat, allocator=allocator + self._grid, dims.EdgeDim, dims.KDim, dtype=vpfloat, allocator=allocator ) """ Declared as z_hydro_corr in ICON. Used for computation of horizontal pressure gradient over steep slope. """ self.rayleigh_damping_factor = data_alloc.zero_field( - self._grid, dims.KDim, dtype=ta.wpfloat, allocator=allocator + self._grid, dims.KDim, dtype=wpfloat, allocator=allocator ) """ Declared as z_raylfac in ICON. """ self.interpolated_fourth_order_divdamp_factor = data_alloc.zero_field( - self._grid, dims.KDim, dtype=ta.wpfloat, allocator=allocator + self._grid, dims.KDim, dtype=wpfloat, allocator=allocator ) """ Declared as enh_divdamp_fac in ICON. """ self.reduced_fourth_order_divdamp_coeff_at_nest_boundary = data_alloc.zero_field( - self._grid, dims.KDim, dtype=ta.wpfloat, allocator=allocator + self._grid, dims.KDim, dtype=wpfloat, allocator=allocator ) """ Declared as bdy_divdamp in ICON. """ self.fourth_order_divdamp_scaling_coeff = data_alloc.zero_field( - self._grid, dims.KDim, dtype=ta.wpfloat, allocator=allocator + self._grid, dims.KDim, dtype=wpfloat, allocator=allocator ) """ Declared as scal_divdamp in ICON. @@ -1008,8 +1005,8 @@ def time_step( diagnostic_state_nh: dycore_states.DiagnosticStateNonHydro, prognostic_states: common_utils.TimeStepPair[prognostics.PrognosticState], prep_adv: dycore_states.PrepAdvection, - second_order_divdamp_factor: float, - dtime: float, + second_order_divdamp_factor: wpfloat, + dtime: wpfloat, ndyn_substeps_var: int, at_initial_timestep: bool, lprep_adv: bool, @@ -1083,13 +1080,12 @@ def time_step( theta_v_new=prognostic_states.next.theta_v, ) - # flake8: noqa: C901 def run_predictor_step( self, diagnostic_state_nh: dycore_states.DiagnosticStateNonHydro, prognostic_states: common_utils.TimeStepPair[prognostics.PrognosticState], z_fields: IntermediateFields, - dtime: float, + dtime: wpfloat, at_initial_timestep: bool, at_first_substep: bool, ): @@ -1260,9 +1256,9 @@ def run_corrector_step( diagnostic_state_nh: dycore_states.DiagnosticStateNonHydro, prognostic_states: common_utils.TimeStepPair[prognostics.PrognosticState], z_fields: IntermediateFields, - second_order_divdamp_factor: float, + second_order_divdamp_factor: wpfloat, prep_adv: dycore_states.PrepAdvection, - dtime: float, + dtime: wpfloat, ndyn_substeps_var: int, lprep_adv: bool, at_first_substep: bool, @@ -1273,23 +1269,42 @@ def run_corrector_step( f"second_order_divdamp_factor = {second_order_divdamp_factor}, at_first_substep = {at_first_substep}, at_last_substep = {at_last_substep} " ) + ndyn_substeps_var_wp = wpfloat(ndyn_substeps_var) # Inverse value of ndyn_substeps for tracer advection precomputations - r_nsubsteps = 1.0 / ndyn_substeps_var + r_nsubsteps = wpfloat(1.0) / ndyn_substeps_var_wp + + second_order_divdamp_factor_wp = wpfloat(second_order_divdamp_factor) # scaling factor for second-order divergence damping: second_order_divdamp_factor_from_sfc_to_divdamp_z*delta_x**2 # delta_x**2 is approximated by the mean cell area # Coefficient for reduced fourth-order divergence d - second_order_divdamp_scaling_coeff = ( - second_order_divdamp_factor * self._grid.global_properties.mean_cell_area + + second_order_divdamp_scaling_coeff = second_order_divdamp_factor_wp * wpfloat( + self._grid.global_properties.mean_cell_area ) - self._calculate_divdamp_fields( - interpolated_fourth_order_divdamp_factor=self.interpolated_fourth_order_divdamp_factor, - fourth_order_divdamp_scaling_coeff=self.fourth_order_divdamp_scaling_coeff, - reduced_fourth_order_divdamp_coeff_at_nest_boundary=self.reduced_fourth_order_divdamp_coeff_at_nest_boundary, - second_order_divdamp_factor=second_order_divdamp_factor, + dycore_utils._calculate_divdamp_fields( + self.interpolated_fourth_order_divdamp_factor, + gtx.int32(self._config.divdamp_order), + wpfloat(self._grid.global_properties.mean_cell_area), + second_order_divdamp_factor_wp, + self._config.max_nudging_coefficient, + constants.WP_EPS, + out=( + self.fourth_order_divdamp_scaling_coeff, + self.reduced_fourth_order_divdamp_coeff_at_nest_boundary, + ), ) + # TODO(pstark): Find and solve bug that appears when running with the compiled self._calculate_divdamp_fields in combination with single precision. + # self._calculate_divdamp_fields( + # interpolated_fourth_order_divdamp_factor=self.interpolated_fourth_order_divdamp_factor, + # fourth_order_divdamp_scaling_coeff=self.fourth_order_divdamp_scaling_coeff, + # reduced_fourth_order_divdamp_coeff_at_nest_boundary=self.reduced_fourth_order_divdamp_coeff_at_nest_boundary, + # second_order_divdamp_factor=second_order_divdamp_factor_wp, + # ) + + log.debug("corrector run velocity advection") self.velocity_advection.run_corrector_step( diagnostic_state=diagnostic_state_nh, @@ -1323,13 +1338,14 @@ def run_corrector_step( log.debug("corrector: start stencil apply_divergence_damping_and_update_vn") apply_2nd_order_divergence_damping = ( self._config.divdamp_order == dycore_states.DivergenceDampingOrder.COMBINED - and second_order_divdamp_scaling_coeff > 1.0e-6 + and second_order_divdamp_scaling_coeff > wpfloat(1.0e-6) ) apply_4th_order_divergence_damping = ( self._config.divdamp_order == dycore_states.DivergenceDampingOrder.FOURTH_ORDER or ( self._config.divdamp_order == dycore_states.DivergenceDampingOrder.COMBINED - and second_order_divdamp_factor <= (4.0 * self._config.fourth_order_divdamp_factor) + and second_order_divdamp_factor_wp + <= (wpfloat(4.0) * self._config.fourth_order_divdamp_factor) ) ) @@ -1397,7 +1413,7 @@ def run_corrector_step( rayleigh_damping_factor=self.rayleigh_damping_factor, lprep_adv=lprep_adv, r_nsubsteps=r_nsubsteps, - ndyn_substeps_var=float(ndyn_substeps_var), + ndyn_substeps_var=ndyn_substeps_var_wp, dtime=dtime, at_first_substep=at_first_substep, at_last_substep=at_last_substep, @@ -1409,7 +1425,7 @@ def run_corrector_step( "corrector step sets prep_adv.dynamical_vertical_mass_flux_at_cells_on_half_levels to zero" ) self._init_cell_kdim_field_with_zero_wp( - field_with_zero_wp=prep_adv.dynamical_vertical_mass_flux_at_cells_on_half_levels, + field_with_zero_wp=prep_adv.dynamical_vertical_mass_flux_at_cells_on_half_levels ) self._update_mass_flux_weighted( rho_ic=diagnostic_state_nh.rho_at_cells_on_half_levels, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro_stencils.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro_stencils.py index bdb18e294c..8a73b39b09 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro_stencils.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro_stencils.py @@ -9,13 +9,7 @@ from gt4py.next.experimental import concat_where from icon4py.model.atmosphere.dycore.dycore_utils import ( - _broadcast_zero_to_three_edge_kdim_fields_wp, -) -from icon4py.model.atmosphere.dycore.stencils.compute_contravariant_correction import ( - _compute_contravariant_correction, -) -from icon4py.model.atmosphere.dycore.stencils.compute_horizontal_kinetic_energy import ( - _compute_horizontal_kinetic_energy, + _broadcast_zero_to_three_edge_kdim_fields_2wp1vp, ) from icon4py.model.atmosphere.dycore.stencils.compute_perturbation_of_rho_and_theta import ( _compute_perturbation_of_rho_and_theta, @@ -26,26 +20,23 @@ from icon4py.model.atmosphere.dycore.stencils.compute_virtual_potential_temperatures_and_pressure_gradient import ( _compute_virtual_potential_temperatures_and_pressure_gradient, ) -from icon4py.model.atmosphere.dycore.stencils.extrapolate_at_top import _extrapolate_at_top from icon4py.model.atmosphere.dycore.stencils.init_cell_kdim_field_with_zero_wp import ( - _init_cell_kdim_field_with_zero_wp, -) -from icon4py.model.atmosphere.dycore.stencils.interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges import ( - _interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges, + _init_cell_kdim_field_with_zero_vp, ) from icon4py.model.atmosphere.dycore.stencils.update_density_exner_wind import ( _update_density_exner_wind, ) from icon4py.model.atmosphere.dycore.stencils.update_wind import _update_wind from icon4py.model.common import dimension as dims, field_type_aliases as fa +from icon4py.model.common.type_alias import vpfloat, wpfloat @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def init_test_fields( - z_rho_e: fa.EdgeKField[float], - z_theta_v_e: fa.EdgeKField[float], - z_dwdz_dd: fa.CellKField[float], - z_graddiv_vn: fa.EdgeKField[float], + z_rho_e: fa.EdgeKField[wpfloat], + z_theta_v_e: fa.EdgeKField[wpfloat], + z_dwdz_dd: fa.CellKField[vpfloat], + z_graddiv_vn: fa.EdgeKField[vpfloat], edges_start: gtx.int32, edges_end: gtx.int32, cells_start: gtx.int32, @@ -53,11 +44,11 @@ def init_test_fields( vertical_start: gtx.int32, vertical_end: gtx.int32, ): - _broadcast_zero_to_three_edge_kdim_fields_wp( + _broadcast_zero_to_three_edge_kdim_fields_2wp1vp( out=(z_rho_e, z_theta_v_e, z_graddiv_vn), domain={dims.EdgeDim: (edges_start, edges_end), dims.KDim: (vertical_start, vertical_end)}, ) - _init_cell_kdim_field_with_zero_wp( + _init_cell_kdim_field_with_zero_vp( out=z_dwdz_dd, domain={dims.CellDim: (cells_start, cells_end), dims.KDim: (vertical_start, vertical_end)}, ) @@ -65,28 +56,28 @@ def init_test_fields( @gtx.field_operator def _compute_pressure_gradient_and_perturbed_rho_and_potential_temperatures( - rho: fa.CellKField[float], - z_rth_pr_1: fa.CellKField[float], - z_rth_pr_2: fa.CellKField[float], - rho_ref_mc: fa.CellKField[float], - theta_v: fa.CellKField[float], - theta_ref_mc: fa.CellKField[float], - rho_ic: fa.CellKField[float], - wgtfac_c: fa.CellKField[float], - vwind_expl_wgt: fa.CellField[float], - exner_pr: fa.CellKField[float], - d_exner_dz_ref_ic: fa.CellKField[float], - ddqz_z_half: fa.CellKField[float], - z_theta_v_pr_ic: fa.CellKField[float], - theta_v_ic: fa.CellKField[float], - z_th_ddz_exner_c: fa.CellKField[float], + rho: fa.CellKField[wpfloat], + z_rth_pr_1: fa.CellKField[vpfloat], + z_rth_pr_2: fa.CellKField[vpfloat], + rho_ref_mc: fa.CellKField[vpfloat], + theta_v: fa.CellKField[wpfloat], + theta_ref_mc: fa.CellKField[vpfloat], + rho_ic: fa.CellKField[wpfloat], + wgtfac_c: fa.CellKField[vpfloat], + vwind_expl_wgt: fa.CellField[wpfloat], + exner_pr: fa.CellKField[wpfloat], + d_exner_dz_ref_ic: fa.CellKField[vpfloat], + ddqz_z_half: fa.CellKField[vpfloat], + z_theta_v_pr_ic: fa.CellKField[vpfloat], + theta_v_ic: fa.CellKField[wpfloat], + z_th_ddz_exner_c: fa.CellKField[vpfloat], ) -> tuple[ - fa.CellKField[float], - fa.CellKField[float], - fa.CellKField[float], - fa.CellKField[float], - fa.CellKField[float], - fa.CellKField[float], + fa.CellKField[vpfloat], + fa.CellKField[vpfloat], + fa.CellKField[wpfloat], + fa.CellKField[vpfloat], + fa.CellKField[wpfloat], + fa.CellKField[vpfloat], ]: (z_rth_pr_1, z_rth_pr_2) = concat_where( dims.KDim == 0, @@ -119,121 +110,18 @@ def _compute_pressure_gradient_and_perturbed_rho_and_potential_temperatures( return z_rth_pr_1, z_rth_pr_2, rho_ic, z_theta_v_pr_ic, theta_v_ic, z_th_ddz_exner_c -@gtx.field_operator -def _predictor_stencils_35_36( - vn: fa.EdgeKField[float], - ddxn_z_full: fa.EdgeKField[float], - ddxt_z_full: fa.EdgeKField[float], - vt: fa.EdgeKField[float], - z_w_concorr_me: fa.EdgeKField[float], - wgtfac_e: fa.EdgeKField[float], - vn_ie: fa.EdgeKField[float], - z_vt_ie: fa.EdgeKField[float], - z_kin_hor_e: fa.EdgeKField[float], - k_field: fa.KField[gtx.int32], - nflatlev_startindex: gtx.int32, -) -> tuple[ - fa.EdgeKField[float], - fa.EdgeKField[float], - fa.EdgeKField[float], - fa.EdgeKField[float], -]: - z_w_concorr_me = concat_where( - dims.KDim >= nflatlev_startindex, - _compute_contravariant_correction(vn, ddxn_z_full, ddxt_z_full, vt), - z_w_concorr_me, - ) - (vn_ie, z_vt_ie, z_kin_hor_e) = concat_where( - dims.KDim >= 1, - _interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges(wgtfac_e, vn, vt), - (vn_ie, z_vt_ie, z_kin_hor_e), - ) - return z_w_concorr_me, vn_ie, z_vt_ie, z_kin_hor_e - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def predictor_stencils_35_36( - vn: fa.EdgeKField[float], - ddxn_z_full: fa.EdgeKField[float], - ddxt_z_full: fa.EdgeKField[float], - vt: fa.EdgeKField[float], - z_w_concorr_me: fa.EdgeKField[float], - wgtfac_e: fa.EdgeKField[float], - vn_ie: fa.EdgeKField[float], - z_vt_ie: fa.EdgeKField[float], - z_kin_hor_e: fa.EdgeKField[float], - k_field: fa.KField[gtx.int32], - nflatlev_startindex: gtx.int32, - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -): - _predictor_stencils_35_36( - vn, - ddxn_z_full, - ddxt_z_full, - vt, - z_w_concorr_me, - wgtfac_e, - vn_ie, - z_vt_ie, - z_kin_hor_e, - k_field, - nflatlev_startindex, - out=(z_w_concorr_me, vn_ie, z_vt_ie, z_kin_hor_e), - domain={ - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end), - }, - ) - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def predictor_stencils_37_38( - vn: fa.EdgeKField[float], - vt: fa.EdgeKField[float], - vn_ie: fa.EdgeKField[float], - z_vt_ie: fa.EdgeKField[float], - z_kin_hor_e: fa.EdgeKField[float], - wgtfacq_e_dsl: fa.EdgeKField[float], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -): - _compute_horizontal_kinetic_energy( - vn, - vt, - out=(vn_ie, z_vt_ie, z_kin_hor_e), - domain={ - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_start + 1), - }, - ) - _extrapolate_at_top( - vn, - wgtfacq_e_dsl, - out=vn_ie, - domain={ - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_end - 1, vertical_end), - }, - ) - - @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def stencils_61_62( - rho_now: fa.CellKField[float], - grf_tend_rho: fa.CellKField[float], - theta_v_now: fa.CellKField[float], - grf_tend_thv: fa.CellKField[float], - w_now: fa.CellKField[float], - grf_tend_w: fa.CellKField[float], - rho_new: fa.CellKField[float], - exner_new: fa.CellKField[float], - w_new: fa.CellKField[float], - dtime: float, + rho_now: fa.CellKField[wpfloat], + grf_tend_rho: fa.CellKField[wpfloat], + theta_v_now: fa.CellKField[wpfloat], + grf_tend_thv: fa.CellKField[wpfloat], + w_now: fa.CellKField[wpfloat], + grf_tend_w: fa.CellKField[wpfloat], + rho_new: fa.CellKField[wpfloat], + exner_new: fa.CellKField[wpfloat], + w_new: fa.CellKField[wpfloat], + dtime: wpfloat, horizontal_start: gtx.int32, horizontal_end: gtx.int32, vertical_start: gtx.int32, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_advection_in_vertical_momentum_equation.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_advection_in_vertical_momentum_equation.py index f6b0533dbe..cb3a86bf80 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_advection_in_vertical_momentum_equation.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_advection_in_vertical_momentum_equation.py @@ -194,8 +194,8 @@ def _compute_contravariant_corrected_w_and_cfl( def _compute_advective_vertical_wind_tendency( vertical_wind_advective_tendency: fa.CellKField[ta.vpfloat], w: fa.CellKField[ta.wpfloat], - horizontal_advection_of_w_at_edges_on_half_levels: fa.EdgeKField[ta.wpfloat], - contravariant_corrected_w_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], + horizontal_advection_of_w_at_edges_on_half_levels: fa.EdgeKField[ta.vpfloat], + contravariant_corrected_w_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], cfl_clipping: fa.CellKField[bool], coeff1_dwdz: fa.CellKField[ta.vpfloat], coeff2_dwdz: fa.CellKField[ta.vpfloat], @@ -453,7 +453,7 @@ def _interpolate_contravariant_correction_to_cells_on_half_levels( def _compute_contravariant_correction_and_advection_in_vertical_momentum_equation( vertical_wind_advective_tendency: fa.CellKField[ta.vpfloat], w: fa.CellKField[ta.wpfloat], - horizontal_advection_of_w_at_edges_on_half_levels: fa.EdgeKField[ta.wpfloat], + horizontal_advection_of_w_at_edges_on_half_levels: fa.EdgeKField[ta.vpfloat], contravariant_correction_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], coeff1_dwdz: fa.CellKField[ta.vpfloat], coeff2_dwdz: fa.CellKField[ta.vpfloat], @@ -539,7 +539,7 @@ def compute_contravariant_correction_and_advection_in_vertical_momentum_equation contravariant_corrected_w_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], vertical_cfl: fa.CellKField[ta.vpfloat], w: fa.CellKField[ta.wpfloat], - horizontal_advection_of_w_at_edges_on_half_levels: fa.EdgeKField[ta.wpfloat], + horizontal_advection_of_w_at_edges_on_half_levels: fa.EdgeKField[ta.vpfloat], contravariant_correction_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], coeff1_dwdz: fa.CellKField[ta.vpfloat], coeff2_dwdz: fa.CellKField[ta.vpfloat], diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 435b607a77..f40406b7ea 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -39,7 +39,7 @@ _init_two_cell_kdim_fields_with_zero_vp, ) from icon4py.model.atmosphere.dycore.stencils.interpolate_to_surface import _interpolate_to_surface -from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import Koff from icon4py.model.common.interpolation.stencils.interpolate_cell_field_to_half_levels_vp import ( _interpolate_cell_field_to_half_levels_vp, @@ -56,13 +56,13 @@ @gtx.field_operator def _calculate_pressure_buoyancy_acceleration_at_cells_on_half_levels( - exner_w_explicit_weight_parameter: fa.CellField[ta.wpfloat], - theta_v_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - perturbed_exner_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], - ddqz_z_half: fa.CellKField[ta.wpfloat], - perturbed_theta_v_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - ddz_of_reference_exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], -) -> fa.CellKField[ta.wpfloat]: + exner_w_explicit_weight_parameter: fa.CellField[wpfloat], + theta_v_at_cells_on_half_levels: fa.CellKField[wpfloat], + perturbed_exner_at_cells_on_model_levels: fa.CellKField[wpfloat], + ddqz_z_half: fa.CellKField[vpfloat], + perturbed_theta_v_at_cells_on_half_levels: fa.CellKField[vpfloat], + ddz_of_reference_exner_at_cells_on_half_levels: fa.CellKField[vpfloat], +) -> fa.CellKField[wpfloat]: return exner_w_explicit_weight_parameter * theta_v_at_cells_on_half_levels * ( perturbed_exner_at_cells_on_model_levels(Koff[-1]) - perturbed_exner_at_cells_on_model_levels @@ -74,31 +74,31 @@ def _calculate_pressure_buoyancy_acceleration_at_cells_on_half_levels( @gtx.field_operator def _compute_perturbed_quantities_and_interpolation( - current_rho: fa.CellKField[ta.wpfloat], - reference_rho_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], - current_theta_v: fa.CellKField[ta.wpfloat], - reference_theta_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], - wgtfac_c: fa.CellKField[ta.vpfloat], - exner_w_explicit_weight_parameter: fa.CellField[ta.wpfloat], - perturbed_exner_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], - ddz_of_reference_exner_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - ddqz_z_half: fa.CellKField[ta.vpfloat], - pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat], - theta_v_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], + current_rho: fa.CellKField[wpfloat], + reference_rho_at_cells_on_model_levels: fa.CellKField[vpfloat], + current_theta_v: fa.CellKField[wpfloat], + reference_theta_at_cells_on_model_levels: fa.CellKField[vpfloat], + wgtfac_c: fa.CellKField[vpfloat], + exner_w_explicit_weight_parameter: fa.CellField[wpfloat], + perturbed_exner_at_cells_on_model_levels: fa.CellKField[wpfloat], + ddz_of_reference_exner_at_cells_on_half_levels: fa.CellKField[vpfloat], + ddqz_z_half: fa.CellKField[vpfloat], + pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[vpfloat], + rho_at_cells_on_half_levels: fa.CellKField[wpfloat], + exner_at_cells_on_half_levels: fa.CellKField[vpfloat], + temporal_extrapolation_of_perturbed_exner: fa.CellKField[vpfloat], + theta_v_at_cells_on_half_levels: fa.CellKField[wpfloat], igradp_method: gtx.int32, nflatlev: gtx.int32, ) -> tuple[ - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.wpfloat], + fa.CellKField[vpfloat], + fa.CellKField[vpfloat], + fa.CellKField[vpfloat], + fa.CellKField[wpfloat], + fa.CellKField[wpfloat], + fa.CellKField[vpfloat], + fa.CellKField[vpfloat], + fa.CellKField[wpfloat], ]: exner_at_cells_on_half_levels = ( concat_where( @@ -135,7 +135,7 @@ def _compute_perturbed_quantities_and_interpolation( _interpolate_cell_field_to_half_levels_vp( wgtfac_c=wgtfac_c, interpolant=perturbed_theta_v_at_cells_on_model_levels ), - broadcast(0.0, (dims.CellDim, dims.KDim)), + broadcast(vpfloat(0.0), (dims.CellDim, dims.KDim)), ) theta_v_at_cells_on_half_levels = concat_where( @@ -175,12 +175,12 @@ def _compute_perturbed_quantities_and_interpolation( @gtx.field_operator def _surface_computations( - wgtfacq_c: fa.CellKField[ta.wpfloat], - exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], + wgtfacq_c: fa.CellKField[vpfloat], + exner_at_cells_on_half_levels: fa.CellKField[vpfloat], igradp_method: gtx.int32, ) -> tuple[ - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.vpfloat], + fa.CellKField[vpfloat], + fa.CellKField[vpfloat], ]: temporal_extrapolation_of_perturbed_exner = _init_cell_kdim_field_with_zero_wp() @@ -278,33 +278,33 @@ def _set_theta_v_and_exner_on_surface_level( @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_perturbed_quantities_and_interpolation( - temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat], - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[ta.vpfloat], - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[ta.vpfloat], - perturbed_exner_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], - exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - perturbed_rho_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - perturbed_theta_v_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - theta_v_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - current_rho: fa.CellKField[ta.wpfloat], - reference_rho_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - current_theta_v: fa.CellKField[ta.wpfloat], - reference_theta_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - reference_theta_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - wgtfacq_c: fa.CellKField[ta.vpfloat], - wgtfac_c: fa.CellKField[ta.vpfloat], - exner_w_explicit_weight_parameter: fa.CellField[ta.wpfloat], - ddz_of_reference_exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - ddqz_z_half: fa.CellKField[ta.vpfloat], - pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - time_extrapolation_parameter_for_exner: fa.CellKField[ta.vpfloat], - current_exner: fa.CellKField[ta.wpfloat], - reference_exner_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - inv_ddqz_z_full: fa.CellKField[ta.wpfloat], - d2dexdz2_fac1_mc: fa.CellKField[ta.vpfloat], - d2dexdz2_fac2_mc: fa.CellKField[ta.vpfloat], + temporal_extrapolation_of_perturbed_exner: fa.CellKField[vpfloat], + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[vpfloat], + d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[vpfloat], + perturbed_exner_at_cells_on_model_levels: fa.CellKField[wpfloat], + exner_at_cells_on_half_levels: fa.CellKField[vpfloat], + perturbed_rho_at_cells_on_model_levels: fa.CellKField[vpfloat], + perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[vpfloat], + rho_at_cells_on_half_levels: fa.CellKField[wpfloat], + perturbed_theta_v_at_cells_on_half_levels: fa.CellKField[vpfloat], + theta_v_at_cells_on_half_levels: fa.CellKField[wpfloat], + current_rho: fa.CellKField[wpfloat], + reference_rho_at_cells_on_model_levels: fa.CellKField[vpfloat], + current_theta_v: fa.CellKField[wpfloat], + reference_theta_at_cells_on_model_levels: fa.CellKField[vpfloat], + reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat], + wgtfacq_c: fa.CellKField[vpfloat], + wgtfac_c: fa.CellKField[vpfloat], + exner_w_explicit_weight_parameter: fa.CellField[wpfloat], + ddz_of_reference_exner_at_cells_on_half_levels: fa.CellKField[vpfloat], + ddqz_z_half: fa.CellKField[vpfloat], + pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[vpfloat], + time_extrapolation_parameter_for_exner: fa.CellKField[vpfloat], + current_exner: fa.CellKField[wpfloat], + reference_exner_at_cells_on_model_levels: fa.CellKField[vpfloat], + inv_ddqz_z_full: fa.CellKField[vpfloat], + d2dexdz2_fac1_mc: fa.CellKField[vpfloat], + d2dexdz2_fac2_mc: fa.CellKField[vpfloat], igradp_method: gtx.int32, nflatlev: gtx.int32, nflat_gradp: gtx.int32, @@ -497,26 +497,26 @@ def compute_perturbed_quantities_and_interpolation( @gtx.field_operator def _interpolate_rho_theta_v_to_half_levels_and_compute_pressure_buoyancy_acceleration( - w: fa.CellKField[ta.wpfloat], - contravariant_correction_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - current_rho: fa.CellKField[ta.wpfloat], - next_rho: fa.CellKField[ta.wpfloat], - current_theta_v: fa.CellKField[ta.wpfloat], - next_theta_v: fa.CellKField[ta.wpfloat], - perturbed_exner_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], - reference_theta_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - ddz_of_reference_exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - ddqz_z_half: fa.CellKField[ta.vpfloat], - wgtfac_c: fa.CellKField[ta.vpfloat], - exner_w_explicit_weight_parameter: fa.CellField[ta.wpfloat], - dtime: ta.wpfloat, - rhotheta_explicit_weight_parameter: ta.wpfloat, - rhotheta_implicit_weight_parameter: ta.wpfloat, + w: fa.CellKField[wpfloat], + contravariant_correction_at_cells_on_half_levels: fa.CellKField[vpfloat], + current_rho: fa.CellKField[wpfloat], + next_rho: fa.CellKField[wpfloat], + current_theta_v: fa.CellKField[wpfloat], + next_theta_v: fa.CellKField[wpfloat], + perturbed_exner_at_cells_on_model_levels: fa.CellKField[wpfloat], + reference_theta_at_cells_on_model_levels: fa.CellKField[vpfloat], + ddz_of_reference_exner_at_cells_on_half_levels: fa.CellKField[vpfloat], + ddqz_z_half: fa.CellKField[vpfloat], + wgtfac_c: fa.CellKField[vpfloat], + exner_w_explicit_weight_parameter: fa.CellField[wpfloat], + dtime: wpfloat, + rhotheta_explicit_weight_parameter: wpfloat, + rhotheta_implicit_weight_parameter: wpfloat, ) -> tuple[ - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.vpfloat], + fa.CellKField[wpfloat], + fa.CellKField[vpfloat], + fa.CellKField[wpfloat], + fa.CellKField[vpfloat], ]: ( contravariant_correction_at_cells_on_half_levels_wp, @@ -607,25 +607,25 @@ def _interpolate_rho_theta_v_to_half_levels_and_compute_pressure_buoyancy_accele @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def interpolate_rho_theta_v_to_half_levels_and_compute_pressure_buoyancy_acceleration( - rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - perturbed_theta_v_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - theta_v_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - w: fa.CellKField[ta.wpfloat], - contravariant_correction_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - current_rho: fa.CellKField[ta.wpfloat], - next_rho: fa.CellKField[ta.wpfloat], - current_theta_v: fa.CellKField[ta.wpfloat], - next_theta_v: fa.CellKField[ta.wpfloat], - perturbed_exner_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], - reference_theta_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], - ddz_of_reference_exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - ddqz_z_half: fa.CellKField[ta.vpfloat], - wgtfac_c: fa.CellKField[ta.vpfloat], - exner_w_explicit_weight_parameter: fa.CellField[ta.wpfloat], - dtime: ta.wpfloat, - rhotheta_explicit_weight_parameter: ta.wpfloat, - rhotheta_implicit_weight_parameter: ta.wpfloat, + rho_at_cells_on_half_levels: fa.CellKField[wpfloat], + perturbed_theta_v_at_cells_on_half_levels: fa.CellKField[vpfloat], + theta_v_at_cells_on_half_levels: fa.CellKField[wpfloat], + pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[vpfloat], + w: fa.CellKField[wpfloat], + contravariant_correction_at_cells_on_half_levels: fa.CellKField[vpfloat], + current_rho: fa.CellKField[wpfloat], + next_rho: fa.CellKField[wpfloat], + current_theta_v: fa.CellKField[wpfloat], + next_theta_v: fa.CellKField[wpfloat], + perturbed_exner_at_cells_on_model_levels: fa.CellKField[wpfloat], + reference_theta_at_cells_on_model_levels: fa.CellKField[vpfloat], + ddz_of_reference_exner_at_cells_on_half_levels: fa.CellKField[vpfloat], + ddqz_z_half: fa.CellKField[vpfloat], + wgtfac_c: fa.CellKField[vpfloat], + exner_w_explicit_weight_parameter: fa.CellField[wpfloat], + dtime: wpfloat, + rhotheta_explicit_weight_parameter: wpfloat, + rhotheta_implicit_weight_parameter: wpfloat, horizontal_start: gtx.int32, horizontal_end: gtx.int32, vertical_start: gtx.int32, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_virtual_potential_temperatures_and_pressure_gradient.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_virtual_potential_temperatures_and_pressure_gradient.py index 6d7063dd45..1eb943b31c 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_virtual_potential_temperatures_and_pressure_gradient.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_virtual_potential_temperatures_and_pressure_gradient.py @@ -8,7 +8,7 @@ import gt4py.next as gtx from gt4py.next import astype -from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import Koff from icon4py.model.common.interpolation.stencils.interpolate_cell_field_to_half_levels_vp import ( _interpolate_cell_field_to_half_levels_vp, @@ -21,17 +21,17 @@ @gtx.field_operator def _compute_virtual_potential_temperatures_and_pressure_gradient( - wgtfac_c: fa.CellKField[ta.vpfloat], - z_rth_pr_2: fa.CellKField[ta.vpfloat], - theta_v: fa.CellKField[ta.wpfloat], - vwind_expl_wgt: fa.CellField[ta.wpfloat], - exner_pr: fa.CellKField[ta.wpfloat], - d_exner_dz_ref_ic: fa.CellKField[ta.vpfloat], - ddqz_z_half: fa.CellKField[ta.vpfloat], + wgtfac_c: fa.CellKField[vpfloat], + z_rth_pr_2: fa.CellKField[vpfloat], + theta_v: fa.CellKField[wpfloat], + vwind_expl_wgt: fa.CellField[wpfloat], + exner_pr: fa.CellKField[wpfloat], + d_exner_dz_ref_ic: fa.CellKField[vpfloat], + ddqz_z_half: fa.CellKField[vpfloat], ) -> tuple[ - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.vpfloat], + fa.CellKField[vpfloat], + fa.CellKField[wpfloat], + fa.CellKField[vpfloat], ]: """Formerly known as _mo_solve_nonhydro_stencil_09.""" wgtfac_c_wp, ddqz_z_half_wp = astype((wgtfac_c, ddqz_z_half), wpfloat) @@ -50,16 +50,16 @@ def _compute_virtual_potential_temperatures_and_pressure_gradient( @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_virtual_potential_temperatures_and_pressure_gradient( - wgtfac_c: fa.CellKField[ta.vpfloat], - z_rth_pr_2: fa.CellKField[ta.vpfloat], - theta_v: fa.CellKField[ta.wpfloat], - vwind_expl_wgt: fa.CellField[ta.wpfloat], - exner_pr: fa.CellKField[ta.wpfloat], - d_exner_dz_ref_ic: fa.CellKField[ta.vpfloat], - ddqz_z_half: fa.CellKField[ta.vpfloat], - z_theta_v_pr_ic: fa.CellKField[ta.vpfloat], - theta_v_ic: fa.CellKField[ta.wpfloat], - z_th_ddz_exner_c: fa.CellKField[ta.vpfloat], + wgtfac_c: fa.CellKField[vpfloat], + z_rth_pr_2: fa.CellKField[vpfloat], + theta_v: fa.CellKField[wpfloat], + vwind_expl_wgt: fa.CellField[wpfloat], + exner_pr: fa.CellKField[wpfloat], + d_exner_dz_ref_ic: fa.CellKField[vpfloat], + ddqz_z_half: fa.CellKField[vpfloat], + z_theta_v_pr_ic: fa.CellKField[vpfloat], + theta_v_ic: fa.CellKField[wpfloat], + z_th_ddz_exner_c: fa.CellKField[vpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, vertical_start: gtx.int32, @@ -83,12 +83,12 @@ def compute_virtual_potential_temperatures_and_pressure_gradient( @gtx.field_operator def _compute_virtual_potential_temperatures( - wgtfac_c: fa.CellKField[ta.vpfloat], - z_rth_pr_2: fa.CellKField[ta.vpfloat], - theta_v: fa.CellKField[ta.wpfloat], + wgtfac_c: fa.CellKField[vpfloat], + z_rth_pr_2: fa.CellKField[vpfloat], + theta_v: fa.CellKField[wpfloat], ) -> tuple[ - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.wpfloat], + fa.CellKField[vpfloat], + fa.CellKField[wpfloat], ]: wgtfac_c_wp = astype(wgtfac_c, wpfloat) @@ -101,13 +101,13 @@ def _compute_virtual_potential_temperatures( @gtx.field_operator def _compute_pressure_gradient( - vwind_expl_wgt: fa.CellField[ta.wpfloat], - theta_v_ic: fa.CellKField[ta.wpfloat], - z_theta_v_pr_ic: fa.CellKField[ta.wpfloat], - exner_pr: fa.CellKField[ta.wpfloat], - d_exner_dz_ref_ic: fa.CellKField[ta.vpfloat], - ddqz_z_half: fa.CellKField[ta.vpfloat], -) -> fa.CellKField[ta.vpfloat]: + vwind_expl_wgt: fa.CellField[wpfloat], + theta_v_ic: fa.CellKField[wpfloat], + z_theta_v_pr_ic: fa.CellKField[vpfloat], + exner_pr: fa.CellKField[wpfloat], + d_exner_dz_ref_ic: fa.CellKField[vpfloat], + ddqz_z_half: fa.CellKField[vpfloat], +) -> fa.CellKField[vpfloat]: ddqz_z_half_wp = astype(ddqz_z_half, wpfloat) z_th_ddz_exner_c_wp = vwind_expl_wgt * theta_v_ic * ( exner_pr(Koff[-1]) - exner_pr diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_cell_kdim_field_with_zero_wp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_cell_kdim_field_with_zero_wp.py index 6439b9248c..f5b84a3910 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_cell_kdim_field_with_zero_wp.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_cell_kdim_field_with_zero_wp.py @@ -9,7 +9,13 @@ from gt4py.next import broadcast from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.type_alias import wpfloat +from icon4py.model.common.type_alias import vpfloat, wpfloat + + +@gtx.field_operator +def _init_cell_kdim_field_with_zero_vp() -> fa.CellKField[vpfloat]: + """Formerly known as _mo_solve_nonhydro_stencil_57 or _mo_solve_nonhydro_stencil_64.""" + return broadcast(vpfloat("0.0"), (dims.CellDim, dims.KDim)) @gtx.field_operator diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/solve_tridiagonal_matrix_for_w_forward_sweep.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/solve_tridiagonal_matrix_for_w_forward_sweep.py index c5a20860e5..714a0afa4d 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/solve_tridiagonal_matrix_for_w_forward_sweep.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/solve_tridiagonal_matrix_for_w_forward_sweep.py @@ -18,11 +18,11 @@ forward=True, init=( vpfloat("0.0"), - 0.0, + wpfloat("0.0"), ), # boundary condition for upper tridiagonal element and w at model top ) def tridiagonal_forward_sweep_for_w( - state_kminus1: tuple[vpfloat, float], + state_kminus1: tuple[vpfloat, wpfloat], a: vpfloat, b: vpfloat, c: vpfloat, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/vertically_implicit_dycore_solver.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/vertically_implicit_dycore_solver.py index 462bb4c220..a55de17b09 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/vertically_implicit_dycore_solver.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/vertically_implicit_dycore_solver.py @@ -48,12 +48,7 @@ from icon4py.model.atmosphere.dycore.stencils.update_mass_volume_flux import ( _update_mass_volume_flux, ) -from icon4py.model.common import ( - constants, - dimension as dims, - field_type_aliases as fa, - type_alias as ta, -) +from icon4py.model.common import constants, dimension as dims, field_type_aliases as fa from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -83,8 +78,8 @@ def _interpolate_contravariant_correction_from_edges_on_model_levels_to_cells_on @gtx.field_operator def _set_surface_boundary_condition_for_computation_of_w( - contravariant_correction_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], -) -> fa.CellKField[ta.wpfloat]: + contravariant_correction_at_cells_on_half_levels: fa.CellKField[vpfloat], +) -> fa.CellKField[wpfloat]: return astype(contravariant_correction_at_cells_on_half_levels, wpfloat) @@ -218,34 +213,34 @@ def solve_w( @gtx.field_operator def _vertically_implicit_solver_at_predictor_step( next_w: fa.CellKField[ - ta.wpfloat + wpfloat ], # necessary input because the last vertical level is set outside this field operator - geofac_div: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], - mass_flux_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat], - theta_v_flux_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat], - predictor_vertical_wind_advective_tendency: fa.CellKField[ta.vpfloat], - pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - contravariant_correction_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - exner_w_explicit_weight_parameter: fa.CellField[ta.wpfloat], - current_exner: fa.CellKField[ta.wpfloat], - current_rho: fa.CellKField[ta.wpfloat], - current_theta_v: fa.CellKField[ta.wpfloat], - current_w: fa.CellKField[ta.wpfloat], - inv_ddqz_z_full: fa.CellKField[ta.vpfloat], - exner_w_implicit_weight_parameter: fa.CellField[ta.wpfloat], - theta_v_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - perturbed_exner_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], - exner_tendency_due_to_slow_physics: fa.CellKField[ta.vpfloat], - rho_iau_increment: fa.CellKField[ta.vpfloat], - exner_iau_increment: fa.CellKField[ta.vpfloat], - ddqz_z_half: fa.CellKField[ta.vpfloat], - exner_dynamical_increment: fa.CellKField[ta.wpfloat], - dwdz_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - rayleigh_damping_factor: fa.KField[ta.wpfloat], - reference_exner_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - iau_wgt_dyn: ta.wpfloat, - dtime: ta.wpfloat, + geofac_div: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], wpfloat], + mass_flux_at_edges_on_model_levels: fa.EdgeKField[wpfloat], + theta_v_flux_at_edges_on_model_levels: fa.EdgeKField[wpfloat], + predictor_vertical_wind_advective_tendency: fa.CellKField[vpfloat], + pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[vpfloat], + rho_at_cells_on_half_levels: fa.CellKField[wpfloat], + contravariant_correction_at_cells_on_half_levels: fa.CellKField[vpfloat], + exner_w_explicit_weight_parameter: fa.CellField[wpfloat], + current_exner: fa.CellKField[wpfloat], + current_rho: fa.CellKField[wpfloat], + current_theta_v: fa.CellKField[wpfloat], + current_w: fa.CellKField[wpfloat], + inv_ddqz_z_full: fa.CellKField[vpfloat], + exner_w_implicit_weight_parameter: fa.CellField[wpfloat], + theta_v_at_cells_on_half_levels: fa.CellKField[wpfloat], + perturbed_exner_at_cells_on_model_levels: fa.CellKField[wpfloat], + exner_tendency_due_to_slow_physics: fa.CellKField[vpfloat], + rho_iau_increment: fa.CellKField[vpfloat], + exner_iau_increment: fa.CellKField[vpfloat], + ddqz_z_half: fa.CellKField[vpfloat], + exner_dynamical_increment: fa.CellKField[vpfloat], + dwdz_at_cells_on_model_levels: fa.CellKField[vpfloat], + rayleigh_damping_factor: fa.KField[wpfloat], + reference_exner_at_cells_on_model_levels: fa.CellKField[vpfloat], + iau_wgt_dyn: wpfloat, + dtime: wpfloat, rayleigh_type: gtx.int32, divdamp_type: gtx.int32, is_iau_active: bool, @@ -254,12 +249,12 @@ def _vertically_implicit_solver_at_predictor_step( kstart_moist: gtx.int32, n_lev: gtx.int32, ) -> tuple[ - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], + fa.CellKField[wpfloat], + fa.CellKField[vpfloat], + fa.CellKField[vpfloat], + fa.CellKField[wpfloat], + fa.CellKField[wpfloat], + fa.CellKField[wpfloat], ]: divergence_of_mass, divergence_of_theta_v = _compute_divergence_of_fluxes_of_rho_and_theta( geofac_div=geofac_div, @@ -403,40 +398,40 @@ def _vertically_implicit_solver_at_predictor_step( @gtx.program def vertically_implicit_solver_at_predictor_step( - contravariant_correction_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - next_w: fa.CellKField[ta.wpfloat], - next_rho: fa.CellKField[ta.wpfloat], - next_exner: fa.CellKField[ta.wpfloat], - next_theta_v: fa.CellKField[ta.wpfloat], - dwdz_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - exner_dynamical_increment: fa.CellKField[ta.vpfloat], - geofac_div: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], - mass_flux_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat], - theta_v_flux_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat], - predictor_vertical_wind_advective_tendency: fa.CellKField[ta.vpfloat], - pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - contravariant_correction_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], - exner_w_explicit_weight_parameter: fa.CellField[ta.wpfloat], - current_exner: fa.CellKField[ta.wpfloat], - current_rho: fa.CellKField[ta.wpfloat], - current_theta_v: fa.CellKField[ta.wpfloat], - current_w: fa.CellKField[ta.wpfloat], - inv_ddqz_z_full: fa.CellKField[ta.vpfloat], - exner_w_implicit_weight_parameter: fa.CellField[ta.wpfloat], - theta_v_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - perturbed_exner_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], - exner_tendency_due_to_slow_physics: fa.CellKField[ta.vpfloat], - rho_iau_increment: fa.CellKField[ta.vpfloat], - exner_iau_increment: fa.CellKField[ta.vpfloat], - ddqz_z_half: fa.CellKField[ta.vpfloat], - rayleigh_damping_factor: fa.KField[ta.wpfloat], - reference_exner_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], + contravariant_correction_at_cells_on_half_levels: fa.CellKField[vpfloat], + next_w: fa.CellKField[wpfloat], + next_rho: fa.CellKField[wpfloat], + next_exner: fa.CellKField[wpfloat], + next_theta_v: fa.CellKField[wpfloat], + dwdz_at_cells_on_model_levels: fa.CellKField[vpfloat], + exner_dynamical_increment: fa.CellKField[vpfloat], + geofac_div: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], wpfloat], + mass_flux_at_edges_on_model_levels: fa.EdgeKField[wpfloat], + theta_v_flux_at_edges_on_model_levels: fa.EdgeKField[wpfloat], + predictor_vertical_wind_advective_tendency: fa.CellKField[vpfloat], + pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[vpfloat], + rho_at_cells_on_half_levels: fa.CellKField[wpfloat], + contravariant_correction_at_edges_on_model_levels: fa.EdgeKField[vpfloat], + exner_w_explicit_weight_parameter: fa.CellField[wpfloat], + current_exner: fa.CellKField[wpfloat], + current_rho: fa.CellKField[wpfloat], + current_theta_v: fa.CellKField[wpfloat], + current_w: fa.CellKField[wpfloat], + inv_ddqz_z_full: fa.CellKField[vpfloat], + exner_w_implicit_weight_parameter: fa.CellField[wpfloat], + theta_v_at_cells_on_half_levels: fa.CellKField[wpfloat], + perturbed_exner_at_cells_on_model_levels: fa.CellKField[wpfloat], + exner_tendency_due_to_slow_physics: fa.CellKField[vpfloat], + rho_iau_increment: fa.CellKField[vpfloat], + exner_iau_increment: fa.CellKField[vpfloat], + ddqz_z_half: fa.CellKField[vpfloat], + rayleigh_damping_factor: fa.KField[wpfloat], + reference_exner_at_cells_on_model_levels: fa.CellKField[vpfloat], e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], wpfloat], wgtfac_c: fa.CellKField[vpfloat], wgtfacq_c: fa.CellKField[vpfloat], - iau_wgt_dyn: ta.wpfloat, - dtime: ta.wpfloat, + iau_wgt_dyn: wpfloat, + dtime: wpfloat, is_iau_active: bool, rayleigh_type: gtx.int32, divdamp_type: gtx.int32, @@ -527,41 +522,41 @@ def vertically_implicit_solver_at_predictor_step( @gtx.field_operator def _vertically_implicit_solver_at_corrector_step( next_w: fa.CellKField[ - ta.wpfloat + wpfloat ], # necessary input because the last vertical level is set outside this field operator - dynamical_vertical_mass_flux_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - dynamical_vertical_volumetric_flux_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - exner_dynamical_increment: fa.CellKField[ta.wpfloat], - geofac_div: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], - mass_flux_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat], - theta_v_flux_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat], - predictor_vertical_wind_advective_tendency: fa.CellKField[ta.vpfloat], - corrector_vertical_wind_advective_tendency: fa.CellKField[ta.vpfloat], - pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - contravariant_correction_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - exner_w_explicit_weight_parameter: fa.CellField[ta.wpfloat], - current_exner: fa.CellKField[ta.wpfloat], - current_rho: fa.CellKField[ta.wpfloat], - current_theta_v: fa.CellKField[ta.wpfloat], - current_w: fa.CellKField[ta.wpfloat], - inv_ddqz_z_full: fa.CellKField[ta.vpfloat], - exner_w_implicit_weight_parameter: fa.CellField[ta.wpfloat], - theta_v_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - perturbed_exner_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], - exner_tendency_due_to_slow_physics: fa.CellKField[ta.vpfloat], - rho_iau_increment: fa.CellKField[ta.vpfloat], - exner_iau_increment: fa.CellKField[ta.vpfloat], - ddqz_z_half: fa.CellKField[ta.vpfloat], - rayleigh_damping_factor: fa.KField[ta.wpfloat], - reference_exner_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - advection_explicit_weight_parameter: ta.wpfloat, - advection_implicit_weight_parameter: ta.wpfloat, + dynamical_vertical_mass_flux_at_cells_on_half_levels: fa.CellKField[wpfloat], + dynamical_vertical_volumetric_flux_at_cells_on_half_levels: fa.CellKField[wpfloat], + exner_dynamical_increment: fa.CellKField[vpfloat], + geofac_div: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], wpfloat], + mass_flux_at_edges_on_model_levels: fa.EdgeKField[wpfloat], + theta_v_flux_at_edges_on_model_levels: fa.EdgeKField[wpfloat], + predictor_vertical_wind_advective_tendency: fa.CellKField[vpfloat], + corrector_vertical_wind_advective_tendency: fa.CellKField[vpfloat], + pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[vpfloat], + rho_at_cells_on_half_levels: fa.CellKField[wpfloat], + contravariant_correction_at_cells_on_half_levels: fa.CellKField[vpfloat], + exner_w_explicit_weight_parameter: fa.CellField[wpfloat], + current_exner: fa.CellKField[wpfloat], + current_rho: fa.CellKField[wpfloat], + current_theta_v: fa.CellKField[wpfloat], + current_w: fa.CellKField[wpfloat], + inv_ddqz_z_full: fa.CellKField[vpfloat], + exner_w_implicit_weight_parameter: fa.CellField[wpfloat], + theta_v_at_cells_on_half_levels: fa.CellKField[wpfloat], + perturbed_exner_at_cells_on_model_levels: fa.CellKField[wpfloat], + exner_tendency_due_to_slow_physics: fa.CellKField[vpfloat], + rho_iau_increment: fa.CellKField[vpfloat], + exner_iau_increment: fa.CellKField[vpfloat], + ddqz_z_half: fa.CellKField[vpfloat], + rayleigh_damping_factor: fa.KField[wpfloat], + reference_exner_at_cells_on_model_levels: fa.CellKField[vpfloat], + advection_explicit_weight_parameter: wpfloat, + advection_implicit_weight_parameter: wpfloat, lprep_adv: bool, - r_nsubsteps: ta.wpfloat, - ndyn_substeps_var: ta.wpfloat, - iau_wgt_dyn: ta.wpfloat, - dtime: ta.wpfloat, + r_nsubsteps: wpfloat, + ndyn_substeps_var: wpfloat, + iau_wgt_dyn: wpfloat, + dtime: wpfloat, is_iau_active: bool, rayleigh_type: gtx.int32, at_first_substep: bool, @@ -570,13 +565,13 @@ def _vertically_implicit_solver_at_corrector_step( kstart_moist: gtx.int32, n_lev: gtx.int32, ) -> tuple[ - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], + fa.CellKField[wpfloat], + fa.CellKField[vpfloat], + fa.CellKField[vpfloat], + fa.CellKField[wpfloat], + fa.CellKField[wpfloat], + fa.CellKField[wpfloat], + fa.CellKField[wpfloat], ]: divergence_of_mass, divergence_of_theta_v = _compute_divergence_of_fluxes_of_rho_and_theta( geofac_div=geofac_div, @@ -743,43 +738,43 @@ def _vertically_implicit_solver_at_corrector_step( @gtx.program def vertically_implicit_solver_at_corrector_step( - next_w: fa.CellKField[ta.wpfloat], - next_rho: fa.CellKField[ta.wpfloat], - next_exner: fa.CellKField[ta.wpfloat], - next_theta_v: fa.CellKField[ta.wpfloat], - dynamical_vertical_mass_flux_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - dynamical_vertical_volumetric_flux_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - exner_dynamical_increment: fa.CellKField[ta.wpfloat], - geofac_div: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], - mass_flux_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat], - theta_v_flux_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat], - predictor_vertical_wind_advective_tendency: fa.CellKField[ta.vpfloat], - corrector_vertical_wind_advective_tendency: fa.CellKField[ta.vpfloat], - pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - contravariant_correction_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - exner_w_explicit_weight_parameter: fa.CellField[ta.wpfloat], - current_exner: fa.CellKField[ta.wpfloat], - current_rho: fa.CellKField[ta.wpfloat], - current_theta_v: fa.CellKField[ta.wpfloat], - current_w: fa.CellKField[ta.wpfloat], - inv_ddqz_z_full: fa.CellKField[ta.vpfloat], - exner_w_implicit_weight_parameter: fa.CellField[ta.wpfloat], - theta_v_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - perturbed_exner_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], - exner_tendency_due_to_slow_physics: fa.CellKField[ta.vpfloat], - rho_iau_increment: fa.CellKField[ta.vpfloat], - exner_iau_increment: fa.CellKField[ta.vpfloat], - ddqz_z_half: fa.CellKField[ta.vpfloat], - rayleigh_damping_factor: fa.KField[ta.wpfloat], - reference_exner_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - advection_explicit_weight_parameter: ta.wpfloat, - advection_implicit_weight_parameter: ta.wpfloat, + next_w: fa.CellKField[wpfloat], + next_rho: fa.CellKField[wpfloat], + next_exner: fa.CellKField[wpfloat], + next_theta_v: fa.CellKField[wpfloat], + dynamical_vertical_mass_flux_at_cells_on_half_levels: fa.CellKField[wpfloat], + dynamical_vertical_volumetric_flux_at_cells_on_half_levels: fa.CellKField[wpfloat], + exner_dynamical_increment: fa.CellKField[vpfloat], + geofac_div: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], wpfloat], + mass_flux_at_edges_on_model_levels: fa.EdgeKField[wpfloat], + theta_v_flux_at_edges_on_model_levels: fa.EdgeKField[wpfloat], + predictor_vertical_wind_advective_tendency: fa.CellKField[vpfloat], + corrector_vertical_wind_advective_tendency: fa.CellKField[vpfloat], + pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[vpfloat], + rho_at_cells_on_half_levels: fa.CellKField[wpfloat], + contravariant_correction_at_cells_on_half_levels: fa.CellKField[vpfloat], + exner_w_explicit_weight_parameter: fa.CellField[wpfloat], + current_exner: fa.CellKField[wpfloat], + current_rho: fa.CellKField[wpfloat], + current_theta_v: fa.CellKField[wpfloat], + current_w: fa.CellKField[wpfloat], + inv_ddqz_z_full: fa.CellKField[vpfloat], + exner_w_implicit_weight_parameter: fa.CellField[wpfloat], + theta_v_at_cells_on_half_levels: fa.CellKField[wpfloat], + perturbed_exner_at_cells_on_model_levels: fa.CellKField[wpfloat], + exner_tendency_due_to_slow_physics: fa.CellKField[vpfloat], + rho_iau_increment: fa.CellKField[vpfloat], + exner_iau_increment: fa.CellKField[vpfloat], + ddqz_z_half: fa.CellKField[vpfloat], + rayleigh_damping_factor: fa.KField[wpfloat], + reference_exner_at_cells_on_model_levels: fa.CellKField[vpfloat], + advection_explicit_weight_parameter: wpfloat, + advection_implicit_weight_parameter: wpfloat, lprep_adv: bool, - r_nsubsteps: ta.wpfloat, - ndyn_substeps_var: ta.wpfloat, - iau_wgt_dyn: ta.wpfloat, - dtime: ta.wpfloat, + r_nsubsteps: wpfloat, + ndyn_substeps_var: wpfloat, + iau_wgt_dyn: wpfloat, + dtime: wpfloat, is_iau_active: bool, rayleigh_type: gtx.int32, at_first_substep: bool, 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 3d1c54dfad..fa703721bf 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 @@ -62,8 +61,8 @@ def __init__( self.edge_params = edge_params self.c_owner_mask = owner_mask - self.cfl_w_limit: float = 0.65 - self.scalfac_exdiff: float = 0.05 + self.cfl_w_limit: ta.vpfloat = ta.vpfloat(0.65) + self.scalfac_exdiff: ta.wpfloat = ta.wpfloat(0.05) self._allocate_local_fields(model_backends.get_allocator(backend)) self._determine_local_domains() @@ -233,6 +232,21 @@ def _determine_local_domains(self): self._end_cell_local = self.grid.end_index(cell_domain(h_grid.Zone.LOCAL)) self._end_cell_halo = self.grid.end_index(cell_domain(h_grid.Zone.HALO)) + def _get_max_vertical_cfl(self): + # Reductions should be performed on flat, contiguous arrays for best cupy performance + # as otherwise cupy won't use cub optimized kernels. + max_xp = ( + self.vertical_cfl.ndarray[ + self._start_cell_lateral_boundary_level_4 : self._end_cell_halo, : + ] + .ravel(order="K") + .max() + ) + if self.vertical_cfl.array_ns.__name__ == "cupy": + return ta.vpfloat(max_xp.get()) + else: + return max_xp + def run_predictor_step( self, skip_compute_predictor_vertical_advection: bool, @@ -289,15 +303,8 @@ def run_predictor_step( skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection, ) - # Reductions should be performed on flat, contiguous arrays for best cupy performance - # as otherwise cupy won't use cub optimized kernels. - max_vertical_cfl = float( - self.vertical_cfl.array_ns.max( - self.vertical_cfl.ndarray[ - self._start_cell_lateral_boundary_level_4 : self._end_cell_halo, : - ].ravel(order="K") - ) - ) + max_vertical_cfl = self._get_max_vertical_cfl() + diagnostic_state.max_vertical_cfl = max(max_vertical_cfl, diagnostic_state.max_vertical_cfl) apply_extra_diffusion_on_vn = max_vertical_cfl > cfl_w_limit * dtime self._compute_advection_in_horizontal_momentum_equation( @@ -314,9 +321,11 @@ def run_predictor_step( ) def _scale_factors_by_dtime(self, dtime): - scaled_cfl_w_limit = self.cfl_w_limit / dtime - scalfac_exdiff = self.scalfac_exdiff / (dtime * (0.85 - scaled_cfl_w_limit * dtime)) - return scaled_cfl_w_limit, scalfac_exdiff + scaled_cfl_w_limit = gtx.astype(self.cfl_w_limit, ta.wpfloat) / dtime + scalfac_exdiff = self.scalfac_exdiff / ( + dtime * (ta.wpfloat(0.85) - scaled_cfl_w_limit * dtime) + ) + return gtx.astype(scaled_cfl_w_limit, ta.vpfloat), scalfac_exdiff def run_corrector_step( self, @@ -356,15 +365,8 @@ def run_corrector_step( dtime=dtime, ) - # Reductions should be performed on flat, contiguous arrays for best cupy performance - # as otherwise cupy won't use cub optimized kernels. - max_vertical_cfl = float( - self.vertical_cfl.array_ns.max( - self.vertical_cfl.ndarray[ - self._start_cell_lateral_boundary_level_4 : self._end_cell_halo, : - ].ravel(order="K") - ) - ) + max_vertical_cfl = self._get_max_vertical_cfl() + diagnostic_state.max_vertical_cfl = max(max_vertical_cfl, diagnostic_state.max_vertical_cfl) apply_extra_diffusion_on_vn = max_vertical_cfl > cfl_w_limit * dtime self._compute_advection_in_horizontal_momentum_equation( 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 e2711f3c82..8667604aef 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 @@ -25,6 +25,7 @@ from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import horizontal as h_grid, vertical as v_grid from icon4py.model.common.math import smagorinsky +from icon4py.model.common.type_alias import vpfloat, wpfloat from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import definitions, test_utils @@ -91,7 +92,7 @@ def test_validate_divdamp_fields_against_savepoint_values( )( fourth_order_divdamp_scaling_coeff, config.max_nudging_coefficient, - constants.DBL_EPS, + constants.WP_EPS, out=reduced_fourth_order_divdamp_coeff_at_nest_boundary, offset_provider={}, ) @@ -190,7 +191,7 @@ def test_nonhydro_predictor_step( rayleigh_damping_height=damping_height, ) vertical_params = utils.create_vertical_params(vertical_config, grid_savepoint) - dtime = sp.get_metadata("dtime").get("dtime") + dtime = sp.dtime() diagnostic_state_nh = utils.construct_diagnostics(sp, icon_grid, backend) @@ -519,7 +520,7 @@ def test_nonhydro_corrector_step( rayleigh_damping_height=damping_height, ) vertical_params = utils.create_vertical_params(vertical_config, grid_savepoint) - dtime = init_savepoint.get_metadata("dtime").get("dtime") + dtime = init_savepoint.dtime() lprep_adv = init_savepoint.get_metadata("prep_adv").get("prep_adv") prep_adv = dycore_states.PrepAdvection( vn_traj=init_savepoint.vn_traj(), @@ -729,7 +730,7 @@ def test_run_solve_nonhydro_single_step( rayleigh_damping_height=damping_height, ) vertical_params = utils.create_vertical_params(vertical_config, grid_savepoint) - dtime = sp.get_metadata("dtime").get("dtime") + dtime = sp.dtime() lprep_adv = sp.get_metadata("prep_adv").get("prep_adv") prep_adv = dycore_states.PrepAdvection( vn_traj=sp.vn_traj(), @@ -857,7 +858,7 @@ def test_run_solve_nonhydro_multi_step( rayleigh_damping_height=damping_height, ) vertical_params = utils.create_vertical_params(vertical_config, grid_savepoint) - dtime = sp.get_metadata("dtime").get("dtime") + dtime = sp.dtime() lprep_adv = sp.get_metadata("prep_adv").get("prep_adv") prep_adv = dycore_states.PrepAdvection( vn_traj=sp.vn_traj(), @@ -1258,7 +1259,7 @@ def test_interpolate_rho_theta_v_to_half_levels_and_compute_pressure_buoyancy_ac sp_ref = savepoint_compute_edge_diagnostics_for_dycore_and_update_vn_init sp_exit = savepoint_nonhydro_exit - dtime = sp_init.get_metadata("dtime").get("dtime") + dtime = sp_init.dtime() current_rho = sp_init.rho_now() next_rho = sp_init.rho_new() @@ -1273,7 +1274,7 @@ def test_interpolate_rho_theta_v_to_half_levels_and_compute_pressure_buoyancy_ac rhotheta_implicit_weight_parameter = sp_init.wgt_nnew_rth() perturbed_theta_v_at_cells_on_half_levels = data_alloc.zero_field( - icon_grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, allocator=backend + icon_grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, allocator=backend, dtype=vpfloat ) pressure_buoyancy_acceleration_at_cells_on_half_levels = data_alloc.zero_field( icon_grid, dims.CellDim, dims.KDim, allocator=backend @@ -1526,7 +1527,7 @@ def test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn( ipeidx_dsl=metrics_savepoint.pg_edgeidx_dsl(), pg_exdist=metrics_savepoint.pg_exdist(), inv_dual_edge_length=grid_savepoint.inv_dual_edge_length(), - dtime=savepoint_nonhydro_init.get_metadata("dtime").get("dtime"), + dtime=savepoint_nonhydro_init.dtime(), iau_wgt_dyn=iau_wgt_dyn, is_iau_active=is_iau_active, limited_area=grid_savepoint.get_metadata("limited_area").get("limited_area"), @@ -1678,7 +1679,7 @@ def test_apply_divergence_damping_and_update_vn( geofac_grdiv=interpolation_savepoint.geofac_grdiv(), advection_explicit_weight_parameter=savepoint_nonhydro_init.wgt_nnow_vel(), advection_implicit_weight_parameter=savepoint_nonhydro_init.wgt_nnew_vel(), - dtime=savepoint_nonhydro_init.get_metadata("dtime").get("dtime"), + dtime=savepoint_nonhydro_init.dtime(), iau_wgt_dyn=iau_wgt_dyn, is_iau_active=is_iau_active, limited_area=grid_savepoint.get_metadata("limited_area").get("limited_area"), @@ -2145,7 +2146,7 @@ def test_vertically_implicit_solver_at_predictor_step( wgtfac_c=metrics_savepoint.wgtfac_c(), wgtfacq_c=metrics_savepoint.wgtfacq_c_dsl(), iau_wgt_dyn=iau_wgt_dyn, - dtime=savepoint_nonhydro_init.get_metadata("dtime").get("dtime"), + dtime=savepoint_nonhydro_init.dtime(), is_iau_active=is_iau_active, rayleigh_type=config.rayleigh_type, divdamp_type=divdamp_type, @@ -2352,7 +2353,7 @@ def test_vertically_implicit_solver_at_corrector_step( r_nsubsteps=r_nsubsteps, ndyn_substeps_var=float(ndyn_substeps), iau_wgt_dyn=iau_wgt_dyn, - dtime=savepoint_nonhydro_init.get_metadata("dtime").get("dtime"), + dtime=savepoint_nonhydro_init.dtime(), is_iau_active=is_iau_active, rayleigh_type=config.rayleigh_type, at_first_substep=at_first_substep, 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 c2df34eb8c..f01c902121 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 @@ -31,13 +31,18 @@ vertical as v_grid, ) from icon4py.model.common.states import prognostic_state as prognostics +from icon4py.model.common.type_alias import vpfloat from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import definitions, serialbox, test_utils +from icon4py.model.testing.test_utils import vp_eps, wp_eps from .. import utils from ..fixtures import * # noqa: F403 +atol_2eps = 2 * vp_eps # for double ≈ 4.44e-16, for single ≈ 2.38e-7 +rtol_8eps = 8 * vp_eps # for double ≈ 1.78e-15, for single ≈ 9.54e-7 + log = logging.getLogger(__name__) @@ -49,6 +54,8 @@ def _compare_cfl( horizontal_end: int, vertical_start: int, vertical_end: int, + rtol: vpfloat = rtol_8eps, + atol: vpfloat = atol_2eps, ) -> None: cfl_clipping_mask = np.where(np.abs(vertical_cfl) > 0.0, True, False) assert ( @@ -56,7 +63,10 @@ def _compare_cfl( == icon_result_cfl_clipping[horizontal_start:horizontal_end, vertical_start:vertical_end] ).all() - assert vertical_cfl[horizontal_start:horizontal_end, :].max() == icon_result_max_vcfl_dyn + assert ( + np.abs(vertical_cfl[horizontal_start:horizontal_end, :].max() - icon_result_max_vcfl_dyn) + <= atol + rtol * icon_result_max_vcfl_dyn + ) def create_vertical_params( @@ -137,7 +147,7 @@ def test_scale_factors_by_dtime( damping_height, backend, ): - dtime = savepoint_velocity_init.get_metadata("dtime").get("dtime") + dtime = savepoint_velocity_init.dtime() interpolation_state = utils.construct_interpolation_state(interpolation_savepoint) metric_state_nonhydro = utils.construct_metric_state(metrics_savepoint, grid_savepoint) vertical_config = v_grid.VerticalGridConfig( @@ -202,7 +212,7 @@ def test_velocity_predictor_step( caplog.set_level(logging.WARN) init_savepoint = savepoint_velocity_init vn_only = init_savepoint.vn_only() - dtime = init_savepoint.get_metadata("dtime").get("dtime") + dtime = init_savepoint.dtime() diagnostic_state = dycore_states.DiagnosticStateNonHydro( max_vertical_cfl=0.0, @@ -357,7 +367,7 @@ def test_velocity_corrector_step( ): init_savepoint = savepoint_velocity_init vn_only = init_savepoint.vn_only() - dtime = init_savepoint.get_metadata("dtime").get("dtime") + dtime = init_savepoint.dtime() assert not vn_only @@ -646,7 +656,7 @@ def test_compute_contravariant_correction_and_advection_in_vertical_momentum_equ end_index_of_damping_layer = grid_savepoint.nrdmax() - dtime = savepoint_velocity_init.get_metadata("dtime").get("dtime") + dtime = savepoint_velocity_init.dtime() cell_domain = h_grid.domain(dims.CellDim) start_cell_nudging_for_vertical_wind_advective_tendency = icon_grid.start_index( cell_domain(h_grid.Zone.NUDGING) @@ -752,6 +762,7 @@ def test_compute_contravariant_correction_and_advection_in_vertical_momentum_equ (definitions.Experiments.EXCLAIM_APE, "2000-01-01T00:00:02.000", "2000-01-01T00:00:02.000"), ], ) +@pytest.mark.single_precision_ready @pytest.mark.parametrize("istep_init, istep_exit", [(2, 2)]) def test_compute_advection_in_vertical_momentum_equation( experiment, @@ -799,7 +810,7 @@ def test_compute_advection_in_vertical_momentum_equation( end_index_of_damping_layer = grid_savepoint.nrdmax() - dtime = savepoint_velocity_init.get_metadata("dtime").get("dtime") + dtime = savepoint_velocity_init.dtime() cell_domain = h_grid.domain(dims.CellDim) start_cell_nudging_for_vertical_wind_advective_tendency = icon_grid.start_index( cell_domain(h_grid.Zone.NUDGING) @@ -851,22 +862,17 @@ def test_compute_advection_in_vertical_momentum_equation( assert test_utils.dallclose( icon_result_z_w_con_c_full.asnumpy(), contravariant_corrected_w_at_cells_on_model_levels.asnumpy(), - rtol=1.0e-15, - atol=1.0e-15, - ) - assert test_utils.dallclose( - icon_result_ddt_w_adv.asnumpy()[ - start_cell_nudging_for_vertical_wind_advective_tendency:end_cell_local_for_vertical_wind_advective_tendency, - :, - ], - vertical_wind_advective_tendency.asnumpy()[ - start_cell_nudging_for_vertical_wind_advective_tendency:end_cell_local_for_vertical_wind_advective_tendency, - :, - ], - rtol=1.0e-15, - atol=1.0e-15, + rtol=rtol_8eps, + atol=atol_2eps, ) + start_idx = start_cell_nudging_for_vertical_wind_advective_tendency + end_idx = end_cell_local_for_vertical_wind_advective_tendency + fortran_res = icon_result_ddt_w_adv[start_idx:end_idx, :].asnumpy() + icon4py_res = vertical_wind_advective_tendency[start_idx:end_idx, :].asnumpy() + + assert test_utils.dallclose(fortran_res, icon4py_res, rtol=rtol_8eps, atol=atol_2eps) + # TODO(OngChia): currently direct comparison of vcfl_dsl is not possible because it is not properly updated in icon run _compare_cfl( vertical_cfl.asnumpy(), @@ -892,6 +898,7 @@ def test_compute_advection_in_vertical_momentum_equation( (definitions.Experiments.EXCLAIM_APE, "2000-01-01T00:00:02.000", "2000-01-01T00:00:02.000"), ], ) +@pytest.mark.single_precision_ready @pytest.mark.parametrize("istep_init, istep_exit", [(1, 1), (2, 2)]) def test_compute_advection_in_horizontal_momentum_equation( experiment, @@ -931,7 +938,7 @@ def test_compute_advection_in_horizontal_momentum_equation( start_edge_nudging_level_2 = icon_grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) end_edge_local = icon_grid.end_index(edge_domain(h_grid.Zone.LOCAL)) - dtime = savepoint_velocity_init.get_metadata("dtime").get("dtime") + dtime = savepoint_velocity_init.dtime() end_index_of_damping_layer = grid_savepoint.nrdmax() icon_result_ddt_vn_apc = savepoint_velocity_exit.ddt_vn_apc_pc(istep_exit - 1) @@ -980,6 +987,6 @@ def test_compute_advection_in_horizontal_momentum_equation( assert test_utils.dallclose( icon_result_ddt_vn_apc.asnumpy(), normal_wind_advective_tendency.asnumpy(), - rtol=1.0e-15, - atol=1.0e-15, + rtol=rtol_8eps, + atol=atol_2eps, ) 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 d9f8b5bbae..3e9da5fad5 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 @@ -95,7 +95,7 @@ def test_run_solve_nonhydro_single_step( rayleigh_damping_height=damping_height, ) vertical_params = utils.create_vertical_params(vertical_config, grid_savepoint) - dtime = savepoint_nonhydro_init.get_metadata("dtime").get("dtime") + dtime = savepoint_nonhydro_init.dtime() lprep_adv = savepoint_nonhydro_init.get_metadata("prep_adv").get("prep_adv") prep_adv = dycore_states.PrepAdvection( vn_traj=savepoint_nonhydro_init.vn_traj(), 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 d49922254d..1ea5d411b5 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py @@ -6,6 +6,8 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +import itertools + import gt4py.next as gtx import numpy as np import pytest @@ -19,6 +21,7 @@ from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.testing import definitions divergence_damp_order = DivergenceDampingOrder() @@ -28,6 +31,23 @@ class TestApplyDivergenceDampingAndUpdateVn(test_helpers.StencilTest): PROGRAM = apply_divergence_damping_and_update_vn OUTPUTS = ("next_vn",) + STATIC_PARAMS = { + test_helpers.StandardStaticVariants.NONE: (), + test_helpers.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + "is_iau_active", + "limited_area", + ), + test_helpers.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + "is_iau_active", + "limited_area", + ), + } @staticmethod def reference( @@ -153,7 +173,10 @@ def reference( return dict(next_vn=next_vn) - @pytest.fixture(params=[True, False]) + @pytest.fixture( + params=[{"limited_area": la} for la in [True, False]], + ids=lambda param: f"limited_area[{param['limited_area']}]", + ) def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: current_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) horizontal_mask_for_3d_divdamp = data_alloc.random_field(grid, dims.EdgeDim) @@ -202,7 +225,7 @@ def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: and (second_order_divdamp_factor <= (4.0 * fourth_order_divdamp_factor)) ) - limited_area = request.param + limited_area = request.param["limited_area"] edge_domain = h_grid.domain(dims.EdgeDim) start_edge_nudging_level_2 = grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) @@ -240,3 +263,43 @@ def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: vertical_start=0, vertical_end=grid.num_levels, ) + + +@pytest.mark.continuous_benchmarking +class TestApplyDivergenceDampingAndUpdateVnContinuousBenchmarking( + TestApplyDivergenceDampingAndUpdateVn +): + @pytest.fixture( + params=[ + {"limited_area": la, "divdamp_order": do, "is_iau_active": ia} + for la, do, ia in itertools.product( + [True, False], + [ + DivergenceDampingOrder.SECOND_ORDER, + DivergenceDampingOrder.FOURTH_ORDER, + DivergenceDampingOrder.COMBINED, + ], + [True, False], + ) + ], + ids=lambda param: f"limited_area[{param['limited_area']}]__divdamp_order[{param['divdamp_order']}]__is_iau_active[{param['is_iau_active']}]", + ) + def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: + base_data = TestApplyDivergenceDampingAndUpdateVn.input_data.__wrapped__( + self, request, grid + ) + base_data["is_iau_active"] = False + fourth_order_divdamp_factor = 0.004 + second_order_divdamp_factor = 0.032 + divdamp_order = request.param["divdamp_order"] + base_data["second_order_divdamp_scaling_coeff"] = 34497.62082646618 # for icon-ch1(_medium) + base_data["apply_2nd_order_divergence_damping"] = ( + divdamp_order == divergence_damp_order.COMBINED + ) and (base_data["second_order_divdamp_scaling_coeff"] > 1.0e-6) + base_data["apply_4th_order_divergence_damping"] = ( + divdamp_order == divergence_damp_order.FOURTH_ORDER + ) or ( + (divdamp_order == divergence_damp_order.COMBINED) + and (second_order_divdamp_factor <= (4.0 * fourth_order_divdamp_factor)) + ) + return base_data diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py index 4340ee9ea4..70c6a4d051 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py @@ -138,6 +138,21 @@ def _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_without_levelm class TestFusedVelocityAdvectionStencilsHMomentum(stencil_tests.StencilTest): PROGRAM = compute_advection_in_horizontal_momentum_equation OUTPUTS = ("normal_wind_advective_tendency",) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "end_index_of_damping_layer", + "vertical_start", + "vertical_end", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "end_index_of_damping_layer", + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( @@ -227,8 +242,13 @@ def reference( return dict(normal_wind_advective_tendency=normal_wind_advective_tendency) - @pytest.fixture - def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: + @pytest.fixture( + params=[{"apply_extra_diffusion_on_vn": value} for value in [True, False]], + ids=lambda param: f"apply_extra_diffusion_on_vn[{param['apply_extra_diffusion_on_vn']}]", + ) + def input_data( + self, request: pytest.FixtureRequest, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: normal_wind_advective_tendency = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) horizontal_kinetic_energy_at_edges_on_model_levels = data_alloc.random_field( @@ -257,7 +277,7 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala scalfac_exdiff = 0.6 dtime = 2.0 cfl_w_limit = 0.65 / dtime - apply_extra_diffusion_on_vn = True + apply_extra_diffusion_on_vn = request.param["apply_extra_diffusion_on_vn"] end_index_of_damping_layer = 5 edge_domain = h_grid.domain(dims.EdgeDim) 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 a98978eeda..22714168dc 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 @@ -19,7 +19,7 @@ from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import stencil_tests +from icon4py.model.testing import definitions, stencil_tests from .test_add_interpolated_horizontal_advection_of_w import ( add_interpolated_horizontal_advection_of_w_numpy, @@ -276,6 +276,21 @@ class TestFusedVelocityAdvectionStencilVMomentum(stencil_tests.StencilTest): "contravariant_corrected_w_at_cells_on_model_levels", "vertical_cfl", ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + "end_index_of_damping_layer", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + "end_index_of_damping_layer", + ), + } @staticmethod def reference( @@ -467,6 +482,17 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala ) +@pytest.mark.continuous_benchmarking +class TestFusedVelocityAdvectionStencilVMomentumContinuousBenchmarking( + TestFusedVelocityAdvectionStencilVMomentum +): + @pytest.fixture + def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: + base_data = TestFusedVelocityAdvectionStencilVMomentum.input_data.__wrapped__(self, grid) + base_data["end_index_of_damping_layer"] = 12 + return base_data + + @pytest.mark.embedded_remap_error class TestFusedVelocityAdvectionStencilVMomentumAndContravariant(stencil_tests.StencilTest): PROGRAM = compute_contravariant_correction_and_advection_in_vertical_momentum_equation @@ -476,6 +502,9 @@ class TestFusedVelocityAdvectionStencilVMomentumAndContravariant(stencil_tests.S "contravariant_corrected_w_at_cells_on_model_levels", "vertical_cfl", ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: (), # For now compile time variants triger error in gt4py + } @staticmethod def reference( @@ -603,7 +632,8 @@ def reference( ) @pytest.fixture( - params=[{"skip_compute_predictor_vertical_advection": value} for value in [True, False]] + params=[{"skip_compute_predictor_vertical_advection": value} for value in [True, False]], + 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 @@ -679,3 +709,10 @@ def input_data( vertical_start=vertical_start, vertical_end=vertical_end, ) + + +@pytest.mark.continuous_benchmarking +class TestFusedVelocityAdvectionStencilVMomentumAndContravariantContinuousBenchmarking( + TestFusedVelocityAdvectionStencilVMomentumAndContravariant +): + pass 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 11a52a0744..5f16f32404 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 @@ -15,18 +15,20 @@ from icon4py.model.atmosphere.dycore.stencils.compute_horizontal_velocity_quantities import ( compute_averaged_vn_and_fluxes_and_prepare_tracer_advection, ) -from icon4py.model.common import dimension as dims, type_alias as ta +from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils -from icon4py.model.testing import stencil_tests +from icon4py.model.common.type_alias import precision, wpfloat +from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest from .test_accumulate_prep_adv_fields import accumulate_prep_adv_fields_numpy from .test_compute_mass_flux import compute_mass_flux_numpy from .test_spatially_average_flux_or_velocity import spatially_average_flux_or_velocity_numpy +@pytest.mark.single_precision_ready @pytest.mark.embedded_remap_error -class TestComputeAveragedVnAndFluxesAndPrepareTracerAdvection(stencil_tests.StencilTest): +class TestComputeAveragedVnAndFluxesAndPrepareTracerAdvection(StencilTest): PROGRAM = compute_averaged_vn_and_fluxes_and_prepare_tracer_advection OUTPUTS = ( "spatially_averaged_vn", @@ -35,6 +37,29 @@ class TestComputeAveragedVnAndFluxesAndPrepareTracerAdvection(stencil_tests.Sten "substep_and_spatially_averaged_vn", "substep_averaged_mass_flux", ) + STATIC_PARAMS = { + # StandardStaticVariants.NONE: (), + StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + "prepare_advection", + "at_first_substep", + ), + # StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + # "vertical_start", + # "vertical_end", + # "limited_area", + # ), + } + # TODO(pstark): rm this again: + FIND_RTOL = True + if precision == "single": + # RTOL = 3e-2 + # ATOL = 1e-4 + RTOL = 1e-1 + ATOL = 1e-2 @staticmethod def reference( @@ -51,7 +76,7 @@ def reference( theta_v_at_edges_on_model_levels: np.ndarray, prepare_advection: bool, at_first_substep: bool, - r_nsubsteps: ta.wpfloat, + r_nsubsteps: wpfloat, horizontal_start: int, horizontal_end: int, **kwargs: Any, @@ -136,41 +161,39 @@ def reference( @pytest.fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - spatially_averaged_vn = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - mass_fl_e = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - z_theta_v_fl_e = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) - - substep_and_spatially_averaged_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - substep_averaged_mass_flux = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - e_flx_avg = data_alloc.random_field(grid, dims.EdgeDim, dims.E2C2EODim) - vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - z_rho_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - ddqz_z_full_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - z_theta_v_e = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - prepare_advection = True - at_first_substep = True - r_nsubsteps = 0.2 + random_fields = data_alloc.get_random_fields( + grid, + [ + "substep_and_spatially_averaged_vn", + "substep_averaged_mass_flux", + "e_flx_avg", + "vn", + "rho_at_edges_on_model_levels", + "ddqz_z_full_e", + "theta_v_at_edges_on_model_levels", + ], + ) - edge_domain = h_grid.domain(dims.EdgeDim) - horizontal_start = grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) - horizontal_end = grid.end_index(edge_domain(h_grid.Zone.LOCAL)) + zero_fields = data_alloc.get_zero_fields( + grid, + [ + "spatially_averaged_vn", + "mass_flux_at_edges_on_model_levels", + "theta_v_flux_at_edges_on_model_levels", + ], + ) - return dict( - spatially_averaged_vn=spatially_averaged_vn, - mass_flux_at_edges_on_model_levels=mass_fl_e, - theta_v_flux_at_edges_on_model_levels=z_theta_v_fl_e, - substep_and_spatially_averaged_vn=substep_and_spatially_averaged_vn, - substep_averaged_mass_flux=substep_averaged_mass_flux, - e_flx_avg=e_flx_avg, - vn=vn, - rho_at_edges_on_model_levels=z_rho_e, - ddqz_z_full_e=ddqz_z_full_e, - theta_v_at_edges_on_model_levels=z_theta_v_e, - prepare_advection=prepare_advection, - at_first_substep=at_first_substep, - r_nsubsteps=r_nsubsteps, - horizontal_start=horizontal_start, - horizontal_end=horizontal_end, - vertical_start=0, - vertical_end=grid.num_levels, + edge_domain = h_grid.domain(dims.EdgeDim) + return ( + random_fields + | zero_fields + | dict( + prepare_advection=True, + at_first_substep=True, + r_nsubsteps=wpfloat(0.2), + horizontal_start=grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)), + horizontal_end=grid.end_index(edge_domain(h_grid.Zone.LOCAL)), + vertical_start=0, + vertical_end=grid.num_levels, + ) ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py index 914bb23c45..be12bc3ee1 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py @@ -16,7 +16,7 @@ from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import stencil_tests +from icon4py.model.testing import definitions, stencil_tests from .test_compute_contravariant_correction import compute_contravariant_correction_numpy from .test_compute_horizontal_advection_term_for_vertical_velocity import ( @@ -154,6 +154,21 @@ class TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContrava "contravariant_correction_at_edges_on_model_levels", "horizontal_advection_of_w_at_edges_on_half_levels", ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + "nflatlev", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + "nflatlev", + ), + } @classmethod def reference( @@ -266,7 +281,8 @@ def reference( ) @pytest.fixture( - params=[{"skip_compute_predictor_vertical_advection": value} for value in [True, False]] + params=[{"skip_compute_predictor_vertical_advection": value} for value in [True, False]], + 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 @@ -334,3 +350,31 @@ def input_data( vertical_start=vertical_start, vertical_end=vertical_end, ) + + +@pytest.mark.continuous_benchmarking +class TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContravariantCorrectionContinuousBenchmarking( + TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContravariantCorrection +): + @pytest.fixture( + params=[{"skip_compute_predictor_vertical_advection": value} for value in [True, False]], + 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 + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + base_data = TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContravariantCorrection.input_data.__wrapped__( + self, grid, request + ) + base_data["skip_compute_predictor_vertical_advection"] = request.param[ + "skip_compute_predictor_vertical_advection" + ] + base_data["nflatlev"] = 6 + edge_domain = h_grid.domain(dims.EdgeDim) + base_data["horizontal_start"] = grid.start_index( + edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5) + ) + base_data["horizontal_end"] = grid.end_index(edge_domain(h_grid.Zone.HALO_LEVEL_2)) + base_data["vertical_start"] = 0 + base_data["vertical_end"] = grid.num_levels + 1 + return base_data diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py index 79217ba75d..8e49c48599 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py @@ -18,7 +18,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest def compute_hydrostatic_correction_term_numpy( @@ -86,11 +86,23 @@ def _apply_index_field( return z_hydro_corr -@pytest.mark.skip_value_error @pytest.mark.uses_as_offset class TestComputeHydrostaticCorrectionTerm(StencilTest): OUTPUTS = ("z_hydro_corr",) PROGRAM = compute_hydrostatic_correction_term + STATIC_PARAMS = { + StandardStaticVariants.NONE: (), + StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( @@ -155,3 +167,10 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala vertical_start=0, vertical_end=gtx.int32(grid.num_levels), ) + + +@pytest.mark.continuous_benchmarking +class TestComputeHydrostaticCorrectionTermContinuousBenchmarking( + TestComputeHydrostaticCorrectionTerm +): + pass diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py index a5415499d8..6e9401fd06 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py @@ -32,7 +32,7 @@ from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import stencil_tests +from icon4py.model.testing import definitions, stencil_tests from .test_compute_approx_of_2nd_vertical_derivative_of_exner import ( compute_approx_of_2nd_vertical_derivative_of_exner_numpy, @@ -80,6 +80,28 @@ class TestComputePerturbedQuantitiesAndInterpolation(stencil_tests.StencilTest): "pressure_buoyancy_acceleration_at_cells_on_half_levels", "d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels", ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "igradp_method", + "nflatlev", + "nflat_gradp", + "start_cell_lateral_boundary", + "start_cell_lateral_boundary_level_3", + "start_cell_halo_level_2", + "end_cell_halo", + "end_cell_halo_level_2", + "model_top", + "surface_level", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "igradp_method", + "nflatlev", + "nflat_gradp", + "model_top", + "surface_level", + ), + } @staticmethod def reference( @@ -439,3 +461,17 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala model_top=0, surface_level=grid.num_levels + 1, ) + + +@pytest.mark.continuous_benchmarking +class TestComputePerturbedQuantitiesAndInterpolationContinuousBenchmarking( + TestComputePerturbedQuantitiesAndInterpolation +): + @pytest.fixture + def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: + base_data = TestComputePerturbedQuantitiesAndInterpolation.input_data.__wrapped__( + self, grid + ) + base_data["nflatlev"] = 6 + base_data["nflat_gradp"] = 35 + return base_data diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py index 350d380792..6a3a75b80e 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py @@ -12,8 +12,6 @@ import numpy as np import pytest -import icon4py.model.common.type_alias as ta -import icon4py.model.testing.stencil_tests as test_helpers from icon4py.model.atmosphere.dycore.dycore_states import ( HorizontalPressureDiscretizationType, RhoThetaAdvectionType, @@ -23,7 +21,9 @@ ) from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid +from icon4py.model.common.type_alias import precision, vpfloat, wpfloat from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.testing import stencil_tests rhotheta_avd_type = RhoThetaAdvectionType() @@ -121,10 +121,10 @@ def compute_theta_rho_face_value_by_miura_scheme_numpy( return rho_at_edges_on_model_levels, theta_v_at_edges_on_model_levels +@pytest.mark.single_precision_ready @pytest.mark.embedded_remap_error -@pytest.mark.skip_value_error @pytest.mark.uses_as_offset -class TestComputeThetaRhoPressureGradientAndUpdateVn(test_helpers.StencilTest): +class TestComputeThetaRhoPressureGradientAndUpdateVn(stencil_tests.StencilTest): PROGRAM = compute_theta_rho_face_values_and_pressure_gradient_and_update_vn OUTPUTS = ( "rho_at_edges_on_model_levels", @@ -132,6 +132,36 @@ class TestComputeThetaRhoPressureGradientAndUpdateVn(test_helpers.StencilTest): "horizontal_pressure_gradient", "next_vn", ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "iau_wgt_dyn", + "is_iau_active", + "limited_area", + "start_edge_lateral_boundary", + "start_edge_lateral_boundary_level_7", + "start_edge_nudging_level_2", + "end_edge_nudging", + "end_edge_halo", + "nflatlev", + "nflat_gradp", + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "is_iau_active", + "limited_area", + "nflatlev", + "nflat_gradp", + "vertical_start", + "vertical_end", + ), + } + if precision == "single": + RTOL = 3e-2 + ATOL = 1e-2 @staticmethod def reference( @@ -169,8 +199,8 @@ def reference( ipeidx_dsl: np.ndarray, pg_exdist: np.ndarray, inv_dual_edge_length: np.ndarray, - dtime: ta.wpfloat, - iau_wgt_dyn: ta.wpfloat, + dtime: wpfloat, + iau_wgt_dyn: wpfloat, is_iau_active: gtx.int32, limited_area: gtx.int32, start_edge_lateral_boundary: gtx.int32, @@ -268,7 +298,10 @@ def reference( perturbed_rho_at_cells_on_model_levels=perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, ), - (rho_at_edges_on_model_levels, theta_v_at_edges_on_model_levels), + ( + np.zeros_like(rho_at_edges_on_model_levels), + np.zeros_like(theta_v_at_edges_on_model_levels), + ), ) # Remaining computations at edge points @@ -406,134 +439,95 @@ def at_neighbor(i: int) -> np.ndarray: @pytest.fixture def input_data(self, grid: base.Grid) -> dict: - geofac_grg_x = data_alloc.random_field(grid, dims.CellDim, dims.C2E2CODim) - geofac_grg_y = data_alloc.random_field(grid, dims.CellDim, dims.C2E2CODim) - current_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - tangential_wind = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - pos_on_tplane_e_x = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - pos_on_tplane_e_y = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - primal_normal_cell_x = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - dual_normal_cell_x = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - primal_normal_cell_y = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - dual_normal_cell_y = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - reference_rho_at_edges_on_model_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim - ) - reference_theta_at_edges_on_model_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim - ) - perturbed_rho_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim - ) - perturbed_theta_v_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim - ) - ddxn_z_full = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - c_lin_e = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim) - temporal_extrapolation_of_perturbed_exner = data_alloc.random_field( - grid, dims.CellDim, dims.KDim - ) - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim - ) - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( - data_alloc.random_field(grid, dims.CellDim, dims.KDim) + random_fields = data_alloc.get_random_fields( + grid, + [ + "geofac_grg_x", + "geofac_grg_y", + "current_vn", + "tangential_wind", + "pos_on_tplane_e_x", + "pos_on_tplane_e_y", + "primal_normal_cell_x", + "dual_normal_cell_x", + "primal_normal_cell_y", + "dual_normal_cell_y", + "reference_rho_at_edges_on_model_levels", + "reference_theta_at_edges_on_model_levels", + "perturbed_rho_at_cells_on_model_levels", + "perturbed_theta_v_at_cells_on_model_levels", + "ddxn_z_full", + "c_lin_e", + "ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels", + "d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels", + "hydrostatic_correction_on_lowest_level", + "zdiff_gradp", + "pg_exdist", + "inv_dual_edge_length", + "predictor_normal_wind_advective_tendency", + "normal_wind_tendency_due_to_slow_physics_process", + "normal_wind_iau_increment", + "grf_tend_vn", + ], ) - hydrostatic_correction_on_lowest_level = data_alloc.random_field(grid, dims.EdgeDim) - zdiff_gradp = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim, dims.KDim) - ipeidx_dsl = data_alloc.random_mask(grid, dims.EdgeDim, dims.KDim) - pg_exdist = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - inv_dual_edge_length = data_alloc.random_field(grid, dims.EdgeDim) - predictor_normal_wind_advective_tendency = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim - ) - normal_wind_tendency_due_to_slow_physics_process = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim - ) - normal_wind_iau_increment = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - grf_tend_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - next_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - theta_v_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - horizontal_pressure_gradient = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - rho_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - - ikoffset = data_alloc.zero_field( - grid, dims.EdgeDim, dims.E2CDim, dims.KDim, dtype=gtx.int32 + + zero_fields = data_alloc.get_zero_fields( + grid, + [ + "next_vn", + "theta_v_at_edges_on_model_levels", + "horizontal_pressure_gradient", + "rho_at_edges_on_model_levels", + ], ) - rng = np.random.default_rng() - k_levels = grid.num_levels - - for k in range(k_levels): - # construct offsets that reach all k-levels except the last (because we are using the entries of this field with `+1`) - ikoffset.ndarray[:, :, k] = rng.integers( # type: ignore[index] - low=0 - k, - high=k_levels - k - 1, - size=(ikoffset.shape[0], ikoffset.shape[1]), - ) - dtime = 0.9 - iau_wgt_dyn = 1.0 - is_iau_active = True - limited_area = True edge_domain = h_grid.domain(dims.EdgeDim) - - start_edge_lateral_boundary = grid.end_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY)) - start_edge_lateral_boundary_level_7 = grid.start_index( - edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_7) + return ( + random_fields + | zero_fields + | dict( + temporal_extrapolation_of_perturbed_exner=data_alloc.random_field( + grid, dims.CellDim, dims.KDim + ), + ipeidx_dsl=data_alloc.random_mask(grid, dims.EdgeDim, dims.KDim), + ikoffset=data_alloc.random_ikoffset(grid, dims.EdgeDim, dims.E2CDim, dims.KDim), + dtime=wpfloat(0.9), + iau_wgt_dyn=wpfloat(1.0), + is_iau_active=True, + limited_area=True, + start_edge_lateral_boundary=grid.end_index( + edge_domain(h_grid.Zone.LATERAL_BOUNDARY) + ), + start_edge_lateral_boundary_level_7=grid.start_index( + edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_7) + ), + start_edge_nudging_level_2=grid.start_index( + edge_domain(h_grid.Zone.NUDGING_LEVEL_2) + ), + end_edge_nudging=grid.end_index(edge_domain(h_grid.Zone.NUDGING)), + end_edge_halo=grid.end_index(edge_domain(h_grid.Zone.HALO)), + nflatlev=4, + nflat_gradp=27, + horizontal_start=0, + horizontal_end=grid.num_edges, + vertical_start=0, + vertical_end=grid.num_levels, + ) ) - start_edge_nudging_level_2 = grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) - end_edge_nudging = grid.end_index(edge_domain(h_grid.Zone.NUDGING)) - end_edge_halo = grid.end_index(edge_domain(h_grid.Zone.HALO)) - nflatlev = 4 - nflat_gradp = 27 - return dict( - rho_at_edges_on_model_levels=rho_at_edges_on_model_levels, - theta_v_at_edges_on_model_levels=theta_v_at_edges_on_model_levels, - horizontal_pressure_gradient=horizontal_pressure_gradient, - next_vn=next_vn, - current_vn=current_vn, - tangential_wind=tangential_wind, - reference_rho_at_edges_on_model_levels=reference_rho_at_edges_on_model_levels, - reference_theta_at_edges_on_model_levels=reference_theta_at_edges_on_model_levels, - perturbed_rho_at_cells_on_model_levels=perturbed_rho_at_cells_on_model_levels, - perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, - temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - hydrostatic_correction_on_lowest_level=hydrostatic_correction_on_lowest_level, - predictor_normal_wind_advective_tendency=predictor_normal_wind_advective_tendency, - normal_wind_tendency_due_to_slow_physics_process=normal_wind_tendency_due_to_slow_physics_process, - normal_wind_iau_increment=normal_wind_iau_increment, - grf_tend_vn=grf_tend_vn, - geofac_grg_x=geofac_grg_x, - geofac_grg_y=geofac_grg_y, - pos_on_tplane_e_x=pos_on_tplane_e_x, - pos_on_tplane_e_y=pos_on_tplane_e_y, - primal_normal_cell_x=primal_normal_cell_x, - dual_normal_cell_x=dual_normal_cell_x, - primal_normal_cell_y=primal_normal_cell_y, - dual_normal_cell_y=dual_normal_cell_y, - ddxn_z_full=ddxn_z_full, - c_lin_e=c_lin_e, - ikoffset=ikoffset, - zdiff_gradp=zdiff_gradp, - ipeidx_dsl=ipeidx_dsl, - pg_exdist=pg_exdist, - inv_dual_edge_length=inv_dual_edge_length, - dtime=dtime, - iau_wgt_dyn=iau_wgt_dyn, - is_iau_active=is_iau_active, - limited_area=limited_area, - nflatlev=nflatlev, - nflat_gradp=nflat_gradp, - start_edge_lateral_boundary=start_edge_lateral_boundary, - start_edge_lateral_boundary_level_7=start_edge_lateral_boundary_level_7, - start_edge_nudging_level_2=start_edge_nudging_level_2, - end_edge_nudging=end_edge_nudging, - end_edge_halo=end_edge_halo, - horizontal_start=0, - horizontal_end=grid.num_edges, - vertical_start=0, - vertical_end=grid.num_levels, + +@pytest.mark.continuous_benchmarking +class TestComputeThetaRhoPressureGradientAndUpdateVnContinuousBenchmarking( + TestComputeThetaRhoPressureGradientAndUpdateVn +): + @pytest.fixture + def input_data(self, grid: base.Grid) -> dict: + base_data = TestComputeThetaRhoPressureGradientAndUpdateVn.input_data.__wrapped__( + self, grid ) + base_data["is_iau_active"] = False + base_data["limited_area"] = grid.limited_area + base_data["nflatlev"] = 5 + base_data["nflat_gradp"] = 34 + base_data["start_edge_lateral_boundary"] = 0 + return base_data diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_dycore_utils.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_dycore_utils.py index bde56aba25..b478fc4656 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_dycore_utils.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_dycore_utils.py @@ -37,7 +37,7 @@ def fourth_order_divdamp_scaling_coeff_for_order_24_numpy( def calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary_numpy( coeff: float, field: np.ndarray ) -> np.ndarray: - return 0.75 / (coeff + constants.DBL_EPS) * np.abs(field) + return 0.75 / (coeff + constants.WP_EPS) * np.abs(field) def test_calculate_fourth_order_divdamp_scaling_coeff_order_24( @@ -50,14 +50,14 @@ def test_calculate_fourth_order_divdamp_scaling_coeff_order_24( interpolated_fourth_order_divdamp_factor = data_alloc.random_field( grid, dims.KDim, allocator=backend ) - out = data_alloc.random_field(grid, dims.KDim, allocator=backend) + fourth_order_divdamp_scaling_coeff = data_alloc.random_field(grid, dims.KDim, allocator=backend) dycore_utils._calculate_fourth_order_divdamp_scaling_coeff.with_backend(backend)( interpolated_fourth_order_divdamp_factor=interpolated_fourth_order_divdamp_factor, second_order_divdamp_factor=second_order_divdamp_factor, divdamp_order=divdamp_order, mean_cell_area=mean_cell_area, - out=out, + out=fourth_order_divdamp_scaling_coeff, offset_provider={}, ) @@ -66,7 +66,7 @@ def test_calculate_fourth_order_divdamp_scaling_coeff_order_24( second_order_divdamp_factor, mean_cell_area, ) - assert test_utils.dallclose(ref, out.asnumpy()) + assert test_utils.dallclose(ref, fourth_order_divdamp_scaling_coeff.asnumpy()) def test_calculate_fourth_order_divdamp_scaling_coeff_any_order( @@ -79,18 +79,18 @@ def test_calculate_fourth_order_divdamp_scaling_coeff_any_order( interpolated_fourth_order_divdamp_factor = data_alloc.random_field( grid, dims.KDim, allocator=backend ) - out = data_alloc.random_field(grid, dims.KDim, allocator=backend) + fourth_order_divdamp_scaling_coeff = data_alloc.random_field(grid, dims.KDim, allocator=backend) dycore_utils._calculate_fourth_order_divdamp_scaling_coeff.with_backend(backend)( interpolated_fourth_order_divdamp_factor=interpolated_fourth_order_divdamp_factor, second_order_divdamp_factor=second_order_divdamp_factor, divdamp_order=divdamp_order, mean_cell_area=mean_cell_area, - out=out, + out=fourth_order_divdamp_scaling_coeff, offset_provider={}, ) enhanced_factor = -interpolated_fourth_order_divdamp_factor.asnumpy() * mean_cell_area**2 - assert test_utils.dallclose(enhanced_factor, out.asnumpy()) + assert test_utils.dallclose(enhanced_factor, fourth_order_divdamp_scaling_coeff.asnumpy()) def test_calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary( @@ -98,13 +98,21 @@ def test_calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary( ) -> None: grid = simple_grid.simple_grid(allocator=backend) fourth_order_divdamp_scaling_coeff = data_alloc.random_field(grid, dims.KDim, allocator=backend) - out = data_alloc.zero_field(grid, dims.KDim, allocator=backend) + reduced_fourth_order_divdamp_coeff_at_nest_boundary = data_alloc.zero_field( + grid, dims.KDim, allocator=backend + ) coeff = 0.3 dycore_utils._calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary.with_backend( backend - )(fourth_order_divdamp_scaling_coeff, coeff, constants.DBL_EPS, out=out, offset_provider={}) + )( + fourth_order_divdamp_scaling_coeff, + coeff, + constants.WP_EPS, + out=reduced_fourth_order_divdamp_coeff_at_nest_boundary, + offset_provider={}, + ) assert test_utils.dallclose( - out.asnumpy(), + reduced_fourth_order_divdamp_coeff_at_nest_boundary.asnumpy(), calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary_numpy( coeff, fourth_order_divdamp_scaling_coeff.asnumpy() ), @@ -139,7 +147,7 @@ def test_calculate_divdamp_fields(backend: gtx_typing.Backend) -> None: mean_cell_area, second_order_divdamp_factor, max_nudging_coefficient, - constants.DBL_EPS, + constants.WP_EPS, out=( fourth_order_divdamp_scaling_coeff, reduced_fourth_order_divdamp_coeff_at_nest_boundary, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py index 658499bc80..943d0419d0 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py @@ -19,12 +19,25 @@ from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import wpfloat from icon4py.model.common.utils.data_allocation import zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest class TestInitCellKdimFieldWithZeroWp(StencilTest): PROGRAM = init_cell_kdim_field_with_zero_wp OUTPUTS = ("field_with_zero_wp",) + STATIC_PARAMS = { + StandardStaticVariants.NONE: (), + StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( @@ -46,3 +59,8 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala vertical_start=0, vertical_end=gtx.int32(grid.num_levels), ) + + +@pytest.mark.continuous_benchmarking +class TestInitCellKdimFieldWithZeroWpContinuousBenchmarking(TestInitCellKdimFieldWithZeroWp): + pass diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py index b3fb185c76..58e876c9da 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py @@ -44,6 +44,19 @@ class TestInterpolateRhoThetaVToHalfLevelsAndComputePressureBuoyancyAcceleration "theta_v_at_cells_on_half_levels", "pressure_buoyancy_acceleration_at_cells_on_half_levels", ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( @@ -236,3 +249,10 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala vertical_start=1, vertical_end=grid.num_levels, ) + + +@pytest.mark.continuous_benchmarking +class TestInterpolateRhoThetaVToHalfLevelsAndComputePressureBuoyancyAccelerationContinuousBenchmarking( + TestInterpolateRhoThetaVToHalfLevelsAndComputePressureBuoyancyAcceleration +): + pass diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py index c9297d94b7..0013e7e1da 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py @@ -19,12 +19,25 @@ from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest class TestUpdateMassFluxWeighted(StencilTest): PROGRAM = update_mass_flux_weighted OUTPUTS = ("mass_flx_ic",) + STATIC_PARAMS = { + StandardStaticVariants.NONE: (), + StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( @@ -71,3 +84,8 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala vertical_start=0, vertical_end=gtx.int32(grid.num_levels), ) + + +@pytest.mark.continuous_benchmarking +class TestUpdateMassFluxWeightedContinuousBenchmarking(TestUpdateMassFluxWeighted): + pass diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py index c43c2de7be..edaf12d1b9 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py @@ -5,6 +5,7 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +import itertools from typing import Any import gt4py.next as gtx @@ -17,6 +18,7 @@ from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils +from icon4py.model.common.type_alias import precision, vpfloat, wpfloat from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests @@ -50,6 +52,7 @@ from .test_update_mass_volume_flux import update_mass_volume_flux_numpy +@pytest.mark.single_precision_ready @pytest.mark.uses_concat_where class TestVerticallyImplicitSolverAtCorrectorStep(stencil_tests.StencilTest): PROGRAM = vertically_implicit_solver_at_corrector_step @@ -62,6 +65,36 @@ class TestVerticallyImplicitSolverAtCorrectorStep(stencil_tests.StencilTest): "dynamical_vertical_volumetric_flux_at_cells_on_half_levels", "exner_dynamical_increment", ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "start_cell_index_nudging", + "end_cell_index_local", + "end_index_of_damping_layer", + "kstart_moist", + "vertical_start_index_model_top", + "vertical_end_index_model_surface", + "at_first_substep", + "at_last_substep", + "lprep_adv", + "is_iau_active", + "rayleigh_type", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "end_index_of_damping_layer", + "kstart_moist", + "vertical_start_index_model_top", + "vertical_end_index_model_surface", + "at_first_substep", + "at_last_substep", + "lprep_adv", + "is_iau_active", + "rayleigh_type", + ), + } + if precision == "single": + RTOL = 0.5 + ATOL = 0.1 @staticmethod def reference( @@ -370,127 +403,267 @@ def reference( exner_dynamical_increment=exner_dynamical_increment, ) - @pytest.fixture - def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) - mass_flux_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - theta_v_flux_at_edges_on_model_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim - ) - current_w = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - predictor_vertical_wind_advective_tendency = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} - ) - corrector_vertical_wind_advective_tendency = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} - ) - pressure_buoyancy_acceleration_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim - ) - rho_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, low=1.0e-5 - ) - contravariant_correction_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} - ) - exner_w_explicit_weight_parameter = data_alloc.random_field(grid, dims.CellDim) - current_exner = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - current_rho = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - current_theta_v = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - inv_ddqz_z_full = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - exner_w_implicit_weight_parameter = data_alloc.random_field(grid, dims.CellDim) - theta_v_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, low=1.0e-5 - ) - perturbed_exner_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim + # @pytest.fixture( + # params=[ + # {"at_first_substep": afs, "at_last_substep": als, "lprep_adv": la} + # for afs, als, la in itertools.product(*([(True, False)] * 3)) + # if not (afs and als) + # ], + # ids=lambda p: ( + # f"at_first_substep[{p['at_first_substep']}]__" + # f"at_last_substep[{p['at_last_substep']}]__" + # f"lprep_adv[{p['lprep_adv']}]" + # ), + # ) + # def input_data( + # self, request: pytest.FixtureRequest, grid: base.Grid + # ) -> dict[str, gtx.Field | state_utils.ScalarType]: + # geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) + # mass_flux_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) + # theta_v_flux_at_edges_on_model_levels = data_alloc.random_field( + # grid, dims.EdgeDim, dims.KDim + # ) + # current_w = data_alloc.random_field(grid, dims.CellDim, dims.KDim) + # predictor_vertical_wind_advective_tendency = data_alloc.random_field( + # grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + # ) + # corrector_vertical_wind_advective_tendency = data_alloc.random_field( + # grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + # ) + # pressure_buoyancy_acceleration_at_cells_on_half_levels = data_alloc.random_field( + # grid, dims.CellDim, dims.KDim + # ) + # rho_at_cells_on_half_levels = data_alloc.random_field( + # grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, low=1.0e-5 + # ) + # contravariant_correction_at_cells_on_half_levels = data_alloc.random_field( + # grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + # ) + # exner_w_explicit_weight_parameter = data_alloc.random_field(grid, dims.CellDim) + # current_exner = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) + # current_rho = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) + # current_theta_v = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) + # inv_ddqz_z_full = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) + # exner_w_implicit_weight_parameter = data_alloc.random_field(grid, dims.CellDim) + # theta_v_at_cells_on_half_levels = data_alloc.random_field( + # grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, low=1.0e-5 + # ) + # perturbed_exner_at_cells_on_model_levels = data_alloc.random_field( + # grid, dims.CellDim, dims.KDim + # ) + # exner_tendency_due_to_slow_physics = data_alloc.random_field(grid, dims.CellDim, dims.KDim) + # rho_iau_increment = data_alloc.random_field(grid, dims.CellDim, dims.KDim) + # exner_iau_increment = data_alloc.random_field(grid, dims.CellDim, dims.KDim) + # ddqz_z_half = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) + # rayleigh_damping_factor = data_alloc.random_field(grid, dims.KDim) + # reference_exner_at_cells_on_model_levels = data_alloc.random_field( + # grid, dims.CellDim, dims.KDim, low=1.0e-5 + # ) + + # next_w = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) + # next_rho = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) + # next_exner = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) + # next_theta_v = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) + # exner_dynamical_increment = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + # dynamical_vertical_mass_flux_at_cells_on_half_levels = data_alloc.zero_field( + # grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + # ) + # dynamical_vertical_volumetric_flux_at_cells_on_half_levels = data_alloc.zero_field( + # grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + # ) + + # lprep_adv = request.param["lprep_adv"] + # r_nsubsteps = 0.5 + # is_iau_active = True + # at_first_substep = request.param["at_first_substep"] + # rayleigh_type = 2 + # end_index_of_damping_layer = 3 + # at_last_substep = request.param["at_last_substep"] + # kstart_moist = 1 + # dtime = 0.001 + # veladv_offctr = 0.25 + # advection_explicit_weight_parameter = 0.5 - veladv_offctr + # advection_implicit_weight_parameter = 0.5 + veladv_offctr + # iau_wgt_dyn = 1.0 + # ndyn_substeps_var = 0.5 + + # cell_domain = h_grid.domain(dims.CellDim) + # start_cell_nudging = grid.start_index(cell_domain(h_grid.Zone.NUDGING)) + # end_cell_local = grid.end_index(cell_domain(h_grid.Zone.LOCAL)) + + # return dict( + # next_w=next_w, + # next_rho=next_rho, + # next_exner=next_exner, + # next_theta_v=next_theta_v, + # dynamical_vertical_mass_flux_at_cells_on_half_levels=dynamical_vertical_mass_flux_at_cells_on_half_levels, + # dynamical_vertical_volumetric_flux_at_cells_on_half_levels=dynamical_vertical_volumetric_flux_at_cells_on_half_levels, + # exner_dynamical_increment=exner_dynamical_increment, + # geofac_div=geofac_div, + # mass_flux_at_edges_on_model_levels=mass_flux_at_edges_on_model_levels, + # theta_v_flux_at_edges_on_model_levels=theta_v_flux_at_edges_on_model_levels, + # predictor_vertical_wind_advective_tendency=predictor_vertical_wind_advective_tendency, + # corrector_vertical_wind_advective_tendency=corrector_vertical_wind_advective_tendency, + # pressure_buoyancy_acceleration_at_cells_on_half_levels=pressure_buoyancy_acceleration_at_cells_on_half_levels, + # rho_at_cells_on_half_levels=rho_at_cells_on_half_levels, + # contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, + # exner_w_explicit_weight_parameter=exner_w_explicit_weight_parameter, + # current_exner=current_exner, + # current_rho=current_rho, + # current_theta_v=current_theta_v, + # current_w=current_w, + # inv_ddqz_z_full=inv_ddqz_z_full, + # exner_w_implicit_weight_parameter=exner_w_implicit_weight_parameter, + # theta_v_at_cells_on_half_levels=theta_v_at_cells_on_half_levels, + # perturbed_exner_at_cells_on_model_levels=perturbed_exner_at_cells_on_model_levels, + # exner_tendency_due_to_slow_physics=exner_tendency_due_to_slow_physics, + # rho_iau_increment=rho_iau_increment, + # exner_iau_increment=exner_iau_increment, + # ddqz_z_half=ddqz_z_half, + # rayleigh_damping_factor=rayleigh_damping_factor, + # reference_exner_at_cells_on_model_levels=reference_exner_at_cells_on_model_levels, + # advection_explicit_weight_parameter=advection_explicit_weight_parameter, + # advection_implicit_weight_parameter=advection_implicit_weight_parameter, + # lprep_adv=lprep_adv, + # r_nsubsteps=r_nsubsteps, + # ndyn_substeps_var=ndyn_substeps_var, + # iau_wgt_dyn=iau_wgt_dyn, + # dtime=dtime, + # is_iau_active=is_iau_active, + # rayleigh_type=rayleigh_type, + # at_first_substep=at_first_substep, + # at_last_substep=at_last_substep, + # end_index_of_damping_layer=end_index_of_damping_layer, + # kstart_moist=kstart_moist, + # start_cell_index_nudging=start_cell_nudging, + # end_cell_index_local=end_cell_local, + # vertical_start_index_model_top=gtx.int32(0), + # vertical_end_index_model_surface=gtx.int32(grid.num_levels + 1), + # ) + + + @pytest.fixture( + params=[ + {"at_first_substep": afs, "at_last_substep": als, "lprep_adv": la} + for afs, als, la in itertools.product(*([(True, False)] * 3)) + if not (afs and als) + ], + ids=lambda p: ( + f"at_first_substep[{p['at_first_substep']}]__" + f"at_last_substep[{p['at_last_substep']}]__" + f"lprep_adv[{p['lprep_adv']}]" + ), + ) + def input_data( + self, request: pytest.FixtureRequest, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + random_fields = data_alloc.get_random_fields( + grid, + [ + "geofac_div", + "mass_flux_at_edges_on_model_levels", + "theta_v_flux_at_edges_on_model_levels", + "predictor_vertical_wind_advective_tendency", + "corrector_vertical_wind_advective_tendency", + "pressure_buoyancy_acceleration_at_cells_on_half_levels", + "contravariant_correction_at_cells_on_half_levels", + "exner_w_explicit_weight_parameter", + "exner_w_implicit_weight_parameter", + "perturbed_exner_at_cells_on_model_levels", + "exner_tendency_due_to_slow_physics", + "rho_iau_increment", + "exner_iau_increment", + "rayleigh_damping_factor", + ], ) - exner_tendency_due_to_slow_physics = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - rho_iau_increment = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - exner_iau_increment = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - ddqz_z_half = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - rayleigh_damping_factor = data_alloc.random_field(grid, dims.KDim) - reference_exner_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, low=1.0e-5 + + random_fields_lowe = data_alloc.get_random_fields( + grid, + [ + "rho_at_cells_on_half_levels", + "current_exner", + "current_rho", + "current_theta_v", + "inv_ddqz_z_full", + "theta_v_at_cells_on_half_levels", + "ddqz_z_half", + "reference_exner_at_cells_on_model_levels", + ], + low=1e-05, ) - next_w = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) - next_rho = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) - next_exner = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) - next_theta_v = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) - exner_dynamical_increment = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - dynamical_vertical_mass_flux_at_cells_on_half_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + constant_fields_e = data_alloc.get_const_fields( + grid, ["next_rho", "next_exner", "next_theta_v"], 1e-05 ) - dynamical_vertical_volumetric_flux_at_cells_on_half_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + + zero_fields = data_alloc.get_zero_fields( + grid, + [ + "next_w", + "exner_dynamical_increment", + "dynamical_vertical_mass_flux_at_cells_on_half_levels", + "dynamical_vertical_volumetric_flux_at_cells_on_half_levels", + ], ) - lprep_adv = True - r_nsubsteps = 0.5 - is_iau_active = True - at_first_substep = True - rayleigh_type = 2 - end_index_of_damping_layer = 3 - at_last_substep = True - kstart_moist = 1 - dtime = 0.001 veladv_offctr = 0.25 - advection_explicit_weight_parameter = 0.5 - veladv_offctr - advection_implicit_weight_parameter = 0.5 + veladv_offctr - iau_wgt_dyn = 1.0 - ndyn_substeps_var = 0.5 cell_domain = h_grid.domain(dims.CellDim) - start_cell_nudging = grid.start_index(cell_domain(h_grid.Zone.NUDGING)) - end_cell_local = grid.end_index(cell_domain(h_grid.Zone.LOCAL)) + return ( + random_fields + | random_fields_lowe + | constant_fields_e + | zero_fields + | dict( + current_w=data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + lprep_adv=request.param["lprep_adv"], + r_nsubsteps=wpfloat(0.5), + is_iau_active=True, + at_first_substep=request.param["at_first_substep"], + rayleigh_type=2, + end_index_of_damping_layer=3, + at_last_substep=request.param["at_last_substep"], + kstart_moist=1, + dtime=wpfloat(0.001), + advection_explicit_weight_parameter=wpfloat(0.5 - veladv_offctr), + advection_implicit_weight_parameter=wpfloat(0.5 + veladv_offctr), + iau_wgt_dyn=wpfloat(1.0), + ndyn_substeps_var=wpfloat(0.5), + start_cell_index_nudging=grid.start_index(cell_domain(h_grid.Zone.NUDGING)), + end_cell_index_local=grid.end_index(cell_domain(h_grid.Zone.LOCAL)), + vertical_start_index_model_top=gtx.int32(0), + vertical_end_index_model_surface=gtx.int32(grid.num_levels + 1), + ) + ) - return dict( - next_w=next_w, - next_rho=next_rho, - next_exner=next_exner, - next_theta_v=next_theta_v, - dynamical_vertical_mass_flux_at_cells_on_half_levels=dynamical_vertical_mass_flux_at_cells_on_half_levels, - dynamical_vertical_volumetric_flux_at_cells_on_half_levels=dynamical_vertical_volumetric_flux_at_cells_on_half_levels, - exner_dynamical_increment=exner_dynamical_increment, - geofac_div=geofac_div, - mass_flux_at_edges_on_model_levels=mass_flux_at_edges_on_model_levels, - theta_v_flux_at_edges_on_model_levels=theta_v_flux_at_edges_on_model_levels, - predictor_vertical_wind_advective_tendency=predictor_vertical_wind_advective_tendency, - corrector_vertical_wind_advective_tendency=corrector_vertical_wind_advective_tendency, - pressure_buoyancy_acceleration_at_cells_on_half_levels=pressure_buoyancy_acceleration_at_cells_on_half_levels, - rho_at_cells_on_half_levels=rho_at_cells_on_half_levels, - contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, - exner_w_explicit_weight_parameter=exner_w_explicit_weight_parameter, - current_exner=current_exner, - current_rho=current_rho, - current_theta_v=current_theta_v, - current_w=current_w, - inv_ddqz_z_full=inv_ddqz_z_full, - exner_w_implicit_weight_parameter=exner_w_implicit_weight_parameter, - theta_v_at_cells_on_half_levels=theta_v_at_cells_on_half_levels, - perturbed_exner_at_cells_on_model_levels=perturbed_exner_at_cells_on_model_levels, - exner_tendency_due_to_slow_physics=exner_tendency_due_to_slow_physics, - rho_iau_increment=rho_iau_increment, - exner_iau_increment=exner_iau_increment, - ddqz_z_half=ddqz_z_half, - rayleigh_damping_factor=rayleigh_damping_factor, - reference_exner_at_cells_on_model_levels=reference_exner_at_cells_on_model_levels, - advection_explicit_weight_parameter=advection_explicit_weight_parameter, - advection_implicit_weight_parameter=advection_implicit_weight_parameter, - lprep_adv=lprep_adv, - r_nsubsteps=r_nsubsteps, - ndyn_substeps_var=ndyn_substeps_var, - iau_wgt_dyn=iau_wgt_dyn, - dtime=dtime, - is_iau_active=is_iau_active, - rayleigh_type=rayleigh_type, - at_first_substep=at_first_substep, - at_last_substep=at_last_substep, - end_index_of_damping_layer=end_index_of_damping_layer, - kstart_moist=kstart_moist, - start_cell_index_nudging=start_cell_nudging, - end_cell_index_local=end_cell_local, - vertical_start_index_model_top=gtx.int32(0), - vertical_end_index_model_surface=gtx.int32(grid.num_levels + 1), + + +@pytest.mark.continuous_benchmarking +class TestVerticallyImplicitSolverAtCorrectorStepContinuousBenchmarking( + TestVerticallyImplicitSolverAtCorrectorStep +): + @pytest.fixture( + params=[ + {"at_first_substep": afs, "at_last_substep": als, "lprep_adv": la} + for afs, als, la in itertools.product(*([(True, False)] * 3)) + if not (afs and als) + ], + ids=lambda p: ( + f"at_first_substep[{p['at_first_substep']}]__" + f"at_last_substep[{p['at_last_substep']}]__" + f"lprep_adv[{p['lprep_adv']}]" + ), + ) + def input_data( + self, request: pytest.FixtureRequest, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + base_data = TestVerticallyImplicitSolverAtCorrectorStep.input_data.__wrapped__( + self, request, grid ) + base_data["at_first_substep"] = request.param["at_first_substep"] + base_data["at_last_substep"] = request.param["at_last_substep"] + base_data["lprep_adv"] = request.param["lprep_adv"] + base_data["is_iau_active"] = False + base_data["end_index_of_damping_layer"] = 13 + base_data["kstart_moist"] = 0 + return base_data 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 b4866a0f75..7f267a0375 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 @@ -18,7 +18,7 @@ from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import stencil_tests +from icon4py.model.testing import definitions, stencil_tests from .test_add_analysis_increments_from_data_assimilation import ( add_analysis_increments_from_data_assimilation_numpy, @@ -53,6 +53,7 @@ ) +@pytest.mark.single_precision_ready @pytest.mark.uses_concat_where class TestVerticallyImplicitSolverAtPredictorStep(stencil_tests.StencilTest): PROGRAM = vertically_implicit_solver_at_predictor_step @@ -64,6 +65,38 @@ class TestVerticallyImplicitSolverAtPredictorStep(stencil_tests.StencilTest): "dwdz_at_cells_on_model_levels", "exner_dynamical_increment", ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "start_cell_index_nudging", + "end_cell_index_local", + "start_cell_index_lateral_lvl3", + "end_cell_index_halo_lvl1", + "end_index_of_damping_layer", + "kstart_moist", + "flat_level_index_plus1", + "vertical_start_index_model_top", + "vertical_end_index_model_surface", + "divdamp_type", + "rayleigh_type", + "is_iau_active", + "at_first_substep", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "end_index_of_damping_layer", + "kstart_moist", + "flat_level_index_plus1", + "vertical_start_index_model_top", + "vertical_end_index_model_surface", + "divdamp_type", + "rayleigh_type", + "is_iau_active", + "at_first_substep", + ), + } + if ta.precision == "single": + RTOL = 1e-1 + ATOL = 1e-2 @staticmethod def reference( @@ -366,126 +399,248 @@ def reference( exner_dynamical_increment=exner_dynamical_increment, ) - @pytest.fixture - def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) - mass_flux_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) - theta_v_flux_at_edges_on_model_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim - ) - predictor_vertical_wind_advective_tendency = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} - ) - pressure_buoyancy_acceleration_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim - ) - rho_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, low=1.0e-5 - ) - contravariant_correction_at_cells_on_half_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} - ) - contravariant_correction_at_edges_on_model_levels = data_alloc.random_field( - grid, dims.EdgeDim, dims.KDim - ) - exner_w_explicit_weight_parameter = data_alloc.random_field(grid, dims.CellDim) - current_exner = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - current_rho = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - current_theta_v = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - current_w = data_alloc.random_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) - inv_ddqz_z_full = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - exner_w_implicit_weight_parameter = data_alloc.random_field(grid, dims.CellDim) - theta_v_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, low=1.0e-5 + # @pytest.fixture( + # params=[{"at_first_substep": value} for value in [True, False]], + # ids=lambda param: f"at_first_substep[{param['at_first_substep']}]", + # ) + # def input_data( + # self, request: pytest.FixtureRequest, grid: base.Grid + # ) -> dict[str, gtx.Field | state_utils.ScalarType]: + # geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) + # mass_flux_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) + # theta_v_flux_at_edges_on_model_levels = data_alloc.random_field( + # grid, dims.EdgeDim, dims.KDim + # ) + # predictor_vertical_wind_advective_tendency = data_alloc.random_field( + # grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + # ) + # pressure_buoyancy_acceleration_at_cells_on_half_levels = data_alloc.random_field( + # grid, dims.CellDim, dims.KDim + # ) + # rho_at_cells_on_half_levels = data_alloc.random_field( + # grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, low=1.0e-5 + # ) + # contravariant_correction_at_cells_on_half_levels = data_alloc.zero_field( + # grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} + # ) + # contravariant_correction_at_edges_on_model_levels = data_alloc.random_field( + # grid, dims.EdgeDim, dims.KDim + # ) + # exner_w_explicit_weight_parameter = data_alloc.random_field(grid, dims.CellDim) + # current_exner = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) + # current_rho = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) + # current_theta_v = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) + # current_w = data_alloc.random_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) + # inv_ddqz_z_full = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) + # exner_w_implicit_weight_parameter = data_alloc.random_field(grid, dims.CellDim) + # theta_v_at_cells_on_half_levels = data_alloc.random_field( + # grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, low=1.0e-5 + # ) + # perturbed_exner_at_cells_on_model_levels = data_alloc.random_field( + # grid, dims.CellDim, dims.KDim + # ) + # exner_tendency_due_to_slow_physics = data_alloc.random_field(grid, dims.CellDim, dims.KDim) + # rho_iau_increment = data_alloc.random_field(grid, dims.CellDim, dims.KDim) + # exner_iau_increment = data_alloc.random_field(grid, dims.CellDim, dims.KDim) + # ddqz_z_half = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) + # rayleigh_damping_factor = data_alloc.random_field(grid, dims.KDim) + # reference_exner_at_cells_on_model_levels = data_alloc.random_field( + # grid, dims.CellDim, dims.KDim, low=1.0e-5 + # ) + # e_bln_c_s = data_alloc.random_field( + # grid, dims.CellDim, dims.C2EDim, low=1.0e-5, high=0.99999 + # ) + # wgtfac_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5, high=0.99999) + # wgtfacq_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5, high=0.99999) + + # next_w = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) + # next_rho = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) + # next_exner = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) + # next_theta_v = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) + # dwdz_at_cells_on_model_levels = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + # exner_dynamical_increment = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) + + # is_iau_active = True + # at_first_substep = request.param["at_first_substep"] + # rayleigh_type = 2 + # divdamp_type = 3 + # end_index_of_damping_layer = 3 + # kstart_moist = 1 + # flat_level_index_plus1 = 3 + # dtime = 0.001 + # iau_wgt_dyn = 1.0 + + # cell_domain = h_grid.domain(dims.CellDim) + # start_cell_nudging = grid.start_index(cell_domain(h_grid.Zone.NUDGING)) + # end_cell_local = grid.end_index(cell_domain(h_grid.Zone.LOCAL)) + # start_cell_index_lateral_lvl3 = grid.start_index( + # cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) + # ) + # end_cell_index_halo_lvl1 = grid.end_index(cell_domain(h_grid.Zone.HALO)) + + # return dict( + # contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, + # next_w=next_w, + # next_rho=next_rho, + # next_exner=next_exner, + # next_theta_v=next_theta_v, + # dwdz_at_cells_on_model_levels=dwdz_at_cells_on_model_levels, + # exner_dynamical_increment=exner_dynamical_increment, + # geofac_div=geofac_div, + # mass_flux_at_edges_on_model_levels=mass_flux_at_edges_on_model_levels, + # theta_v_flux_at_edges_on_model_levels=theta_v_flux_at_edges_on_model_levels, + # predictor_vertical_wind_advective_tendency=predictor_vertical_wind_advective_tendency, + # pressure_buoyancy_acceleration_at_cells_on_half_levels=pressure_buoyancy_acceleration_at_cells_on_half_levels, + # rho_at_cells_on_half_levels=rho_at_cells_on_half_levels, + # contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, + # exner_w_explicit_weight_parameter=exner_w_explicit_weight_parameter, + # current_exner=current_exner, + # current_rho=current_rho, + # current_theta_v=current_theta_v, + # current_w=current_w, + # inv_ddqz_z_full=inv_ddqz_z_full, + # exner_w_implicit_weight_parameter=exner_w_implicit_weight_parameter, + # theta_v_at_cells_on_half_levels=theta_v_at_cells_on_half_levels, + # perturbed_exner_at_cells_on_model_levels=perturbed_exner_at_cells_on_model_levels, + # exner_tendency_due_to_slow_physics=exner_tendency_due_to_slow_physics, + # rho_iau_increment=rho_iau_increment, + # exner_iau_increment=exner_iau_increment, + # ddqz_z_half=ddqz_z_half, + # rayleigh_damping_factor=rayleigh_damping_factor, + # reference_exner_at_cells_on_model_levels=reference_exner_at_cells_on_model_levels, + # e_bln_c_s=e_bln_c_s, + # wgtfac_c=wgtfac_c, + # wgtfacq_c=wgtfacq_c, + # iau_wgt_dyn=iau_wgt_dyn, + # dtime=dtime, + # is_iau_active=is_iau_active, + # rayleigh_type=rayleigh_type, + # divdamp_type=divdamp_type, + # at_first_substep=at_first_substep, + # end_index_of_damping_layer=end_index_of_damping_layer, + # kstart_moist=kstart_moist, + # flat_level_index_plus1=flat_level_index_plus1, + # start_cell_index_nudging=start_cell_nudging, + # end_cell_index_local=end_cell_local, + # start_cell_index_lateral_lvl3=start_cell_index_lateral_lvl3, + # end_cell_index_halo_lvl1=end_cell_index_halo_lvl1, + # vertical_start_index_model_top=gtx.int32(0), + # vertical_end_index_model_surface=gtx.int32(grid.num_levels + 1), + # ) + + @pytest.fixture( + params=[{"at_first_substep": value} for value in [True, False]], + ids=lambda param: f"at_first_substep[{param['at_first_substep']}]", + ) + def input_data( + self, request: pytest.FixtureRequest, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + random_fields = data_alloc.get_random_fields( + grid, + [ + "geofac_div", + "mass_flux_at_edges_on_model_levels", + "theta_v_flux_at_edges_on_model_levels", + "pressure_buoyancy_acceleration_at_cells_on_half_levels", + "contravariant_correction_at_edges_on_model_levels", + "exner_w_explicit_weight_parameter", + "exner_w_implicit_weight_parameter", + "perturbed_exner_at_cells_on_model_levels", + "exner_tendency_due_to_slow_physics", + "rho_iau_increment", + "exner_iau_increment", + "rayleigh_damping_factor", + "predictor_vertical_wind_advective_tendency", + "contravariant_correction_at_cells_on_half_levels", + "current_w", + ], ) - perturbed_exner_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim + + random_with_low = data_alloc.get_random_fields( + grid, + [ + "current_exner", + "current_rho", + "rho_at_cells_on_half_levels", + "inv_ddqz_z_full", + "ddqz_z_half", + "reference_exner_at_cells_on_model_levels", + ], + low=1.0e-5, + ) # "current_theta_v" + + random_low_and_high = data_alloc.get_random_fields( + grid, ["e_bln_c_s", "wgtfac_c", "wgtfacq_c"], low=1.0e-5, high=0.99999 ) - exner_tendency_due_to_slow_physics = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - rho_iau_increment = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - exner_iau_increment = data_alloc.random_field(grid, dims.CellDim, dims.KDim) - ddqz_z_half = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5) - rayleigh_damping_factor = data_alloc.random_field(grid, dims.KDim) - reference_exner_at_cells_on_model_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, low=1.0e-5 + + # out vars + zero_fields = data_alloc.get_zero_fields( + grid, + [ + "dwdz_at_cells_on_model_levels", + "exner_dynamical_increment", + "next_rho", + "next_exner", + "next_w", + ], ) - e_bln_c_s = data_alloc.random_field( - grid, dims.CellDim, dims.C2EDim, low=1.0e-5, high=0.99999 + + # theta_fields = data_alloc.get_random_fields(grid, ["next_theta_v", "current_theta_v", "theta_v_at_cells_on_half_levels"], low=1.0e-5) + theta_fields = data_alloc.get_random_fields( + grid, + ["next_theta_v", "current_theta_v", "theta_v_at_cells_on_half_levels"], + low=250, + high=320, ) - wgtfac_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5, high=0.99999) - wgtfacq_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5, high=0.99999) - - next_w = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) - next_rho = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) - next_exner = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) - next_theta_v = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) - dwdz_at_cells_on_model_levels = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - exner_dynamical_increment = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) - - is_iau_active = True - at_first_substep = True - rayleigh_type = 2 - divdamp_type = 3 - end_index_of_damping_layer = 3 - kstart_moist = 1 - flat_level_index_plus1 = 3 - dtime = 0.001 - iau_wgt_dyn = 1.0 + # TODO(pstark): add rho and exner fields that are random deviations from the hydrostatic atmosphere to work with more realistic magnitudes (and potentially have a better handle for analysing error propagation) cell_domain = h_grid.domain(dims.CellDim) - start_cell_nudging = grid.start_index(cell_domain(h_grid.Zone.NUDGING)) - end_cell_local = grid.end_index(cell_domain(h_grid.Zone.LOCAL)) - start_cell_index_lateral_lvl3 = grid.start_index( - cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) + + # return random_fields | random_with_extend | random_with_extend_and_low | random_with_low | random_low_and_high | constant_fields | zero_fields | \ + return ( + random_fields + | random_with_low + | random_low_and_high + | zero_fields + | theta_fields + | dict( + is_iau_active=True, + at_first_substep=request.param["at_first_substep"], + rayleigh_type=2, + divdamp_type=3, + end_index_of_damping_layer=3, + kstart_moist=1, + flat_level_index_plus1=3, + dtime=ta.wpfloat(0.001), + iau_wgt_dyn=ta.wpfloat(1.0), + start_cell_index_nudging=grid.start_index(cell_domain(h_grid.Zone.NUDGING)), + end_cell_index_local=grid.end_index(cell_domain(h_grid.Zone.LOCAL)), + start_cell_index_lateral_lvl3=grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) + ), + end_cell_index_halo_lvl1=grid.end_index(cell_domain(h_grid.Zone.HALO)), + vertical_start_index_model_top=gtx.int32(0), + vertical_end_index_model_surface=gtx.int32(grid.num_levels + 1), + ) ) - end_cell_index_halo_lvl1 = grid.end_index(cell_domain(h_grid.Zone.HALO)) - return dict( - contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, - next_w=next_w, - next_rho=next_rho, - next_exner=next_exner, - next_theta_v=next_theta_v, - dwdz_at_cells_on_model_levels=dwdz_at_cells_on_model_levels, - exner_dynamical_increment=exner_dynamical_increment, - geofac_div=geofac_div, - mass_flux_at_edges_on_model_levels=mass_flux_at_edges_on_model_levels, - theta_v_flux_at_edges_on_model_levels=theta_v_flux_at_edges_on_model_levels, - predictor_vertical_wind_advective_tendency=predictor_vertical_wind_advective_tendency, - pressure_buoyancy_acceleration_at_cells_on_half_levels=pressure_buoyancy_acceleration_at_cells_on_half_levels, - rho_at_cells_on_half_levels=rho_at_cells_on_half_levels, - contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, - exner_w_explicit_weight_parameter=exner_w_explicit_weight_parameter, - current_exner=current_exner, - current_rho=current_rho, - current_theta_v=current_theta_v, - current_w=current_w, - inv_ddqz_z_full=inv_ddqz_z_full, - exner_w_implicit_weight_parameter=exner_w_implicit_weight_parameter, - theta_v_at_cells_on_half_levels=theta_v_at_cells_on_half_levels, - perturbed_exner_at_cells_on_model_levels=perturbed_exner_at_cells_on_model_levels, - exner_tendency_due_to_slow_physics=exner_tendency_due_to_slow_physics, - rho_iau_increment=rho_iau_increment, - exner_iau_increment=exner_iau_increment, - ddqz_z_half=ddqz_z_half, - rayleigh_damping_factor=rayleigh_damping_factor, - reference_exner_at_cells_on_model_levels=reference_exner_at_cells_on_model_levels, - e_bln_c_s=e_bln_c_s, - wgtfac_c=wgtfac_c, - wgtfacq_c=wgtfacq_c, - iau_wgt_dyn=iau_wgt_dyn, - dtime=dtime, - is_iau_active=is_iau_active, - rayleigh_type=rayleigh_type, - divdamp_type=divdamp_type, - at_first_substep=at_first_substep, - end_index_of_damping_layer=end_index_of_damping_layer, - kstart_moist=kstart_moist, - flat_level_index_plus1=flat_level_index_plus1, - start_cell_index_nudging=start_cell_nudging, - end_cell_index_local=end_cell_local, - start_cell_index_lateral_lvl3=start_cell_index_lateral_lvl3, - end_cell_index_halo_lvl1=end_cell_index_halo_lvl1, - vertical_start_index_model_top=gtx.int32(0), - vertical_end_index_model_surface=gtx.int32(grid.num_levels + 1), +@pytest.mark.continuous_benchmarking +class TestVerticallyImplicitSolverAtPredictorStepContinuousBenchmarking( + TestVerticallyImplicitSolverAtPredictorStep +): + @pytest.fixture( + params=[{"at_first_substep": value} for value in [True, False]], + ids=lambda param: f"at_first_substep[{param['at_first_substep']}]", + ) + def input_data( + self, request: pytest.FixtureRequest, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + base_data = TestVerticallyImplicitSolverAtPredictorStep.input_data.__wrapped__( + self, request, grid ) + base_data["at_first_substep"] = request.param["at_first_substep"] + base_data["is_iau_active"] = False + base_data["divdamp_type"] = 32 + base_data["end_index_of_damping_layer"] = 13 + base_data["kstart_moist"] = 0 + return base_data diff --git a/model/common/src/icon4py/model/common/constants.py b/model/common/src/icon4py/model/common/constants.py index 6e8fb138fc..0d6fdfa96f 100644 --- a/model/common/src/icon4py/model/common/constants.py +++ b/model/common/src/icon4py/model/common/constants.py @@ -6,123 +6,119 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -import sys from typing import Final from gt4py.eve import utils as eve_utils +from numpy import finfo as float_info -from icon4py.model.common import type_alias as ta +from icon4py.model.common.type_alias import vpfloat, wpfloat #: Gas constant for dry air [J/K/kg], called 'rd' in ICON (mo_physical_constants.f90), #: see https://glossary.ametsoc.org/wiki/Gas_constant. -GAS_CONSTANT_DRY_AIR: Final[ta.wpfloat] = 287.04 -RD: Final[ta.wpfloat] = GAS_CONSTANT_DRY_AIR +GAS_CONSTANT_DRY_AIR: Final[wpfloat] = wpfloat(287.04) +RD: Final[wpfloat] = GAS_CONSTANT_DRY_AIR #: Specific heat capacity of dry air at constant pressure [J/K/kg] -SPECIFIC_HEAT_CAPACITY_PRESSURE_DRY_AIR: Final[ta.wpfloat] = 1004.64 +SPECIFIC_HEAT_CAPACITY_PRESSURE_DRY_AIR: Final[wpfloat] = wpfloat(1004.64) CPD = SPECIFIC_HEAT_CAPACITY_PRESSURE_DRY_AIR #: [J/K/kg] specific heat capacity at constant volume -SPECIFIC_HEAT_CAPACITY_VOLUME_DRY_AIR: Final[ta.wpfloat] = CPD - RD -CVD: Final[ta.wpfloat] = SPECIFIC_HEAT_CAPACITY_VOLUME_DRY_AIR -CVD_O_RD: Final[ta.wpfloat] = CVD / RD -RD_O_CPD: Final[ta.wpfloat] = RD / CPD -CPD_O_RD: Final[ta.wpfloat] = CPD / RD -RD_O_CVD: Final[ta.wpfloat] = RD / CVD +SPECIFIC_HEAT_CAPACITY_VOLUME_DRY_AIR: Final[wpfloat] = CPD - RD +CVD: Final[wpfloat] = SPECIFIC_HEAT_CAPACITY_VOLUME_DRY_AIR +CVD_O_RD: Final[wpfloat] = CVD / RD +RD_O_CPD: Final[wpfloat] = RD / CPD +CPD_O_RD: Final[wpfloat] = CPD / RD +RD_O_CVD: Final[wpfloat] = RD / CVD #: Gas constant for water vapor [J/K/kg], rv in ICON. -GAS_CONSTANT_WATER_VAPOR: Final[ta.wpfloat] = 461.51 -RV: Final[ta.wpfloat] = GAS_CONSTANT_WATER_VAPOR +GAS_CONSTANT_WATER_VAPOR: Final[wpfloat] = wpfloat(461.51) +RV: Final[wpfloat] = GAS_CONSTANT_WATER_VAPOR #: Specific heat capacity of water vapor at constant pressure [J/K/kg] -SPECIFIC_HEAT_CAPACITY_PRESSURE_WATER_VAPOR: Final[ta.wpfloat] = 1869.46 +SPECIFIC_HEAT_CAPACITY_PRESSURE_WATER_VAPOR: Final[wpfloat] = wpfloat(1869.46) CPV = SPECIFIC_HEAT_CAPACITY_PRESSURE_WATER_VAPOR #: Specific heat capacity of water vapor at constant volume [J/K/kg] -SPECIFIC_HEAT_CAPACITY_VOLUME_WATER_VAPOR: Final[ta.wpfloat] = CPV - RV +SPECIFIC_HEAT_CAPACITY_VOLUME_WATER_VAPOR: Final[wpfloat] = CPV - RV CVV = SPECIFIC_HEAT_CAPACITY_VOLUME_WATER_VAPOR #: cp_dry_air / cp_liquid_water - 1 -_RCPL: Final[ta.wpfloat] = 3.1733 +_RCPL: Final[wpfloat] = wpfloat(3.1733) #: Specific heat capacity of liquid water [J/K/kg]. Originally expressed as clw in ICON. -SPECIFIC_HEAT_CAPACITY_LIQUID_WATER: Final[ta.wpfloat] = (_RCPL + 1.0) * CPD +SPECIFIC_HEAT_CAPACITY_LIQUID_WATER: Final[wpfloat] = (_RCPL + wpfloat(1.0)) * CPD CPL = SPECIFIC_HEAT_CAPACITY_LIQUID_WATER #: density of liquid water. Originally expressed as rhow in ICON. [kg/m3] -WATER_DENSITY: Final[ta.wpfloat] = 1.000e3 +WATER_DENSITY: Final[wpfloat] = wpfloat(1.000e3) #: specific heat capacity of ice. Originally expressed as ci in ICON. [J/K/kg] -SPECIFIC_HEAT_CAPACITY_ICE: Final[ta.wpfloat] = 2108.0 +SPECIFIC_HEAT_CAPACITY_ICE: Final[wpfloat] = wpfloat(2108.0) #: Melting temperature of ice/snow [K]. Originally expressed as tmelt in ICON. -MELTING_TEMPERATURE: Final[ta.wpfloat] = 273.15 +MELTING_TEMPERATURE: Final[wpfloat] = wpfloat(273.15) #: Latent heat of vaporisation for water [J/kg]. Originally expressed as alv in ICON. -LATENT_HEAT_FOR_VAPORISATION: Final[ta.wpfloat] = 2.5008e6 +LATENT_HEAT_FOR_VAPORISATION: Final[wpfloat] = wpfloat(2.5008e6) #: Latent heat of sublimation for water [J/kg]. Originally expressed as als in ICON. -LATENT_HEAT_FOR_SUBLIMATION: Final[ta.wpfloat] = 2.8345e6 +LATENT_HEAT_FOR_SUBLIMATION: Final[wpfloat] = wpfloat(2.8345e6) #: Latent heat of fusion for water [J/kg]. Originally expressed as alf in ICON. -LATENT_HEAT_FOR_FUSION: Final[ta.wpfloat] = ( - LATENT_HEAT_FOR_SUBLIMATION - LATENT_HEAT_FOR_VAPORISATION -) +LATENT_HEAT_FOR_FUSION: Final[wpfloat] = LATENT_HEAT_FOR_SUBLIMATION - LATENT_HEAT_FOR_VAPORISATION #: Triple point of water at 611hPa [K] -WATER_TRIPLE_POINT_TEMPERATURE: Final[ta.wpfloat] = 273.16 +WATER_TRIPLE_POINT_TEMPERATURE: Final[wpfloat] = wpfloat(273.16) #: RV/RD - 1, tvmpc1 in ICON. -RV_O_RD_MINUS_1: Final[ta.wpfloat] = GAS_CONSTANT_WATER_VAPOR / GAS_CONSTANT_DRY_AIR - 1.0 -TVMPC1: Final[ta.wpfloat] = RV_O_RD_MINUS_1 +RV_O_RD_MINUS_1: Final[wpfloat] = GAS_CONSTANT_WATER_VAPOR / GAS_CONSTANT_DRY_AIR - wpfloat(1.0) +TVMPC1: Final[wpfloat] = RV_O_RD_MINUS_1 #: Av. gravitational acceleration [m/s^2] -GRAVITATIONAL_ACCELERATION: Final[ta.wpfloat] = 9.80665 -GRAV: Final[ta.wpfloat] = GRAVITATIONAL_ACCELERATION -GRAV_O_RD: Final[ta.wpfloat] = GRAV / RD -GRAV_O_CPD: Final[ta.wpfloat] = GRAV / CPD +GRAVITATIONAL_ACCELERATION: Final[wpfloat] = wpfloat(9.80665) +GRAV: Final[wpfloat] = GRAVITATIONAL_ACCELERATION +GRAV_O_RD: Final[wpfloat] = GRAV / RD +GRAV_O_CPD: Final[wpfloat] = GRAV / CPD #: reference pressure for Exner function [Pa] -REFERENCE_PRESSURE: Final[ta.wpfloat] = 100000.0 -P0REF: Final[ta.wpfloat] = REFERENCE_PRESSURE -RD_O_P0REF: Final[ta.wpfloat] = RD / P0REF +REFERENCE_PRESSURE: Final[wpfloat] = wpfloat(100000.0) +P0REF: Final[wpfloat] = REFERENCE_PRESSURE +RD_O_P0REF: Final[wpfloat] = RD / P0REF #: sea level pressure [Pa] -SEA_LEVEL_PRESSURE: Final[ta.wpfloat] = 101325.0 -P0SL_BG: Final[ta.wpfloat] = SEA_LEVEL_PRESSURE +SEA_LEVEL_PRESSURE: Final[wpfloat] = wpfloat(101325.0) +P0SL_BG: Final[wpfloat] = SEA_LEVEL_PRESSURE # average earth radius in [m] -EARTH_RADIUS: Final[float] = 6.371229e6 +EARTH_RADIUS: Final[wpfloat] = wpfloat(6.371229e6) #: Earth angular velocity [rad/s] -EARTH_ANGULAR_VELOCITY: Final[ta.wpfloat] = 7.29212e-5 +EARTH_ANGULAR_VELOCITY: Final[wpfloat] = wpfloat(7.29212e-5) #: sea level temperature for reference atmosphere [K] -SEA_LEVEL_TEMPERATURE: Final[ta.wpfloat] = 288.15 -T0SL_BG: Final[ta.wpfloat] = SEA_LEVEL_TEMPERATURE +SEA_LEVEL_TEMPERATURE: Final[wpfloat] = wpfloat(288.15) +T0SL_BG: Final[wpfloat] = SEA_LEVEL_TEMPERATURE #: difference between sea level temperature and asymptotic stratospheric temperature -DELTA_TEMPERATURE: Final[ta.wpfloat] = 75.0 -DEL_T_BG: Final[ta.wpfloat] = DELTA_TEMPERATURE +DELTA_TEMPERATURE: Final[wpfloat] = wpfloat(75.0) +DEL_T_BG: Final[wpfloat] = DELTA_TEMPERATURE #: height scale for reference atmosphere [m], defined in mo_vertical_grid #: scale height [m] -HEIGHT_SCALE_FOR_REFERENCE_ATMOSPHERE = 10000.0 -_H_SCAL_BG: Final[ta.wpfloat] = HEIGHT_SCALE_FOR_REFERENCE_ATMOSPHERE +HEIGHT_SCALE_FOR_REFERENCE_ATMOSPHERE = wpfloat(10000.0) +_H_SCAL_BG: Final[wpfloat] = HEIGHT_SCALE_FOR_REFERENCE_ATMOSPHERE # Math constants -DBL_EPS = sys.float_info.epsilon # EPSILON(1._wp) +WP_EPS = float_info(wpfloat).eps # EPSILON(1._wp) +VP_EPS = float_info(vpfloat).eps # Implementation constants #: default dynamics to physics time step ratio -DEFAULT_DYNAMICS_TO_PHYSICS_TIMESTEP_RATIO: Final[float] = 5.0 +DEFAULT_DYNAMICS_TO_PHYSICS_TIMESTEP_RATIO: Final[wpfloat] = wpfloat(5.0) -#: average earth radius in [m] -EARTH_RADIUS: Final[ta.wpfloat] = 6.371229e6 - -class PhysicsConstants(eve_utils.FrozenNamespace[ta.wpfloat]): +class PhysicsConstants(eve_utils.FrozenNamespace[wpfloat]): """ Constants used in gt4py stencils. """ @@ -150,7 +146,7 @@ class PhysicsConstants(eve_utils.FrozenNamespace[ta.wpfloat]): grav_o_cpd = GRAV_O_CPD grav_o_rd = GRAV_O_RD p0ref = REFERENCE_PRESSURE - eps = DBL_EPS + eps = WP_EPS class RayleighType(eve_utils.FrozenNamespace[int]): diff --git a/model/common/src/icon4py/model/common/grid/geometry_stencils.py b/model/common/src/icon4py/model/common/grid/geometry_stencils.py index 90addf46df..ba7599b932 100644 --- a/model/common/src/icon4py/model/common/grid/geometry_stencils.py +++ b/model/common/src/icon4py/model/common/grid/geometry_stencils.py @@ -12,7 +12,7 @@ from gt4py import next as gtx from gt4py.next import sin, where -from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import E2C, E2C2V, E2V, EdgeDim from icon4py.model.common.math.helpers import ( arc_length_on_edges, @@ -22,15 +22,16 @@ normalize_cartesian_vector_on_edges, zonal_and_meridional_components_on_edges, ) +from icon4py.model.common.type_alias import wpfloat from icon4py.model.common.utils import data_allocation as alloc @gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) def cartesian_coordinates_of_edge_tangent( - vertex_lat: fa.VertexField[ta.wpfloat], - vertex_lon: fa.VertexField[ta.wpfloat], - edge_orientation: fa.EdgeField[ta.wpfloat], -) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: + vertex_lat: fa.VertexField[wpfloat], + vertex_lon: fa.VertexField[wpfloat], + edge_orientation: fa.EdgeField[wpfloat], +) -> tuple[fa.EdgeField[wpfloat], fa.EdgeField[wpfloat], fa.EdgeField[wpfloat]]: """ Compute normalized cartesian vector tangential to an edge. @@ -58,15 +59,15 @@ def cartesian_coordinates_of_edge_tangent( @gtx.field_operator def cartesian_coordinates_of_edge_normal( - edge_lat: fa.EdgeField[ta.wpfloat], - edge_lon: fa.EdgeField[ta.wpfloat], - edge_tangent_x: fa.EdgeField[ta.wpfloat], - edge_tangent_y: fa.EdgeField[ta.wpfloat], - edge_tangent_z: fa.EdgeField[ta.wpfloat], + edge_lat: fa.EdgeField[wpfloat], + edge_lon: fa.EdgeField[wpfloat], + edge_tangent_x: fa.EdgeField[wpfloat], + edge_tangent_y: fa.EdgeField[wpfloat], + edge_tangent_z: fa.EdgeField[wpfloat], ) -> tuple[ - fa.EdgeField[ta.wpfloat], - fa.EdgeField[ta.wpfloat], - fa.EdgeField[ta.wpfloat], + fa.EdgeField[wpfloat], + fa.EdgeField[wpfloat], + fa.EdgeField[wpfloat], ]: """ Compute the normal to the edge tangent vector. @@ -95,18 +96,18 @@ def cartesian_coordinates_of_edge_normal( @gtx.field_operator def cartesian_coordinates_edge_tangent_and_normal( - vertex_lat: fa.VertexField[ta.wpfloat], - vertex_lon: fa.VertexField[ta.wpfloat], - edge_lat: fa.EdgeField[ta.wpfloat], - edge_lon: fa.EdgeField[ta.wpfloat], - edge_orientation: fa.EdgeField[ta.wpfloat], + vertex_lat: fa.VertexField[wpfloat], + vertex_lon: fa.VertexField[wpfloat], + edge_lat: fa.EdgeField[wpfloat], + edge_lon: fa.EdgeField[wpfloat], + edge_orientation: fa.EdgeField[wpfloat], ) -> tuple[ - fa.EdgeField[ta.wpfloat], - fa.EdgeField[ta.wpfloat], - fa.EdgeField[ta.wpfloat], - fa.EdgeField[ta.wpfloat], - fa.EdgeField[ta.wpfloat], - fa.EdgeField[ta.wpfloat], + fa.EdgeField[wpfloat], + fa.EdgeField[wpfloat], + fa.EdgeField[wpfloat], + fa.EdgeField[wpfloat], + fa.EdgeField[wpfloat], + fa.EdgeField[wpfloat], ]: """Compute normalized cartesian vectors of edge tangent and edge normal.""" tangent_x, tangent_y, tangent_z = cartesian_coordinates_of_edge_tangent( @@ -125,17 +126,17 @@ def cartesian_coordinates_edge_tangent_and_normal( @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_cartesian_coordinates_of_edge_tangent_and_normal( - vertex_lat: fa.VertexField[ta.wpfloat], - vertex_lon: fa.VertexField[ta.wpfloat], - edge_lat: fa.EdgeField[ta.wpfloat], - edge_lon: fa.EdgeField[ta.wpfloat], - edge_orientation: fa.EdgeField[ta.wpfloat], - tangent_x: fa.EdgeField[ta.wpfloat], - tangent_y: fa.EdgeField[ta.wpfloat], - tangent_z: fa.EdgeField[ta.wpfloat], - normal_x: fa.EdgeField[ta.wpfloat], - normal_y: fa.EdgeField[ta.wpfloat], - normal_z: fa.EdgeField[ta.wpfloat], + vertex_lat: fa.VertexField[wpfloat], + vertex_lon: fa.VertexField[wpfloat], + edge_lat: fa.EdgeField[wpfloat], + edge_lon: fa.EdgeField[wpfloat], + edge_orientation: fa.EdgeField[wpfloat], + tangent_x: fa.EdgeField[wpfloat], + tangent_y: fa.EdgeField[wpfloat], + tangent_z: fa.EdgeField[wpfloat], + normal_x: fa.EdgeField[wpfloat], + normal_y: fa.EdgeField[wpfloat], + normal_z: fa.EdgeField[wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, ): @@ -152,20 +153,20 @@ def compute_cartesian_coordinates_of_edge_tangent_and_normal( @gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) def zonal_and_meridional_component_of_edge_field_at_vertex( - vertex_lat: fa.VertexField[ta.wpfloat], - vertex_lon: fa.VertexField[ta.wpfloat], - x: fa.EdgeField[ta.wpfloat], - y: fa.EdgeField[ta.wpfloat], - z: fa.EdgeField[ta.wpfloat], + vertex_lat: fa.VertexField[wpfloat], + vertex_lon: fa.VertexField[wpfloat], + x: fa.EdgeField[wpfloat], + y: fa.EdgeField[wpfloat], + z: fa.EdgeField[wpfloat], ) -> tuple[ - fa.EdgeField[ta.wpfloat], - fa.EdgeField[ta.wpfloat], - fa.EdgeField[ta.wpfloat], - fa.EdgeField[ta.wpfloat], - fa.EdgeField[ta.wpfloat], - fa.EdgeField[ta.wpfloat], - fa.EdgeField[ta.wpfloat], - fa.EdgeField[ta.wpfloat], + fa.EdgeField[wpfloat], + fa.EdgeField[wpfloat], + fa.EdgeField[wpfloat], + fa.EdgeField[wpfloat], + fa.EdgeField[wpfloat], + fa.EdgeField[wpfloat], + fa.EdgeField[wpfloat], + fa.EdgeField[wpfloat], ]: """ Compute the zonal (u) an meridional (v) component of a cartesian vector (x, y, z) at the vertex position (lat, lon). @@ -223,19 +224,19 @@ def zonal_and_meridional_component_of_edge_field_at_vertex( @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_zonal_and_meridional_component_of_edge_field_at_vertex( - vertex_lat: fa.VertexField[ta.wpfloat], - vertex_lon: fa.VertexField[ta.wpfloat], - x: fa.EdgeField[ta.wpfloat], - y: fa.EdgeField[ta.wpfloat], - z: fa.EdgeField[ta.wpfloat], - u_vertex_1: fa.EdgeField[ta.wpfloat], - v_vertex_1: fa.EdgeField[ta.wpfloat], - u_vertex_2: fa.EdgeField[ta.wpfloat], - v_vertex_2: fa.EdgeField[ta.wpfloat], - u_vertex_3: fa.EdgeField[ta.wpfloat], - v_vertex_3: fa.EdgeField[ta.wpfloat], - u_vertex_4: fa.EdgeField[ta.wpfloat], - v_vertex_4: fa.EdgeField[ta.wpfloat], + vertex_lat: fa.VertexField[wpfloat], + vertex_lon: fa.VertexField[wpfloat], + x: fa.EdgeField[wpfloat], + y: fa.EdgeField[wpfloat], + z: fa.EdgeField[wpfloat], + u_vertex_1: fa.EdgeField[wpfloat], + v_vertex_1: fa.EdgeField[wpfloat], + u_vertex_2: fa.EdgeField[wpfloat], + v_vertex_2: fa.EdgeField[wpfloat], + u_vertex_3: fa.EdgeField[wpfloat], + v_vertex_3: fa.EdgeField[wpfloat], + u_vertex_4: fa.EdgeField[wpfloat], + v_vertex_4: fa.EdgeField[wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, ): @@ -261,16 +262,16 @@ def compute_zonal_and_meridional_component_of_edge_field_at_vertex( @gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) def zonal_and_meridional_component_of_edge_field_at_cell_center( - cell_lat: fa.CellField[ta.wpfloat], - cell_lon: fa.CellField[ta.wpfloat], - x: fa.EdgeField[ta.wpfloat], - y: fa.EdgeField[ta.wpfloat], - z: fa.EdgeField[ta.wpfloat], + cell_lat: fa.CellField[wpfloat], + cell_lon: fa.CellField[wpfloat], + x: fa.EdgeField[wpfloat], + y: fa.EdgeField[wpfloat], + z: fa.EdgeField[wpfloat], ) -> tuple[ - fa.EdgeField[ta.wpfloat], - fa.EdgeField[ta.wpfloat], - fa.EdgeField[ta.wpfloat], - fa.EdgeField[ta.wpfloat], + fa.EdgeField[wpfloat], + fa.EdgeField[wpfloat], + fa.EdgeField[wpfloat], + fa.EdgeField[wpfloat], ]: """ Compute zonal (U) and meridional (V) component of a vector (x, y, z) at cell centers (lat, lon) @@ -307,15 +308,15 @@ def zonal_and_meridional_component_of_edge_field_at_cell_center( @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_zonal_and_meridional_component_of_edge_field_at_cell_center( - cell_lat: fa.CellField[ta.wpfloat], - cell_lon: fa.CellField[ta.wpfloat], - x: fa.EdgeField[ta.wpfloat], - y: fa.EdgeField[ta.wpfloat], - z: fa.EdgeField[ta.wpfloat], - u_cell_1: fa.EdgeField[ta.wpfloat], - v_cell_1: fa.EdgeField[ta.wpfloat], - u_cell_2: fa.EdgeField[ta.wpfloat], - v_cell_2: fa.EdgeField[ta.wpfloat], + cell_lat: fa.CellField[wpfloat], + cell_lon: fa.CellField[wpfloat], + x: fa.EdgeField[wpfloat], + y: fa.EdgeField[wpfloat], + z: fa.EdgeField[wpfloat], + u_cell_1: fa.EdgeField[wpfloat], + v_cell_1: fa.EdgeField[wpfloat], + u_cell_2: fa.EdgeField[wpfloat], + v_cell_2: fa.EdgeField[wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, ): @@ -337,12 +338,12 @@ def compute_zonal_and_meridional_component_of_edge_field_at_cell_center( @gtx.field_operator def cell_center_arc_distance( - lat_neighbor_0: fa.EdgeField[ta.wpfloat], - lon_neighbor_0: fa.EdgeField[ta.wpfloat], - lat_neighbor_1: fa.EdgeField[ta.wpfloat], - lon_neighbor_1: fa.EdgeField[ta.wpfloat], - radius: ta.wpfloat, -) -> fa.EdgeField[ta.wpfloat]: + lat_neighbor_0: fa.EdgeField[wpfloat], + lon_neighbor_0: fa.EdgeField[wpfloat], + lat_neighbor_1: fa.EdgeField[wpfloat], + lon_neighbor_1: fa.EdgeField[wpfloat], + radius: wpfloat, +) -> fa.EdgeField[wpfloat]: """ Compute the distance between to cell centers. @@ -368,10 +369,10 @@ def cell_center_arc_distance( @gtx.field_operator def arc_distance_of_far_edges_in_diamond( - vertex_lat: fa.VertexField[ta.wpfloat], - vertex_lon: fa.VertexField[ta.wpfloat], - radius: ta.wpfloat, -) -> fa.EdgeField[ta.wpfloat]: + vertex_lat: fa.VertexField[wpfloat], + vertex_lon: fa.VertexField[wpfloat], + radius: wpfloat, +) -> fa.EdgeField[wpfloat]: """ Compute the arc length between the "far" vertices of an edge. @@ -412,10 +413,10 @@ def arc_distance_of_far_edges_in_diamond( @gtx.field_operator def edge_length( - vertex_lat: fa.VertexField[ta.wpfloat], - vertex_lon: fa.VertexField[ta.wpfloat], - radius: ta.wpfloat, -) -> fa.EdgeField[ta.wpfloat]: + vertex_lat: fa.VertexField[wpfloat], + vertex_lon: fa.VertexField[wpfloat], + radius: wpfloat, +) -> fa.EdgeField[wpfloat]: """ Compute the arc length of an edge. @@ -446,10 +447,10 @@ def edge_length( @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_edge_length( - vertex_lat: fa.VertexField[ta.wpfloat], - vertex_lon: fa.VertexField[ta.wpfloat], - radius: ta.wpfloat, - length: fa.EdgeField[ta.wpfloat], + vertex_lat: fa.VertexField[wpfloat], + vertex_lon: fa.VertexField[wpfloat], + radius: wpfloat, + length: fa.EdgeField[wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, ): @@ -464,12 +465,12 @@ def compute_edge_length( @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_cell_center_arc_distance( - edge_neighbor_0_lat: fa.EdgeField[ta.wpfloat], - edge_neighbor_0_lon: fa.EdgeField[ta.wpfloat], - edge_neighbor_1_lat: fa.EdgeField[ta.wpfloat], - edge_neighbor_1_lon: fa.EdgeField[ta.wpfloat], - radius: ta.wpfloat, - dual_edge_length: fa.EdgeField[ta.wpfloat], + edge_neighbor_0_lat: fa.EdgeField[wpfloat], + edge_neighbor_0_lon: fa.EdgeField[wpfloat], + edge_neighbor_1_lat: fa.EdgeField[wpfloat], + edge_neighbor_1_lon: fa.EdgeField[wpfloat], + radius: wpfloat, + dual_edge_length: fa.EdgeField[wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, ): @@ -486,10 +487,10 @@ def compute_cell_center_arc_distance( @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_arc_distance_of_far_edges_in_diamond( - vertex_lat: fa.VertexField[ta.wpfloat], - vertex_lon: fa.VertexField[ta.wpfloat], - radius: ta.wpfloat, - far_vertex_distance: fa.EdgeField[ta.wpfloat], + vertex_lat: fa.VertexField[wpfloat], + vertex_lon: fa.VertexField[wpfloat], + radius: wpfloat, + far_vertex_distance: fa.EdgeField[wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, ): @@ -505,9 +506,9 @@ def compute_arc_distance_of_far_edges_in_diamond( @gtx.field_operator def edge_area( owner_mask: fa.EdgeField[bool], - primal_edge_length: fa.EdgeField[ta.wpfloat], - dual_edge_length: fa.EdgeField[ta.wpfloat], -) -> fa.EdgeField[ta.wpfloat]: + primal_edge_length: fa.EdgeField[wpfloat], + dual_edge_length: fa.EdgeField[wpfloat], +) -> fa.EdgeField[wpfloat]: """ Compute the area spanned by an edge and the its dual edge Args: @@ -519,15 +520,15 @@ def edge_area( area """ - return where(owner_mask, primal_edge_length * dual_edge_length, 0.0) + return where(owner_mask, primal_edge_length * dual_edge_length, wpfloat(0.0)) @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_edge_area( owner_mask: fa.EdgeField[bool], - primal_edge_length: fa.EdgeField[ta.wpfloat], - dual_edge_length: fa.EdgeField[ta.wpfloat], - area: fa.EdgeField[ta.wpfloat], + primal_edge_length: fa.EdgeField[wpfloat], + dual_edge_length: fa.EdgeField[wpfloat], + area: fa.EdgeField[wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, ): @@ -542,9 +543,9 @@ def compute_edge_area( @gtx.field_operator def coriolis_parameter_on_edges( - edge_center_lat: fa.EdgeField[ta.wpfloat], - angular_velocity: ta.wpfloat, -) -> fa.EdgeField[ta.wpfloat]: + edge_center_lat: fa.EdgeField[wpfloat], + angular_velocity: wpfloat, +) -> fa.EdgeField[wpfloat]: """ Compute the coriolis force on edges. Args: @@ -554,14 +555,14 @@ def coriolis_parameter_on_edges( Returns: coriolis parameter """ - return 2.0 * angular_velocity * sin(edge_center_lat) + return wpfloat(2.0) * angular_velocity * sin(edge_center_lat) @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_coriolis_parameter_on_edges( - edge_center_lat: fa.EdgeField[ta.wpfloat], - angular_velocity: ta.wpfloat, - coriolis_parameter: fa.EdgeField[ta.wpfloat], + edge_center_lat: fa.EdgeField[wpfloat], + angular_velocity: wpfloat, + coriolis_parameter: fa.EdgeField[wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, ): diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py index 76ddd675a9..17e34d729c 100644 --- a/model/common/src/icon4py/model/common/grid/vertical.py +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -19,9 +19,9 @@ import numpy as np import icon4py.model.common.states.metadata as data -import icon4py.model.common.type_alias as ta from icon4py.model.common import dimension as dims, exceptions, field_type_aliases as fa from icon4py.model.common.grid import topography as topo +from icon4py.model.common.type_alias import wpfloat from icon4py.model.common.utils import data_allocation as data_alloc @@ -84,40 +84,40 @@ class VerticalGridConfig: #: Number of full levels. num_levels: int #: Defined as max_lay_thckn in ICON namelist mo_sleve_nml. Maximum thickness of grid cells below top_height_limit_for_maximal_layer_thickness. - maximal_layer_thickness: Final[ta.wpfloat] = 25000.0 + maximal_layer_thickness: Final[wpfloat] = 25000.0 #: Defined as htop_thcknlimit in ICON namelist mo_sleve_nml. Height below which thickness of grid cells must not exceed maximal_layer_thickness. - top_height_limit_for_maximal_layer_thickness: Final[ta.wpfloat] = 15000.0 + top_height_limit_for_maximal_layer_thickness: Final[wpfloat] = 15000.0 #: Defined as min_lay_thckn in ICON namelist mo_sleve_nml. Thickness of lowest level grid cells. - lowest_layer_thickness: Final[ta.wpfloat] = 50.0 + lowest_layer_thickness: Final[wpfloat] = 50.0 #: Model top height in ICON namelist mo_sleve_nml. - model_top_height: Final[ta.wpfloat] = 23500.0 + model_top_height: Final[wpfloat] = 23500.0 #: Defined in ICON namelist mo_sleve_nml. Height above which coordinate surfaces are flat - flat_height: Final[ta.wpfloat] = 16000.0 + flat_height: Final[wpfloat] = 16000.0 #: Defined as stretch_fac in ICON namelist mo_sleve_nml. Scaling factor for stretching/squeezing the model layer distribution. - stretch_factor: Final[ta.wpfloat] = 1.0 + stretch_factor: Final[wpfloat] = 1.0 #: Defined as damp_height in ICON namelist nonhydrostatic_nml. Height [m] at which Rayleigh damping of vertical wind starts. - rayleigh_damping_height: Final[ta.wpfloat] = 45000.0 + rayleigh_damping_height: Final[wpfloat] = 45000.0 #: Defined in ICON namelist nonhydrostatic_nml. Height [m] above which moist physics and advection of cloud and precipitation variables are turned off. - htop_moist_proc: Final[ta.wpfloat] = 22500.0 + htop_moist_proc: Final[wpfloat] = 22500.0 #: file name containing vct_a and vct_b table file_path: pathlib.Path | None = None # Parameters for setting up the decay function of the topographic signal for # SLEVE. Default values from mo_sleve_nml. #: Decay scale for large-scale topography component - SLEVE_decay_scale_1: Final[ta.wpfloat] = 4000.0 + SLEVE_decay_scale_1: Final[wpfloat] = 4000.0 #: Decay scale for small-scale topography component - SLEVE_decay_scale_2: Final[ta.wpfloat] = 2500.0 + SLEVE_decay_scale_2: Final[wpfloat] = 2500.0 #: Exponent for decay function - SLEVE_decay_exponent: Final[ta.wpfloat] = 1.2 + SLEVE_decay_exponent: Final[wpfloat] = 1.2 #: minimum absolute layer thickness 1 for SLEVE coordinates - SLEVE_minimum_layer_thickness_1: Final[ta.wpfloat] = 100.0 + SLEVE_minimum_layer_thickness_1: Final[wpfloat] = 100.0 #: minimum absolute layer thickness 2 for SLEVE coordinates - SLEVE_minimum_layer_thickness_2: Final[ta.wpfloat] = 500.0 + SLEVE_minimum_layer_thickness_2: Final[wpfloat] = 500.0 #: minimum relative layer thickness for nominal thicknesses <= SLEVE_minimum_layer_thickness_1 - SLEVE_minimum_relative_layer_thickness_1: Final[ta.wpfloat] = 1.0 / 3.0 + SLEVE_minimum_relative_layer_thickness_1: Final[wpfloat] = 1.0 / 3.0 #: minimum relative layer thickness for a nominal thickness of SLEVE_minimum_layer_thickness_2 - SLEVE_minimum_relative_layer_thickness_2: Final[ta.wpfloat] = 0.5 + SLEVE_minimum_relative_layer_thickness_2: Final[wpfloat] = 0.5 @dataclasses.dataclass(frozen=True) @@ -132,10 +132,10 @@ class VerticalGrid: """ config: VerticalGridConfig - vct_a: dataclasses.InitVar[fa.KField[ta.wpfloat]] - vct_b: dataclasses.InitVar[fa.KField[ta.wpfloat]] - _vct_a: fa.KField[ta.wpfloat] = dataclasses.field(init=False) - _vct_b: fa.KField[ta.wpfloat] = dataclasses.field(init=False) + vct_a: dataclasses.InitVar[fa.KField[wpfloat]] + vct_b: dataclasses.InitVar[fa.KField[wpfloat]] + _vct_a: fa.KField[wpfloat] = dataclasses.field(init=False) + _vct_b: fa.KField[wpfloat] = dataclasses.field(init=False) _end_index_of_damping_layer: Final[gtx.int32] = dataclasses.field(init=False) _start_index_for_moist_physics: Final[gtx.int32] = dataclasses.field(init=False) _end_index_of_flat_layer: Final[gtx.int32] = dataclasses.field(init=False) @@ -182,7 +182,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) ), ] @@ -225,8 +225,8 @@ def _bottom_level(self, domain: Domain) -> int: return self.size(domain.dim) @property - def interface_physical_height(self) -> fa.KField[ta.wpfloat]: - return self._vct_a + def interface_physical_height(self) -> fa.KField[wpfloat]: + return gtx.astype(self._vct_a, wpfloat) @functools.cached_property def kstart_moist(self) -> gtx.int32: @@ -263,7 +263,7 @@ def size(self, dim: gtx.Dimension) -> int: @classmethod def _determine_start_level_of_moist_physics( - cls, vct_a: np.ndarray, top_moist_threshold: ta.wpfloat, nshift_total: int = 0 + cls, vct_a: np.ndarray, top_moist_threshold: wpfloat, nshift_total: int = 0 ) -> gtx.int32: n_levels = vct_a.shape[0] interface_height = 0.5 * (vct_a[: n_levels - 1 - nshift_total] + vct_a[1 + nshift_total :]) @@ -271,7 +271,7 @@ def _determine_start_level_of_moist_physics( @classmethod def _determine_damping_height_index( - cls, vct_a: np.ndarray, damping_height: ta.wpfloat + cls, vct_a: np.ndarray, damping_height: wpfloat ) -> gtx.int32: assert damping_height >= 0.0, "Damping height must be positive." return ( @@ -282,7 +282,7 @@ def _determine_damping_height_index( @classmethod def _determine_end_index_of_flat_layers( - cls, vct_a: np.ndarray, flat_height: ta.wpfloat + cls, vct_a: np.ndarray, flat_height: wpfloat ) -> gtx.int32: assert flat_height >= 0.0, "Flat surface height must be positive." return ( @@ -313,16 +313,16 @@ def _read_vct_a_and_vct_b_from_file( Returns: one dimensional vct_a and vct_b arrays. """ num_levels_plus_one = num_levels + 1 - vct_a = np.zeros(num_levels_plus_one, dtype=ta.wpfloat) - vct_b = np.zeros(num_levels_plus_one, dtype=ta.wpfloat) + vct_a = np.zeros(num_levels_plus_one, dtype=wpfloat) + vct_b = np.zeros(num_levels_plus_one, dtype=wpfloat) try: with file_path.open() as vertical_grid_file: # skip the first line that contains titles vertical_grid_file.readline() for k in range(num_levels_plus_one): grid_content = vertical_grid_file.readline().split() - vct_a[k] = ta.wpfloat(grid_content[1]) - vct_b[k] = ta.wpfloat(grid_content[2]) + vct_a[k] = wpfloat(grid_content[1]) + vct_b[k] = wpfloat(grid_content[2]) except OSError as err: raise FileNotFoundError( f"Vertical coord table file {file_path} could not be read." @@ -389,8 +389,8 @@ def _compute_vct_a_and_vct_b( # noqa: PLR0912 [too-many-branches] 2.0 / math.pi * np.arccos( - ta.wpfloat(vertical_config.num_levels - 1) ** vertical_config.stretch_factor - / ta.wpfloat(vertical_config.num_levels) ** vertical_config.stretch_factor + wpfloat(vertical_config.num_levels - 1) ** vertical_config.stretch_factor + / wpfloat(vertical_config.num_levels) ** vertical_config.stretch_factor ) ) @@ -400,9 +400,9 @@ def _compute_vct_a_and_vct_b( # noqa: PLR0912 [too-many-branches] 2.0 / math.pi * np.arccos( - np.arange(vertical_config.num_levels + 1, dtype=ta.wpfloat) + np.arange(vertical_config.num_levels + 1, dtype=wpfloat) ** vertical_config.stretch_factor - / ta.wpfloat(vertical_config.num_levels) ** vertical_config.stretch_factor + / wpfloat(vertical_config.num_levels) ** vertical_config.stretch_factor ) ) ** vct_a_exponential_factor @@ -417,7 +417,7 @@ def _compute_vct_a_and_vct_b( # noqa: PLR0912 [too-many-branches] lowest_level_exceeding_limit = np.max( np.where(layer_thickness > vertical_config.maximal_layer_thickness) ) - modified_vct_a = np.zeros(num_levels_plus_one, dtype=ta.wpfloat) + modified_vct_a = np.zeros(num_levels_plus_one, dtype=wpfloat) lowest_level_unmodified_thickness = 0 shifted_levels = 0 for k in range(vertical_config.num_levels - 1, -1, -1): @@ -443,13 +443,13 @@ def _compute_vct_a_and_vct_b( # noqa: PLR0912 [too-many-branches] else ( vct_a[0] - modified_vct_a[lowest_level_unmodified_thickness] - - ta.wpfloat(lowest_level_unmodified_thickness) + - wpfloat(lowest_level_unmodified_thickness) * vertical_config.maximal_layer_thickness ) / ( modified_vct_a[0] - modified_vct_a[lowest_level_unmodified_thickness] - - ta.wpfloat(lowest_level_unmodified_thickness) + - wpfloat(lowest_level_unmodified_thickness) * vertical_config.maximal_layer_thickness ) ) @@ -509,11 +509,8 @@ def _compute_vct_a_and_vct_b( # noqa: PLR0912 [too-many-branches] else: vct_a = ( vertical_config.model_top_height - * ( - ta.wpfloat(vertical_config.num_levels) - - np.arange(num_levels_plus_one, dtype=ta.wpfloat) - ) - / ta.wpfloat(vertical_config.num_levels) + * (wpfloat(vertical_config.num_levels) - np.arange(num_levels_plus_one, dtype=wpfloat)) + / wpfloat(vertical_config.num_levels) ) vct_b = np.exp(-vct_a / 5000.0) @@ -561,10 +558,10 @@ def _compute_SLEVE_coordinate_from_vcta_and_topography( geofac_n2s: data_alloc.NDArray, c2e2co: data_alloc.NDArray, nflatlev: int, - model_top_height: ta.wpfloat, - SLEVE_decay_scale_1: ta.wpfloat, - SLEVE_decay_exponent: ta.wpfloat, - SLEVE_decay_scale_2: ta.wpfloat, + model_top_height: wpfloat, + SLEVE_decay_scale_1: wpfloat, + SLEVE_decay_exponent: wpfloat, + SLEVE_decay_scale_2: wpfloat, array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ @@ -581,9 +578,9 @@ def _compute_SLEVE_coordinate_from_vcta_and_topography( def _decay_func( vct_a: data_alloc.NDArray, - model_top_height: ta.wpfloat, - decay_scale: ta.wpfloat, - decay_exponent: ta.wpfloat, + model_top_height: wpfloat, + decay_scale: wpfloat, + decay_exponent: wpfloat, ) -> data_alloc.NDArray: return array_ns.sinh( (model_top_height / decay_scale) ** decay_exponent @@ -597,7 +594,7 @@ def _decay_func( c2e2co=c2e2co, ) - vertical_coordinate = array_ns.zeros((num_cells, num_levels + 1), dtype=ta.wpfloat) + vertical_coordinate = array_ns.zeros((num_cells, num_levels + 1), dtype=wpfloat) vertical_coordinate[:, num_levels] = topography # Small-scale topography (i.e. full topo - smooth topo) @@ -632,11 +629,11 @@ def _decay_func( def _check_and_correct_layer_thickness( vertical_coordinate: data_alloc.NDArray, vct_a: data_alloc.NDArray, - SLEVE_minimum_layer_thickness_1: ta.wpfloat, - SLEVE_minimum_relative_layer_thickness_1: ta.wpfloat, - SLEVE_minimum_layer_thickness_2: ta.wpfloat, - SLEVE_minimum_relative_layer_thickness_2: ta.wpfloat, - lowest_layer_thickness: ta.wpfloat, + SLEVE_minimum_layer_thickness_1: wpfloat, + SLEVE_minimum_relative_layer_thickness_1: wpfloat, + SLEVE_minimum_layer_thickness_2: wpfloat, + SLEVE_minimum_relative_layer_thickness_2: wpfloat, + lowest_layer_thickness: wpfloat, array_ns: ModuleType = np, ) -> data_alloc.NDArray: num_cells = vertical_coordinate.shape[0] @@ -740,15 +737,15 @@ def compute_vertical_coordinate( geofac_n2s: data_alloc.NDArray, c2e2co: data_alloc.NDArray, nflatlev: int, - model_top_height: ta.wpfloat, - SLEVE_decay_scale_1: ta.wpfloat, - SLEVE_decay_exponent: ta.wpfloat, - SLEVE_decay_scale_2: ta.wpfloat, - SLEVE_minimum_layer_thickness_1: ta.wpfloat, - SLEVE_minimum_relative_layer_thickness_1: ta.wpfloat, - SLEVE_minimum_layer_thickness_2: ta.wpfloat, - SLEVE_minimum_relative_layer_thickness_2: ta.wpfloat, - lowest_layer_thickness: ta.wpfloat, + model_top_height: wpfloat, + SLEVE_decay_scale_1: wpfloat, + SLEVE_decay_exponent: wpfloat, + SLEVE_decay_scale_2: wpfloat, + SLEVE_minimum_layer_thickness_1: wpfloat, + SLEVE_minimum_relative_layer_thickness_1: wpfloat, + SLEVE_minimum_layer_thickness_2: wpfloat, + SLEVE_minimum_relative_layer_thickness_2: wpfloat, + lowest_layer_thickness: wpfloat, array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ diff --git a/model/common/src/icon4py/model/common/math/helpers.py b/model/common/src/icon4py/model/common/math/helpers.py index 8d47cf0a16..919d5746ff 100644 --- a/model/common/src/icon4py/model/common/math/helpers.py +++ b/model/common/src/icon4py/model/common/math/helpers.py @@ -9,7 +9,7 @@ from gt4py import next as gtx from gt4py.next import arccos, cos, sin, sqrt, where -from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import E2C, E2V, Koff from icon4py.model.common.type_alias import wpfloat @@ -29,7 +29,7 @@ def average_level_plus1_on_cells( Returns: Field[Dims[CellDim, dims.KDim], wpfloat] full level field """ - return 0.5 * (half_level_field + half_level_field(Koff[1])) + return wpfloat(0.5) * (half_level_field + half_level_field(Koff[1])) @gtx.field_operator @@ -47,7 +47,7 @@ def average_level_plus1_on_edges( Returns: fa.EdgeKField[wpfloat] full level field """ - return 0.5 * (half_level_field + half_level_field(Koff[1])) + return wpfloat(0.5) * (half_level_field + half_level_field(Koff[1])) @gtx.field_operator @@ -100,8 +100,8 @@ def _grad_fd_tang( @gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) def geographical_to_cartesian_on_cells( - lat: fa.CellField[ta.wpfloat], lon: fa.CellField[ta.wpfloat] -) -> tuple[fa.CellField[ta.wpfloat], fa.CellField[ta.wpfloat], fa.CellField[ta.wpfloat]]: + lat: fa.CellField[wpfloat], lon: fa.CellField[wpfloat] +) -> tuple[fa.CellField[wpfloat], fa.CellField[wpfloat], fa.CellField[wpfloat]]: """ Convert geographical (lat, lon) coordinates to cartesian coordinates on the unit sphere. @@ -123,8 +123,8 @@ def geographical_to_cartesian_on_cells( @gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) def geographical_to_cartesian_on_edges( - lat: fa.EdgeField[ta.wpfloat], lon: fa.EdgeField[ta.wpfloat] -) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: + lat: fa.EdgeField[wpfloat], lon: fa.EdgeField[wpfloat] +) -> tuple[fa.EdgeField[wpfloat], fa.EdgeField[wpfloat], fa.EdgeField[wpfloat]]: """ Convert geographical (lat, lon) coordinates to cartesian coordinates on the unit sphere. @@ -146,8 +146,8 @@ def geographical_to_cartesian_on_edges( @gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) def geographical_to_cartesian_on_vertices( - lat: fa.VertexField[ta.wpfloat], lon: fa.VertexField[ta.wpfloat] -) -> tuple[fa.VertexField[ta.wpfloat], fa.VertexField[ta.wpfloat], fa.VertexField[ta.wpfloat]]: + lat: fa.VertexField[wpfloat], lon: fa.VertexField[wpfloat] +) -> tuple[fa.VertexField[wpfloat], fa.VertexField[wpfloat], fa.VertexField[wpfloat]]: """ Convert geographical (lat, lon) coordinates to cartesian coordinates on the unit sphere. @@ -169,52 +169,52 @@ def geographical_to_cartesian_on_vertices( @gtx.field_operator def dot_product_on_edges( - x1: fa.EdgeField[ta.wpfloat], - x2: fa.EdgeField[ta.wpfloat], - y1: fa.EdgeField[ta.wpfloat], - y2: fa.EdgeField[ta.wpfloat], - z1: fa.EdgeField[ta.wpfloat], - z2: fa.EdgeField[ta.wpfloat], -) -> fa.EdgeField[ta.wpfloat]: + x1: fa.EdgeField[wpfloat], + x2: fa.EdgeField[wpfloat], + y1: fa.EdgeField[wpfloat], + y2: fa.EdgeField[wpfloat], + z1: fa.EdgeField[wpfloat], + z2: fa.EdgeField[wpfloat], +) -> fa.EdgeField[wpfloat]: """Compute dot product of cartesian vectors (x1, y1, z1) * (x2, y2, z2)""" return x1 * x2 + y1 * y2 + z1 * z2 @gtx.field_operator def dot_product_on_cells( - x1: fa.CellField[ta.wpfloat], - x2: fa.CellField[ta.wpfloat], - y1: fa.CellField[ta.wpfloat], - y2: fa.CellField[ta.wpfloat], - z1: fa.CellField[ta.wpfloat], - z2: fa.CellField[ta.wpfloat], -) -> fa.CellField[ta.wpfloat]: + x1: fa.CellField[wpfloat], + x2: fa.CellField[wpfloat], + y1: fa.CellField[wpfloat], + y2: fa.CellField[wpfloat], + z1: fa.CellField[wpfloat], + z2: fa.CellField[wpfloat], +) -> fa.CellField[wpfloat]: """Compute dot product of cartesian vectors (x1, y1, z1) * (x2, y2, z2)""" return x1 * x2 + y1 * y2 + z1 * z2 @gtx.field_operator def dot_product_on_vertices( - x1: fa.VertexField[ta.wpfloat], - x2: fa.VertexField[ta.wpfloat], - y1: fa.VertexField[ta.wpfloat], - y2: fa.VertexField[ta.wpfloat], - z1: fa.VertexField[ta.wpfloat], - z2: fa.VertexField[ta.wpfloat], -) -> fa.VertexField[ta.wpfloat]: + x1: fa.VertexField[wpfloat], + x2: fa.VertexField[wpfloat], + y1: fa.VertexField[wpfloat], + y2: fa.VertexField[wpfloat], + z1: fa.VertexField[wpfloat], + z2: fa.VertexField[wpfloat], +) -> fa.VertexField[wpfloat]: """Compute dot product of cartesian vectors (x1, y1, z1) * (x2, y2, z2)""" return x1 * x2 + y1 * y2 + z1 * z2 @gtx.field_operator def cross_product_on_edges( - x1: fa.EdgeField[ta.wpfloat], - x2: fa.EdgeField[ta.wpfloat], - y1: fa.EdgeField[ta.wpfloat], - y2: fa.EdgeField[ta.wpfloat], - z1: fa.EdgeField[ta.wpfloat], - z2: fa.EdgeField[ta.wpfloat], -) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: + x1: fa.EdgeField[wpfloat], + x2: fa.EdgeField[wpfloat], + y1: fa.EdgeField[wpfloat], + y2: fa.EdgeField[wpfloat], + z1: fa.EdgeField[wpfloat], + z2: fa.EdgeField[wpfloat], +) -> tuple[fa.EdgeField[wpfloat], fa.EdgeField[wpfloat], fa.EdgeField[wpfloat]]: """Compute cross product of cartesian vectors (x1, y1, z1) x (x2, y2, z2)""" x = y1 * z2 - z1 * y2 y = z1 * x2 - x1 * z2 @@ -224,8 +224,8 @@ def cross_product_on_edges( @gtx.field_operator def norm2_on_edges( - x: fa.EdgeField[ta.wpfloat], y: fa.EdgeField[ta.wpfloat], z: fa.EdgeField[ta.wpfloat] -) -> fa.EdgeField[ta.wpfloat]: + x: fa.EdgeField[wpfloat], y: fa.EdgeField[wpfloat], z: fa.EdgeField[wpfloat] +) -> fa.EdgeField[wpfloat]: """ Compute 2 norm of a cartesian vector (x, y, z) Args: @@ -242,8 +242,8 @@ def norm2_on_edges( @gtx.field_operator def norm2_on_cells( - x: fa.CellField[ta.wpfloat], y: fa.CellField[ta.wpfloat], z: fa.CellField[ta.wpfloat] -) -> fa.CellField[ta.wpfloat]: + x: fa.CellField[wpfloat], y: fa.CellField[wpfloat], z: fa.CellField[wpfloat] +) -> fa.CellField[wpfloat]: """ Compute 2 norm of a cartesian vector (x, y, z) Args: @@ -260,8 +260,8 @@ def norm2_on_cells( @gtx.field_operator def norm2_on_vertices( - x: fa.VertexField[ta.wpfloat], y: fa.VertexField[ta.wpfloat], z: fa.VertexField[ta.wpfloat] -) -> fa.VertexField[ta.wpfloat]: + x: fa.VertexField[wpfloat], y: fa.VertexField[wpfloat], z: fa.VertexField[wpfloat] +) -> fa.VertexField[wpfloat]: """ Compute 2 norm of a cartesian vector (x, y, z) Args: @@ -278,8 +278,8 @@ def norm2_on_vertices( @gtx.field_operator def normalize_cartesian_vector_on_edges( - v_x: fa.EdgeField[ta.wpfloat], v_y: fa.EdgeField[ta.wpfloat], v_z: fa.EdgeField[ta.wpfloat] -) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: + v_x: fa.EdgeField[wpfloat], v_y: fa.EdgeField[wpfloat], v_z: fa.EdgeField[wpfloat] +) -> tuple[fa.EdgeField[wpfloat], fa.EdgeField[wpfloat], fa.EdgeField[wpfloat]]: """ Normalize a cartesian vector. @@ -297,7 +297,7 @@ def normalize_cartesian_vector_on_edges( @gtx.field_operator -def invert_edge_field(f: fa.EdgeField[ta.wpfloat]) -> fa.EdgeField[ta.wpfloat]: +def invert_edge_field(f: fa.EdgeField[wpfloat]) -> fa.EdgeField[wpfloat]: """ Invert values. Args: @@ -306,13 +306,13 @@ def invert_edge_field(f: fa.EdgeField[ta.wpfloat]) -> fa.EdgeField[ta.wpfloat]: Returns: 1/f where f is not zero. """ - return where(f != 0.0, 1.0 / f, f) + return where(f != wpfloat(0.0), wpfloat(1.0) / f, f) @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_inverse_on_edges( - f: fa.EdgeField[ta.wpfloat], - f_inverse: fa.EdgeField[ta.wpfloat], + f: fa.EdgeField[wpfloat], + f_inverse: fa.EdgeField[wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, ): @@ -321,12 +321,12 @@ def compute_inverse_on_edges( @gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) def zonal_and_meridional_components_on_cells( - lat: fa.CellField[ta.wpfloat], - lon: fa.CellField[ta.wpfloat], - x: fa.CellField[ta.wpfloat], - y: fa.CellField[ta.wpfloat], - z: fa.CellField[ta.wpfloat], -) -> tuple[fa.CellField[ta.wpfloat], fa.CellField[ta.wpfloat]]: + lat: fa.CellField[wpfloat], + lon: fa.CellField[wpfloat], + x: fa.CellField[wpfloat], + y: fa.CellField[wpfloat], + z: fa.CellField[wpfloat], +) -> tuple[fa.CellField[wpfloat], fa.CellField[wpfloat]]: """ Compute normalized zonal and meridional components of a cartesian vector (x, y, z) at point (lat, lon) @@ -355,12 +355,12 @@ def zonal_and_meridional_components_on_cells( @gtx.field_operator def zonal_and_meridional_components_on_edges( - lat: fa.EdgeField[ta.wpfloat], - lon: fa.EdgeField[ta.wpfloat], - x: fa.EdgeField[ta.wpfloat], - y: fa.EdgeField[ta.wpfloat], - z: fa.EdgeField[ta.wpfloat], -) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: + lat: fa.EdgeField[wpfloat], + lon: fa.EdgeField[wpfloat], + x: fa.EdgeField[wpfloat], + y: fa.EdgeField[wpfloat], + z: fa.EdgeField[wpfloat], +) -> tuple[fa.EdgeField[wpfloat], fa.EdgeField[wpfloat]]: """ Compute the zonal and meridional component of a vector (x, y, z) at position (lat, lon) @@ -389,13 +389,13 @@ def zonal_and_meridional_components_on_edges( @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_zonal_and_meridional_components_on_edges( - lat: fa.EdgeField[ta.wpfloat], - lon: fa.EdgeField[ta.wpfloat], - x: fa.EdgeField[ta.wpfloat], - y: fa.EdgeField[ta.wpfloat], - z: fa.EdgeField[ta.wpfloat], - u: fa.EdgeField[ta.wpfloat], - v: fa.EdgeField[ta.wpfloat], + lat: fa.EdgeField[wpfloat], + lon: fa.EdgeField[wpfloat], + x: fa.EdgeField[wpfloat], + y: fa.EdgeField[wpfloat], + z: fa.EdgeField[wpfloat], + u: fa.EdgeField[wpfloat], + v: fa.EdgeField[wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, ): @@ -406,11 +406,11 @@ def compute_zonal_and_meridional_components_on_edges( @gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) def cartesian_coordinates_from_zonal_and_meridional_components_on_edges( - lat: fa.EdgeField[ta.wpfloat], - lon: fa.EdgeField[ta.wpfloat], - u: fa.EdgeField[ta.wpfloat], - v: fa.EdgeField[ta.wpfloat], -) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: + lat: fa.EdgeField[wpfloat], + lon: fa.EdgeField[wpfloat], + u: fa.EdgeField[wpfloat], + v: fa.EdgeField[wpfloat], +) -> tuple[fa.EdgeField[wpfloat], fa.EdgeField[wpfloat], fa.EdgeField[wpfloat]]: """ Compute cartesian coordinates from zonal an meridional components at position (lat, lon) Args: @@ -438,13 +438,13 @@ def cartesian_coordinates_from_zonal_and_meridional_components_on_edges( @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_cartesian_coordinates_from_zonal_and_meridional_components_on_edges( - edge_lat: fa.EdgeField[ta.wpfloat], - edge_lon: fa.EdgeField[ta.wpfloat], - u: fa.EdgeField[ta.wpfloat], - v: fa.EdgeField[ta.wpfloat], - x: fa.EdgeField[ta.wpfloat], - y: fa.EdgeField[ta.wpfloat], - z: fa.EdgeField[ta.wpfloat], + edge_lat: fa.EdgeField[wpfloat], + edge_lon: fa.EdgeField[wpfloat], + u: fa.EdgeField[wpfloat], + v: fa.EdgeField[wpfloat], + x: fa.EdgeField[wpfloat], + y: fa.EdgeField[wpfloat], + z: fa.EdgeField[wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, ): @@ -460,11 +460,11 @@ def compute_cartesian_coordinates_from_zonal_and_meridional_components_on_edges( @gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) def cartesian_coordinates_from_zonal_and_meridional_components_on_cells( - lat: fa.CellField[ta.wpfloat], - lon: fa.CellField[ta.wpfloat], - u: fa.CellField[ta.wpfloat], - v: fa.CellField[ta.wpfloat], -) -> tuple[fa.CellField[ta.wpfloat], fa.CellField[ta.wpfloat], fa.CellField[ta.wpfloat]]: + lat: fa.CellField[wpfloat], + lon: fa.CellField[wpfloat], + u: fa.CellField[wpfloat], + v: fa.CellField[wpfloat], +) -> tuple[fa.CellField[wpfloat], fa.CellField[wpfloat], fa.CellField[wpfloat]]: """ Compute cartesian coordinates form zonal an meridonal components at position (lat, lon) Args: @@ -492,13 +492,13 @@ def cartesian_coordinates_from_zonal_and_meridional_components_on_cells( @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_cartesian_coordinates_from_zonal_and_meridional_components_on_cells( - cell_lat: fa.CellField[ta.wpfloat], - cell_lon: fa.CellField[ta.wpfloat], - u: fa.CellField[ta.wpfloat], - v: fa.CellField[ta.wpfloat], - x: fa.CellField[ta.wpfloat], - y: fa.CellField[ta.wpfloat], - z: fa.CellField[ta.wpfloat], + cell_lat: fa.CellField[wpfloat], + cell_lon: fa.CellField[wpfloat], + u: fa.CellField[wpfloat], + v: fa.CellField[wpfloat], + x: fa.CellField[wpfloat], + y: fa.CellField[wpfloat], + z: fa.CellField[wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, ): @@ -514,13 +514,13 @@ def compute_cartesian_coordinates_from_zonal_and_meridional_components_on_cells( @gtx.field_operator def arc_length_on_edges( - x0: fa.EdgeField[ta.wpfloat], - x1: fa.EdgeField[ta.wpfloat], - y0: fa.EdgeField[ta.wpfloat], - y1: fa.EdgeField[ta.wpfloat], - z0: fa.EdgeField[ta.wpfloat], - z1: fa.EdgeField[ta.wpfloat], - radius: ta.wpfloat, + x0: fa.EdgeField[wpfloat], + x1: fa.EdgeField[wpfloat], + y0: fa.EdgeField[wpfloat], + y1: fa.EdgeField[wpfloat], + z0: fa.EdgeField[wpfloat], + z1: fa.EdgeField[wpfloat], + radius: wpfloat, ): """ Compute the arc length between two points on the sphere. diff --git a/model/common/src/icon4py/model/common/math/smagorinsky.py b/model/common/src/icon4py/model/common/math/smagorinsky.py index 1698d5443d..6288a17c69 100644 --- a/model/common/src/icon4py/model/common/math/smagorinsky.py +++ b/model/common/src/icon4py/model/common/math/smagorinsky.py @@ -10,20 +10,21 @@ from icon4py.model.common import field_type_aliases as fa from icon4py.model.common.dimension import KDim, Koff +from icon4py.model.common.type_alias import wpfloat @gtx.field_operator def _en_smag_fac_for_zero_nshift( - vect_a: fa.KField[float], - hdiff_smag_fac: float, - hdiff_smag_fac2: float, - hdiff_smag_fac3: float, - hdiff_smag_fac4: float, - hdiff_smag_z: float, - hdiff_smag_z2: float, - hdiff_smag_z3: float, - hdiff_smag_z4: float, -) -> fa.KField[float]: + vect_a: fa.KField[wpfloat], + hdiff_smag_fac: wpfloat, + hdiff_smag_fac2: wpfloat, + hdiff_smag_fac3: wpfloat, + hdiff_smag_fac4: wpfloat, + hdiff_smag_z: wpfloat, + hdiff_smag_z2: wpfloat, + hdiff_smag_z3: wpfloat, + hdiff_smag_z4: wpfloat, +) -> fa.KField[wpfloat]: dz21 = hdiff_smag_z2 - hdiff_smag_z alin = (hdiff_smag_fac2 - hdiff_smag_fac) / dz21 df32 = hdiff_smag_fac3 - hdiff_smag_fac2 @@ -33,8 +34,8 @@ def _en_smag_fac_for_zero_nshift( bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32)) aqdr = df32 / dz32 - bqdr * dz32 - zf = 0.5 * (vect_a + vect_a(Koff[1])) - zero = broadcast(0.0, (KDim,)) + zf = wpfloat(0.5) * (vect_a + vect_a(Koff[1])) + zero = broadcast(wpfloat(0.0), (KDim,)) dzlin = minimum(broadcast(dz21, (KDim,)), maximum(zero, zf - hdiff_smag_z)) dzqdr = minimum(broadcast(dz42, (KDim,)), maximum(zero, zf - hdiff_smag_z2)) @@ -44,16 +45,16 @@ def _en_smag_fac_for_zero_nshift( @gtx.program def en_smag_fac_for_zero_nshift( - vect_a: fa.KField[float], - hdiff_smag_fac: float, - hdiff_smag_fac2: float, - hdiff_smag_fac3: float, - hdiff_smag_fac4: float, - hdiff_smag_z: float, - hdiff_smag_z2: float, - hdiff_smag_z3: float, - hdiff_smag_z4: float, - enh_smag_fac: fa.KField[float], + vect_a: fa.KField[wpfloat], + hdiff_smag_fac: wpfloat, + hdiff_smag_fac2: wpfloat, + hdiff_smag_fac3: wpfloat, + hdiff_smag_fac4: wpfloat, + hdiff_smag_z: wpfloat, + hdiff_smag_z2: wpfloat, + hdiff_smag_z3: wpfloat, + hdiff_smag_z4: wpfloat, + enh_smag_fac: fa.KField[wpfloat], ): _en_smag_fac_for_zero_nshift( vect_a, diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 6fd36b09ba..7b1e643727 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -71,7 +71,7 @@ def _compute_ddqz_z_half( def compute_ddqz_z_half( z_ifc: fa.CellKField[wpfloat], z_mc: fa.CellKField[wpfloat], - ddqz_z_half: fa.CellKField[wpfloat], + ddqz_z_half: fa.CellKField[vpfloat], nlev: gtx.int32, horizontal_start: gtx.int32, horizontal_end: gtx.int32, @@ -537,7 +537,7 @@ def compute_exner_exfac( @gtx.program def compute_wgtfac_e( wgtfac_c: fa.CellKField[wpfloat], - c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], float], + c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], wpfloat], wgtfac_e: fa.EdgeKField[wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, @@ -683,7 +683,7 @@ def _compute_pressure_gradient_downward_extrapolation_mask_distance( @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_pressure_gradient_downward_extrapolation_mask_distance( z_mc: fa.CellKField[wpfloat], - c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], float], + c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], wpfloat], topography: fa.CellField[wpfloat], e_owner_mask: fa.EdgeField[bool], flat_idx_max: fa.EdgeField[gtx.int32], diff --git a/model/common/src/icon4py/model/common/model_options.py b/model/common/src/icon4py/model/common/model_options.py index 04187a79c0..5a71283403 100644 --- a/model/common/src/icon4py/model/common/model_options.py +++ b/model/common/src/icon4py/model/common/model_options.py @@ -13,6 +13,7 @@ import gt4py.next as gtx import gt4py.next.typing as gtx_typing from gt4py.next import backend as gtx_backend +from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations from icon4py.model.common import model_backends @@ -27,6 +28,28 @@ def dict_values_to_list(d: dict[str, Any]) -> dict[str, list]: def get_dace_options( program_name: str, **backend_descriptor: Any ) -> model_backends.BackendDescriptor: + optimization_args = backend_descriptor.get("optimization_args", {}) + optimization_hooks = optimization_args.get("optimization_hooks", {}) + if program_name in [ + "vertically_implicit_solver_at_corrector_step", + "vertically_implicit_solver_at_predictor_step", + ]: + if gtx_transformations.GT4PyAutoOptHook.TopLevelDataFlowStep not in optimization_hooks: + # TODO(iomaganaris): Enable this for CPU once the issue with the strides of memlets from the nested SDFG + # to a global Access Node is resolved. + # Enable pass that removes access node (next_w) copies for vertically implicit solver programs + if backend_descriptor["device"] == model_backends.GPU: + optimization_hooks[gtx_transformations.GT4PyAutoOptHook.TopLevelDataFlowStep] = ( + lambda sdfg: sdfg.apply_transformations_repeated( + gtx_transformations.RemoveAccessNodeCopies(), + validate=False, + validate_all=False, + ) + ) + if optimization_hooks: + optimization_args["optimization_hooks"] = optimization_hooks + if optimization_args: + backend_descriptor["optimization_args"] = optimization_args return backend_descriptor diff --git a/model/common/src/icon4py/model/common/orchestration/dtypes.py b/model/common/src/icon4py/model/common/orchestration/dtypes.py index 0d15baadae..59b8182911 100644 --- a/model/common/src/icon4py/model/common/orchestration/dtypes.py +++ b/model/common/src/icon4py/model/common/orchestration/dtypes.py @@ -10,7 +10,7 @@ from typing import Final import dace -from gt4py.next import Field, common, int32, int64 +from gt4py.next import Field, common, int32, int64, float32, float64 from icon4py.model.common import type_alias @@ -21,17 +21,18 @@ KDim_sym = dace.symbol("KDim_sym") ICON4PY_PRIMITIVE_DTYPES: Final = ( - type_alias.wpfloat, - type_alias.vpfloat, + float32, + float64, float, bool, int32, int64, int, ) + DACE_PRIMITIVE_DTYPES: Final = ( + dace.float32, dace.float64, - dace.float64 if type_alias.precision == "double" else dace.float32, dace.float64, dace.bool, dace.int32, @@ -39,11 +40,9 @@ dace.int64 if sys.maxsize > 2**32 else dace.int32, ) - def stride_symbol_name_from_field(cls: type, field_name: str, stride: int) -> str: return f"{cls.__name__}_{field_name}_s{stride}_sym" - def gt4py_dim_to_dace_symbol(dim: common.Dimension) -> dace.symbol: # See dims.global_dimensions.values() # TODO(kotsaloscv): generalize this diff --git a/model/common/src/icon4py/model/common/type_alias.py b/model/common/src/icon4py/model/common/type_alias.py index 02b933bfa3..87c5857ae3 100644 --- a/model/common/src/icon4py/model/common/type_alias.py +++ b/model/common/src/icon4py/model/common/type_alias.py @@ -12,26 +12,31 @@ import gt4py.next as gtx -DEFAULT_PRECISION = "double" +DEFAULT_PRECISION = "single" # TODO(pstark): reset default to double wpfloat: TypeAlias = gtx.float64 -vpfloat: type[gtx.float32] | type[gtx.float64] = wpfloat +vpfloat: TypeAlias = wpfloat precision = os.environ.get("FLOAT_PRECISION", DEFAULT_PRECISION).lower() -def set_precision(new_precision: Literal["double", "mixed"]) -> None: +def set_precision(new_precision: Literal["double", "mixed", "single"]) -> None: global precision # noqa: PLW0603 [global-statement] - global vpfloat # noqa: PLW0603 [global-statement] + global vpfloat, wpfloat # noqa: PLW0603 [global-statement] precision = new_precision.lower() match precision: case "double": + wpfloat = gtx.float64 vpfloat = wpfloat case "mixed": + wpfloat = gtx.float64 vpfloat = gtx.float32 + case "single": + wpfloat = gtx.float32 + vpfloat = wpfloat case _: - raise ValueError("Only 'double' and 'mixed' precision are supported.") + raise ValueError("Only 'double', 'mixed' and 'single' precision are supported.") set_precision(precision) 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 375c5b3471..99c063c709 100644 --- a/model/common/src/icon4py/model/common/utils/data_allocation.py +++ b/model/common/src/icon4py/model/common/utils/data_allocation.py @@ -8,7 +8,9 @@ from __future__ import annotations +import json import logging as log +from pathlib import Path from types import ModuleType from typing import TYPE_CHECKING, TypeAlias @@ -18,7 +20,8 @@ from gt4py import next as gtx from gt4py.next import allocators as gtx_allocators -from icon4py.model.common import type_alias as ta +from icon4py.model.common import dimension as dims # noqa: F401 +from icon4py.model.common.type_alias import vpfloat, wpfloat # noqa: F401 from icon4py.model.common.utils import device_utils @@ -80,44 +83,43 @@ def as_field( def random_field( grid: grid_base.Grid, - *dims: gtx.Dimension, + *dimensions: gtx.Dimension, low: float = -1.0, high: float = 1.0, - dtype: npt.DTypeLike | None = None, + dtype: npt.DTypeLike = wpfloat, extend: dict[gtx.Dimension, int] | None = None, allocator: gtx_allocators.FieldBufferAllocationUtil | None = None, ) -> gtx.Field: arr = np.random.default_rng().uniform( - low=low, high=high, size=_shape(grid, *dims, extend=extend) + low=low, high=high, size=_shape(grid, *dimensions, extend=extend) ) - if dtype: - arr = arr.astype(dtype) - return gtx.as_field(dims, arr, allocator=allocator) + arr = arr.astype(dtype) + return gtx.as_field(dimensions, arr, allocator=allocator) def random_sign( grid: grid_base.Grid, - *dims: gtx.Dimension, + *dimensions: gtx.Dimension, dtype: npt.DTypeLike | None = None, extend: dict[gtx.Dimension, int] | None = None, allocator: gtx_allocators.FieldBufferAllocationUtil | None = None, ) -> gtx.Field: """Generate a random field with values -1 or 1.""" - arr = np.random.default_rng().choice([-1, 1], size=_shape(grid, *dims, extend=extend)) + arr = np.random.default_rng().choice([-1, 1], size=_shape(grid, *dimensions, extend=extend)) if dtype: arr = arr.astype(dtype) - return gtx.as_field(dims, arr, allocator=allocator) + return gtx.as_field(dimensions, arr, allocator=allocator) def random_mask( grid: grid_base.Grid, - *dims: gtx.Dimension, + *dimensions: gtx.Dimension, dtype: npt.DTypeLike | None = None, extend: dict[gtx.Dimension, int] | None = None, allocator: gtx_typing.Backend | None = None, ) -> gtx.Field: rng = np.random.default_rng() - shape = _shape(grid, *dims, extend=extend) + shape = _shape(grid, *dimensions, extend=extend) arr = np.full(shape, False).flatten() num_true = int(arr.size * 0.5) arr[:num_true] = True @@ -125,41 +127,75 @@ def random_mask( arr = np.reshape(arr, newshape=shape) if dtype: arr = arr.astype(dtype) - return gtx.as_field(dims, arr, allocator=allocator) + return gtx.as_field(dimensions, arr, allocator=allocator) + + +def random_ikoffset( + grid: grid_base.Grid, + *dimensions: gtx.Dimension, + dtype=gtx.int32, + extend: dict[gtx.Dimension, int] | None = None, + allocator: gtx_allocators.FieldBufferAllocationUtil | None = None +): + ikoffset = empty_field(grid, *dimensions, dtype=dtype, extend=extend, allocator=allocator) + rng=np.random.default_rng() + for k in range(grid.num_levels): + # construct offsets that reach all k-levels except the last (because we are using the entries of this field with `+1`) + ikoffset.ndarray[:, :, k] = rng.integers( # type: ignore[index] + low=0 - k, + high=grid.num_levels - k - 1, + size=(ikoffset.shape[0], ikoffset.shape[1]), + ) + return ikoffset + + +def empty_field( + grid: grid_base.Grid, + *dimensions: gtx.Dimension, + dtype=wpfloat, + extend: dict[gtx.Dimension, int] | None = None, + allocator: gtx_allocators.FieldBufferAllocationUtil | None = None, +) -> gtx.Field: + field_domain = { + dim: (0, stop) for dim, stop in zip(dimensions, _shape(grid, *dimensions, extend=extend)) + } + return gtx.constructors.empty(field_domain, dtype=dtype, allocator=allocator) def zero_field( grid: grid_base.Grid, - *dims: gtx.Dimension, - dtype=ta.wpfloat, + *dimensions: gtx.Dimension, + dtype=wpfloat, extend: dict[gtx.Dimension, int] | None = None, allocator: gtx_allocators.FieldBufferAllocationUtil | None = None, ) -> gtx.Field: - field_domain = {dim: (0, stop) for dim, stop in zip(dims, _shape(grid, *dims, extend=extend))} + field_domain = { + dim: (0, stop) for dim, stop in zip(dimensions, _shape(grid, *dimensions, extend=extend)) + } return gtx.constructors.zeros(field_domain, dtype=dtype, allocator=allocator) def constant_field( grid: grid_base.Grid, value: float, - *dims: gtx.Dimension, - dtype=ta.wpfloat, + *dimensions: gtx.Dimension, + dtype=wpfloat, + extend: dict[gtx.Dimension, int] | None = None, allocator: gtx_allocators.FieldBufferAllocationUtil | None = None, ) -> gtx.Field: - return gtx.as_field( - dims, - value * np.ones(shape=tuple(map(lambda x: grid.size[x], dims)), dtype=dtype), - allocator=allocator, - ) + field_domain = { + dim: (0, stop) for dim, stop in zip(dimensions, _shape(grid, *dimensions, extend=extend)) + } + return gtx.constructors.full(field_domain, value, dtype=dtype, allocator=allocator) def _shape( grid: grid_base.Grid, - *dims: gtx.Dimension, + *dimensions: gtx.Dimension, extend: dict[gtx.Dimension, int] | None = None, ) -> tuple[int, ...]: extend = extend or {} - return tuple(grid.size[dim] + extend.get(dim, 0) for dim in dims) + return tuple(grid.size[dim] + extend.get(dim, 0) for dim in dimensions) def index_field( @@ -172,3 +208,43 @@ 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) + + +# load variable properties (dtype, field dimensions and extend) from json +p = Path(__file__).resolve().parent / "variable_properties.json" +with p.open("r", encoding="utf-8") as fh: + _variable_properties = json.load(fh) + + +def make_fields(varnames, gen_fct, *args, **kwargs) -> dict[str, gtx.Field]: + fields = {} + for var in varnames: + dtype, dimensions, extend = eval(_variable_properties[var]) + this_field = gen_fct(*args, *dimensions, dtype=dtype, extend=extend, **kwargs) + fields[var] = this_field + return fields + + +def get_random_fields( + grid: grid_base.Grid, + varnames: list, + low: float = -1.0, + high: float = 1.0, + allocator: gtx_allocators.FieldBufferAllocationUtil | None = None, +) -> dict[str, gtx.Field]: + return make_fields(varnames, random_field, grid, low=low, high=high, allocator=allocator) + + +def get_zero_fields( + grid, varnames: list, allocator: gtx_allocators.FieldBufferAllocationUtil | None = None +) -> dict[str, gtx.Field]: + return make_fields(varnames, zero_field, grid, allocator=allocator) + + +def get_const_fields( + grid: grid_base.Grid, + varnames: list, + value, + allocator: gtx_allocators.FieldBufferAllocationUtil | None = None, +) -> dict[str, gtx.Field]: + return make_fields(varnames, constant_field, grid, value, allocator=allocator) diff --git a/model/common/src/icon4py/model/common/utils/variable_properties.json b/model/common/src/icon4py/model/common/utils/variable_properties.json new file mode 100644 index 0000000000..0dffac1f1b --- /dev/null +++ b/model/common/src/icon4py/model/common/utils/variable_properties.json @@ -0,0 +1,189 @@ +{ + "PROPERTIES": "[dtype, dims, extend]", + "area": "[wpfloat, (dims.CellDim,), None]", + "area_edge": "[wpfloat, (dims.EdgeDim,), None]", + "c_intp": "[wpfloat, (dims.VertexDim, dims.V2CDim,), None]", + "c_lin_e": "[wpfloat, (dims.EdgeDim, dims.E2CDim,), None]", + "coeff1_dwdz": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "coeff2_dwdz": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "coeff_gradekin": "[wpfloat, (dims.EdgeDim, dims.E2CDim,), None]", + "contravariant_corrected_w_at_cells_on_model_levels": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "contravariant_correction_at_cells_on_half_levels": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "contravariant_correction_at_edges_on_model_levels": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "coriolis_frequency": "[wpfloat, (dims.EdgeDim,), None]", + "corrector_normal_wind_advective_tendency": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "corrector_vertical_wind_advective_tendency": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "current_exner": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "current_rho": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "current_theta_v": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "current_vn": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "current_w": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "d2dexdz2_fac1_mc": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "d2dexdz2_fac2_mc": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "d_exner_dz_ref_ic": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "ddqz_z_full_e": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "ddqz_z_half": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "ddt_exner_phy": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "ddt_vn_apc": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "ddt_w_adv": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "ddt_w_adv_ntl1": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "ddt_w_adv_ntl2": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "ddxn_z_full": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "ddxt_z_full": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "ddz_of_reference_exner_at_cells_on_half_levels": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "divdamp_field": "[wpfloat, (dims.KDim,), None]", + "dual_normal_cell_1": "[wpfloat, (dims.EdgeDim, dims.E2CDim,), None]", + "dual_normal_cell_2": "[wpfloat, (dims.EdgeDim, dims.E2CDim,), None]", + "dual_normal_cell_x": "[wpfloat, (dims.EdgeDim, dims.E2CDim,), None]", + "dual_normal_cell_y": "[wpfloat, (dims.EdgeDim, dims.E2CDim,), None]", + "dwdz_at_cells_on_model_levels": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "dynamical_vertical_mass_flux_at_cells_on_half_levels": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "dynamical_vertical_volumetric_flux_at_cells_on_half_levels": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "e_bln_c_s": "[wpfloat, (dims.CellDim, dims.C2EDim,), None]", + "e_flx_avg": "[wpfloat, (dims.EdgeDim, dims.E2C2EODim,), None]", + "exner": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "exner_at_cells_on_half_levels": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "exner_dynamical_increment": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "exner_iau_increment": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "exner_incr": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "exner_new": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "exner_nnow": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "exner_now": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "exner_pr": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "exner_ref_mc": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "exner_tendency_due_to_slow_physics": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "exner_w_explicit_weight_parameter": "[wpfloat, (dims.CellDim,), None]", + "exner_w_implicit_weight_parameter": "[wpfloat, (dims.CellDim,), None]", + "fourth_order_divdamp_scaling_coeff": "[wpfloat, (dims.KDim,), None]", + "geofac_div": "[wpfloat, (dims.CellDim, dims.C2EDim,), None]", + "geofac_grdiv": "[wpfloat, (dims.EdgeDim, dims.E2C2EODim,), None]", + "geofac_grg_x": "[wpfloat, (dims.CellDim, dims.C2E2CODim,), None]", + "geofac_grg_y": "[wpfloat, (dims.CellDim, dims.C2E2CODim,), None]", + "geofac_n2s": "[wpfloat, (dims.CellDim, dims.C2E2CODim,), None]", + "geofac_rot": "[wpfloat, (dims.VertexDim, dims.V2EDim,), None]", + "grf_tend_vn": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "hmask_dd3d": "[wpfloat, (dims.EdgeDim,), None]", + "horizontal_advection_of_w_at_edges_on_half_levels": "[wpfloat, (dims.EdgeDim, dims.KDim,), {dims.KDim: 1}]", + "horizontal_gradient_of_normal_wind_divergence": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "horizontal_kinetic_energy_at_edges_on_model_levels": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "horizontal_mask_for_3d_divdamp": "[wpfloat, (dims.EdgeDim,), None]", + "horizontal_pressure_gradient": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "hydrostatic_correction_on_lowest_level": "[wpfloat, (dims.EdgeDim,), None]", + "ikoffset": "[gtx.int32, (dims.EdgeDim, dims.E2CDim, dims.KDim), None]", + "interpolant": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "interpolated_fourth_order_divdamp_factor": "[wpfloat, (dims.KDim,), None]", + "interpolation": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "inv_ddqz_z_full": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "inv_dual_edge_length": "[wpfloat, (dims.EdgeDim,), None]", + "inv_primal_edge_length": "[wpfloat, (dims.EdgeDim,), None]", + "mass_fl_e": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "mass_flux_at_edges_on_model_levels": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "mass_flx_ic": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "next_exner": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "next_rho": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "next_theta_v": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "next_vn": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "next_w": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "normal_wind_advective_tendency": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "normal_wind_iau_increment": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "normal_wind_tendency_due_to_slow_physics_process": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "nudgecoeff_e": "[wpfloat, (dims.EdgeDim,), None]", + "p_vn": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "p_vt": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "perturbed_exner_at_cells_on_model_levels": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "perturbed_rho_at_cells_on_model_levels": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "perturbed_theta_v_at_cells_on_half_levels": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "perturbed_theta_v_at_cells_on_model_levels": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "pg_exdist": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "pos_on_tplane_e_1": "[wpfloat, (dims.EdgeDim, dims.E2CDim,), None]", + "pos_on_tplane_e_2": "[wpfloat, (dims.EdgeDim, dims.E2CDim,), None]", + "pos_on_tplane_e_x": "[wpfloat, (dims.EdgeDim, dims.E2CDim,), None]", + "pos_on_tplane_e_y": "[wpfloat, (dims.EdgeDim, dims.E2CDim,), None]", + "predictor_normal_wind_advective_tendency": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "predictor_vertical_wind_advective_tendency": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "pressure_buoyancy_acceleration_at_cells_on_half_levels": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "primal_normal_cell_1": "[wpfloat, (dims.EdgeDim, dims.E2CDim,), None]", + "primal_normal_cell_2": "[wpfloat, (dims.EdgeDim, dims.E2CDim,), None]", + "primal_normal_cell_x": "[wpfloat, (dims.EdgeDim, dims.E2CDim,), None]", + "primal_normal_cell_y": "[wpfloat, (dims.EdgeDim, dims.E2CDim,), None]", + "rayleigh_damping_factor": "[wpfloat, (dims.KDim,), None]", + "rbf_vec_coeff_e": "[wpfloat, (dims.EdgeDim, dims.E2C2EDim,), None]", + "reduced_fourth_order_divdamp_coeff_at_nest_boundary": "[wpfloat, (dims.KDim,), None]", + "reference_exner_at_cells_on_model_levels": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "reference_rho_at_cells_on_model_levels": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "reference_rho_at_edges_on_model_levels": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "reference_theta_at_cells_on_half_levels": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "reference_theta_at_cells_on_model_levels": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "reference_theta_at_edges_on_model_levels": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "rho": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "rho_at_cells_on_half_levels": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "rho_at_edges_on_model_levels": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "rho_iau_increment": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "rho_ic": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "rho_incr": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "rho_new": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "rho_nnow": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "rho_now": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "rho_ref_me": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "rho_var": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "scalfac_dd3d": "[wpfloat, (dims.KDim,), None]", + "scaling_factor_for_3d_divdamp": "[wpfloat, (dims.KDim,), None]", + "spatially_averaged_vn": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "substep_and_spatially_averaged_vn": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "substep_averaged_mass_flux": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "tangent_orientation": "[wpfloat, (dims.EdgeDim,), None]", + "tangential_wind": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "tangential_wind_on_half_levels": "[wpfloat, (dims.EdgeDim, dims.KDim,), {dims.KDim: 1}]", + "temporal_extrapolation_of_perturbed_exner": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "theta_now": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "theta_ref_mc": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "theta_ref_me": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "theta_v": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "theta_v_at_cells_on_half_levels": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "theta_v_at_edges_on_model_levels": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "theta_v_flux_at_edges_on_model_levels": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "theta_v_ic": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "theta_v_new": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "theta_v_nnow": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "theta_v_now": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "theta_var": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "time_extrapolation_parameter_for_exner": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "vertical_cfl": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "vertical_wind_advective_tendency": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "vn": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "vn_on_half_levels": "[wpfloat, (dims.EdgeDim, dims.KDim,), {dims.KDim: 1}]", + "vol_flx_ic": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "vt": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "vwind_expl_wgt": "[wpfloat, (dims.CellDim,), None]", + "vwind_impl_wgt": "[wpfloat, (dims.CellDim,), None]", + "w": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "w_concorr_c": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "w_nnow": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "wgtfac_c": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "wgtfac_e": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "wgtfacq_c": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "wgtfacq_e": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "z_alpha": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "z_beta": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "z_contr_w_fl_l": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "z_dwdz_dd": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "z_exner_expl": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "z_flxdiv_mass": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "z_flxdiv_theta": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "z_graddiv_vn": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "z_hydro_corr": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "z_q": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "z_rho_e": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "z_rho_expl": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "z_th_ddz_exner_c": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "z_theta_v_e": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "z_theta_v_fl_e": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "z_theta_v_pr_ic": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "z_v_grad_w": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "z_vt_ie": "[wpfloat, (dims.EdgeDim, dims.KDim,), None]", + "z_w_con_c_full": "[wpfloat, (dims.CellDim, dims.KDim,), None]", + "z_w_expl": "[wpfloat, (dims.CellDim, dims.KDim,), {dims.KDim: 1}]", + "zdiff_gradp": "[wpfloat, (dims.EdgeDim, dims.E2CDim, dims.KDim), None]" +} diff --git a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index f6e5034fa9..2ddde5d342 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -18,13 +18,25 @@ ) from icon4py.model.common.type_alias import wpfloat from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest -@pytest.mark.skip_value_error class TestMoIntpRbfRbfVecInterpolVertex(StencilTest): PROGRAM = mo_intp_rbf_rbf_vec_interpol_vertex OUTPUTS = ("p_u_out", "p_v_out") + STATIC_PARAMS = { + StandardStaticVariants.NONE: (), + StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( @@ -62,3 +74,8 @@ def input_data(self, grid: base.Grid) -> dict: vertical_start=0, vertical_end=gtx.int32(grid.num_levels), ) + + +@pytest.mark.continuous_benchmarking +class TestMoIntpRbfRbfVecInterpolVertexContinuousBenchmarking(TestMoIntpRbfRbfVecInterpolVertex): + pass diff --git a/model/driver/src/icon4py/model/driver/icon4py_configuration.py b/model/driver/src/icon4py/model/driver/icon4py_configuration.py index 5cf7ba1424..5089f6b5d0 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_configuration.py +++ b/model/driver/src/icon4py/model/driver/icon4py_configuration.py @@ -143,8 +143,8 @@ def _mch_ch_r04b09_config(): def _jablownoski_Williamson_config(): icon_run_config = Icon4pyRunConfig( - dtime=datetime.timedelta(seconds=300.0), - end_date=datetime.datetime(1, 1, 1, 0, 30, 0), + dtime=datetime.timedelta(seconds=30.0), + end_date=datetime.datetime(1, 1, 1, 2, 0, 0), apply_initial_stabilization=False, n_substeps=5, backend=backend, diff --git a/model/driver/src/icon4py/model/driver/icon4py_driver.py b/model/driver/src/icon4py/model/driver/icon4py_driver.py index d54768ee7b..8ead989783 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_driver.py +++ b/model/driver/src/icon4py/model/driver/icon4py_driver.py @@ -26,6 +26,7 @@ diagnostic_state as diagnostics, prognostic_state as prognostics, ) +from icon4py.model.common.type_alias import wpfloat from icon4py.model.common.utils import device_utils from icon4py.model.driver import ( icon4py_configuration as driver_config, @@ -54,9 +55,9 @@ def __init__( self._n_time_steps: int = int( (self.run_config.end_date - self.run_config.start_date) / self.run_config.dtime ) - self.dtime_in_seconds: float = self.run_config.dtime.total_seconds() + self.dtime_in_seconds: wpfloat = wpfloat(self.run_config.dtime.total_seconds()) self._n_substeps_var: int = self.run_config.n_substeps - self._substep_timestep: float = float(self.dtime_in_seconds / self._n_substeps_var) + self._substep_timestep: wpfloat = self.dtime_in_seconds / wpfloat(self._n_substeps_var) self._validate_config() @@ -116,7 +117,7 @@ def time_integration( solve_nonhydro_diagnostic_state: dycore_states.DiagnosticStateNonHydro, prognostic_states: common_utils.TimeStepPair[prognostics.PrognosticState], prep_adv: dycore_states.PrepAdvection, - second_order_divdamp_factor: float, + second_order_divdamp_factor: wpfloat, do_prep_adv: bool, profiling: driver_config.ProfilingConfig | None = None, ): @@ -202,7 +203,7 @@ def _integrate_one_time_step( solve_nonhydro_diagnostic_state: dycore_states.DiagnosticStateNonHydro, prognostic_states: common_utils.TimeStepPair[prognostics.PrognosticState], prep_adv: dycore_states.PrepAdvection, - second_order_divdamp_factor: float, + second_order_divdamp_factor: wpfloat, do_prep_adv: bool, ): # TODO(OngChia): Add update_spinup_damping here to compute second_order_divdamp_factor @@ -270,7 +271,7 @@ def _do_dyn_substepping( solve_nonhydro_diagnostic_state: dycore_states.DiagnosticStateNonHydro, prognostic_states: common_utils.TimeStepPair[prognostics.PrognosticState], prep_adv: dycore_states.PrepAdvection, - second_order_divdamp_factor: float, + second_order_divdamp_factor: wpfloat, do_prep_adv: bool, ): # TODO(OngChia): compute airmass for prognostic_state here @@ -333,7 +334,7 @@ class DriverParams(NamedTuple): second_order_divdamp_factor: Second order divergence damping factor. """ - second_order_divdamp_factor: float + second_order_divdamp_factor: wpfloat def initialize( @@ -527,9 +528,10 @@ def initialize( ) @click.option( "--enable_profiling", - is_flag=True, - default=False, - help="Enable detailed profiling with GT4Py metrics.", + is_flag=False, + flag_value="gt4py_metrics.json", + default="", + help="Enable detailed profiling with GT4Py metrics. Can be a flag (--enable_profiling) or provide a filename (--enable_profiling='gt4py_metrics.json').", ) @click.option( "--icon4py_driver_backend", @@ -607,7 +609,9 @@ def icon4py_driver( ds.prep_advection_prognostic, dp.second_order_divdamp_factor, do_prep_adv=False, - profiling=driver_config.ProfilingConfig() if enable_profiling else None, + profiling=driver_config.ProfilingConfig(gt4py_metrics_output_file=enable_profiling) + if enable_profiling + else None, ) log.info("time loop: DONE") diff --git a/model/testing/src/icon4py/model/testing/pytest_hooks.py b/model/testing/src/icon4py/model/testing/pytest_hooks.py index 670f7ff775..83eb88fe73 100644 --- a/model/testing/src/icon4py/model/testing/pytest_hooks.py +++ b/model/testing/src/icon4py/model/testing/pytest_hooks.py @@ -35,10 +35,9 @@ def pytest_configure(config): "markers", "level(name): marks test as unit or integration tests, mostly applicable where both are available", ) - - # Check if the --enable-mixed-precision option is set and set the environment variable accordingly - if config.getoption("--enable-mixed-precision"): - os.environ["FLOAT_PRECISION"] = "mixed" + config.addinivalue_line( + "markers", "single_precision_ready: intended to run if single precision is selected" + ) # Handle datatest options: --datatest-only and --datatest-skip if m_option := config.getoption("-m", []): @@ -49,6 +48,9 @@ def pytest_configure(config): if config.getoption("--datatest-skip"): config.option.markexpr = " and ".join(["not datatest", *m_option]) + if os.environ.get("FLOAT_PRECISION", "double").lower() == "single": + config.option.markexpr = " and ".join(["single_precision_ready", *m_option]) + def pytest_addoption(parser: pytest.Parser): """Add custom commandline options for pytest.""" @@ -83,14 +85,6 @@ def pytest_addoption(parser: pytest.Parser): help="Grid to use.", ) - with contextlib.suppress(ValueError): - parser.addoption( - "--enable-mixed-precision", - action="store_true", - help="Switch unit tests from double to mixed-precision", - default=False, - ) - with contextlib.suppress(ValueError): parser.addoption( "--level", @@ -102,6 +96,23 @@ def pytest_addoption(parser: pytest.Parser): def pytest_collection_modifyitems(config, items): + """Modify collected test items based on command line options.""" + test_grid = config.getoption("--grid") + for item in items: + if (marker := item.get_closest_marker("continuous_benchmarking")) is not None: + if test_grid is not None and not test_grid.startswith("icon_benchmark"): + item.add_marker( + pytest.mark.skip( + reason="Continuous benchmarking tests only run with --grid icon_benchmark_{regional,global}." + ) + ) + continue + if not config.getoption("--benchmark-only"): + item.add_marker( + pytest.mark.benchmark_only( + reason="Continuous benchmarking tests shouldn't be verified." + ) + ) test_level = config.getoption("--level") if test_level == "any": return @@ -161,3 +172,56 @@ def pytest_benchmark_update_json(output_json): for bench in output_json["benchmarks"]: bench["fullname"] = _name_from_fullname(bench["fullname"]) + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_makereport(item, call): + """ + Gather GT4Py timer metrics from benchmark fixture and add them to the test report. + """ + outcome = yield + report = outcome.get_result() + if call.when == "call": + benchmark = item.funcargs.get("benchmark", None) + if benchmark and hasattr(benchmark, "extra_info"): + info = benchmark.extra_info.get("gtx_metrics", None) + if info: + filtered_benchmark_name = benchmark.name.split("test_Test")[-1] + # Combine the benchmark name in a readable form with the gtx_metrics data + report.sections.append(("benchmark-extra", tuple([filtered_benchmark_name, info]))) + + +def pytest_terminal_summary(terminalreporter, exitstatus, config): + """ + Add a custom section to the terminal summary with GT4Py timer metrics from benchmarks. + """ + # Gather gtx_metrics + benchmark_gtx_metrics = [] + for outcome in ("passed", "failed", "skipped"): + all_reports = terminalreporter.stats.get(outcome, []) + for report in all_reports: + for secname, info in getattr(report, "sections", []): + if secname == "benchmark-extra": + benchmark_gtx_metrics.append(info) + # Calculate the maximum length of benchmark names for formatting + max_name_len = 0 + for benchmark_name, _ in benchmark_gtx_metrics: + max_name_len = max(len(benchmark_name), max_name_len) + # Print the GT4Py timer report table + if benchmark_gtx_metrics: + terminalreporter.ensure_newline() + header = f"{'Benchmark Name':<{max_name_len}} | {'Mean (s)':>10} | {'Median (s)':>10} | {'Std Dev':>10} | {'Runs':>4}" + title = " GT4Py Timer Report " + sep_len = max(0, len(header) - len(title)) + left = sep_len // 2 + right = sep_len - left + terminalreporter.line("-" * left + title + "-" * right, bold=True, blue=True) + terminalreporter.line(header) + terminalreporter.line("-" * len(header), blue=True) + import numpy as np + + for benchmark_name, gtx_metrics in benchmark_gtx_metrics: + terminalreporter.line( + f"{benchmark_name:<{max_name_len}} | {np.mean(gtx_metrics):>10.8f} | {np.median(gtx_metrics):>10.8f} | {np.std(gtx_metrics):>10.8f} | {len(gtx_metrics):>4}" + ) + terminalreporter.line("-" * len(header), blue=True) diff --git a/model/testing/src/icon4py/model/testing/serialbox.py b/model/testing/src/icon4py/model/testing/serialbox.py index aca51de9ab..335d293e8d 100644 --- a/model/testing/src/icon4py/model/testing/serialbox.py +++ b/model/testing/src/icon4py/model/testing/serialbox.py @@ -17,9 +17,10 @@ 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 from icon4py.model.common.grid import base, horizontal as h_grid, icon from icon4py.model.common.states import prognostic_state +from icon4py.model.common.type_alias import vpfloat, wpfloat from icon4py.model.common.utils import data_allocation as data_alloc @@ -56,7 +57,7 @@ def __init__( self.backend = backend self.xp = data_alloc.import_array_ns(self.backend) - def optionally_registered(*dims, dtype=type_alias.wpfloat): + def optionally_registered(*dims, dtype=wpfloat): def decorator(func): @functools.wraps(func) def wrapper(self, *args, **kwargs): @@ -84,15 +85,17 @@ def wrapper(self, *args, **kwargs): def log_meta_info(self): self.log.info(self.savepoint.metainfo) - def _get_field(self, name, *dimensions, dtype=float): + def _get_field(self, name, *dimensions, dtype=wpfloat): buffer = np.squeeze(self.serializer.read(name, self.savepoint).astype(dtype)) buffer = self._reduce_to_dim_size(buffer, dimensions) self.log.debug(f"{name} {buffer.shape}") return gtx.as_field(dimensions, buffer, allocator=self.backend) - def _get_field_component(self, name: str, level: int, dims: tuple[gtx.Dimension, gtx]): - buffer = self.serializer.read(name, self.savepoint).astype(float) + def _get_field_component( + self, name: str, level: int, dims: tuple[gtx.Dimension, gtx], dtype=wpfloat + ): + buffer = self.serializer.read(name, self.savepoint).astype(dtype) buffer = np.squeeze(buffer)[:, :, level] buffer = self._reduce_to_dim_size(buffer, dims) self.log.debug(f"{name} {buffer.shape}") @@ -105,7 +108,7 @@ def _reduce_to_dim_size(self, buffer, dimensions): ) return buffer[tuple(map(slice, buffer_size))] - def _get_field_from_ndarray(self, ar, *dimensions, dtype=float): + def _get_field_from_ndarray(self, ar, *dimensions, dtype=wpfloat): ar = self._reduce_to_dim_size(ar, dimensions) return gtx.as_field(dimensions, ar, allocator=self.backend, dtype=dtype) @@ -113,6 +116,15 @@ def get_metadata(self, *names): metadata = self.savepoint.metainfo.to_dict() return {n: metadata[n] for n in names if n in metadata} + def dtime(self, dtype=wpfloat): + metadata = self.savepoint.metainfo.to_dict() + try: + return dtype(metadata["dtime"]) + except KeyError as e: + raise RuntimeError( + "Invalid call to dtime() for a static savepoint. No time information in metadata." + ) from e + def _read_int32_shift1(self, name: str): """ Read a start indices field. @@ -294,7 +306,7 @@ def edge_center_lon(self): return self._get_field("edges_center_lon", dims.EdgeDim) def mean_cell_area(self): - return self.serializer.read("mean_cell_area", self.savepoint).astype(float)[0] + return self.serializer.read("mean_cell_area", self.savepoint).astype(wpfloat)[0] def edge_areas(self): return self._get_field("edge_areas", dims.EdgeDim) @@ -578,7 +590,7 @@ def geofac_grdiv(self): return self._get_field("geofac_grdiv", dims.EdgeDim, dims.E2C2EODim) def geofac_grg(self): - grg = np.squeeze(self.serializer.read("geofac_grg", self.savepoint)) + grg = np.squeeze(self.serializer.read("geofac_grg", self.savepoint).astype(wpfloat)) num_cells = self.sizes[dims.CellDim] return gtx.as_field( (dims.CellDim, dims.C2E2CODim), grg[:num_cells, :, 0], allocator=self.backend @@ -612,7 +624,7 @@ def rbf_vec_coeff_e(self): def rbf_vec_coeff_c1(self): dimensions = (dims.CellDim, dims.C2E2C2EDim) buffer = np.squeeze( - self.serializer.read("rbf_vec_coeff_c1", self.savepoint).astype(float) + self.serializer.read("rbf_vec_coeff_c1", self.savepoint).astype(wpfloat) ).transpose() buffer = self._reduce_to_dim_size(buffer, dimensions) return gtx.as_field(dimensions, buffer, allocator=self.backend) @@ -621,7 +633,7 @@ def rbf_vec_coeff_c1(self): def rbf_vec_coeff_c2(self): dimensions = (dims.CellDim, dims.C2E2C2EDim) buffer = np.squeeze( - self.serializer.read("rbf_vec_coeff_c2", self.savepoint).astype(float) + self.serializer.read("rbf_vec_coeff_c2", self.savepoint).astype(wpfloat) ).transpose() buffer = self._reduce_to_dim_size(buffer, dimensions) return gtx.as_field(dimensions, buffer, allocator=self.backend) @@ -647,25 +659,25 @@ def bdy_halo_c(self): return self._get_field("bdy_halo_c", dims.CellDim, dtype=bool) def d2dexdz2_fac1_mc(self): - return self._get_field("d2dexdz2_fac1_mc", dims.CellDim, dims.KDim) + return self._get_field("d2dexdz2_fac1_mc", dims.CellDim, dims.KDim, dtype=vpfloat) def d2dexdz2_fac2_mc(self): - return self._get_field("d2dexdz2_fac2_mc", dims.CellDim, dims.KDim) + return self._get_field("d2dexdz2_fac2_mc", dims.CellDim, dims.KDim, dtype=vpfloat) def d_exner_dz_ref_ic(self): - return self._get_field("d_exner_dz_ref_ic", dims.CellDim, dims.KDim) + return self._get_field("d_exner_dz_ref_ic", dims.CellDim, dims.KDim, dtype=vpfloat) def exner_exfac(self): - return self._get_field("exner_exfac", dims.CellDim, dims.KDim) + return self._get_field("exner_exfac", dims.CellDim, dims.KDim, dtype=vpfloat) def exner_ref_mc(self): - return self._get_field("exner_ref_mc", dims.CellDim, dims.KDim) + return self._get_field("exner_ref_mc", dims.CellDim, dims.KDim, dtype=vpfloat) def hmask_dd3d(self): return self._get_field("hmask_dd3d", dims.EdgeDim) def inv_ddqz_z_full(self): - return self._get_field("inv_ddqz_z_full", dims.CellDim, dims.KDim) + return self._get_field("inv_ddqz_z_full", dims.CellDim, dims.KDim, dtype=vpfloat) @IconSavepoint.optionally_registered(dims.CellDim, dims.KDim) def ddqz_z_full(self): @@ -675,7 +687,7 @@ def mask_prog_halo_c(self): return self._get_field("mask_prog_halo_c", dims.CellDim, dtype=bool) def pg_exdist(self): - return self._get_field("pg_exdist_dsl", dims.EdgeDim, dims.KDim) + return self._get_field("pg_exdist_dsl", dims.EdgeDim, dims.KDim, dtype=vpfloat) def pg_edgeidx_dsl(self): return self._get_field("pg_edgeidx_dsl", dims.EdgeDim, dims.KDim, dtype=bool) @@ -684,16 +696,16 @@ def rayleigh_w(self): return self._get_field("rayleigh_w", dims.KDim) def rho_ref_mc(self): - return self._get_field("rho_ref_mc", dims.CellDim, dims.KDim) + return self._get_field("rho_ref_mc", dims.CellDim, dims.KDim, dtype=vpfloat) def rho_ref_me(self): - return self._get_field("rho_ref_me", dims.EdgeDim, dims.KDim) + return self._get_field("rho_ref_me", dims.EdgeDim, dims.KDim, dtype=vpfloat) def scalfac_dd3d(self): return self._get_field("scalfac_dd3d", dims.KDim) def theta_ref_ic(self): - return self._get_field("theta_ref_ic", dims.CellDim, dims.KDim) + return self._get_field("theta_ref_ic", dims.CellDim, dims.KDim, dtype=vpfloat) def z_ifc(self): return self._get_field("z_ifc", dims.CellDim, dims.KDim) @@ -702,7 +714,7 @@ def z_mc(self): return self._get_field("z_mc", dims.CellDim, dims.KDim) def theta_ref_me(self): - return self._get_field("theta_ref_me", dims.EdgeDim, dims.KDim) + return self._get_field("theta_ref_me", dims.EdgeDim, dims.KDim, dtype=vpfloat) def vwind_expl_wgt(self): return self._get_field("vwind_expl_wgt", dims.CellDim) @@ -711,10 +723,12 @@ def vwind_impl_wgt(self): return self._get_field("vwind_impl_wgt", dims.CellDim) def wgtfacq_c_dsl(self): - return self._get_field("wgtfacq_c_dsl", dims.CellDim, dims.KDim) + return self._get_field("wgtfacq_c_dsl", dims.CellDim, dims.KDim, dtype=vpfloat) def zdiff_gradp(self): - return self._get_field("zdiff_gradp_dsl", dims.EdgeDim, dims.E2CDim, dims.KDim) + return self._get_field( + "zdiff_gradp_dsl", dims.EdgeDim, dims.E2CDim, dims.KDim, dtype=vpfloat + ) def vertoffset_gradp(self): return self._get_field( @@ -722,44 +736,44 @@ def vertoffset_gradp(self): ) def coeff1_dwdz(self): - return self._get_field("coeff1_dwdz", dims.CellDim, dims.KDim) + return self._get_field("coeff1_dwdz", dims.CellDim, dims.KDim, dtype=vpfloat) def coeff2_dwdz(self): - return self._get_field("coeff2_dwdz", dims.CellDim, dims.KDim) + return self._get_field("coeff2_dwdz", dims.CellDim, dims.KDim, dtype=vpfloat) def coeff_gradekin(self): - return self._get_field("coeff_gradekin", dims.EdgeDim, dims.E2CDim) + return self._get_field("coeff_gradekin", dims.EdgeDim, dims.E2CDim, dtype=vpfloat) def ddqz_z_full_e(self): - return self._get_field("ddqz_z_full_e", dims.EdgeDim, dims.KDim) + return self._get_field("ddqz_z_full_e", dims.EdgeDim, dims.KDim, dtype=vpfloat) def ddqz_z_half(self): - return self._get_field("ddqz_z_half", dims.CellDim, dims.KDim) + return self._get_field("ddqz_z_half", dims.CellDim, dims.KDim, dtype=vpfloat) def ddxn_z_full(self): - return self._get_field("ddxn_z_full", dims.EdgeDim, dims.KDim) + return self._get_field("ddxn_z_full", dims.EdgeDim, dims.KDim, dtype=vpfloat) def ddxt_z_full(self): - return self._get_field("ddxt_z_full", dims.EdgeDim, dims.KDim) + return self._get_field("ddxt_z_full", dims.EdgeDim, dims.KDim, dtype=vpfloat) @IconSavepoint.optionally_registered(dims.CellDim, dims.KDim, dtype=gtx.bool) def mask_hdiff(self): return self._get_field("mask_hdiff", dims.CellDim, dims.KDim, dtype=bool) def theta_ref_mc(self): - return self._get_field("theta_ref_mc", dims.CellDim, dims.KDim) + return self._get_field("theta_ref_mc", dims.CellDim, dims.KDim, dtype=vpfloat) def wgtfac_c(self): - return self._get_field("wgtfac_c", dims.CellDim, dims.KDim) + return self._get_field("wgtfac_c", dims.CellDim, dims.KDim, dtype=vpfloat) def wgtfac_e(self): - return self._get_field("wgtfac_e", dims.EdgeDim, dims.KDim) + return self._get_field("wgtfac_e", dims.EdgeDim, dims.KDim, dtype=vpfloat) def wgtfacq_e_dsl(self, k_level): - ar = np.squeeze(self.serializer.read("wgtfacq_e", self.savepoint)) + ar = np.squeeze(self.serializer.read("wgtfacq_e", self.savepoint).astype(vpfloat)) 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) + ar = np.pad(ar[:, ::-1], ((0, 0), (k, 0)), "constant", constant_values=(vpfloat(0.0),)) + return self._get_field_from_ndarray(ar, dims.EdgeDim, dims.KDim, dtype=vpfloat) @IconSavepoint.optionally_registered(dims.CellDim, dims.KDim) def zd_diffcoef(self): @@ -830,19 +844,19 @@ def tracer(self, ntracer: int): class IconDiffusionInitSavepoint(IconSavepoint): @IconSavepoint.optionally_registered(dims.CellDim, dims.KDim) def hdef_ic(self): - return self._get_field("hdef_ic", dims.CellDim, dims.KDim) + return self._get_field("hdef_ic", dims.CellDim, dims.KDim, dtype=vpfloat) - @IconSavepoint.optionally_registered(dims.CellDim, dims.KDim) + @IconSavepoint.optionally_registered(dims.CellDim, dims.KDim, dtype=vpfloat) def div_ic(self): - return self._get_field("div_ic", dims.CellDim, dims.KDim) + return self._get_field("div_ic", dims.CellDim, dims.KDim, dtype=vpfloat) - @IconSavepoint.optionally_registered(dims.CellDim, dims.KDim) + @IconSavepoint.optionally_registered(dims.CellDim, dims.KDim, dtype=vpfloat) def dwdx(self): - return self._get_field("dwdx", dims.CellDim, dims.KDim) + return self._get_field("dwdx", dims.CellDim, dims.KDim, dtype=vpfloat) - @IconSavepoint.optionally_registered(dims.CellDim, dims.KDim) + @IconSavepoint.optionally_registered(dims.CellDim, dims.KDim, dtype=vpfloat) def dwdy(self): - return self._get_field("dwdy", dims.CellDim, dims.KDim) + return self._get_field("dwdy", dims.CellDim, dims.KDim, dtype=vpfloat) def vn(self): return self._get_field("vn", dims.EdgeDim, dims.KDim) @@ -857,34 +871,34 @@ def exner(self): return self._get_field("exner", dims.CellDim, dims.KDim) def diff_multfac_smag(self): - return np.squeeze(self.serializer.read("diff_multfac_smag", self.savepoint)) + return np.squeeze(self.serializer.read("diff_multfac_smag", self.savepoint).astype(vpfloat)) def enh_smag_fac(self): - return np.squeeze(self.serializer.read("enh_smag_fac", self.savepoint)) + return np.squeeze(self.serializer.read("enh_smag_fac", self.savepoint).astype(vpfloat)) def smag_limit(self): - return np.squeeze(self.serializer.read("smag_limit", self.savepoint)) + return np.squeeze(self.serializer.read("smag_limit", self.savepoint).astype(vpfloat)) def diff_multfac_n2w(self): - return np.squeeze(self.serializer.read("diff_multfac_n2w", self.savepoint)) + return np.squeeze(self.serializer.read("diff_multfac_n2w", self.savepoint).astype(wpfloat)) - def nudgezone_diff(self) -> int: - return self.serializer.read("nudgezone_diff", self.savepoint)[0] + def nudgezone_diff(self): + return self.serializer.read("nudgezone_diff", self.savepoint).astype(vpfloat)[0] - def bdy_diff(self) -> int: - return self.serializer.read("bdy_diff", self.savepoint)[0] + def bdy_diff(self): + return self.serializer.read("bdy_diff", self.savepoint).astype(vpfloat)[0] - def fac_bdydiff_v(self) -> int: - return self.serializer.read("fac_bdydiff_v", self.savepoint)[0] + def fac_bdydiff_v(self): + return self.serializer.read("fac_bdydiff_v", self.savepoint).astype(wpfloat)[0] def smag_offset(self): - return self.serializer.read("smag_offset", self.savepoint)[0] + return self.serializer.read("smag_offset", self.savepoint).astype(vpfloat)[0] def diff_multfac_w(self): - return self.serializer.read("diff_multfac_w", self.savepoint)[0] + return self.serializer.read("diff_multfac_w", self.savepoint).astype(wpfloat)[0] def diff_multfac_vn(self): - return self.serializer.read("diff_multfac_vn", self.savepoint) + return self.serializer.read("diff_multfac_vn", self.savepoint).astype(wpfloat) def rho(self): return self._get_field("rho", dims.CellDim, dims.KDim) @@ -910,19 +924,19 @@ def w(self): return self._get_field("w", dims.CellDim, dims.KDim) def dwdx(self): - return self._get_field("dwdx", dims.CellDim, dims.KDim) + return self._get_field("dwdx", dims.CellDim, dims.KDim, dtype=vpfloat) def dwdy(self): - return self._get_field("dwdy", dims.CellDim, dims.KDim) + return self._get_field("dwdy", dims.CellDim, dims.KDim, dtype=vpfloat) def exner(self): return self._get_field("exner", dims.CellDim, dims.KDim) def div_ic(self): - return self._get_field("div_ic", dims.CellDim, dims.KDim) + return self._get_field("div_ic", dims.CellDim, dims.KDim, dtype=vpfloat) def hdef_ic(self): - return self._get_field("hdef_ic", dims.CellDim, dims.KDim) + return self._get_field("hdef_ic", dims.CellDim, dims.KDim, dtype=vpfloat) class IconNonHydroInitSavepoint(IconSavepoint): @@ -933,22 +947,22 @@ def z_kin_hor_e(self): return self._get_field("z_kin_hor_e", dims.EdgeDim, dims.KDim) def vn_ie(self): - return self._get_field("vn_ie", dims.EdgeDim, dims.KDim) + return self._get_field("vn_ie", dims.EdgeDim, dims.KDim, dtype=vpfloat) def vt(self): - return self._get_field("vt", dims.EdgeDim, dims.KDim) + return self._get_field("vt", dims.EdgeDim, dims.KDim, dtype=vpfloat) def bdy_divdamp(self): return self._get_field("bdy_divdamp", dims.KDim) def divdamp_fac_o2(self): - return self.serializer.read("divdamp_fac_o2", self.savepoint).astype(float)[0] + return self.serializer.read("divdamp_fac_o2", self.savepoint).astype(wpfloat)[0] def ddt_exner_phy(self): - return self._get_field("ddt_exner_phy", dims.CellDim, dims.KDim) + return self._get_field("ddt_exner_phy", dims.CellDim, dims.KDim, dtype=vpfloat) def ddt_vn_phy(self): - return self._get_field("ddt_vn_phy", dims.EdgeDim, dims.KDim) + return self._get_field("ddt_vn_phy", dims.EdgeDim, dims.KDim, dtype=vpfloat) def exner_now(self): return self._get_field("exner_now", dims.CellDim, dims.KDim) @@ -981,7 +995,7 @@ def grf_tend_vn(self): return self._get_field("grf_tend_vn", dims.EdgeDim, dims.KDim) def w_concorr_c(self): - return self._get_field("w_concorr_c", dims.CellDim, dims.KDim) + return self._get_field("w_concorr_c", dims.CellDim, dims.KDim, dtype=vpfloat) def ddt_vn_apc_pc(self, ntnd): return self._get_field_component("ddt_vn_apc_pc", ntnd, (dims.EdgeDim, dims.KDim)) @@ -1004,26 +1018,26 @@ def mass_flx_ic(self): def rho_ic(self): return self._get_field("rho_ic", dims.CellDim, dims.KDim) - @IconSavepoint.optionally_registered() + @IconSavepoint.optionally_registered(dtype=vpfloat) def rho_incr(self): - return self._get_field("rho_incr", dims.CellDim, dims.KDim) + return self._get_field("rho_incr", dims.CellDim, dims.KDim, dtype=vpfloat) - @IconSavepoint.optionally_registered() + @IconSavepoint.optionally_registered(dtype=vpfloat) def exner_incr(self): - return self._get_field("exner_incr", dims.CellDim, dims.KDim) + return self._get_field("exner_incr", dims.CellDim, dims.KDim, dtype=vpfloat) - @IconSavepoint.optionally_registered() + @IconSavepoint.optionally_registered(dtype=vpfloat) def vn_incr(self): - return self._get_field("vn_incr", dims.EdgeDim, dims.KDim) + return self._get_field("vn_incr", dims.EdgeDim, dims.KDim, dtype=vpfloat) - @IconSavepoint.optionally_registered() + @IconSavepoint.optionally_registered(dtype=vpfloat) def exner_dyn_incr(self): - return self._get_field("exner_dyn_incr", dims.CellDim, dims.KDim) + return self._get_field("exner_dyn_incr", dims.CellDim, dims.KDim, dtype=vpfloat) - def scal_divdamp_o2(self) -> float: - return self.serializer.read("scal_divdamp_o2", self.savepoint)[0] + def scal_divdamp_o2(self) -> wpfloat: + return self.serializer.read("scal_divdamp_o2", self.savepoint).astype(wpfloat)[0] - def scal_divdamp(self) -> fa.KField[float]: + def scal_divdamp(self) -> fa.KField[wpfloat]: return self._get_field("scal_divdamp", dims.KDim) def theta_v_ic(self): @@ -1068,17 +1082,17 @@ def z_contr_w_fl_l(self): def z_q(self): return self._get_field("z_q", dims.CellDim, dims.KDim) - def wgt_nnow_rth(self) -> float: - return self.serializer.read("wgt_nnow_rth", self.savepoint)[0] + def wgt_nnow_rth(self) -> wpfloat: + return self.serializer.read("wgt_nnow_rth", self.savepoint).astype(wpfloat)[0] - def wgt_nnew_rth(self) -> float: - return self.serializer.read("wgt_nnew_rth", self.savepoint)[0] + def wgt_nnew_rth(self) -> wpfloat: + return self.serializer.read("wgt_nnew_rth", self.savepoint).astype(wpfloat)[0] - def wgt_nnow_vel(self) -> float: - return self.serializer.read("wgt_nnow_vel", self.savepoint)[0] + def wgt_nnow_vel(self) -> wpfloat: + return self.serializer.read("wgt_nnow_vel", self.savepoint).astype(wpfloat)[0] - def wgt_nnew_vel(self) -> float: - return self.serializer.read("wgt_nnew_vel", self.savepoint)[0] + def wgt_nnew_vel(self) -> wpfloat: + return self.serializer.read("wgt_nnew_vel", self.savepoint).astype(wpfloat)[0] def w_now(self): return self._get_field("w_now", dims.CellDim, dims.KDim) @@ -1101,7 +1115,7 @@ def vn(self): return self._get_field("vn_now", dims.EdgeDim, dims.KDim) def vt(self): - return self._get_field("vt", dims.EdgeDim, dims.KDim) + return self._get_field("vt", dims.EdgeDim, dims.KDim, dtype=vpfloat) def z_rth_pr(self, ind: TwoIndex): return self._get_field_component("z_rth_pr", ind, (dims.CellDim, dims.KDim)) @@ -1125,10 +1139,10 @@ def ddt_vn_apc_ntl(self, ntnd): return self._get_field_component("ddt_vn_apc_pc", ntnd, (dims.EdgeDim, dims.KDim)) def ddt_vn_phy(self): - return self._get_field("ddt_vn_phy", dims.EdgeDim, dims.KDim) + return self._get_field("ddt_vn_phy", dims.EdgeDim, dims.KDim, dtype=vpfloat) def vn_incr(self): - return self._get_field("vn_now", dims.EdgeDim, dims.KDim) + return self._get_field("vn_now", dims.EdgeDim, dims.KDim, dtype=vpfloat) def bdy_divdamp(self): return self._get_field("bdy_divdamp", dims.KDim) @@ -1184,7 +1198,7 @@ def rho_ic(self): return self._get_field("rho_ic", dims.CellDim, dims.KDim) def w_concorr_c(self): - return self._get_field("w_concorr_c", dims.CellDim, dims.KDim) + return self._get_field("w_concorr_c", dims.CellDim, dims.KDim, dtype=vpfloat) def exner_nnow(self): return self._get_field("exner_now", dims.CellDim, dims.KDim) @@ -1220,15 +1234,15 @@ def exner_pr(self): return self._get_field("exner_pr", dims.CellDim, dims.KDim) def ddt_exner_phy(self): - return self._get_field("ddt_exner_phy", dims.CellDim, dims.KDim) + return self._get_field("ddt_exner_phy", dims.CellDim, dims.KDim, dtype=vpfloat) @IconSavepoint.optionally_registered() def rho_incr(self): - return self._get_field("rho_now", dims.CellDim, dims.KDim) + return self._get_field("rho_now", dims.CellDim, dims.KDim, dtype=vpfloat) - @IconSavepoint.optionally_registered() + @IconSavepoint.optionally_registered(dtype=vpfloat) def exner_incr(self): - return self._get_field("exner_now", dims.CellDim, dims.KDim) + return self._get_field("exner_now", dims.CellDim, dims.KDim, dtype=vpfloat) def z_raylfac(self): return self._get_field("z_raylfac", dims.KDim) @@ -1245,9 +1259,9 @@ def theta_v(self): def z_dwdz_dd(self): return self._get_field("z_dwdz_dd", dims.CellDim, dims.KDim) - @IconSavepoint.optionally_registered() + @IconSavepoint.optionally_registered(dtype=vpfloat) def exner_dyn_incr(self): - return self._get_field("exner_dyn_incr", dims.CellDim, dims.KDim) + return self._get_field("exner_dyn_incr", dims.CellDim, dims.KDim, dtype=vpfloat) def mass_flx_ic(self): return self._get_field("mass_flx_ic", dims.CellDim, dims.KDim) @@ -1267,7 +1281,7 @@ def vn(self): return self._get_field("vn", dims.EdgeDim, dims.KDim) def vt(self): - return self._get_field("vt", dims.EdgeDim, dims.KDim) + return self._get_field("vt", dims.EdgeDim, dims.KDim, dtype=vpfloat) def z_rho_e(self): return self._get_field("z_rho_e", dims.EdgeDim, dims.KDim) @@ -1279,7 +1293,7 @@ def z_vt_ie(self): return self._get_field("z_vt_ie", dims.EdgeDim, dims.KDim) def vn_ie(self): - return self._get_field("vn_ie", dims.EdgeDim, dims.KDim) + return self._get_field("vn_ie", dims.EdgeDim, dims.KDim, dtype=vpfloat) def mass_fl_e(self): return self._get_field("mass_fl_e", dims.EdgeDim, dims.KDim) @@ -1302,7 +1316,7 @@ def z_graddiv_vn(self): return self._get_field("z_graddiv_vn", dims.EdgeDim, dims.KDim) def vt(self): - return self._get_field("vt", dims.EdgeDim, dims.KDim) + return self._get_field("vt", dims.EdgeDim, dims.KDim, dtype=vpfloat) def mass_fl_e(self): return self._get_field("mass_fl_e", dims.EdgeDim, dims.KDim) @@ -1311,7 +1325,7 @@ def z_theta_v_fl_e(self): return self._get_field("z_theta_v_fl_e", dims.EdgeDim, dims.KDim) def vn_ie(self): - return self._get_field("vn_ie", dims.EdgeDim, dims.KDim) + return self._get_field("vn_ie", dims.EdgeDim, dims.KDim, dtype=vpfloat) def z_vt_ie(self): return self._get_field("z_vt_ie", dims.EdgeDim, dims.KDim) @@ -1403,7 +1417,7 @@ def vn_traj(self): return self._get_field("vn_traj", dims.EdgeDim, dims.KDim) def exner_dyn_incr(self): - return self._get_field("exner_dyn_incr", dims.CellDim, dims.KDim) + return self._get_field("exner_dyn_incr", dims.CellDim, dims.KDim, dtype=vpfloat) def z_exner_ic(self): return self._get_field("z_exner_ic", dims.CellDim, dims.KDim) @@ -1430,7 +1444,7 @@ def z_theta_v_pr_ic(self): return self._get_field("z_theta_v_pr_ic", dims.CellDim, dims.KDim) def vt(self): - return self._get_field("vt", dims.EdgeDim, dims.KDim) + return self._get_field("vt", dims.EdgeDim, dims.KDim, dtype=vpfloat) def z_flxdiv_mass(self): return self._get_field("z_flxdiv_mass", dims.CellDim, dims.KDim) @@ -1445,7 +1459,7 @@ def z_contr_w_fl_l(self): return self._get_field("z_contr_w_fl_l", dims.CellDim, dims.KDim) def vn_ie(self): - return self._get_field("vn_ie", dims.EdgeDim, dims.KDim) + return self._get_field("vn_ie", dims.EdgeDim, dims.KDim, dtype=vpfloat) def z_vt_ie(self): return self._get_field("z_vt_ie", dims.EdgeDim, dims.KDim) @@ -1454,7 +1468,7 @@ def z_w_concorr_me(self): return self._get_field("z_w_concorr_me", dims.EdgeDim, dims.KDim) def w_concorr_c(self): - return self._get_field("w_concorr_c", dims.CellDim, dims.KDim) + return self._get_field("w_concorr_c", dims.CellDim, dims.KDim, dtype=vpfloat) def z_theta_v_fl_e(self): return self._get_field("z_theta_v_fl_e", dims.EdgeDim, dims.KDim) @@ -1493,17 +1507,17 @@ def exner_new(self): class IconVelocityInitSavepoint(IconSavepoint): - def cfl_w_limit(self) -> float: - return self.serializer.read("cfl_w_limit", self.savepoint)[0] + def cfl_w_limit(self) -> vpfloat: + return self.serializer.read("cfl_w_limit", self.savepoint).astype(vpfloat)[0] def vn_only(self) -> bool: return bool(self.serializer.read("vn_only", self.savepoint)[0]) def max_vcfl_dyn(self): - return self.serializer.read("max_vcfl_dyn", self.savepoint)[0] + return self.serializer.read("max_vcfl_dyn", self.savepoint).astype(vpfloat)[0] - def scalfac_exdiff(self) -> float: - return self.serializer.read("scalfac_exdiff", self.savepoint)[0] + def scalfac_exdiff(self) -> wpfloat: + return self.serializer.read("scalfac_exdiff", self.savepoint).astype(wpfloat)[0] def ddt_vn_apc_pc(self, ntnd: TimeIndex): return self._get_field_component("ddt_vn_apc_pc", ntnd, (dims.EdgeDim, dims.KDim)) @@ -1515,10 +1529,10 @@ def vn(self): return self._get_field("vn", dims.EdgeDim, dims.KDim) def vn_ie(self): - return self._get_field("vn_ie", dims.EdgeDim, dims.KDim) + return self._get_field("vn_ie", dims.EdgeDim, dims.KDim, dtype=vpfloat) def vt(self): - return self._get_field("vt", dims.EdgeDim, dims.KDim) + return self._get_field("vt", dims.EdgeDim, dims.KDim, dtype=vpfloat) def w(self): return self._get_field("w", dims.CellDim, dims.KDim) @@ -1533,7 +1547,7 @@ def z_w_concorr_me(self): return self._get_field("z_w_concorr_me", dims.EdgeDim, dims.KDim) def w_concorr_c(self): - return self._get_field("w_concorr_c", dims.CellDim, dims.KDim) + return self._get_field("w_concorr_c", dims.CellDim, dims.KDim, dtype=vpfloat) def lvn_only(self) -> bool: return bool(self.serializer.read("vn_only", self.savepoint)[0]) @@ -1544,7 +1558,7 @@ def z_w_con_c_full(self): class IconVelocityExitSavepoint(IconSavepoint): def max_vcfl_dyn(self): - return self.serializer.read("max_vcfl_dyn", self.savepoint)[0] + return self.serializer.read("max_vcfl_dyn", self.savepoint).astype(vpfloat)[0] def ddt_vn_apc_pc(self, ntnd: TimeIndex): return self._get_field_component("ddt_vn_apc_pc", ntnd, (dims.EdgeDim, dims.KDim)) @@ -1559,13 +1573,13 @@ def w(self): return self._get_field("w", dims.CellDim, dims.KDim) def vt(self): - return self._get_field("vt", dims.EdgeDim, dims.KDim) + return self._get_field("vt", dims.EdgeDim, dims.KDim, dtype=vpfloat) def vn_ie(self): - return self._get_field("vn_ie", dims.EdgeDim, dims.KDim) + return self._get_field("vn_ie", dims.EdgeDim, dims.KDim, dtype=vpfloat) def w_concorr_c(self): - return self._get_field("w_concorr_c", dims.CellDim, dims.KDim) + return self._get_field("w_concorr_c", dims.CellDim, dims.KDim, dtype=vpfloat) def z_vt_ie(self): return self._get_field("z_vt_ie", dims.EdgeDim, dims.KDim) @@ -1730,7 +1744,7 @@ def qnc(self): return self._get_field("qnc", dims.CellDim) def dtime(self): - return self.serializer.read("dtime", self.savepoint)[0] + return self.serializer.read("dtime", self.savepoint).astype(wpfloat)[0] class IconSatadExitSavepoint(IconSavepoint): diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index 489381db13..750a8fc041 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -9,6 +9,7 @@ from __future__ import annotations import dataclasses +import warnings from collections.abc import Callable, Sequence from typing import Any, ClassVar @@ -16,13 +17,19 @@ import numpy as np import pytest from gt4py import eve -from gt4py.next import constructors, typing as gtx_typing +from gt4py.next import ( + config as gtx_config, + constructors, + metrics as gtx_metrics, + typing as gtx_typing, +) # TODO(havogt): import will disappear after FieldOperators support `.compile` from gt4py.next.ffront.decorator import FieldOperator from icon4py.model.common import model_backends, model_options from icon4py.model.common.grid import base +from icon4py.model.common.type_alias import precision from icon4py.model.common.utils import device_utils @@ -77,29 +84,55 @@ def test_and_benchmark( grid: base.Grid, _properly_allocated_input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], _configured_program: Callable[..., None], + request: pytest.FixtureRequest, ) -> None: - reference_outputs = self.reference( - _ConnectivityConceptFixer( - grid # TODO(havogt): pass as keyword argument (needs fixes in some tests) - ), - **{ - k: v.asnumpy() if isinstance(v, gtx.Field) else v - for k, v in _properly_allocated_input_data.items() - }, - ) - - _configured_program(**_properly_allocated_input_data, offset_provider=grid.connectivities) - self._verify_stencil_test( - input_data=_properly_allocated_input_data, reference_outputs=reference_outputs - ) + benchmark_only_option = request.config.getoption("benchmark_only") + benchmark_only_mark = request.node.get_closest_marker("benchmark_only") is not None + if (not benchmark_only_option) and (not benchmark_only_mark): + reference_outputs = self.reference( + _ConnectivityConceptFixer( + grid # TODO(havogt): pass as keyword argument (needs fixes in some tests) + ), + **{ + k: v.asnumpy() if isinstance(v, gtx.Field) else v + for k, v in _properly_allocated_input_data.items() + }, + ) + + _configured_program(**_properly_allocated_input_data, offset_provider=grid.connectivities) + self._verify_stencil_test( + input_data=_properly_allocated_input_data, reference_outputs=reference_outputs + ) if benchmark is not None and benchmark.enabled: + warmup_enabled = request.config.getoption("benchmark_warmup") + if warmup_enabled: + print("[WARNING] Benchmark warmup enabled, GT4Py timers include warmup iterations.") + # Clean up GT4Py metrics from previous runs + if gtx_config.COLLECT_METRICS_LEVEL > 0: + gtx_metrics.sources.clear() + benchmark( _configured_program, **_properly_allocated_input_data, offset_provider=grid.connectivities, ) + # Collect GT4Py runtime metrics if enabled + if gtx_config.COLLECT_METRICS_LEVEL > 0: + assert ( + len(gtx_metrics.sources) == 1 + ), "Expected exactly one entry in gtx_metrics.sources" + # Store GT4Py metrics in benchmark.extra_info + metrics_data = gtx_metrics.sources + key = next(iter(metrics_data)) + compute_samples = metrics_data[key].metrics["compute"].samples + # emprically exclude first few iterations run for warmup + initial_program_iterations_to_skip = 2 + benchmark.extra_info["gtx_metrics"] = compute_samples[ + initial_program_iterations_to_skip: + ] + class StencilTest: """ @@ -125,6 +158,37 @@ class StencilTest: OUTPUTS: ClassVar[tuple[str | Output, ...]] STATIC_PARAMS: ClassVar[dict[str, Sequence[str]] | None] = None + # standard tolerances of np.testing.assert_allclose: + RTOL = 1e-7 + ATOL = 0.0 + + def try_allclose( + self, + actual, + desired, + rtol=RTOL, + atol=ATOL, + equal_nan=True, + err_msg="Verification failed for ", + ): + try: + np.testing.assert_allclose( + actual, + desired, + rtol=rtol, + atol=atol, + equal_nan=equal_nan, + err_msg=err_msg, + ) + except AssertionError as e: + if precision != "single": + raise e + else: + # Because these stencil_tests are ran with input fields in unrealistic ranges the single precision version is not required to pass them. The tolerances can be changed s.t. they would pass but they are not very meaningful. + warnings.warn( + "As expected the stencil test did not pass for single: " + str(e), UserWarning + ) + reference: ClassVar[Callable[..., dict[str, np.ndarray | tuple[np.ndarray, ...]]]] @pytest.fixture @@ -186,18 +250,23 @@ def _verify_stencil_test( input_data_name = input_data[name] # for mypy if isinstance(input_data_name, tuple): for i_out_field, out_field in enumerate(input_data_name): - np.testing.assert_allclose( + self.try_allclose( out_field.asnumpy()[gtslice], reference_outputs[name][i_out_field][refslice], + rtol=self.RTOL, + atol=self.ATOL, equal_nan=True, err_msg=f"Verification failed for '{name}[{i_out_field}]'", ) else: reference_outputs_name = reference_outputs[name] # for mypy assert isinstance(reference_outputs_name, np.ndarray) - np.testing.assert_allclose( + + self.try_allclose( input_data_name.asnumpy()[gtslice], reference_outputs_name[refslice], + rtol=self.RTOL, + atol=self.ATOL, equal_nan=True, err_msg=f"Verification failed for '{name}'", ) @@ -216,11 +285,33 @@ def static_variant(request: pytest.FixtureRequest) -> Sequence[str]: def __init_subclass__(cls, **kwargs: Any) -> None: super().__init_subclass__(**kwargs) - setattr(cls, f"test_{cls.__name__}", test_and_benchmark) + pytest_prefix = "test_" + + # only add `test_and_benchmark` to direct subclasses of `StencilTest`. Inherited tests can use the same function. + if cls.__base__ == StencilTest: + setattr(cls, f"{pytest_prefix}{cls.__name__}", test_and_benchmark) # decorate `static_variant` with parametrized fixtures, since the # parametrization is only available in the concrete subclass definition - if cls.STATIC_PARAMS is None: + # Check if cls.static_variant has changes in inherited subclasses and update it accordingly + if hasattr(cls.static_variant, "_pytestfixturefunction"): + base_static = getattr(cls.__base__, "STATIC_PARAMS", None) + if base_static != cls.STATIC_PARAMS: + if cls.STATIC_PARAMS is None: + cls.static_variant = staticmethod(pytest.fixture(lambda: ())) # type: ignore[method-assign, assignment] + else: + # copy of `static_variant` + def _new_static_variant(request: pytest.FixtureRequest) -> Sequence[str]: + _, variant = request.param + return () if variant is None else variant + + # replace with new parametrized fixture + cls.static_variant = staticmethod( # type: ignore[method-assign] + pytest.fixture( + params=cls.STATIC_PARAMS.items(), scope="class", ids=lambda p: p[0] + )(_new_static_variant) + ) + elif cls.STATIC_PARAMS is None: # not parametrized, return an empty tuple cls.static_variant = staticmethod(pytest.fixture(lambda: ())) # type: ignore[method-assign, assignment] # we override with a non-parametrized function else: diff --git a/model/testing/src/icon4py/model/testing/test_utils.py b/model/testing/src/icon4py/model/testing/test_utils.py index 8d23824591..f50a2f54bb 100644 --- a/model/testing/src/icon4py/model/testing/test_utils.py +++ b/model/testing/src/icon4py/model/testing/test_utils.py @@ -16,13 +16,21 @@ from typing_extensions import Buffer from icon4py.model.common import model_options +from icon4py.model.common.constants import VP_EPS, WP_EPS +from icon4py.model.common.type_alias import vpfloat + + +wp_eps = WP_EPS # to enable to set tolerances with eps dependance +vp_eps = VP_EPS + +tol_big = 5e3 * vp_eps # for double ≈ 1.11e-12 def dallclose( a: npt.ArrayLike, b: npt.ArrayLike, - rtol: float = 1.0e-12, - atol: float = 0.0, + rtol: vpfloat = tol_big, + atol: vpfloat = vp_eps, equal_nan: bool = False, ) -> bool: return np.allclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) diff --git a/noxfile.py b/noxfile.py index 2ff5bb10d1..44e5b564f1 100644 --- a/noxfile.py +++ b/noxfile.py @@ -9,6 +9,7 @@ from __future__ import annotations import os +import pathlib import re from collections.abc import Sequence from datetime import datetime @@ -152,11 +153,49 @@ def test_model( _install_session_venv(session, extras=["fortran", "io", "testing"], groups=["test"]) pytest_args = _selection_to_pytest_args(selection) + + posargs_list = list(session.posargs) + if "--single-precision" in posargs_list: + session.env["FLOAT_PRECISION"] = "single" + posargs_list.remove("--single-precision") with session.chdir(f"model/{subpackage}"): session.run( *f"pytest -sv --benchmark-disable -n {os.environ.get('NUM_PROCESSES', 'auto')}".split(), *pytest_args, - *session.posargs, + *posargs_list, + success_codes=[0, NO_TESTS_COLLECTED_EXIT_CODE], + ) + + +# TODO(pstark): Is this a good way to add the single precision tests? +# Or is a "--single-precision" session flag in test_model() nicer? +@nox.session(python=["3.10", "3.11"]) +def __test_file(session: nox.Session) -> None: + """Run tests for specific pytest file.""" + _install_session_venv(session, extras=["dace", "fortran", "io", "testing"], groups=["test"]) + + testfile = None + filtered_posargs = [] + for arg in session.posargs: + if arg.startswith("--testfile="): + testfile = arg.split("=", 1)[1] + else: + filtered_posargs.append(arg) + + if not testfile: + session.error("Missing required argument: --testfile=/path/to/file") + + if not pathlib.Path(testfile).is_file(): + session.error(f"File '{testfile}' not found.") + + # Use the directory of the test file as working directory + test_dir = pathlib.Path(testfile).parent if pathlib.Path(testfile).parent else "." + test_filename = pathlib.Path(testfile).name + + with session.chdir(test_dir): + session.run( + *f"pytest {test_filename} --benchmark-disable -n {os.environ.get('NUM_PROCESSES', 'auto')}".split(), + *filtered_posargs, success_codes=[0, NO_TESTS_COLLECTED_EXIT_CODE], ) diff --git a/tools/tests/tools/py2fgen/wrappers/test_diffusion_wrapper.py b/tools/tests/tools/py2fgen/wrappers/test_diffusion_wrapper.py index 88f155a0c2..b998bf1704 100644 --- a/tools/tests/tools/py2fgen/wrappers/test_diffusion_wrapper.py +++ b/tools/tests/tools/py2fgen/wrappers/test_diffusion_wrapper.py @@ -115,11 +115,11 @@ def test_diffusion_wrapper_granule_inputs( exner = test_utils.array_to_array_info(savepoint_diffusion_init.exner().ndarray) theta_v = test_utils.array_to_array_info(savepoint_diffusion_init.theta_v().ndarray) rho = test_utils.array_to_array_info(savepoint_diffusion_init.rho().ndarray) - dtime = savepoint_diffusion_init.get_metadata("dtime")["dtime"] + dtime = savepoint_diffusion_init.dtime() # --- Expected objects that form inputs into init and run functions expected_icon_grid = icon_grid - expected_dtime = savepoint_diffusion_init.get_metadata("dtime").get("dtime") + expected_dtime = dtime expected_edge_geometry: grid_states.EdgeParams = grid_savepoint.construct_edge_geometry() expected_cell_geometry: grid_states.CellParams = grid_savepoint.construct_cell_geometry() expected_interpolation_state = diffusion_states.DiffusionInterpolationState( @@ -376,7 +376,7 @@ def test_diffusion_wrapper_single_step( exner = test_utils.array_to_array_info(savepoint_diffusion_init.exner().ndarray) theta_v = test_utils.array_to_array_info(savepoint_diffusion_init.theta_v().ndarray) rho = test_utils.array_to_array_info(savepoint_diffusion_init.rho().ndarray) - dtime = savepoint_diffusion_init.get_metadata("dtime")["dtime"] + dtime = savepoint_diffusion_init.dtime() ffi = cffi.FFI() # Call diffusion_init diff --git a/tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py b/tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py index 5e70372e42..e94f15b31d 100644 --- a/tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py +++ b/tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py @@ -331,7 +331,7 @@ def test_dycore_wrapper_granule_inputs( ) # undo the -1 to go back to Fortran value # other params - dtime = sp.get_metadata("dtime").get("dtime") + dtime = sp.dtime() lprep_adv = sp.get_metadata("prep_adv").get("prep_adv") # metric state parameters @@ -598,7 +598,7 @@ def test_dycore_wrapper_granule_inputs( ), # TODO(): sp.vol_flx_ic(), ) expected_second_order_divdamp_factor = sp.divdamp_fac_o2() - expected_dtime = sp.get_metadata("dtime").get("dtime") + expected_dtime = sp.dtime() expected_lprep_adv = sp.get_metadata("prep_adv").get("prep_adv") expected_at_first_substep = substep_init == 1 expected_at_last_substep = substep_init == ndyn_substeps @@ -885,7 +885,7 @@ def test_granule_solve_nonhydro_single_step_regional( sp_step_exit = savepoint_nonhydro_step_final # other params - dtime = sp.get_metadata("dtime").get("dtime") + dtime = sp.dtime() lprep_adv = sp.get_metadata("prep_adv").get("prep_adv") # solve nh run parameters @@ -1064,7 +1064,7 @@ def test_granule_solve_nonhydro_multi_step_regional( sp_step_exit = savepoint_nonhydro_step_final # other params - dtime = sp.get_metadata("dtime").get("dtime") + dtime = sp.dtime() lprep_adv = sp.get_metadata("prep_adv").get("prep_adv") # solve nh run parameters