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 c1a0180bb3..88398af28a 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 @@ -21,6 +21,10 @@ 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.set_cell_kdim_field_to_zero_wp import ( + set_cell_kdim_field_to_zero_wp, +) + from icon4py.model.atmosphere.dycore.accumulate_prep_adv_fields import ( accumulate_prep_adv_fields, ) @@ -136,8 +140,6 @@ _allocate_indices, _calculate_divdamp_fields, compute_z_raylfac, - set_zero_c_k, - set_zero_e_k, ) from icon4py.model.atmosphere.dycore.update_dynamical_exner_time_increment import ( update_dynamical_exner_time_increment, @@ -244,7 +246,6 @@ def __init__( rayleigh_type: int = 2, rayleigh_coeff: float = 0.05, divdamp_order: int = 24, # the ICON default is 4, - idiv_method: int = 1, is_iau_active: bool = False, iau_wgt_dyn: float = 0.0, divdamp_type: int = 3, @@ -328,9 +329,6 @@ def __init__( #: IAU weight for dynamics fields self.iau_wgt_dyn: float = iau_wgt_dyn - #: from mo_dynamics_nml.f90 - self.idiv_method: int = idiv_method - self._validate() def _validate(self): @@ -345,9 +343,6 @@ def _validate(self): if self.itime_scheme != 4: raise NotImplementedError("itime_scheme can only be 4") - if self.idiv_method != 1: - raise NotImplementedError("idiv_method can only be 1") - if self.divdamp_order != 24: raise NotImplementedError("divdamp_order can only be 24") @@ -709,6 +704,9 @@ def run_predictor_step( start_edge_lb_plus4 = self.grid.get_start_index( EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4 ) + start_edge_local_minus2 = self.grid.get_start_index( + EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2 + ) end_edge_local_minus2 = self.grid.get_end_index( EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2 ) @@ -898,50 +896,26 @@ def run_predictor_step( }, ) if self.config.iadv_rhotheta <= 2: - tmp_0_0 = self.grid.get_start_index(EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2) - offset = 2 if self.config.idiv_method == 1 else 3 - tmp_0_1 = self.grid.get_end_index( - EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - offset - ) - - set_zero_e_k.with_backend(backend)( - field=z_fields.z_rho_e, - horizontal_start=tmp_0_0, - horizontal_end=tmp_0_1, - vertical_start=0, - vertical_end=self.grid.num_levels, - offset_provider={}, - ) - - set_zero_e_k.with_backend(backend)( - field=z_fields.z_theta_v_e, - horizontal_start=tmp_0_0, - horizontal_end=tmp_0_1, + set_two_edge_kdim_fields_to_zero_wp.with_backend(backend)( + edge_kdim_field_to_zero_wp_1=z_fields.z_rho_e, + edge_kdim_field_to_zero_wp_2=z_fields.z_theta_v_e, + horizontal_start=start_edge_local_minus2, + horizontal_end=end_edge_local_minus2, vertical_start=0, vertical_end=self.grid.num_levels, offset_provider={}, ) - # initialize also nest boundary points with zero if self.grid.limited_area: - set_zero_e_k.with_backend(backend)( - field=z_fields.z_rho_e, - horizontal_start=start_edge_lb, - horizontal_end=end_edge_local_minus1, - vertical_start=0, - vertical_end=self.grid.num_levels, - offset_provider={}, - ) - - set_zero_e_k.with_backend(backend)( - field=z_fields.z_theta_v_e, + set_two_edge_kdim_fields_to_zero_wp.with_backend(backend)( + edge_kdim_field_to_zero_wp_1=z_fields.z_rho_e, + edge_kdim_field_to_zero_wp_2=z_fields.z_theta_v_e, horizontal_start=start_edge_lb, horizontal_end=end_edge_local_minus1, vertical_start=0, vertical_end=self.grid.num_levels, offset_provider={}, ) - if self.config.iadv_rhotheta == 2: # Compute upwind-biased values for rho and theta starting from centered differences # Note: the length of the backward trajectory should be 0.5*dtime*(vn,vt) in order to arrive @@ -1134,20 +1108,19 @@ def run_predictor_step( }, ) - if self.config.idiv_method == 1: - compute_mass_flux.with_backend(backend)( - z_rho_e=z_fields.z_rho_e, - z_vn_avg=self.z_vn_avg, - ddqz_z_full_e=self.metric_state_nonhydro.ddqz_z_full_e, - z_theta_v_e=z_fields.z_theta_v_e, - mass_fl_e=diagnostic_state_nh.mass_fl_e, - z_theta_v_fl_e=self.z_theta_v_fl_e, - horizontal_start=start_edge_lb_plus4, - horizontal_end=end_edge_local_minus2, - vertical_start=0, - vertical_end=self.grid.num_levels, - offset_provider={}, - ) + compute_mass_flux.with_backend(backend)( + z_rho_e=z_fields.z_rho_e, + z_vn_avg=self.z_vn_avg, + ddqz_z_full_e=self.metric_state_nonhydro.ddqz_z_full_e, + z_theta_v_e=z_fields.z_theta_v_e, + mass_fl_e=diagnostic_state_nh.mass_fl_e, + z_theta_v_fl_e=self.z_theta_v_fl_e, + horizontal_start=start_edge_lb_plus4, + horizontal_end=end_edge_local_minus2, + vertical_start=0, + vertical_end=self.grid.num_levels, + offset_provider={}, + ) nhsolve_prog.predictor_stencils_35_36.with_backend(backend)( vn=prognostic_state[nnew].vn, @@ -1203,22 +1176,21 @@ def run_predictor_step( }, ) - if self.config.idiv_method == 1: - compute_divergence_of_fluxes_of_rho_and_theta.with_backend(backend)( - geofac_div=self.interpolation_state.geofac_div, - mass_fl_e=diagnostic_state_nh.mass_fl_e, - z_theta_v_fl_e=self.z_theta_v_fl_e, - z_flxdiv_mass=self.z_flxdiv_mass, - z_flxdiv_theta=self.z_flxdiv_theta, - horizontal_start=start_cell_nudging, - horizontal_end=end_cell_local, - vertical_start=0, - vertical_end=self.grid.num_levels, - offset_provider={ - "C2E": self.grid.get_offset_provider("C2E"), - "C2CE": self.grid.get_offset_provider("C2CE"), - }, - ) + compute_divergence_of_fluxes_of_rho_and_theta.with_backend(backend)( + geofac_div=self.interpolation_state.geofac_div, + mass_fl_e=diagnostic_state_nh.mass_fl_e, + z_theta_v_fl_e=self.z_theta_v_fl_e, + z_flxdiv_mass=self.z_flxdiv_mass, + z_flxdiv_theta=self.z_flxdiv_theta, + horizontal_start=start_cell_nudging, + horizontal_end=end_cell_local, + vertical_start=0, + vertical_end=self.grid.num_levels, + offset_provider={ + "C2E": self.grid.get_offset_provider("C2E"), + "C2CE": self.grid.get_offset_provider("C2CE"), + }, + ) nhsolve_prog.stencils_43_44_45_45b.with_backend(backend)( z_w_expl=z_fields.z_w_expl, @@ -1686,15 +1658,41 @@ def run_corrector_step( }, ) - if self.config.idiv_method == 1: - log.debug("corrector: start stencil 32") - compute_mass_flux.with_backend(backend)( - z_rho_e=z_fields.z_rho_e, + log.debug("corrector: start stencil 32") + compute_mass_flux.with_backend(backend)( + z_rho_e=z_fields.z_rho_e, + z_vn_avg=self.z_vn_avg, + ddqz_z_full_e=self.metric_state_nonhydro.ddqz_z_full_e, + z_theta_v_e=z_fields.z_theta_v_e, + mass_fl_e=diagnostic_state_nh.mass_fl_e, + z_theta_v_fl_e=self.z_theta_v_fl_e, + horizontal_start=start_edge_lb_plus4, + horizontal_end=end_edge_local_minus2, # TODO: (halungge) this is actually the second halo line + vertical_start=0, + vertical_end=self.grid.num_levels, + offset_provider={}, + ) + + if lprep_adv: # Preparations for tracer advection + log.debug("corrector: doing prep advection") + if lclean_mflx: + log.debug("corrector: start stencil 33") + set_two_edge_kdim_fields_to_zero_wp.with_backend(backend)( + edge_kdim_field_to_zero_wp_1=prep_adv.vn_traj, + edge_kdim_field_to_zero_wp_2=prep_adv.mass_flx_me, + horizontal_start=start_edge_lb, + horizontal_end=end_edge_end, + vertical_start=0, + vertical_end=self.grid.num_levels, + offset_provider={}, + ) + log.debug(f"corrector: start stencil 34") + accumulate_prep_adv_fields.with_backend(backend)( z_vn_avg=self.z_vn_avg, - ddqz_z_full_e=self.metric_state_nonhydro.ddqz_z_full_e, - z_theta_v_e=z_fields.z_theta_v_e, mass_fl_e=diagnostic_state_nh.mass_fl_e, - z_theta_v_fl_e=self.z_theta_v_fl_e, + vn_traj=prep_adv.vn_traj, + mass_flx_me=prep_adv.mass_flx_me, + r_nsubsteps=r_nsubsteps, horizontal_start=start_edge_lb_plus4, horizontal_end=end_edge_local_minus2, vertical_start=0, @@ -1702,34 +1700,6 @@ def run_corrector_step( offset_provider={}, ) - if lprep_adv: # Preparations for tracer advection - log.debug("corrector: doing prep advection") - if lclean_mflx: - log.debug("corrector: start stencil 33") - set_two_edge_kdim_fields_to_zero_wp.with_backend(backend)( - edge_kdim_field_to_zero_wp_1=prep_adv.vn_traj, - edge_kdim_field_to_zero_wp_2=prep_adv.mass_flx_me, - horizontal_start=start_edge_lb, - horizontal_end=end_edge_end, - vertical_start=0, - vertical_end=self.grid.num_levels, - offset_provider={}, - ) - log.debug(f"corrector: start stencil 34") - accumulate_prep_adv_fields.with_backend(backend)( - z_vn_avg=self.z_vn_avg, - mass_fl_e=diagnostic_state_nh.mass_fl_e, - vn_traj=prep_adv.vn_traj, - mass_flx_me=prep_adv.mass_flx_me, - r_nsubsteps=r_nsubsteps, - horizontal_start=start_edge_lb_plus4, - horizontal_end=end_edge_local_minus2, - vertical_start=0, - vertical_end=self.grid.num_levels, - offset_provider={}, - ) - - if self.config.idiv_method == 1: # verified for e-9 log.debug(f"corrector: start stencile 41") compute_divergence_of_fluxes_of_rho_and_theta.with_backend(backend)( @@ -1961,7 +1931,7 @@ def run_corrector_step( r_nsubsteps=r_nsubsteps, horizontal_start=start_cell_nudging, horizontal_end=end_cell_local, - vertical_start=0, + vertical_start=1, vertical_end=self.grid.num_levels, offset_provider={}, ) @@ -1982,8 +1952,8 @@ def run_corrector_step( if lprep_adv: if lclean_mflx: log.debug(f"corrector set prep_adv.mass_flx_ic to zero") - set_zero_c_k.with_backend(backend)( - field=prep_adv.mass_flx_ic, + set_cell_kdim_field_to_zero_wp.with_backend(backend)( + field_to_zero_wp=prep_adv.mass_flx_ic, horizontal_start=start_cell_lb, horizontal_end=end_cell_nudging, vertical_start=0, 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 db910df114..c1f24445de 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 @@ -67,13 +67,18 @@ from icon4py.model.atmosphere.dycore.set_cell_kdim_field_to_zero_vp import ( _set_cell_kdim_field_to_zero_vp, ) +from icon4py.model.atmosphere.dycore.set_cell_kdim_field_to_zero_wp import ( + _set_cell_kdim_field_to_zero_wp, +) from icon4py.model.atmosphere.dycore.set_lower_boundary_condition_for_w_and_contravariant_correction import ( _set_lower_boundary_condition_for_w_and_contravariant_correction, ) from icon4py.model.atmosphere.dycore.set_theta_v_prime_ic_at_lower_boundary import ( _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.state_utils.utils import ( + _broadcast_zero_to_three_edge_kdim_fields_wp, +) 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 @@ -92,19 +97,12 @@ def init_test_fields( indices_cells_2: int32, nlev: int32, ): - _set_zero_e_k( - out=z_rho_e, - domain={EdgeDim: (indices_edges_1, indices_edges_2), KDim: (0, nlev)}, - ) - _set_zero_e_k( - out=z_theta_v_e, + _broadcast_zero_to_three_edge_kdim_fields_wp( + out=(z_rho_e, z_theta_v_e, z_graddiv_vn), domain={EdgeDim: (indices_edges_1, indices_edges_2), KDim: (0, nlev)}, ) - _set_zero_e_k( - out=z_graddiv_vn, - domain={EdgeDim: (indices_edges_1, indices_edges_2), KDim: (0, nlev)}, - ) - _set_zero_c_k( + + _set_cell_kdim_field_to_zero_wp( out=z_dwdz_dd, domain={CellDim: (indices_cells_1, indices_cells_2), KDim: (0, nlev)}, ) @@ -125,7 +123,7 @@ def _predictor_stencils_2_3( _extrapolate_temporally_exner_pressure(exner_exfac, exner, exner_ref_mc, exner_pr), (z_exner_ex_pr, exner_pr), ) - z_exner_ex_pr = where(k_field == nlev, _set_zero_c_k(), z_exner_ex_pr) + z_exner_ex_pr = where(k_field == nlev, _set_cell_kdim_field_to_zero_wp(), z_exner_ex_pr) return z_exner_ex_pr, exner_pr diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/utils.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/utils.py index ae12c6205f..466fcc0790 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/utils.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/utils.py @@ -22,7 +22,8 @@ maximum, ) -from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.type_alias import wpfloat def indices_field(dim: Dimension, grid, is_halfdim, dtype=int): @@ -57,56 +58,17 @@ def scale_k(field: Field[[KDim], float], factor: float, scaled_field: Field[[KDi @field_operator -def _set_zero_v_k() -> Field[[VertexDim, KDim], float]: - return broadcast(0.0, (VertexDim, KDim)) - - -@program -def set_zero_v_k(field: Field[[VertexDim, KDim], float]): - _set_zero_v_k(out=field) - - -@field_operator -def _set_zero_e_k() -> Field[[EdgeDim, KDim], float]: - return broadcast(0.0, (EdgeDim, KDim)) - - -@program -def set_zero_e_k( - field: Field[[EdgeDim, KDim], float], - horizontal_start: int32, - horizontal_end: int32, - vertical_start: int32, - vertical_end: int32, -): - _set_zero_e_k( - out=field, - domain={ - EdgeDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, - ) - - -@field_operator -def _set_zero_c_k() -> Field[[CellDim, KDim], float]: - return broadcast(0.0, (CellDim, KDim)) - - -@program -def set_zero_c_k( - field: Field[[CellDim, KDim], float], - horizontal_start: int32, - horizontal_end: int32, - vertical_start: int32, - vertical_end: int32, +def _broadcast_zero_to_three_edge_kdim_fields_wp() -> ( + tuple[ + Field[[EdgeDim, KDim], wpfloat], + Field[[EdgeDim, KDim], wpfloat], + Field[[EdgeDim, KDim], wpfloat], + ] ): - _set_zero_c_k( - out=field, - domain={ - CellDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), - }, + return ( + broadcast(wpfloat("0.0"), (EdgeDim, KDim)), + broadcast(wpfloat("0.0"), (EdgeDim, KDim)), + broadcast(wpfloat("0.0"), (EdgeDim, 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 053511ca75..8061ded0af 100644 --- a/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py @@ -358,8 +358,8 @@ def test_nonhydro_predictor_step( # stencil 32 # TODO: @abishekg7 higher tol. assert dallclose( - solve_nonhydro.z_theta_v_fl_e.asnumpy(), - sp_exit.z_theta_v_fl_e().asnumpy(), + solve_nonhydro.z_theta_v_fl_e.asnumpy()[edge_start_lb_plus4:, :], + sp_exit.z_theta_v_fl_e().asnumpy()[edge_start_lb_plus4:, :], atol=1e-9, ) @@ -372,8 +372,8 @@ def test_nonhydro_predictor_step( # stencil 35,36, 37,38 assert dallclose( - solve_nonhydro.intermediate_fields.z_vt_ie.asnumpy(), - sp_exit.z_vt_ie().asnumpy(), + solve_nonhydro.intermediate_fields.z_vt_ie.asnumpy()[edge_start_lb_plus4:, :], + sp_exit.z_vt_ie().asnumpy()[edge_start_lb_plus4:, :], atol=2e-14, ) # stencil 35,36 @@ -384,8 +384,8 @@ def test_nonhydro_predictor_step( ) # stencil 35 assert dallclose( - solve_nonhydro.z_w_concorr_me.asnumpy()[:, nflatlev:], - sp_exit.z_w_concorr_me().asnumpy()[:, nflatlev:], + solve_nonhydro.z_w_concorr_me.asnumpy()[edge_start_lb_plus4:, nflatlev:], + sp_exit.z_w_concorr_me().asnumpy()[edge_start_lb_plus4:, nflatlev:], atol=1e-15, ) diff --git a/model/common/src/icon4py/model/common/test_utils/datatest_utils.py b/model/common/src/icon4py/model/common/test_utils/datatest_utils.py index 80284c464b..d2eba8917a 100644 --- a/model/common/src/icon4py/model/common/test_utils/datatest_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/datatest_utils.py @@ -46,11 +46,11 @@ def get_test_data_root_path() -> Path: GRIDS_PATH = TEST_DATA_ROOT.joinpath("grids") DATA_URIS = { - 1: "https://polybox.ethz.ch/index.php/s/y2IMUx4pOQ6lyZ7/download", + 1: "https://polybox.ethz.ch/index.php/s/xhooaubvGffG8Qy/download", 2: "https://polybox.ethz.ch/index.php/s/YyC5qDJWyC39y7u/download", 4: "https://polybox.ethz.ch/index.php/s/UIHOVJs6FVPpz9V/download", } -DATA_URIS_APE = {1: "https://polybox.ethz.ch/index.php/s/uK3jtrWK90Z4kHC/download"} +DATA_URIS_APE = {1: "https://polybox.ethz.ch/index.php/s/y9WRP1mpPlf2BtM/download"} def get_processor_properties_for_run(run_instance):