From 40268eadc3c999b5b4ef3e69dc4cd7c80a2c0bd5 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 6 Feb 2024 09:22:14 +0100 Subject: [PATCH 1/2] Rename second batch of stencils in mo_solve_nonhydro (#379) This PR renames and cleans up stencils in mo_solve_nonhydro. --- ...ysis_increments_from_data_assimilation.py} | 7 +- ... => apply_4th_order_divergence_damping.py} | 7 +- ...py => apply_rayleigh_damping_mechanism.py} | 7 +- ...=> compute_dwdz_for_divergence_damping.py} | 7 +- ...l_67.py => compute_exner_from_rhotheta.py} | 7 +- ...te_results_for_thermodynamic_variables.py} | 7 +- ...encil_66.py => compute_theta_and_exner.py} | 7 +- .../dycore/nh_solve/solve_nonhydro.py | 106 ++++++++---------- .../dycore/nh_solve/solve_nonhydro_program.py | 12 +- ...iagonal_matrix_for_w_back_substitution.py} | 9 +- ...tridiagonal_matrix_for_w_forward_sweep.py} | 11 +- ...il_61.py => update_densety_exener_wind.py} | 13 +-- ... update_dynamical_exner_time_increment.py} | 7 +- ...ydro_stencil_58.py => update_mass_flux.py} | 7 +- ...cil_65.py => update_mass_flux_weighted.py} | 7 +- ...nhydro_stencil_68.py => update_theta_v.py} | 7 +- ..._nonhydro_stencil_62.py => update_wind.py} | 7 +- ...ysis_increments_from_data_assimilation.py} | 6 +- ...est_apply_4th_order_divergence_damping.py} | 6 +- ... test_apply_rayleigh_damping_mechanism.py} | 6 +- ...st_compute_dwdz_for_divergence_damping.py} | 6 +- ...py => test_compute_exner_from_rhotheta.py} | 6 +- ...te_results_for_thermodynamic_variables.py} | 6 +- ..._66.py => test_compute_theta_and_exner.py} | 6 +- ...iagonal_matrix_for_w_back_substitution.py} | 6 +- ...tridiagonal_matrix_for_w_forward_sweep.py} | 12 +- ....py => test_update_densety_exener_wind.py} | 6 +- ..._update_dynamical_exner_time_increment.py} | 6 +- ...stencil_58.py => test_update_mass_flux.py} | 6 +- ...5.py => test_update_mass_flux_weighted.py} | 6 +- ...o_stencil_68.py => test_update_theta_v.py} | 6 +- ...ydro_stencil_62.py => test_update_wind.py} | 6 +- 32 files changed, 158 insertions(+), 173 deletions(-) rename model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/{mo_solve_nonhydro_stencil_50.py => add_analysis_increments_from_data_assimilation.py} (90%) rename model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/{mo_solve_nonhydro_4th_order_divdamp.py => apply_4th_order_divergence_damping.py} (90%) rename model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/{mo_solve_nonhydro_stencil_54.py => apply_rayleigh_damping_mechanism.py} (90%) rename model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/{mo_solve_nonhydro_stencil_56_63.py => compute_dwdz_for_divergence_damping.py} (90%) rename model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/{mo_solve_nonhydro_stencil_67.py => compute_exner_from_rhotheta.py} (91%) rename model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/{mo_solve_nonhydro_stencil_55.py => compute_results_for_thermodynamic_variables.py} (94%) rename model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/{mo_solve_nonhydro_stencil_66.py => compute_theta_and_exner.py} (92%) rename model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/{mo_solve_nonhydro_stencil_53.py => solve_tridiagonal_matrix_for_w_back_substitution.py} (81%) rename model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/{mo_solve_nonhydro_stencil_52.py => solve_tridiagonal_matrix_for_w_forward_sweep.py} (93%) rename model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/{mo_solve_nonhydro_stencil_61.py => update_densety_exener_wind.py} (88%) rename model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/{mo_solve_nonhydro_stencil_60.py => update_dynamical_exner_time_increment.py} (91%) rename model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/{mo_solve_nonhydro_stencil_58.py => update_mass_flux.py} (93%) rename model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/{mo_solve_nonhydro_stencil_65.py => update_mass_flux_weighted.py} (93%) rename model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/{mo_solve_nonhydro_stencil_68.py => update_theta_v.py} (95%) rename model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/{mo_solve_nonhydro_stencil_62.py => update_wind.py} (93%) rename model/atmosphere/dycore/tests/dycore_stencil_tests/{test_mo_solve_nonhydro_stencil_50.py => test_add_analysis_increments_from_data_assimilation.py} (91%) rename model/atmosphere/dycore/tests/dycore_stencil_tests/{test_mo_solve_nonhydro_4th_order_divdamp.py => test_apply_4th_order_divergence_damping.py} (91%) rename model/atmosphere/dycore/tests/dycore_stencil_tests/{test_mo_solve_nonhydro_stencil_54.py => test_apply_rayleigh_damping_mechanism.py} (91%) rename model/atmosphere/dycore/tests/dycore_stencil_tests/{test_mo_solve_nonhydro_stencil_56_63.py => test_compute_dwdz_for_divergence_damping.py} (91%) rename model/atmosphere/dycore/tests/dycore_stencil_tests/{test_mo_solve_nonhydro_stencil_67.py => test_compute_exner_from_rhotheta.py} (92%) rename model/atmosphere/dycore/tests/dycore_stencil_tests/{test_mo_solve_nonhydro_stencil_55.py => test_compute_results_for_thermodynamic_variables.py} (95%) rename model/atmosphere/dycore/tests/dycore_stencil_tests/{test_mo_solve_nonhydro_stencil_66.py => test_compute_theta_and_exner.py} (93%) rename model/atmosphere/dycore/tests/dycore_stencil_tests/{test_mo_solve_nonhydro_stencil_53.py => test_solve_tridiagonal_matrix_for_w_back_substitution.py} (89%) rename model/atmosphere/dycore/tests/dycore_stencil_tests/{test_mo_solve_nonhydro_stencil_52.py => test_solve_tridiagonal_matrix_for_w_forward_sweep.py} (90%) rename model/atmosphere/dycore/tests/dycore_stencil_tests/{test_mo_solve_nonhydro_stencil_61.py => test_update_densety_exener_wind.py} (94%) rename model/atmosphere/dycore/tests/dycore_stencil_tests/{test_mo_solve_nonhydro_stencil_60.py => test_update_dynamical_exner_time_increment.py} (91%) rename model/atmosphere/dycore/tests/dycore_stencil_tests/{test_mo_solve_nonhydro_stencil_58.py => test_update_mass_flux.py} (93%) rename model/atmosphere/dycore/tests/dycore_stencil_tests/{test_mo_solve_nonhydro_stencil_65.py => test_update_mass_flux_weighted.py} (94%) rename model/atmosphere/dycore/tests/dycore_stencil_tests/{test_mo_solve_nonhydro_stencil_68.py => test_update_theta_v.py} (94%) rename model/atmosphere/dycore/tests/dycore_stencil_tests/{test_mo_solve_nonhydro_stencil_62.py => test_update_wind.py} (91%) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_50.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_analysis_increments_from_data_assimilation.py similarity index 90% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_50.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_analysis_increments_from_data_assimilation.py index d366c532cd..b56d59ff9d 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_50.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_analysis_increments_from_data_assimilation.py @@ -20,13 +20,14 @@ @field_operator -def _mo_solve_nonhydro_stencil_50( +def _add_analysis_increments_from_data_assimilation( z_rho_expl: Field[[CellDim, KDim], wpfloat], z_exner_expl: Field[[CellDim, KDim], wpfloat], rho_incr: Field[[CellDim, KDim], vpfloat], exner_incr: Field[[CellDim, KDim], vpfloat], iau_wgt_dyn: wpfloat, ) -> tuple[Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat]]: + """Formerly known as _mo_solve_nonhydro_stencil_50.""" rho_incr_wp, exner_incr_wp = astype((rho_incr, exner_incr), wpfloat) z_rho_expl_wp = z_rho_expl + iau_wgt_dyn * rho_incr_wp @@ -35,7 +36,7 @@ def _mo_solve_nonhydro_stencil_50( @program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_50( +def add_analysis_increments_from_data_assimilation( z_rho_expl: Field[[CellDim, KDim], wpfloat], z_exner_expl: Field[[CellDim, KDim], wpfloat], rho_incr: Field[[CellDim, KDim], vpfloat], @@ -46,7 +47,7 @@ def mo_solve_nonhydro_stencil_50( vertical_start: int32, vertical_end: int32, ): - _mo_solve_nonhydro_stencil_50( + _add_analysis_increments_from_data_assimilation( z_rho_expl, z_exner_expl, rho_incr, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_4th_order_divdamp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_4th_order_divergence_damping.py similarity index 90% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_4th_order_divdamp.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_4th_order_divergence_damping.py index 95570b6153..5df5ea4b6a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_4th_order_divdamp.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_4th_order_divergence_damping.py @@ -20,11 +20,12 @@ @field_operator -def _mo_solve_nonhydro_4th_order_divdamp( +def _apply_4th_order_divergence_damping( scal_divdamp: Field[[KDim], wpfloat], z_graddiv2_vn: Field[[EdgeDim, KDim], vpfloat], vn: Field[[EdgeDim, KDim], wpfloat], ) -> Field[[EdgeDim, KDim], wpfloat]: + """Formelry known as _mo_solve_nonhydro_4th_order_divdamp.""" z_graddiv2_vn_wp = astype(z_graddiv2_vn, wpfloat) scal_divdamp = broadcast(scal_divdamp, (EdgeDim, KDim)) vn_wp = vn + (scal_divdamp * z_graddiv2_vn_wp) @@ -32,7 +33,7 @@ def _mo_solve_nonhydro_4th_order_divdamp( @program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_4th_order_divdamp( +def apply_4th_order_divergence_damping( scal_divdamp: Field[[KDim], wpfloat], z_graddiv2_vn: Field[[EdgeDim, KDim], vpfloat], vn: Field[[EdgeDim, KDim], wpfloat], @@ -41,7 +42,7 @@ def mo_solve_nonhydro_4th_order_divdamp( vertical_start: int32, vertical_end: int32, ): - _mo_solve_nonhydro_4th_order_divdamp( + _apply_4th_order_divergence_damping( scal_divdamp, z_graddiv2_vn, vn, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_54.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_rayleigh_damping_mechanism.py similarity index 90% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_54.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_rayleigh_damping_mechanism.py index a83159f01f..77b079ea80 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_54.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_rayleigh_damping_mechanism.py @@ -20,18 +20,19 @@ @field_operator -def _mo_solve_nonhydro_stencil_54( +def _apply_rayleigh_damping_mechanism( z_raylfac: Field[[KDim], wpfloat], w_1: Field[[CellDim], wpfloat], w: Field[[CellDim, KDim], wpfloat], ) -> Field[[CellDim, KDim], wpfloat]: + """Formerly known as _mo_solve_nonhydro_stencil_54.""" z_raylfac = broadcast(z_raylfac, (CellDim, KDim)) w_wp = z_raylfac * w + (wpfloat("1.0") - z_raylfac) * w_1 return w_wp @program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_54( +def apply_rayleigh_damping_mechanism( z_raylfac: Field[[KDim], wpfloat], w_1: Field[[CellDim], wpfloat], w: Field[[CellDim, KDim], wpfloat], @@ -40,7 +41,7 @@ def mo_solve_nonhydro_stencil_54( vertical_start: int32, vertical_end: int32, ): - _mo_solve_nonhydro_stencil_54( + _apply_rayleigh_damping_mechanism( z_raylfac, w_1, w, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_dwdz_for_divergence_damping.py similarity index 90% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_dwdz_for_divergence_damping.py index 52bb21f7fd..98ee7c26b9 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_56_63.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_dwdz_for_divergence_damping.py @@ -20,11 +20,12 @@ @field_operator -def _mo_solve_nonhydro_stencil_56_63( +def _compute_dwdz_for_divergence_damping( inv_ddqz_z_full: Field[[CellDim, KDim], vpfloat], w: Field[[CellDim, KDim], wpfloat], w_concorr_c: Field[[CellDim, KDim], vpfloat], ) -> Field[[CellDim, KDim], vpfloat]: + """Formerly known as _mo_solve_nonhydro_stencil_56_63.""" inv_ddqz_z_full_wp = astype(inv_ddqz_z_full, wpfloat) z_dwdz_dd_wp = inv_ddqz_z_full_wp * ( @@ -34,7 +35,7 @@ def _mo_solve_nonhydro_stencil_56_63( @program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_56_63( +def compute_dwdz_for_divergence_damping( inv_ddqz_z_full: Field[[CellDim, KDim], vpfloat], w: Field[[CellDim, KDim], wpfloat], w_concorr_c: Field[[CellDim, KDim], vpfloat], @@ -44,7 +45,7 @@ def mo_solve_nonhydro_stencil_56_63( vertical_start: int32, vertical_end: int32, ): - _mo_solve_nonhydro_stencil_56_63( + _compute_dwdz_for_divergence_damping( inv_ddqz_z_full, w, w_concorr_c, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_exner_from_rhotheta.py similarity index 91% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_exner_from_rhotheta.py index 4604481037..3df296af79 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_67.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_exner_from_rhotheta.py @@ -20,20 +20,21 @@ @field_operator -def _mo_solve_nonhydro_stencil_67( +def _compute_exner_from_rhotheta( rho: Field[[CellDim, KDim], wpfloat], theta_v: Field[[CellDim, KDim], wpfloat], exner: Field[[CellDim, KDim], wpfloat], rd_o_cvd: wpfloat, rd_o_p0ref: wpfloat, ) -> tuple[Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat]]: + """Formerly known as _mo_solve_nonhydro_stencil_67.""" theta_v_wp = exner exner_wp = exp(rd_o_cvd * log(rd_o_p0ref * rho * theta_v_wp)) return theta_v_wp, exner_wp @program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_67( +def compute_exner_from_rhotheta( rho: Field[[CellDim, KDim], wpfloat], theta_v: Field[[CellDim, KDim], wpfloat], exner: Field[[CellDim, KDim], wpfloat], @@ -44,7 +45,7 @@ def mo_solve_nonhydro_stencil_67( vertical_start: int32, vertical_end: int32, ): - _mo_solve_nonhydro_stencil_67( + _compute_exner_from_rhotheta( rho, theta_v, exner, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_results_for_thermodynamic_variables.py similarity index 94% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_results_for_thermodynamic_variables.py index fe8c1753d8..38ac6556d1 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_55.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_results_for_thermodynamic_variables.py @@ -20,7 +20,7 @@ @field_operator -def _mo_solve_nonhydro_stencil_55( +def _compute_results_for_thermodynamic_variables( z_rho_expl: Field[[CellDim, KDim], wpfloat], vwind_impl_wgt: Field[[CellDim], wpfloat], inv_ddqz_z_full: Field[[CellDim, KDim], vpfloat], @@ -40,6 +40,7 @@ def _mo_solve_nonhydro_stencil_55( Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat], ]: + """Formerly known as _mo_solve_nonhydro_stencil_55.""" inv_ddqz_z_full_wp, exner_ref_mc_wp, z_alpha_wp, z_beta_wp = astype( (inv_ddqz_z_full, exner_ref_mc, z_alpha, z_beta), wpfloat ) @@ -62,7 +63,7 @@ def _mo_solve_nonhydro_stencil_55( @program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_55( +def compute_results_for_thermodynamic_variables( z_rho_expl: Field[[CellDim, KDim], wpfloat], vwind_impl_wgt: Field[[CellDim], wpfloat], inv_ddqz_z_full: Field[[CellDim, KDim], vpfloat], @@ -85,7 +86,7 @@ def mo_solve_nonhydro_stencil_55( vertical_start: int32, vertical_end: int32, ): - _mo_solve_nonhydro_stencil_55( + _compute_results_for_thermodynamic_variables( z_rho_expl, vwind_impl_wgt, inv_ddqz_z_full, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_66.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_theta_and_exner.py similarity index 92% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_66.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_theta_and_exner.py index 9fcae5defa..b3337bb508 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_66.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_theta_and_exner.py @@ -20,7 +20,7 @@ @field_operator -def _mo_solve_nonhydro_stencil_66( +def _compute_theta_and_exner( bdy_halo_c: Field[[CellDim], bool], rho: Field[[CellDim, KDim], wpfloat], theta_v: Field[[CellDim, KDim], wpfloat], @@ -28,13 +28,14 @@ def _mo_solve_nonhydro_stencil_66( rd_o_cvd: wpfloat, rd_o_p0ref: wpfloat, ) -> tuple[Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat]]: + """Formelry known as _mo_solve_nonhydro_stencil_66.""" theta_v_wp = where(bdy_halo_c, exner, theta_v) exner_wp = where(bdy_halo_c, exp(rd_o_cvd * log(rd_o_p0ref * rho * exner)), exner) return theta_v_wp, exner_wp @program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_66( +def compute_theta_and_exner( bdy_halo_c: Field[[CellDim], bool], rho: Field[[CellDim, KDim], wpfloat], theta_v: Field[[CellDim, KDim], wpfloat], @@ -46,7 +47,7 @@ def mo_solve_nonhydro_stencil_66( vertical_start: int32, vertical_end: int32, ): - _mo_solve_nonhydro_stencil_66( + _compute_theta_and_exner( bdy_halo_c, rho, theta_v, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py index 5d700c2b89..0307d17065 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py @@ -22,6 +22,9 @@ import icon4py.model.atmosphere.dycore.nh_solve.solve_nonhydro_program as nhsolve_prog import icon4py.model.common.constants as constants from icon4py.model.atmosphere.dycore.accumulate_prep_adv_fields import accumulate_prep_adv_fields +from icon4py.model.atmosphere.dycore.add_analysis_increments_from_data_assimilation import ( + add_analysis_increments_from_data_assimilation, +) from icon4py.model.atmosphere.dycore.add_analysis_increments_to_vn import ( add_analysis_increments_to_vn, ) @@ -37,9 +40,15 @@ from icon4py.model.atmosphere.dycore.apply_2nd_order_divergence_damping import ( apply_2nd_order_divergence_damping, ) +from icon4py.model.atmosphere.dycore.apply_4th_order_divergence_damping import ( + apply_4th_order_divergence_damping, +) from icon4py.model.atmosphere.dycore.apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure import ( apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure, ) +from icon4py.model.atmosphere.dycore.apply_rayleigh_damping_mechanism import ( + apply_rayleigh_damping_mechanism, +) from icon4py.model.atmosphere.dycore.apply_weighted_2nd_and_4th_order_divergence_damping import ( apply_weighted_2nd_and_4th_order_divergence_damping, ) @@ -53,6 +62,10 @@ from icon4py.model.atmosphere.dycore.compute_divergence_of_fluxes_of_rho_and_theta import ( compute_divergence_of_fluxes_of_rho_and_theta, ) +from icon4py.model.atmosphere.dycore.compute_dwdz_for_divergence_damping import ( + compute_dwdz_for_divergence_damping, +) +from icon4py.model.atmosphere.dycore.compute_exner_from_rhotheta import compute_exner_from_rhotheta from icon4py.model.atmosphere.dycore.compute_graddiv2_of_vn import compute_graddiv2_of_vn from icon4py.model.atmosphere.dycore.compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates import ( compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates, @@ -70,9 +83,13 @@ from icon4py.model.atmosphere.dycore.compute_pertubation_of_rho_and_theta import ( compute_pertubation_of_rho_and_theta, ) +from icon4py.model.atmosphere.dycore.compute_results_for_thermodynamic_variables import ( + compute_results_for_thermodynamic_variables, +) from icon4py.model.atmosphere.dycore.compute_rho_virtual_potential_temperatures_and_pressure_gradient import ( compute_rho_virtual_potential_temperatures_and_pressure_gradient, ) +from icon4py.model.atmosphere.dycore.compute_theta_and_exner import compute_theta_and_exner from icon4py.model.atmosphere.dycore.compute_vn_on_lateral_boundary import ( compute_vn_on_lateral_boundary, ) @@ -83,45 +100,6 @@ from icon4py.model.atmosphere.dycore.mo_math_gradients_grad_green_gauss_cell_dsl import ( mo_math_gradients_grad_green_gauss_cell_dsl, ) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_4th_order_divdamp import ( - mo_solve_nonhydro_4th_order_divdamp, -) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_50 import ( - mo_solve_nonhydro_stencil_50, -) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_52 import ( - mo_solve_nonhydro_stencil_52, -) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_53 import ( - mo_solve_nonhydro_stencil_53, -) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_54 import ( - mo_solve_nonhydro_stencil_54, -) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_55 import ( - mo_solve_nonhydro_stencil_55, -) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_56_63 import ( - mo_solve_nonhydro_stencil_56_63, -) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_58 import ( - mo_solve_nonhydro_stencil_58, -) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_60 import ( - mo_solve_nonhydro_stencil_60, -) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_65 import ( - mo_solve_nonhydro_stencil_65, -) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_66 import ( - mo_solve_nonhydro_stencil_66, -) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_67 import ( - mo_solve_nonhydro_stencil_67, -) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_68 import ( - mo_solve_nonhydro_stencil_68, -) from icon4py.model.atmosphere.dycore.set_two_cell_kdim_fields_to_zero_vp import ( set_two_cell_kdim_fields_to_zero_vp, ) @@ -131,6 +109,12 @@ from icon4py.model.atmosphere.dycore.set_two_edge_kdim_fields_to_zero_wp import ( set_two_edge_kdim_fields_to_zero_wp, ) +from icon4py.model.atmosphere.dycore.solve_tridiagonal_matrix_for_w_back_substitution import ( + solve_tridiagonal_matrix_for_w_back_substitution, +) +from icon4py.model.atmosphere.dycore.solve_tridiagonal_matrix_for_w_forward_sweep import ( + solve_tridiagonal_matrix_for_w_forward_sweep, +) from icon4py.model.atmosphere.dycore.state_utils.states import ( DiagnosticStateNonHydro, InterpolationState, @@ -145,6 +129,12 @@ set_zero_c_k, set_zero_e_k, ) +from icon4py.model.atmosphere.dycore.update_dynamical_exner_time_increment import ( + update_dynamical_exner_time_increment, +) +from icon4py.model.atmosphere.dycore.update_mass_flux import update_mass_flux +from icon4py.model.atmosphere.dycore.update_mass_flux_weighted import update_mass_flux_weighted +from icon4py.model.atmosphere.dycore.update_theta_v import update_theta_v from icon4py.model.atmosphere.dycore.velocity.velocity_advection import VelocityAdvection from icon4py.model.common.decomposition.definitions import ExchangeRuntime, SingleNodeExchange from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim @@ -572,7 +562,7 @@ def time_step( start_cell_halo = self.grid.get_start_index(CellDim, HorizontalMarkerIndex.halo(CellDim)) end_cell_end = self.grid.get_end_index(CellDim, HorizontalMarkerIndex.end(CellDim)) if self.grid.limited_area: - mo_solve_nonhydro_stencil_66.with_backend(backend)( + compute_theta_and_exner.with_backend(backend)( bdy_halo_c=self.metric_state_nonhydro.bdy_halo_c, rho=prognostic_state_ls[nnew].rho, theta_v=prognostic_state_ls[nnew].theta_v, @@ -586,7 +576,7 @@ def time_step( offset_provider={}, ) - mo_solve_nonhydro_stencil_67.with_backend(backend)( + compute_exner_from_rhotheta.with_backend(backend)( rho=prognostic_state_ls[nnew].rho, theta_v=prognostic_state_ls[nnew].theta_v, exner=prognostic_state_ls[nnew].exner, @@ -599,7 +589,7 @@ def time_step( offset_provider={}, ) - mo_solve_nonhydro_stencil_68.with_backend(backend)( + update_theta_v.with_backend(backend)( mask_prog_halo_c=self.metric_state_nonhydro.mask_prog_halo_c, rho_now=prognostic_state_ls[nnow].rho, theta_v_now=prognostic_state_ls[nnow].theta_v, @@ -1275,7 +1265,7 @@ def run_predictor_step( ) if self.config.is_iau_active: - mo_solve_nonhydro_stencil_50.with_backend(backend)( + add_analysis_increments_from_data_assimilation.with_backend(backend)( z_fields.z_rho_expl, z_fields.z_exner_expl, diagnostic_state_nh.rho_incr, @@ -1288,7 +1278,7 @@ def run_predictor_step( offset_provider={}, ) - mo_solve_nonhydro_stencil_52.with_backend(backend)( + solve_tridiagonal_matrix_for_w_forward_sweep.with_backend(backend)( vwind_impl_wgt=self.metric_state_nonhydro.vwind_impl_wgt, theta_v_ic=diagnostic_state_nh.theta_v_ic, ddqz_z_half=self.metric_state_nonhydro.ddqz_z_half, @@ -1307,7 +1297,7 @@ def run_predictor_step( offset_provider={"Koff": KDim}, ) - mo_solve_nonhydro_stencil_53.with_backend(backend)( + solve_tridiagonal_matrix_for_w_back_substitution.with_backend(backend)( z_q=z_fields.z_q, w=prognostic_state[nnew].w, horizontal_start=start_cell_nudging, @@ -1318,7 +1308,7 @@ def run_predictor_step( ) if self.config.rayleigh_type == constants.RAYLEIGH_KLEMP: - mo_solve_nonhydro_stencil_54.with_backend(backend)( + apply_rayleigh_damping_mechanism.with_backend(backend)( z_raylfac=self.z_raylfac, w_1=prognostic_state[nnew].w_1, w=prognostic_state[nnew].w, @@ -1331,7 +1321,7 @@ def run_predictor_step( offset_provider={}, ) - mo_solve_nonhydro_stencil_55.with_backend(backend)( + compute_results_for_thermodynamic_variables.with_backend(backend)( z_rho_expl=z_fields.z_rho_expl, vwind_impl_wgt=self.metric_state_nonhydro.vwind_impl_wgt, inv_ddqz_z_full=self.metric_state_nonhydro.inv_ddqz_z_full, @@ -1358,7 +1348,7 @@ def run_predictor_step( # compute dw/dz for divergence damping term if self.config.lhdiff_rcf and self.config.divdamp_type >= 3: - mo_solve_nonhydro_stencil_56_63.with_backend(backend)( + compute_dwdz_for_divergence_damping.with_backend(backend)( inv_ddqz_z_full=self.metric_state_nonhydro.inv_ddqz_z_full, w=prognostic_state[nnew].w, w_concorr_c=diagnostic_state_nh.w_concorr_c, @@ -1403,7 +1393,7 @@ def run_predictor_step( ) if self.config.lhdiff_rcf and self.config.divdamp_type >= 3: - mo_solve_nonhydro_stencil_56_63.with_backend(backend)( + compute_dwdz_for_divergence_damping.with_backend(backend)( inv_ddqz_z_full=self.metric_state_nonhydro.inv_ddqz_z_full, w=prognostic_state[nnew].w, w_concorr_c=diagnostic_state_nh.w_concorr_c, @@ -1633,7 +1623,7 @@ def run_corrector_step( ) else: log.debug("corrector start stencil 4th order divdamp") - mo_solve_nonhydro_4th_order_divdamp.with_backend(backend)( + apply_4th_order_divergence_damping.with_backend(backend)( scal_divdamp=self.scal_divdamp, z_graddiv2_vn=self.z_graddiv2_vn, vn=prognostic_state[nnew].vn, @@ -1843,7 +1833,7 @@ def run_corrector_step( # TODO: this is not tested in green line so far if self.config.is_iau_active: log.debug(f"corrector start stencil 50") - mo_solve_nonhydro_stencil_50( + add_analysis_increments_from_data_assimilation( z_rho_expl=z_fields.z_rho_expl, z_exner_expl=z_fields.z_exner_expl, rho_incr=diagnostic_state_nh.rho_incr, @@ -1856,7 +1846,7 @@ def run_corrector_step( offset_provider={}, ) log.debug(f"corrector start stencil 52") - mo_solve_nonhydro_stencil_52.with_backend(backend)( + solve_tridiagonal_matrix_for_w_forward_sweep.with_backend(backend)( vwind_impl_wgt=self.metric_state_nonhydro.vwind_impl_wgt, theta_v_ic=diagnostic_state_nh.theta_v_ic, ddqz_z_half=self.metric_state_nonhydro.ddqz_z_half, @@ -1875,7 +1865,7 @@ def run_corrector_step( offset_provider={"Koff": KDim}, ) log.debug(f"corrector start stencil 53") - mo_solve_nonhydro_stencil_53.with_backend(backend)( + solve_tridiagonal_matrix_for_w_back_substitution.with_backend(backend)( z_q=z_fields.z_q, w=prognostic_state[nnew].w, horizontal_start=start_cell_nudging, @@ -1887,7 +1877,7 @@ def run_corrector_step( if self.config.rayleigh_type == constants.RAYLEIGH_KLEMP: log.debug(f"corrector start stencil 54") - mo_solve_nonhydro_stencil_54.with_backend(backend)( + apply_rayleigh_damping_mechanism.with_backend(backend)( z_raylfac=self.z_raylfac, w_1=prognostic_state[nnew].w_1, w=prognostic_state[nnew].w, @@ -1900,7 +1890,7 @@ def run_corrector_step( offset_provider={}, ) log.debug(f"corrector start stencil 55") - mo_solve_nonhydro_stencil_55.with_backend(backend)( + compute_results_for_thermodynamic_variables.with_backend(backend)( z_rho_expl=z_fields.z_rho_expl, vwind_impl_wgt=self.metric_state_nonhydro.vwind_impl_wgt, inv_ddqz_z_full=self.metric_state_nonhydro.inv_ddqz_z_full, @@ -1937,7 +1927,7 @@ def run_corrector_step( offset_provider={}, ) log.debug(f"corrector start stencil 58") - mo_solve_nonhydro_stencil_58.with_backend(backend)( + update_mass_flux.with_backend(backend)( z_contr_w_fl_l=z_fields.z_contr_w_fl_l, rho_ic=diagnostic_state_nh.rho_ic, vwind_impl_wgt=self.metric_state_nonhydro.vwind_impl_wgt, @@ -1951,7 +1941,7 @@ def run_corrector_step( offset_provider={}, ) if at_last_substep: - mo_solve_nonhydro_stencil_60.with_backend(backend)( + update_dynamical_exner_time_increment.with_backend(backend)( exner=prognostic_state[nnew].exner, ddt_exner_phy=diagnostic_state_nh.ddt_exner_phy, exner_dyn_incr=diagnostic_state_nh.exner_dyn_incr, @@ -1976,7 +1966,7 @@ def run_corrector_step( offset_provider={}, ) log.debug(f" corrector: start stencil 65") - mo_solve_nonhydro_stencil_65.with_backend(backend)( + update_mass_flux_weighted.with_backend(backend)( rho_ic=diagnostic_state_nh.rho_ic, vwind_expl_wgt=self.metric_state_nonhydro.vwind_expl_wgt, vwind_impl_wgt=self.metric_state_nonhydro.vwind_impl_wgt, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro_program.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro_program.py index 7006183854..db910df114 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro_program.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro_program.py @@ -64,12 +64,6 @@ from icon4py.model.atmosphere.dycore.interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges import ( _interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges, ) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_61 import ( - _mo_solve_nonhydro_stencil_61, -) -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_62 import ( - _mo_solve_nonhydro_stencil_62, -) from icon4py.model.atmosphere.dycore.set_cell_kdim_field_to_zero_vp import ( _set_cell_kdim_field_to_zero_vp, ) @@ -80,6 +74,8 @@ _set_theta_v_prime_ic_at_lower_boundary, ) from icon4py.model.atmosphere.dycore.state_utils.utils import _set_zero_c_k, _set_zero_e_k +from icon4py.model.atmosphere.dycore.update_densety_exener_wind import _update_densety_exener_wind +from icon4py.model.atmosphere.dycore.update_wind import _update_wind from icon4py.model.common.dimension import CEDim, CellDim, ECDim, EdgeDim, KDim @@ -1002,14 +998,14 @@ def _stencils_61_62( ]: (rho_new, exner_new, w_new) = where( (k_field >= int32(0)) & (k_field < nlev), - _mo_solve_nonhydro_stencil_61( + _update_densety_exener_wind( rho_now, grf_tend_rho, theta_v_now, grf_tend_thv, w_now, grf_tend_w, dtime ), (rho_new, exner_new, w_new), ) w_new = where( k_field == nlev, - _mo_solve_nonhydro_stencil_62(w_now, grf_tend_w, dtime), + _update_wind(w_now, grf_tend_w, dtime), w_new, ) return rho_new, exner_new, w_new diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_53.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_tridiagonal_matrix_for_w_back_substitution.py similarity index 81% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_53.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_tridiagonal_matrix_for_w_back_substitution.py index 21598f9f38..919b612434 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_53.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_tridiagonal_matrix_for_w_back_substitution.py @@ -20,12 +20,15 @@ @scan_operator(axis=KDim, forward=False, init=wpfloat("0.0")) -def _mo_solve_nonhydro_stencil_53_scan(w_state: wpfloat, z_q: vpfloat, w: wpfloat) -> wpfloat: +def _solve_tridiagonal_matrix_for_w_back_substitution_scan( + w_state: wpfloat, z_q: vpfloat, w: wpfloat +) -> wpfloat: + """Formerly known as _mo_solve_nonhydro_stencil_53_scan.""" return w + w_state * astype(z_q, wpfloat) @program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_53( +def solve_tridiagonal_matrix_for_w_back_substitution( z_q: Field[[CellDim, KDim], vpfloat], w: Field[[CellDim, KDim], wpfloat], horizontal_start: int32, @@ -33,7 +36,7 @@ def mo_solve_nonhydro_stencil_53( vertical_start: int32, vertical_end: int32, ): - _mo_solve_nonhydro_stencil_53_scan( + _solve_tridiagonal_matrix_for_w_back_substitution_scan( z_q, w, out=w, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_52.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_tridiagonal_matrix_for_w_forward_sweep.py similarity index 93% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_52.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_tridiagonal_matrix_for_w_forward_sweep.py index 31aa2cf0dd..b5390d64d3 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_52.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_tridiagonal_matrix_for_w_forward_sweep.py @@ -39,7 +39,7 @@ def _w( @field_operator -def _mo_solve_nonhydro_stencil_52( +def _solve_tridiagonal_matrix_for_w_forward_sweep( vwind_impl_wgt: Field[[CellDim], wpfloat], theta_v_ic: Field[[CellDim, KDim], wpfloat], ddqz_z_half: Field[[CellDim, KDim], vpfloat], @@ -52,9 +52,8 @@ def _mo_solve_nonhydro_stencil_52( dtime: wpfloat, cpd: wpfloat, ) -> tuple[Field[[CellDim, KDim], vpfloat], Field[[CellDim, KDim], wpfloat]]: - ddqz_z_half_wp, z_alpha_wp, z_beta_wp, z_q_wp = astype( - (ddqz_z_half, z_alpha, z_beta, z_q), wpfloat - ) + """Formerly known as _mo_solve_nonhydro_stencil_52.""" + ddqz_z_half_wp = astype(ddqz_z_half, wpfloat) z_gamma_vp = astype(dtime * cpd * vwind_impl_wgt * theta_v_ic / ddqz_z_half_wp, vpfloat) z_a = (vpfloat("0.0") - z_gamma_vp) * z_beta(Koff[-1]) * z_alpha(Koff[-1]) @@ -69,7 +68,7 @@ def _mo_solve_nonhydro_stencil_52( @program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_52( +def solve_tridiagonal_matrix_for_w_forward_sweep( vwind_impl_wgt: Field[[CellDim], wpfloat], theta_v_ic: Field[[CellDim, KDim], wpfloat], ddqz_z_half: Field[[CellDim, KDim], vpfloat], @@ -86,7 +85,7 @@ def mo_solve_nonhydro_stencil_52( vertical_start: int32, vertical_end: int32, ): - _mo_solve_nonhydro_stencil_52( + _solve_tridiagonal_matrix_for_w_forward_sweep( vwind_impl_wgt, theta_v_ic, ddqz_z_half, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_61.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_densety_exener_wind.py similarity index 88% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_61.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_densety_exener_wind.py index ef87f67c43..1557ce28b2 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_61.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_densety_exener_wind.py @@ -15,15 +15,13 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32 -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_62 import ( - _mo_solve_nonhydro_stencil_62, -) +from icon4py.model.atmosphere.dycore.update_wind import _update_wind from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.type_alias import wpfloat @field_operator -def _mo_solve_nonhydro_stencil_61( +def _update_densety_exener_wind( rho_now: Field[[CellDim, KDim], wpfloat], grf_tend_rho: Field[[CellDim, KDim], wpfloat], theta_v_now: Field[[CellDim, KDim], wpfloat], @@ -36,14 +34,15 @@ def _mo_solve_nonhydro_stencil_61( Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat], ]: + """Formerly known as _mo_solve_nonhydro_stencil_61.""" rho_new_wp = rho_now + dtime * grf_tend_rho exner_new_wp = theta_v_now + dtime * grf_tend_thv - w_new_wp = _mo_solve_nonhydro_stencil_62(w_now=w_now, grf_tend_w=grf_tend_w, dtime=dtime) + w_new_wp = _update_wind(w_now=w_now, grf_tend_w=grf_tend_w, dtime=dtime) return rho_new_wp, exner_new_wp, w_new_wp @program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_61( +def update_densety_exener_wind( rho_now: Field[[CellDim, KDim], wpfloat], grf_tend_rho: Field[[CellDim, KDim], wpfloat], theta_v_now: Field[[CellDim, KDim], wpfloat], @@ -59,7 +58,7 @@ def mo_solve_nonhydro_stencil_61( vertical_start: int32, vertical_end: int32, ): - _mo_solve_nonhydro_stencil_61( + _update_densety_exener_wind( rho_now, grf_tend_rho, theta_v_now, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_dynamical_exner_time_increment.py similarity index 91% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_dynamical_exner_time_increment.py index 8fb0a70674..7b6f9f07eb 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_60.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_dynamical_exner_time_increment.py @@ -20,13 +20,14 @@ @field_operator -def _mo_solve_nonhydro_stencil_60( +def _update_dynamical_exner_time_increment( exner: Field[[CellDim, KDim], wpfloat], ddt_exner_phy: Field[[CellDim, KDim], vpfloat], exner_dyn_incr: Field[[CellDim, KDim], vpfloat], ndyn_substeps_var: wpfloat, dtime: wpfloat, ) -> Field[[CellDim, KDim], vpfloat]: + """Formerly known as _mo_solve_nonhydro_stencil_60.""" exner_dyn_incr_wp, ddt_exner_phy_wp = astype((exner_dyn_incr, ddt_exner_phy), wpfloat) exner_dyn_incr_wp = exner - (exner_dyn_incr_wp + ndyn_substeps_var * dtime * ddt_exner_phy_wp) @@ -34,7 +35,7 @@ def _mo_solve_nonhydro_stencil_60( @program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_60( +def update_dynamical_exner_time_increment( exner: Field[[CellDim, KDim], wpfloat], ddt_exner_phy: Field[[CellDim, KDim], vpfloat], exner_dyn_incr: Field[[CellDim, KDim], vpfloat], @@ -45,7 +46,7 @@ def mo_solve_nonhydro_stencil_60( vertical_start: int32, vertical_end: int32, ): - _mo_solve_nonhydro_stencil_60( + _update_dynamical_exner_time_increment( exner, ddt_exner_phy, exner_dyn_incr, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_flux.py similarity index 93% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_flux.py index 6850f9b68e..9c22b97e86 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_58.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_flux.py @@ -20,7 +20,7 @@ @field_operator -def _mo_solve_nonhydro_stencil_58( +def _update_mass_flux( z_contr_w_fl_l: Field[[CellDim, KDim], wpfloat], rho_ic: Field[[CellDim, KDim], wpfloat], vwind_impl_wgt: Field[[CellDim], wpfloat], @@ -28,12 +28,13 @@ def _mo_solve_nonhydro_stencil_58( mass_flx_ic: Field[[CellDim, KDim], wpfloat], r_nsubsteps: wpfloat, ) -> Field[[CellDim, KDim], wpfloat]: + """Formerly known as _mo_solve_nonhydro_stencil_58.""" mass_flx_ic_wp = mass_flx_ic + (r_nsubsteps * (z_contr_w_fl_l + rho_ic * vwind_impl_wgt * w)) return mass_flx_ic_wp @program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_58( +def update_mass_flux( z_contr_w_fl_l: Field[[CellDim, KDim], wpfloat], rho_ic: Field[[CellDim, KDim], wpfloat], vwind_impl_wgt: Field[[CellDim], wpfloat], @@ -45,7 +46,7 @@ def mo_solve_nonhydro_stencil_58( vertical_start: int32, vertical_end: int32, ): - _mo_solve_nonhydro_stencil_58( + _update_mass_flux( z_contr_w_fl_l, rho_ic, vwind_impl_wgt, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_flux_weighted.py similarity index 93% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_flux_weighted.py index 7824abd725..f9786dd7ef 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_65.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_flux_weighted.py @@ -20,7 +20,7 @@ @field_operator -def _mo_solve_nonhydro_stencil_65( +def _update_mass_flux_weighted( rho_ic: Field[[CellDim, KDim], wpfloat], vwind_expl_wgt: Field[[CellDim], wpfloat], vwind_impl_wgt: Field[[CellDim], wpfloat], @@ -30,6 +30,7 @@ def _mo_solve_nonhydro_stencil_65( mass_flx_ic: Field[[CellDim, KDim], wpfloat], r_nsubsteps: wpfloat, ) -> Field[[CellDim, KDim], wpfloat]: + """Formerly known as _mo_solve_nonhydro_stencil_65.""" w_concorr_c_wp = astype(w_concorr_c, wpfloat) mass_flx_ic_wp = mass_flx_ic + ( @@ -39,7 +40,7 @@ def _mo_solve_nonhydro_stencil_65( @program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_65( +def update_mass_flux_weighted( rho_ic: Field[[CellDim, KDim], wpfloat], vwind_expl_wgt: Field[[CellDim], wpfloat], vwind_impl_wgt: Field[[CellDim], wpfloat], @@ -53,7 +54,7 @@ def mo_solve_nonhydro_stencil_65( vertical_start: int32, vertical_end: int32, ): - _mo_solve_nonhydro_stencil_65( + _update_mass_flux_weighted( rho_ic, vwind_expl_wgt, vwind_impl_wgt, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_v.py similarity index 95% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_v.py index e011e10c78..5a1c203a03 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_68.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_v.py @@ -20,7 +20,7 @@ @field_operator -def _mo_solve_nonhydro_stencil_68( +def _update_theta_v( mask_prog_halo_c: Field[[CellDim], bool], rho_now: Field[[CellDim, KDim], wpfloat], theta_v_now: Field[[CellDim, KDim], wpfloat], @@ -30,6 +30,7 @@ def _mo_solve_nonhydro_stencil_68( theta_v_new: Field[[CellDim, KDim], wpfloat], cvd_o_rd: wpfloat, ) -> Field[[CellDim, KDim], wpfloat]: + """Formerly known as _mo_solve_nonhydro_stencil_68.""" theta_v_new_wp = where( mask_prog_halo_c, rho_now @@ -42,7 +43,7 @@ def _mo_solve_nonhydro_stencil_68( @program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_68( +def update_theta_v( mask_prog_halo_c: Field[[CellDim], bool], rho_now: Field[[CellDim, KDim], wpfloat], theta_v_now: Field[[CellDim, KDim], wpfloat], @@ -56,7 +57,7 @@ def mo_solve_nonhydro_stencil_68( vertical_start: int32, vertical_end: int32, ): - _mo_solve_nonhydro_stencil_68( + _update_theta_v( mask_prog_halo_c, rho_now, theta_v_now, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_62.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_wind.py similarity index 93% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_62.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_wind.py index 7d8caf8544..d42dae1e28 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_62.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_wind.py @@ -20,17 +20,18 @@ @field_operator -def _mo_solve_nonhydro_stencil_62( +def _update_wind( w_now: Field[[CellDim, KDim], wpfloat], grf_tend_w: Field[[CellDim, KDim], wpfloat], dtime: wpfloat, ) -> Field[[CellDim, KDim], wpfloat]: + """Formerly known as _mo_solve_nonhydro_stencil_62.""" w_new_wp = w_now + dtime * grf_tend_w return w_new_wp @program(grid_type=GridType.UNSTRUCTURED) -def mo_solve_nonhydro_stencil_62( +def update_wind( w_now: Field[[CellDim, KDim], wpfloat], grf_tend_w: Field[[CellDim, KDim], wpfloat], w_new: Field[[CellDim, KDim], wpfloat], @@ -40,7 +41,7 @@ def mo_solve_nonhydro_stencil_62( vertical_start: int32, vertical_end: int32, ): - _mo_solve_nonhydro_stencil_62( + _update_wind( w_now, grf_tend_w, dtime, diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_50.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_analysis_increments_from_data_assimilation.py similarity index 91% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_50.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_analysis_increments_from_data_assimilation.py index 3d12c8b990..0fcf919916 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_50.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_add_analysis_increments_from_data_assimilation.py @@ -15,8 +15,8 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_50 import ( - mo_solve_nonhydro_stencil_50, +from icon4py.model.atmosphere.dycore.add_analysis_increments_from_data_assimilation import ( + add_analysis_increments_from_data_assimilation, ) from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field @@ -24,7 +24,7 @@ class TestMoSolveNonhydroStencil50(StencilTest): - PROGRAM = mo_solve_nonhydro_stencil_50 + PROGRAM = add_analysis_increments_from_data_assimilation OUTPUTS = ("z_rho_expl", "z_exner_expl") @staticmethod diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_4th_order_divdamp.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_4th_order_divergence_damping.py similarity index 91% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_4th_order_divdamp.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_4th_order_divergence_damping.py index c0fccfe30a..8a173ff7d0 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_4th_order_divdamp.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_4th_order_divergence_damping.py @@ -15,8 +15,8 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_4th_order_divdamp import ( - mo_solve_nonhydro_4th_order_divdamp, +from icon4py.model.atmosphere.dycore.apply_4th_order_divergence_damping import ( + apply_4th_order_divergence_damping, ) from icon4py.model.common.dimension import EdgeDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field @@ -24,7 +24,7 @@ class TestMoSolveNonhydro4thOrderDivdamp(StencilTest): - PROGRAM = mo_solve_nonhydro_4th_order_divdamp + PROGRAM = apply_4th_order_divergence_damping OUTPUTS = ("vn",) @staticmethod diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_54.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_rayleigh_damping_mechanism.py similarity index 91% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_54.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_rayleigh_damping_mechanism.py index a85c6ebec6..c6fd9f98d2 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_54.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_apply_rayleigh_damping_mechanism.py @@ -15,8 +15,8 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_54 import ( - mo_solve_nonhydro_stencil_54, +from icon4py.model.atmosphere.dycore.apply_rayleigh_damping_mechanism import ( + apply_rayleigh_damping_mechanism, ) from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field @@ -24,7 +24,7 @@ class TestMoSolveNonhydroStencil54(StencilTest): - PROGRAM = mo_solve_nonhydro_stencil_54 + PROGRAM = apply_rayleigh_damping_mechanism OUTPUTS = ("w",) @staticmethod diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_56_63.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_dwdz_for_divergence_damping.py similarity index 91% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_56_63.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_dwdz_for_divergence_damping.py index 4ea424616f..7c7e70a445 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_56_63.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_dwdz_for_divergence_damping.py @@ -15,8 +15,8 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_56_63 import ( - mo_solve_nonhydro_stencil_56_63, +from icon4py.model.atmosphere.dycore.compute_dwdz_for_divergence_damping import ( + compute_dwdz_for_divergence_damping, ) from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field @@ -24,7 +24,7 @@ class TestMoSolveNonhydroStencil5663(StencilTest): - PROGRAM = mo_solve_nonhydro_stencil_56_63 + PROGRAM = compute_dwdz_for_divergence_damping OUTPUTS = ("z_dwdz_dd",) @staticmethod diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_67.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_exner_from_rhotheta.py similarity index 92% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_67.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_exner_from_rhotheta.py index 76ffb6c366..9eebc5d212 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_67.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_exner_from_rhotheta.py @@ -15,16 +15,14 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_67 import ( - mo_solve_nonhydro_stencil_67, -) +from icon4py.model.atmosphere.dycore.compute_exner_from_rhotheta import compute_exner_from_rhotheta from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field from icon4py.model.common.type_alias import wpfloat class TestMoSolveNonhydroStencil67(StencilTest): - PROGRAM = mo_solve_nonhydro_stencil_67 + PROGRAM = compute_exner_from_rhotheta OUTPUTS = ("theta_v", "exner") @staticmethod diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_55.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_results_for_thermodynamic_variables.py similarity index 95% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_55.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_results_for_thermodynamic_variables.py index 3644338599..df0d8714b0 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_55.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_results_for_thermodynamic_variables.py @@ -15,8 +15,8 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_55 import ( - mo_solve_nonhydro_stencil_55, +from icon4py.model.atmosphere.dycore.compute_results_for_thermodynamic_variables import ( + compute_results_for_thermodynamic_variables, ) from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field @@ -24,7 +24,7 @@ class TestMoSolveNonhydroStencil55(StencilTest): - PROGRAM = mo_solve_nonhydro_stencil_55 + PROGRAM = compute_results_for_thermodynamic_variables OUTPUTS = ("rho_new", "exner_new", "theta_v_new") @staticmethod diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_66.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_theta_and_exner.py similarity index 93% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_66.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_theta_and_exner.py index fb21cecce4..4edb90aa0c 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_66.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_theta_and_exner.py @@ -15,16 +15,14 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_66 import ( - mo_solve_nonhydro_stencil_66, -) +from icon4py.model.atmosphere.dycore.compute_theta_and_exner import compute_theta_and_exner from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field, random_mask from icon4py.model.common.type_alias import wpfloat class TestMoSolveNonhydroStencil66(StencilTest): - PROGRAM = mo_solve_nonhydro_stencil_66 + PROGRAM = compute_theta_and_exner OUTPUTS = ("theta_v", "exner") @staticmethod diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_53.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_solve_tridiagonal_matrix_for_w_back_substitution.py similarity index 89% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_53.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_solve_tridiagonal_matrix_for_w_back_substitution.py index a68750d93b..92d6250bf5 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_53.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_solve_tridiagonal_matrix_for_w_back_substitution.py @@ -15,8 +15,8 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_53 import ( - mo_solve_nonhydro_stencil_53, +from icon4py.model.atmosphere.dycore.solve_tridiagonal_matrix_for_w_back_substitution import ( + solve_tridiagonal_matrix_for_w_back_substitution, ) from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field @@ -24,7 +24,7 @@ class TestMoSolveNonhydroStencil53(StencilTest): - PROGRAM = mo_solve_nonhydro_stencil_53 + PROGRAM = solve_tridiagonal_matrix_for_w_back_substitution OUTPUTS = ("w",) @staticmethod diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_52.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_solve_tridiagonal_matrix_for_w_forward_sweep.py similarity index 90% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_52.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_solve_tridiagonal_matrix_for_w_forward_sweep.py index 702efb1521..ebf5df38c0 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_52.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_solve_tridiagonal_matrix_for_w_forward_sweep.py @@ -15,8 +15,8 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.program_processors.runners.gtfn import run_gtfn -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_52 import ( - mo_solve_nonhydro_stencil_52, +from icon4py.model.atmosphere.dycore.solve_tridiagonal_matrix_for_w_forward_sweep import ( + solve_tridiagonal_matrix_for_w_forward_sweep, ) from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.grid.simple import SimpleGrid @@ -24,7 +24,7 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat -def mo_solve_nonhydro_stencil_52_numpy( +def solve_tridiagonal_matrix_for_w_forward_sweep_numpy( vwind_impl_wgt: np.array, theta_v_ic: np.array, ddqz_z_half: np.array, @@ -60,7 +60,7 @@ def mo_solve_nonhydro_stencil_52_numpy( return z_q, w -def test_mo_solve_nonhydro_stencil_52(): +def test_solve_tridiagonal_matrix_for_w_forward_sweep(): grid = SimpleGrid() vwind_impl_wgt = random_field(grid, CellDim, dtype=wpfloat) theta_v_ic = random_field(grid, CellDim, KDim, dtype=wpfloat) @@ -75,7 +75,7 @@ def test_mo_solve_nonhydro_stencil_52(): z_q = random_field(grid, CellDim, KDim, dtype=vpfloat) w = random_field(grid, CellDim, KDim, dtype=wpfloat) - z_q_ref, w_ref = mo_solve_nonhydro_stencil_52_numpy( + z_q_ref, w_ref = solve_tridiagonal_matrix_for_w_forward_sweep_numpy( vwind_impl_wgt.asnumpy(), theta_v_ic.asnumpy(), ddqz_z_half.asnumpy(), @@ -93,7 +93,7 @@ def test_mo_solve_nonhydro_stencil_52(): v_start = int32(1) v_end = int32(grid.num_levels) # TODO we run this test with the C++ backend as the `embedded` backend doesn't handle this pattern - mo_solve_nonhydro_stencil_52.with_backend(run_gtfn)( + solve_tridiagonal_matrix_for_w_forward_sweep.with_backend(run_gtfn)( vwind_impl_wgt=vwind_impl_wgt, theta_v_ic=theta_v_ic, ddqz_z_half=ddqz_z_half, diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_61.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_densety_exener_wind.py similarity index 94% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_61.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_densety_exener_wind.py index 9f98d2bb17..e925d9f7a0 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_61.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_densety_exener_wind.py @@ -15,16 +15,14 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_61 import ( - mo_solve_nonhydro_stencil_61, -) +from icon4py.model.atmosphere.dycore.update_densety_exener_wind import update_densety_exener_wind from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field from icon4py.model.common.type_alias import wpfloat class TestMoSolveNonhydroStencil61(StencilTest): - PROGRAM = mo_solve_nonhydro_stencil_61 + PROGRAM = update_densety_exener_wind OUTPUTS = ("rho_new", "exner_new", "w_new") @staticmethod diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_60.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_dynamical_exner_time_increment.py similarity index 91% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_60.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_dynamical_exner_time_increment.py index b835069c81..e4f0b44fb6 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_60.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_dynamical_exner_time_increment.py @@ -15,8 +15,8 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_60 import ( - mo_solve_nonhydro_stencil_60, +from icon4py.model.atmosphere.dycore.update_dynamical_exner_time_increment import ( + update_dynamical_exner_time_increment, ) from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field @@ -24,7 +24,7 @@ class TestMoSolveNonhydroStencil60(StencilTest): - PROGRAM = mo_solve_nonhydro_stencil_60 + PROGRAM = update_dynamical_exner_time_increment OUTPUTS = ("exner_dyn_incr",) @staticmethod diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_58.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_flux.py similarity index 93% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_58.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_flux.py index 6bee1d2e23..a556776547 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_58.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_flux.py @@ -15,16 +15,14 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_58 import ( - mo_solve_nonhydro_stencil_58, -) +from icon4py.model.atmosphere.dycore.update_mass_flux import update_mass_flux from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field from icon4py.model.common.type_alias import wpfloat class TestMoSolveNonhydroStencil58(StencilTest): - PROGRAM = mo_solve_nonhydro_stencil_58 + PROGRAM = update_mass_flux OUTPUTS = ("mass_flx_ic",) @staticmethod diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_65.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_flux_weighted.py similarity index 94% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_65.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_flux_weighted.py index 9e55f77d2d..4216be6f57 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_65.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_mass_flux_weighted.py @@ -15,16 +15,14 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_65 import ( - mo_solve_nonhydro_stencil_65, -) +from icon4py.model.atmosphere.dycore.update_mass_flux_weighted import update_mass_flux_weighted from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field from icon4py.model.common.type_alias import vpfloat, wpfloat class TestMoSolveNonhydroStencil65(StencilTest): - PROGRAM = mo_solve_nonhydro_stencil_65 + PROGRAM = update_mass_flux_weighted OUTPUTS = ("mass_flx_ic",) @staticmethod diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_68.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_theta_v.py similarity index 94% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_68.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_theta_v.py index 7ceb02d934..88e1d285b0 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_68.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_theta_v.py @@ -15,16 +15,14 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_68 import ( - mo_solve_nonhydro_stencil_68, -) +from icon4py.model.atmosphere.dycore.update_theta_v import update_theta_v from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field, random_mask from icon4py.model.common.type_alias import wpfloat class TestMoSolveNonhydroStencil68(StencilTest): - PROGRAM = mo_solve_nonhydro_stencil_68 + PROGRAM = update_theta_v OUTPUTS = ("theta_v_new",) @staticmethod diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_62.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_wind.py similarity index 91% rename from model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_62.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_wind.py index 9ac70d90ca..9bde15f031 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_62.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_update_wind.py @@ -15,16 +15,14 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_62 import ( - mo_solve_nonhydro_stencil_62, -) +from icon4py.model.atmosphere.dycore.update_wind import update_wind from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field from icon4py.model.common.type_alias import wpfloat class TestMoSolveNonhydroStencil62(StencilTest): - PROGRAM = mo_solve_nonhydro_stencil_62 + PROGRAM = update_wind OUTPUTS = ("w_new",) @staticmethod From a5c8e188f88a4cbbe078f6b526490464c11f5f79 Mon Sep 17 00:00:00 2001 From: Magdalena Date: Tue, 6 Feb 2024 12:50:16 +0100 Subject: [PATCH 2/2] Fix test dependencies on savepoints (#349) (cleanup) Make test_solve_nonhydro.py indenpendent of `VelocityInitSavepoint`. --- .../atmosphere/diffusion/diffusion_utils.py | 4 +- .../diffusion/tests/diffusion_tests/utils.py | 2 +- .../dycore/nh_solve/solve_nonhydro.py | 2 +- .../tests/dycore_tests/test_solve_nonhydro.py | 94 +++++++++---------- .../icon4py/model/common/math/smagorinsky.py | 31 +++++- .../common/test_utils/datatest_fixtures.py | 2 +- .../common/test_utils/serialbox_utils.py | 35 ++++--- .../tests/math_tests/test_smagorinsky.py | 7 +- .../src/icon4py/model/driver/io_utils.py | 4 +- 9 files changed, 108 insertions(+), 73 deletions(-) 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 d5db49b331..62986325e1 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 @@ -19,7 +19,7 @@ from gt4py.next.ffront.fbuiltins import broadcast, int32, minimum from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim -from icon4py.model.common.math.smagorinsky import en_smag_fac_for_zero_nshift +from icon4py.model.common.math.smagorinsky import _en_smag_fac_for_zero_nshift # TODO(Magdalena): fix duplication: duplicated from test testutils/utils.py @@ -115,7 +115,7 @@ def _init_diffusion_local_fields_for_regular_timestemp( ) -> tuple[Field[[KDim], float], Field[[KDim], float], Field[[KDim], float]]: 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( + enh_smag_fac = _en_smag_fac_for_zero_nshift( vect_a, hdiff_smag_fac, hdiff_smag_fac2, diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/utils.py b/model/atmosphere/diffusion/tests/diffusion_tests/utils.py index d1ef9b094f..2edfe1850a 100644 --- a/model/atmosphere/diffusion/tests/diffusion_tests/utils.py +++ b/model/atmosphere/diffusion/tests/diffusion_tests/utils.py @@ -169,7 +169,7 @@ def construct_diagnostics( savepoint: IconDiffusionInitSavepoint, grid_savepoint: IconGridSavepoint, ) -> DiffusionDiagnosticState: - grid = grid_savepoint.construct_icon_grid() + grid = grid_savepoint.construct_icon_grid(on_gpu=False) dwdx = savepoint.dwdx() if savepoint.dwdx() else zero_field(grid, CellDim, KDim) dwdy = savepoint.dwdy() if savepoint.dwdy() else zero_field(grid, CellDim, KDim) return DiffusionDiagnosticState( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py index 0307d17065..17202a8bc0 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py @@ -433,7 +433,7 @@ def init( self.config.divdamp_z2, self.config.divdamp_z3, self.config.divdamp_z4, - out=self.enh_divdamp_fac, + self.enh_divdamp_fac, offset_provider={"Koff": KDim}, ) diff --git a/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py index eee7ac5c91..23c504bc6b 100644 --- a/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py @@ -38,6 +38,7 @@ from icon4py.model.common.states.prognostic_state import PrognosticState from icon4py.model.common.test_utils.datatest_utils import GLOBAL_EXPERIMENT, REGIONAL_EXPERIMENT from icon4py.model.common.test_utils.helpers import dallclose +from icon4py.model.common.test_utils.serialbox_utils import IconNonHydroInitSavepoint from .utils import ( construct_config, @@ -71,7 +72,7 @@ def test_validate_divdamp_fields_against_savepoint_values( config.divdamp_z2, config.divdamp_z3, config.divdamp_z4, - out=enh_divdamp_fac, + enh_divdamp_fac, offset_provider={"Koff": KDim}, ) _calculate_scal_divdamp.with_backend(backend)( @@ -109,7 +110,6 @@ def test_nonhydro_predictor_step( savepoint_nonhydro_init, damping_height, grid_savepoint, - savepoint_velocity_init, metrics_savepoint, interpolation_savepoint, savepoint_nonhydro_exit, @@ -123,15 +123,14 @@ def test_nonhydro_predictor_step( sp_exit = savepoint_nonhydro_exit nonhydro_params = NonHydrostaticParams(config) vertical_params = create_vertical_params(damping_height, grid_savepoint) - sp_v = savepoint_velocity_init - dtime = sp_v.get_metadata("dtime").get("dtime") - recompute = sp_v.get_metadata("recompute").get("recompute") - linit = sp_v.get_metadata("linit").get("linit") + dtime = sp.get_metadata("dtime").get("dtime") + recompute = sp.get_metadata("recompute").get("recompute") + linit = sp.get_metadata("linit").get("linit") nnow = 0 nnew = 1 - diagnostic_state_nh = construct_diagnostics(sp, sp_v) + diagnostic_state_nh = construct_diagnostics(sp) interpolation_state = construct_interpolation_state_for_nonhydro(interpolation_savepoint) metric_state_nonhydro = construct_nh_metric_state(metrics_savepoint, icon_grid.num_levels) @@ -442,29 +441,29 @@ def test_nonhydro_predictor_step( assert dallclose(prognostic_state_nnew.theta_v.asnumpy(), sp_exit.theta_v_new().asnumpy()) -def construct_diagnostics(sp, sp_v): +def construct_diagnostics(init_savepoint: IconNonHydroInitSavepoint): return DiagnosticStateNonHydro( - theta_v_ic=sp.theta_v_ic(), - exner_pr=sp.exner_pr(), - rho_ic=sp.rho_ic(), - ddt_exner_phy=sp.ddt_exner_phy(), - grf_tend_rho=sp.grf_tend_rho(), - grf_tend_thv=sp.grf_tend_thv(), - grf_tend_w=sp.grf_tend_w(), - mass_fl_e=sp.mass_fl_e(), - ddt_vn_phy=sp.ddt_vn_phy(), - grf_tend_vn=sp.grf_tend_vn(), - ddt_vn_apc_ntl1=sp_v.ddt_vn_apc_pc(1), - ddt_vn_apc_ntl2=sp_v.ddt_vn_apc_pc(2), - ddt_w_adv_ntl1=sp_v.ddt_w_adv_pc(1), - ddt_w_adv_ntl2=sp_v.ddt_w_adv_pc(2), - vt=sp_v.vt(), - vn_ie=sp_v.vn_ie(), - w_concorr_c=sp_v.w_concorr_c(), + theta_v_ic=init_savepoint.theta_v_ic(), + exner_pr=init_savepoint.exner_pr(), + rho_ic=init_savepoint.rho_ic(), + ddt_exner_phy=init_savepoint.ddt_exner_phy(), + grf_tend_rho=init_savepoint.grf_tend_rho(), + grf_tend_thv=init_savepoint.grf_tend_thv(), + grf_tend_w=init_savepoint.grf_tend_w(), + mass_fl_e=init_savepoint.mass_fl_e(), + ddt_vn_phy=init_savepoint.ddt_vn_phy(), + grf_tend_vn=init_savepoint.grf_tend_vn(), + ddt_vn_apc_ntl1=init_savepoint.ddt_vn_apc_pc(1), + ddt_vn_apc_ntl2=init_savepoint.ddt_vn_apc_pc(2), + ddt_w_adv_ntl1=init_savepoint.ddt_w_adv_pc(1), + ddt_w_adv_ntl2=init_savepoint.ddt_w_adv_pc(2), + vt=init_savepoint.vt(), + vn_ie=init_savepoint.vn_ie(), + w_concorr_c=init_savepoint.w_concorr_c(), rho_incr=None, # sp.rho_incr(), vn_incr=None, # sp.vn_incr(), exner_incr=None, # sp.exner_incr(), - exner_dyn_incr=sp.exner_dyn_incr(), + exner_dyn_incr=init_savepoint.exner_dyn_incr(), ) @@ -496,7 +495,6 @@ def test_nonhydro_corrector_step( savepoint_nonhydro_init, damping_height, grid_savepoint, - savepoint_velocity_init, metrics_savepoint, interpolation_savepoint, savepoint_nonhydro_exit, @@ -514,10 +512,9 @@ def test_nonhydro_corrector_step( nflatlev=grid_savepoint.nflatlev(), nflat_gradp=grid_savepoint.nflat_gradp(), ) - sp_v = savepoint_velocity_init - dtime = sp_v.get_metadata("dtime").get("dtime") - clean_mflx = sp_v.get_metadata("clean_mflx").get("clean_mflx") - lprep_adv = sp_v.get_metadata("prep_adv").get("prep_adv") + dtime = sp.get_metadata("dtime").get("dtime") + clean_mflx = sp.get_metadata("clean_mflx").get("clean_mflx") + lprep_adv = sp.get_metadata("prep_adv").get("prep_adv") prep_adv = PrepAdvection( vn_traj=sp.vn_traj(), mass_flx_me=sp.mass_flx_me(), mass_flx_ic=sp.mass_flx_ic() ) @@ -525,7 +522,7 @@ def test_nonhydro_corrector_step( nnow = 0 # TODO: @abishekg7 read from serialized data? nnew = 1 - diagnostic_state_nh = construct_diagnostics(sp, sp_v) + diagnostic_state_nh = construct_diagnostics(sp) z_fields = IntermediateFields( z_gradh_exner=sp.z_gradh_exner(), @@ -540,8 +537,8 @@ def test_nonhydro_corrector_step( z_graddiv_vn=sp.z_graddiv_vn(), z_rho_expl=sp.z_rho_expl(), z_dwdz_dd=sp.z_dwdz_dd(), - z_kin_hor_e=sp_v.z_kin_hor_e(), - z_vt_ie=sp_v.z_vt_ie(), + z_kin_hor_e=sp.z_kin_hor_e(), + z_vt_ie=sp.z_vt_ie(), ) divdamp_fac_o2 = sp.divdamp_fac_o2() @@ -683,7 +680,6 @@ def test_run_solve_nonhydro_single_step( savepoint_nonhydro_init, damping_height, grid_savepoint, - savepoint_velocity_init, # TODO (magdalena) this should not be needed in test_solve_nonhydro.py, only for test_velocity_advection.py metrics_savepoint, interpolation_savepoint, savepoint_nonhydro_exit, @@ -697,20 +693,19 @@ def test_run_solve_nonhydro_single_step( sp_step_exit = savepoint_nonhydro_step_exit nonhydro_params = NonHydrostaticParams(config) vertical_params = create_vertical_params(damping_height, grid_savepoint) - sp_v = savepoint_velocity_init - dtime = sp_v.get_metadata("dtime").get("dtime") - lprep_adv = sp_v.get_metadata("prep_adv").get("prep_adv") - clean_mflx = sp_v.get_metadata("clean_mflx").get("clean_mflx") + dtime = sp.get_metadata("dtime").get("dtime") + lprep_adv = sp.get_metadata("prep_adv").get("prep_adv") + clean_mflx = sp.get_metadata("clean_mflx").get("clean_mflx") prep_adv = PrepAdvection( vn_traj=sp.vn_traj(), mass_flx_me=sp.mass_flx_me(), mass_flx_ic=sp.mass_flx_ic() ) nnow = 0 nnew = 1 - recompute = sp_v.get_metadata("recompute").get("recompute") - linit = sp_v.get_metadata("linit").get("linit") + recompute = sp.get_metadata("recompute").get("recompute") + linit = sp.get_metadata("linit").get("linit") - diagnostic_state_nh = construct_diagnostics(sp, sp_v) + diagnostic_state_nh = construct_diagnostics(sp) interpolation_state = construct_interpolation_state_for_nonhydro(interpolation_savepoint) metric_state_nonhydro = construct_nh_metric_state(metrics_savepoint, icon_grid.num_levels) @@ -797,7 +792,7 @@ def test_run_solve_nonhydro_multi_step( savepoint_nonhydro_init, damping_height, grid_savepoint, - savepoint_velocity_init, + vn_only, metrics_savepoint, interpolation_savepoint, savepoint_nonhydro_exit, @@ -810,20 +805,19 @@ def test_run_solve_nonhydro_multi_step( sp_step_exit = savepoint_nonhydro_step_exit nonhydro_params = NonHydrostaticParams(config) vertical_params = create_vertical_params(damping_height, grid_savepoint) - sp_v = savepoint_velocity_init - dtime = sp_v.get_metadata("dtime").get("dtime") - lprep_adv = sp_v.get_metadata("prep_adv").get("prep_adv") - clean_mflx = sp_v.get_metadata("clean_mflx").get("clean_mflx") + dtime = sp.get_metadata("dtime").get("dtime") + lprep_adv = sp.get_metadata("prep_adv").get("prep_adv") + clean_mflx = sp.get_metadata("clean_mflx").get("clean_mflx") prep_adv = PrepAdvection( vn_traj=sp.vn_traj(), mass_flx_me=sp.mass_flx_me(), mass_flx_ic=sp.mass_flx_ic() ) nnow = 0 nnew = 1 - recompute = sp_v.get_metadata("recompute").get("recompute") - linit = sp_v.get_metadata("linit").get("linit") - diagnostic_state_nh = construct_diagnostics(sp, sp_v) + recompute = sp.get_metadata("recompute").get("recompute") + linit = sp.get_metadata("linit").get("linit") + diagnostic_state_nh = construct_diagnostics(sp) prognostic_state_ls = create_prognostic_states(sp) interpolation_state = construct_interpolation_state_for_nonhydro(interpolation_savepoint) diff --git a/model/common/src/icon4py/model/common/math/smagorinsky.py b/model/common/src/icon4py/model/common/math/smagorinsky.py index 6fca59d129..ace495078d 100644 --- a/model/common/src/icon4py/model/common/math/smagorinsky.py +++ b/model/common/src/icon4py/model/common/math/smagorinsky.py @@ -11,14 +11,14 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next import Field, field_operator +from gt4py.next import Field, field_operator, program from gt4py.next.ffront.fbuiltins import broadcast, maximum, minimum from icon4py.model.common.dimension import KDim, Koff @field_operator -def en_smag_fac_for_zero_nshift( +def _en_smag_fac_for_zero_nshift( vect_a: Field[[KDim], float], hdiff_smag_fac: float, hdiff_smag_fac2: float, @@ -45,3 +45,30 @@ def en_smag_fac_for_zero_nshift( dzqdr = minimum(broadcast(dz42, (KDim,)), maximum(zero, zf - hdiff_smag_z2)) enh_smag_fac = hdiff_smag_fac + (dzlin * alin) + dzqdr * (aqdr + dzqdr * bqdr) return enh_smag_fac + + +@program +def en_smag_fac_for_zero_nshift( + vect_a: Field[[KDim], 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: Field[[KDim], float], +): + _en_smag_fac_for_zero_nshift( + vect_a, + hdiff_smag_fac, + hdiff_smag_fac2, + hdiff_smag_fac3, + hdiff_smag_fac4, + hdiff_smag_z, + hdiff_smag_z2, + hdiff_smag_z3, + hdiff_smag_z4, + out=enh_smag_fac, + ) diff --git a/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py b/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py index f449e600b6..2a3f1adcc1 100644 --- a/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py +++ b/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py @@ -103,7 +103,7 @@ def icon_grid(grid_savepoint, experiment): Uses the special grid_savepoint that contains data from p_patch """ - return grid_savepoint.construct_icon_grid() + return grid_savepoint.construct_icon_grid(on_gpu=False) @pytest.fixture diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py index af8a1c22b3..fa7c6062d3 100644 --- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -690,6 +690,18 @@ def construct_prognostics(self) -> PrognosticState: class IconNonHydroInitSavepoint(IconSavepoint): + def z_vt_ie(self): + return self._get_field("z_vt_ie", EdgeDim, KDim) + + def z_kin_hor_e(self): + return self._get_field("z_kin_hor_e", EdgeDim, KDim) + + def vn_ie(self): + return self._get_field("vn_ie", EdgeDim, KDim) + + def vt(self): + return self._get_field("vt", EdgeDim, KDim) + def bdy_divdamp(self): return self._get_field("bdy_divdamp", KDim) @@ -732,21 +744,20 @@ def grf_tend_thv(self): def grf_tend_vn(self): return self._get_field("grf_tend_vn", EdgeDim, KDim) + def w_concorr_c(self): + return self._get_field("w_concorr_c", CellDim, KDim) + + def ddt_vn_apc_pc(self, ntnd): + return self._get_field_component("ddt_vn_apc_pc", ntnd, (EdgeDim, KDim)) + + def ddt_w_adv_pc(self, ntnd): + return self._get_field_component("ddt_w_adv_ntl", ntnd, (CellDim, KDim)) + def ddt_vn_adv_ntl(self, ntl): - buffer = np.squeeze(self.serializer.read("ddt_vn_adv_ntl", self.savepoint).astype(float))[ - :, :, ntl - 1 - ] - dims = (EdgeDim, KDim) - buffer = self._reduce_to_dim_size(buffer, dims) - return as_field(dims, buffer) + return self._get_field_component("ddt_vn_apc_pc", ntl, (EdgeDim, KDim)) def ddt_w_adv_ntl(self, ntl): - buffer = np.squeeze(self.serializer.read("ddt_w_adv_ntl", self.savepoint).astype(float))[ - :, :, ntl - 1 - ] - dims = (CellDim, KDim) - buffer = self._reduce_to_dim_size(buffer, dims) - return as_field(dims, buffer) + return self._get_field_component("ddt_w_adv_ntl", ntl, (CellDim, KDim)) def grf_tend_w(self): return self._get_field("grf_tend_w", CellDim, KDim) diff --git a/model/common/tests/math_tests/test_smagorinsky.py b/model/common/tests/math_tests/test_smagorinsky.py index 7413392932..a5b31f82a7 100644 --- a/model/common/tests/math_tests/test_smagorinsky.py +++ b/model/common/tests/math_tests/test_smagorinsky.py @@ -12,6 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +from gt4py.next.program_processors.runners.roundtrip import backend as roundtrip from icon4py.model.common.dimension import KDim from icon4py.model.common.grid.simple import SimpleGrid @@ -20,6 +21,8 @@ from icon4py.model.common.test_utils.reference_funcs import enhanced_smagorinski_factor_numpy +# TODO (magdalena) stencil does not run on embedded backend, broadcast(0.0, (KDim,)) return scalar? +# TODO (magdalena) run as to StencilTest def test_init_enh_smag_fac(): grid = SimpleGrid() enh_smag_fac = zero_field(grid, KDim) @@ -28,11 +31,11 @@ def test_init_enh_smag_fac(): z = (0.1, 0.2, 0.3, 0.4) enhanced_smag_fac_np = enhanced_smagorinski_factor_numpy(fac, z, a_vec.asnumpy()) - en_smag_fac_for_zero_nshift( + en_smag_fac_for_zero_nshift.with_backend(roundtrip)( a_vec, *fac, *z, - out=enh_smag_fac, + enh_smag_fac, offset_provider={"Koff": KDim}, ) assert np.allclose(enhanced_smag_fac_np, enh_smag_fac.asnumpy()) diff --git a/model/driver/src/icon4py/model/driver/io_utils.py b/model/driver/src/icon4py/model/driver/io_utils.py index d7e7abd750..44c05f1063 100644 --- a/model/driver/src/icon4py/model/driver/io_utils.py +++ b/model/driver/src/icon4py/model/driver/io_utils.py @@ -71,7 +71,7 @@ def read_icon_grid( return ( sb.IconSerialDataProvider("icon_pydycore", str(path.absolute()), False, mpi_rank=rank) .from_savepoint_grid() - .construct_icon_grid() + .construct_icon_grid(on_gpu=False) ) else: raise NotImplementedError(SB_ONLY_MSG) @@ -251,7 +251,7 @@ def read_static_fields( icon_grid = ( sb.IconSerialDataProvider("icon_pydycore", str(path.absolute()), False, mpi_rank=rank) .from_savepoint_grid() - .construct_icon_grid() + .construct_icon_grid(on_gpu=False) ) diffusion_interpolation_state = construct_interpolation_state_for_diffusion( dataprovider.from_interpolation_savepoint()