diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py index 21a177e86f..2de72b691f 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py @@ -154,13 +154,12 @@ def __init__( rayleigh_type: constants.RayleighType = constants.RayleighType.KLEMP, rayleigh_coeff: float = 0.05, divdamp_order: dycore_states.DivergenceDampingOrder = dycore_states.DivergenceDampingOrder.COMBINED, # the ICON default is 4, - is_iau_active: bool = False, - iau_wgt_dyn: float = 0.0, divdamp_type: dycore_states.DivergenceDampingType = dycore_states.DivergenceDampingType.THREE_DIMENSIONAL, divdamp_trans_start: float = 12500.0, divdamp_trans_end: float = 17500.0, l_vert_nested: bool = False, deepatmos_mode: bool = False, + iau_init: bool = False, rhotheta_offctr: float = -0.1, veladv_offctr: float = 0.25, _nudge_max_coeff: float | None = None, # default is set in __init__ @@ -286,11 +285,8 @@ def __init__( #: deep atmosphere mode, originally defined as ldeepatmo in ICON self.deepatmos_mode: bool = deepatmos_mode - #: from mo_initicon_nml.f90/ mo_initicon_config.f90 - #: whether IAU is active at current time - self.is_iau_active: bool = is_iau_active - #: IAU weight for dynamics fields - self.iau_wgt_dyn: float = iau_wgt_dyn + #: incremental analysis init mode, defined as one of the ICON init modes + self.iau_init: bool = iau_init self._validate() @@ -458,10 +454,9 @@ def __init__( "zdiff_gradp": self._metric_state_nonhydro.zdiff_gradp, "pg_exdist": self._metric_state_nonhydro.pg_exdist, "inv_dual_edge_length": self._edge_geometry.inverse_dual_edge_lengths, - "iau_wgt_dyn": self._config.iau_wgt_dyn, - "is_iau_active": self._config.is_iau_active, "limited_area": self._grid.limited_area, }, + variants={"is_iau_active": [False, True] if self._config.iau_init else [False]}, horizontal_sizes={ "start_edge_lateral_boundary": self._start_edge_lateral_boundary, "start_edge_lateral_boundary_level_7": self._start_edge_lateral_boundary_level_7, @@ -491,8 +486,6 @@ def __init__( "geofac_grdiv": self._interpolation_state.geofac_grdiv, "advection_explicit_weight_parameter": self._params.advection_explicit_weight_parameter, "advection_implicit_weight_parameter": self._params.advection_implicit_weight_parameter, - "iau_wgt_dyn": self._config.iau_wgt_dyn, - "is_iau_active": self._config.is_iau_active, "limited_area": self._grid.limited_area, "divdamp_order": gtx.int32(self._config.divdamp_order), "mean_cell_area": self._cell_params.mean_cell_area, @@ -502,6 +495,7 @@ def __init__( variants={ "apply_2nd_order_divergence_damping": [False, True], "apply_4th_order_divergence_damping": [False, True], + "is_iau_active": [False, True] if self._config.iau_init else [False], }, horizontal_sizes={ "horizontal_start": gtx.int32(self._start_edge_nudging_level_2), @@ -574,13 +568,12 @@ def __init__( "e_bln_c_s": self._interpolation_state.e_bln_c_s, "wgtfac_c": self._metric_state_nonhydro.wgtfac_c, "wgtfacq_c": self._metric_state_nonhydro.wgtfacq_c, - "iau_wgt_dyn": self._config.iau_wgt_dyn, - "is_iau_active": self._config.is_iau_active, "rayleigh_type": self._config.rayleigh_type, "divdamp_type": self._config.divdamp_type, }, variants={ "at_first_substep": [False, True], + "is_iau_active": [False, True] if self._config.iau_init else [False], }, horizontal_sizes={ "start_cell_index_nudging": self._start_cell_nudging, @@ -609,14 +602,13 @@ def __init__( "reference_exner_at_cells_on_model_levels": self._metric_state_nonhydro.reference_exner_at_cells_on_model_levels, "advection_explicit_weight_parameter": self._params.advection_explicit_weight_parameter, "advection_implicit_weight_parameter": self._params.advection_implicit_weight_parameter, - "iau_wgt_dyn": self._config.iau_wgt_dyn, - "is_iau_active": self._config.is_iau_active, "rayleigh_type": self._config.rayleigh_type, }, variants={ "at_first_substep": [False, True], "at_last_substep": [False, True], "lprep_adv": [False, True], + "is_iau_active": [False, True] if self._config.iau_init else [False], }, horizontal_sizes={ "start_cell_index_nudging": self._start_cell_nudging, @@ -1008,6 +1000,8 @@ def time_step( lprep_adv: bool, at_first_substep: bool, at_last_substep: bool, + is_iau_active: bool = False, + iau_wgt_dyn: float = 0.0, ): """ Update prognostic variables (prognostic_states.next) after the dynamical process over one substep. @@ -1022,6 +1016,8 @@ def time_step( lprep_adv: Preparation for tracer advection at_first_substep: first substep at_last_substep: last substep + is_iau_active: Incremental analysis update active during dycore step + iau_wgt_dyn: weight scalar for the incremental analysis update """ log.info( f"running timestep: dtime = {dtime}, initial_timestep = {at_initial_timestep}, first_substep = {at_first_substep}, last_substep = {at_last_substep}, prep_adv = {lprep_adv}" @@ -1042,6 +1038,8 @@ def time_step( dtime=dtime, at_initial_timestep=at_initial_timestep, at_first_substep=at_first_substep, + is_iau_active=is_iau_active, + iau_wgt_dyn=iau_wgt_dyn, ) self.run_corrector_step( @@ -1055,6 +1053,8 @@ def time_step( lprep_adv=lprep_adv, at_first_substep=at_first_substep, at_last_substep=at_last_substep, + is_iau_active=is_iau_active, + iau_wgt_dyn=iau_wgt_dyn, ) if self._grid.limited_area: self._compute_exner_from_rhotheta_in_lateral_boundary( @@ -1079,6 +1079,8 @@ def run_predictor_step( dtime: float, at_initial_timestep: bool, at_first_substep: bool, + is_iau_active: bool, + iau_wgt_dyn: float, ): """ Runs the predictor step of the non-hydrostatic solver. @@ -1147,6 +1149,8 @@ def run_predictor_step( normal_wind_tendency_due_to_slow_physics_process=diagnostic_state_nh.normal_wind_tendency_due_to_slow_physics_process, normal_wind_iau_increment=diagnostic_state_nh.normal_wind_iau_increment, grf_tend_vn=diagnostic_state_nh.grf_tend_vn, + is_iau_active=is_iau_active, + iau_wgt_dyn=iau_wgt_dyn, dtime=dtime, ) @@ -1199,6 +1203,8 @@ def run_predictor_step( rayleigh_damping_factor=self._get_rayleigh_damping_factor(dtime), dtime=dtime, at_first_substep=at_first_substep, + is_iau_active=is_iau_active, + iau_wgt_dyn=iau_wgt_dyn, ) if self._grid.limited_area: @@ -1252,6 +1258,8 @@ def run_corrector_step( lprep_adv: bool, at_first_substep: bool, at_last_substep: bool, + is_iau_active: bool, + iau_wgt_dyn: float, ): log.info( f"running corrector step: dtime = {dtime}, prep_adv = {lprep_adv}, " @@ -1326,6 +1334,8 @@ def run_corrector_step( dtime=dtime, apply_2nd_order_divergence_damping=apply_2nd_order_divergence_damping, apply_4th_order_divergence_damping=apply_4th_order_divergence_damping, + is_iau_active=is_iau_active, + iau_wgt_dyn=iau_wgt_dyn, ) log.debug("exchanging prognostic field 'vn'") @@ -1374,6 +1384,8 @@ def run_corrector_step( exner_tendency_due_to_slow_physics=diagnostic_state_nh.exner_tendency_due_to_slow_physics, rho_iau_increment=diagnostic_state_nh.rho_iau_increment, exner_iau_increment=diagnostic_state_nh.exner_iau_increment, + is_iau_active=is_iau_active, + iau_wgt_dyn=iau_wgt_dyn, rayleigh_damping_factor=self._get_rayleigh_damping_factor(dtime), lprep_adv=lprep_adv, r_nsubsteps=r_nsubsteps, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py index 8f681a4ad5..5682511635 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py @@ -169,8 +169,8 @@ def _compute_rho_theta_pgrad_and_update_vn( pg_exdist: fa.EdgeKField[ta.vpfloat], inv_dual_edge_length: fa.EdgeField[ta.wpfloat], dtime: ta.wpfloat, - iau_wgt_dyn: ta.wpfloat, is_iau_active: bool, + iau_wgt_dyn: ta.wpfloat, limited_area: bool, nflatlev: gtx.int32, nflat_gradp: gtx.int32, @@ -299,8 +299,8 @@ def _apply_divergence_damping_and_update_vn( advection_explicit_weight_parameter: ta.wpfloat, advection_implicit_weight_parameter: ta.wpfloat, dtime: ta.wpfloat, - iau_wgt_dyn: ta.wpfloat, is_iau_active: bool, + iau_wgt_dyn: ta.wpfloat, limited_area: bool, apply_2nd_order_divergence_damping: bool, apply_4th_order_divergence_damping: bool, @@ -411,8 +411,8 @@ def compute_rho_theta_pgrad_and_update_vn( pg_exdist: fa.EdgeKField[ta.vpfloat], inv_dual_edge_length: fa.EdgeField[ta.wpfloat], dtime: ta.wpfloat, - iau_wgt_dyn: ta.wpfloat, is_iau_active: bool, + iau_wgt_dyn: ta.wpfloat, limited_area: bool, nflatlev: gtx.int32, nflat_gradp: gtx.int32, @@ -467,8 +467,8 @@ def compute_rho_theta_pgrad_and_update_vn( - pg_exdist: vertical distance between current cell height and neighboring cell height for hydrostatic correction [m] - inv_dual_edge_length: inverse dual edge length [m] - dtime: time step [s] - - iau_wgt_dyn: a scaling factor for iau increment - is_iau_active: option for iau increment analysis + - iau_wgt_dyn: a scaling factor for iau increment - limited_area: option indicating the grid is limited area or not - iadv_rhotheta: advection type for air density and virtual potential temperature (see RhoThetaAdvectionType) - igradp_method: option for pressure gradient computation (see HorizontalPressureDiscretizationType) @@ -516,8 +516,8 @@ def compute_rho_theta_pgrad_and_update_vn( pg_exdist=pg_exdist, inv_dual_edge_length=inv_dual_edge_length, dtime=dtime, - iau_wgt_dyn=iau_wgt_dyn, is_iau_active=is_iau_active, + iau_wgt_dyn=iau_wgt_dyn, limited_area=limited_area, nflatlev=nflatlev, nflat_gradp=nflat_gradp, @@ -561,8 +561,8 @@ def apply_divergence_damping_and_update_vn( advection_explicit_weight_parameter: ta.wpfloat, advection_implicit_weight_parameter: ta.wpfloat, dtime: ta.wpfloat, - iau_wgt_dyn: ta.wpfloat, is_iau_active: bool, + iau_wgt_dyn: ta.wpfloat, limited_area: bool, apply_2nd_order_divergence_damping: bool, apply_4th_order_divergence_damping: bool, @@ -605,8 +605,8 @@ def apply_divergence_damping_and_update_vn( - advection_explicit_weight_parameter: explicitness weight of normal_wind_advective_tendency - advection_implicit_weight_parameter: implicitness weight of normal_wind_advective_tendency - dtime: time step [s] - - iau_wgt_dyn: a scaling factor for iau increment - is_iau_active: option for iau increment analysis + - iau_wgt_dyn: a scaling factor for iau increment - limited_area: option indicating the grid is limited area or not - apply_2nd_order_divergence_damping: scaling factor for second order divergence damping - apply_4th_order_divergence_damping: scaling factor for fourth order divergence damping @@ -640,8 +640,8 @@ def apply_divergence_damping_and_update_vn( advection_explicit_weight_parameter=advection_explicit_weight_parameter, advection_implicit_weight_parameter=advection_implicit_weight_parameter, dtime=dtime, - iau_wgt_dyn=iau_wgt_dyn, is_iau_active=is_iau_active, + iau_wgt_dyn=iau_wgt_dyn, limited_area=limited_area, apply_2nd_order_divergence_damping=apply_2nd_order_divergence_damping, apply_4th_order_divergence_damping=apply_4th_order_divergence_damping, diff --git a/model/atmosphere/dycore/tests/dycore/fixtures.py b/model/atmosphere/dycore/tests/dycore/fixtures.py index 6fc5afa77d..c7131d5146 100644 --- a/model/atmosphere/dycore/tests/dycore/fixtures.py +++ b/model/atmosphere/dycore/tests/dycore/fixtures.py @@ -17,8 +17,10 @@ flat_height, grid_savepoint, htop_moist_proc, + iau_wgt_dyn, icon_grid, interpolation_savepoint, + is_iau_active, istep_exit, istep_init, linit, diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py index ef9bc2eb8f..02165ec8e9 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py @@ -54,7 +54,6 @@ def solve_nonhydro( config = solve_nh.NonHydrostaticConfig( rayleigh_coeff=0.1, divdamp_order=dycore_states.DivergenceDampingOrder.COMBINED, # type: ignore[arg-type] - iau_wgt_dyn=1.0, fourth_order_divdamp_factor=0.004, max_nudging_coefficient=0.375, ) diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index 080ffac929..1b5867deb0 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -167,6 +167,8 @@ def test_nonhydro_predictor_step( model_top_height, stretch_factor, damping_height, + is_iau_active, + iau_wgt_dyn, grid_savepoint, metrics_savepoint, interpolation_savepoint, @@ -229,6 +231,8 @@ def test_nonhydro_predictor_step( dtime=dtime, at_initial_timestep=at_initial_timestep, at_first_substep=at_first_substep, + is_iau_active=is_iau_active, + iau_wgt_dyn=iau_wgt_dyn, ) cell_domain = h_grid.domain(dims.CellDim) @@ -498,6 +502,8 @@ def test_nonhydro_corrector_step( model_top_height, stretch_factor, damping_height, + is_iau_active, + iau_wgt_dyn, grid_savepoint, metrics_savepoint, interpolation_savepoint, @@ -583,6 +589,8 @@ def test_nonhydro_corrector_step( lprep_adv=lprep_adv, at_first_substep=at_first_substep, at_last_substep=at_last_substep, + is_iau_active=is_iau_active, + iau_wgt_dyn=iau_wgt_dyn, ) # stencil 10 @@ -697,6 +705,8 @@ def test_run_solve_nonhydro_single_step( model_top_height, stretch_factor, damping_height, + is_iau_active, + iau_wgt_dyn, grid_savepoint, metrics_savepoint, interpolation_savepoint, @@ -765,6 +775,8 @@ def test_run_solve_nonhydro_single_step( lprep_adv=lprep_adv, at_first_substep=substep_init == 1, at_last_substep=substep_init == ndyn_substeps, + is_iau_active=is_iau_active, + iau_wgt_dyn=iau_wgt_dyn, ) prognostic_state_nnew = prognostic_states.next assert test_utils.dallclose( @@ -827,6 +839,8 @@ def test_run_solve_nonhydro_multi_step( model_top_height, stretch_factor, damping_height, + is_iau_active, + iau_wgt_dyn, grid_savepoint, metrics_savepoint, interpolation_savepoint, @@ -904,6 +918,8 @@ def test_run_solve_nonhydro_multi_step( lprep_adv=lprep_adv, at_first_substep=at_first_substep, at_last_substep=at_last_substep, + is_iau_active=is_iau_active, + iau_wgt_dyn=iau_wgt_dyn, ) if not at_last_substep: @@ -1384,6 +1400,8 @@ def test_compute_rho_theta_pgrad_and_update_vn( model_top_height, stretch_factor, damping_height, + is_iau_active, + iau_wgt_dyn, grid_savepoint, metrics_savepoint, interpolation_savepoint, @@ -1451,8 +1469,6 @@ def test_compute_rho_theta_pgrad_and_update_vn( dual_normal_cell_1 = grid_savepoint.dual_normal_cell_x() dual_normal_cell_2 = grid_savepoint.dual_normal_cell_y() - iau_wgt_dyn = config.iau_wgt_dyn - is_iau_active = config.is_iau_active igradp_method = config.igradp_method z_rho_e_ref = sp_stencil_exit.z_rho_e() @@ -1599,6 +1615,8 @@ def test_apply_divergence_damping_and_update_vn( model_top_height, stretch_factor, damping_height, + is_iau_active, + iau_wgt_dyn, grid_savepoint, metrics_savepoint, interpolation_savepoint, @@ -1636,7 +1654,6 @@ def test_apply_divergence_damping_and_update_vn( allocator=backend, ) - iau_wgt_dyn = config.iau_wgt_dyn divdamp_order = config.divdamp_order second_order_divdamp_scaling_coeff = sp_nh_init.divdamp_fac_o2() * mean_cell_area second_order_divdamp_factor = savepoint_nonhydro_init.divdamp_fac_o2() @@ -1651,7 +1668,6 @@ def test_apply_divergence_damping_and_update_vn( and second_order_divdamp_factor <= (4.0 * config.fourth_order_divdamp_factor) ) ) - is_iau_active = config.is_iau_active vn_ref = sp_nh_exit.vn_new() @@ -1691,8 +1707,8 @@ def test_apply_divergence_damping_and_update_vn( advection_explicit_weight_parameter=savepoint_nonhydro_init.wgt_nnow_vel(), advection_implicit_weight_parameter=savepoint_nonhydro_init.wgt_nnew_vel(), dtime=savepoint_nonhydro_init.get_metadata("dtime").get("dtime"), - iau_wgt_dyn=iau_wgt_dyn, is_iau_active=is_iau_active, + iau_wgt_dyn=iau_wgt_dyn, limited_area=grid_savepoint.get_metadata("limited_area").get("limited_area"), apply_2nd_order_divergence_damping=apply_2nd_order_divergence_damping, apply_4th_order_divergence_damping=apply_4th_order_divergence_damping, @@ -2049,6 +2065,8 @@ def test_vertically_implicit_solver_at_predictor_step( model_top_height, stretch_factor, damping_height, + is_iau_active, + iau_wgt_dyn, grid_savepoint, metrics_savepoint, interpolation_savepoint, @@ -2098,8 +2116,6 @@ def test_vertically_implicit_solver_at_predictor_step( dwdz_at_cells_on_model_levels = sp_stencil_init.z_dwdz_dd() exner_dynamical_increment = sp_stencil_init.exner_dyn_incr() - iau_wgt_dyn = config.iau_wgt_dyn - is_iau_active = config.is_iau_active divdamp_type = config.divdamp_type w_concorr_c_ref = sp_nh_exit.w_concorr_c() @@ -2254,6 +2270,8 @@ def test_vertically_implicit_solver_at_corrector_step( model_top_height, stretch_factor, damping_height, + is_iau_active, + iau_wgt_dyn, grid_savepoint, metrics_savepoint, interpolation_savepoint, @@ -2307,9 +2325,6 @@ def test_vertically_implicit_solver_at_corrector_step( r_nsubsteps = 1.0 / ndyn_substeps kstart_moist = vertical_params.kstart_moist - iau_wgt_dyn = config.iau_wgt_dyn - is_iau_active = config.is_iau_active - w_ref = sp_nh_exit.w_new() rho_ref = sp_nh_exit.rho_new() exner_ref = sp_nh_exit.exner_new() diff --git a/model/atmosphere/dycore/tests/dycore/mpi_tests/test_parallel_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/mpi_tests/test_parallel_solve_nonhydro.py index 6636ee4a24..afad94b36d 100644 --- a/model/atmosphere/dycore/tests/dycore/mpi_tests/test_parallel_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/mpi_tests/test_parallel_solve_nonhydro.py @@ -58,6 +58,8 @@ def test_run_solve_nonhydro_single_step( model_top_height: ta.wpfloat, stretch_factor: ta.wpfloat, damping_height: ta.wpfloat, + is_iau_active: bool, + iau_wgt_dyn: ta.wpfloat, grid_savepoint: serialbox.IconGridSavepoint, metrics_savepoint: serialbox.MetricSavepoint, interpolation_savepoint: serialbox.InterpolationSavepoint, @@ -157,6 +159,8 @@ def test_run_solve_nonhydro_single_step( lprep_adv=lprep_adv, at_first_substep=(substep_init == 1), at_last_substep=(substep_init == ndyn_substeps), + is_iau_active=is_iau_active, + iau_wgt_dyn=iau_wgt_dyn, ) _log.info(f"rank={processor_props.rank}/{processor_props.comm_size}: dycore step run ") diff --git a/model/driver/src/icon4py/model/driver/icon4py_configuration.py b/model/driver/src/icon4py/model/driver/icon4py_configuration.py index cec879df8b..bfad4f9c3d 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_configuration.py +++ b/model/driver/src/icon4py/model/driver/icon4py_configuration.py @@ -90,7 +90,6 @@ def _mch_ch_r04b09_diffusion_config(): def _mch_ch_r04b09_nonhydro_config(): return solve_nh.NonHydrostaticConfig( divdamp_order=dycore_states.DivergenceDampingOrder.COMBINED, - iau_wgt_dyn=1.0, fourth_order_divdamp_factor=0.004, max_nudging_coefficient=0.375, ) diff --git a/model/testing/src/icon4py/model/testing/definitions.py b/model/testing/src/icon4py/model/testing/definitions.py index 41f2f0ff76..deffe01997 100644 --- a/model/testing/src/icon4py/model/testing/definitions.py +++ b/model/testing/src/icon4py/model/testing/definitions.py @@ -313,7 +313,6 @@ def construct_nonhydrostatic_config(experiment: Experiment) -> solve_nh.NonHydro if experiment == Experiments.MCH_CH_R04B09: return solve_nh.NonHydrostaticConfig( divdamp_order=dycore_states.DivergenceDampingOrder.COMBINED, # type: ignore[arg-type] # TODO(havogt): typing in `NonHydrostaticConfig` needs to be fixed - iau_wgt_dyn=1.0, fourth_order_divdamp_factor=0.004, max_nudging_coefficient=0.375, ) diff --git a/model/testing/src/icon4py/model/testing/fixtures/datatest.py b/model/testing/src/icon4py/model/testing/fixtures/datatest.py index 7730656d5a..867a6616f2 100644 --- a/model/testing/src/icon4py/model/testing/fixtures/datatest.py +++ b/model/testing/src/icon4py/model/testing/fixtures/datatest.py @@ -598,3 +598,13 @@ def rayleigh_type() -> int: @pytest.fixture def top_height_limit_for_maximal_layer_thickness() -> float: return 15000.0 + + +@pytest.fixture +def is_iau_active() -> bool: + return False + + +@pytest.fixture +def iau_wgt_dyn() -> float: + return 0.0 diff --git a/tools/src/icon4py/tools/py2fgen/wrappers/dycore_wrapper.py b/tools/src/icon4py/tools/py2fgen/wrappers/dycore_wrapper.py index b529bc4374..197655667e 100644 --- a/tools/src/icon4py/tools/py2fgen/wrappers/dycore_wrapper.py +++ b/tools/src/icon4py/tools/py2fgen/wrappers/dycore_wrapper.py @@ -109,13 +109,12 @@ def solve_nh_init( rayleigh_type: gtx.int32, rayleigh_coeff: gtx.float64, divdamp_order: gtx.int32, - is_iau_active: bool, - iau_wgt_dyn: gtx.float64, divdamp_type: gtx.int32, divdamp_trans_start: gtx.float64, divdamp_trans_end: gtx.float64, l_vert_nested: bool, ldeepatmo: bool, + iau_init: bool, rhotheta_offctr: gtx.float64, veladv_offctr: gtx.float64, nudge_max_coeff: gtx.float64, # note: this is the scaled ICON value, i.e. not the namelist value @@ -165,13 +164,12 @@ def solve_nh_init( rayleigh_type=rayleigh_type, rayleigh_coeff=rayleigh_coeff, divdamp_order=divdamp_order, - is_iau_active=is_iau_active, - iau_wgt_dyn=iau_wgt_dyn, divdamp_type=divdamp_type, divdamp_trans_start=divdamp_trans_start, divdamp_trans_end=divdamp_trans_end, l_vert_nested=l_vert_nested, deepatmos_mode=ldeepatmo, + iau_init=iau_init, rhotheta_offctr=rhotheta_offctr, veladv_offctr=veladv_offctr, max_nudging_coefficient=nudge_max_coeff, @@ -351,6 +349,8 @@ def solve_nh_run( divdamp_fac_o2: gtx.float64, ndyn_substeps_var: gtx.int32, idyn_timestep: gtx.int32, + is_iau_active: bool, + iau_wgt_dyn: gtx.float64, ): if granule is None: raise RuntimeError("SolveNonhydro granule not initialized. Call 'solve_nh_init' first.") @@ -438,6 +438,8 @@ def solve_nh_run( lprep_adv=lprep_adv, at_first_substep=idyn_timestep == 0, at_last_substep=idyn_timestep == (ndyn_substeps_var - 1), + is_iau_active=is_iau_active, + iau_wgt_dyn=iau_wgt_dyn, ) # TODO(havogt): create separate bindings for writing the timers diff --git a/tools/tests/tools/py2fgen/fortran_samples/test_dycore.f90 b/tools/tests/tools/py2fgen/fortran_samples/test_dycore.f90 index fc22596d1c..c94c783652 100644 --- a/tools/tests/tools/py2fgen/fortran_samples/test_dycore.f90 +++ b/tools/tests/tools/py2fgen/fortran_samples/test_dycore.f90 @@ -174,8 +174,6 @@ program solve_nh_simulation integer(c_int), parameter :: rayleigh_type = 1 real(c_double), parameter :: rayleigh_coeff = 0.1 integer(c_int), parameter :: divdamp_order = 24 - logical(c_int), parameter :: is_iau_active = .false. - real(c_double), parameter :: iau_wgt_dyn = 0.5 real(c_double), parameter :: divdamp_fac_o2 = 0.5 integer(c_int), parameter :: divdamp_type = 1 real(c_double), parameter :: divdamp_trans_start = 1000.0 @@ -752,8 +750,6 @@ program solve_nh_simulation rayleigh_type=rayleigh_type, & rayleigh_coeff=rayleigh_coeff, & divdamp_order=divdamp_order, & - is_iau_active=is_iau_active, & - iau_wgt_dyn=iau_wgt_dyn, & divdamp_type=divdamp_type, & divdamp_trans_start=divdamp_trans_start, & divdamp_trans_end=divdamp_trans_end, & diff --git a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.f90 b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.f90 index 322658ce45..62f23c55cb 100644 --- a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.f90 +++ b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.f90 @@ -121,6 +121,8 @@ function solve_nh_run_wrapper(rho_now, & divdamp_fac_o2, & ndyn_substeps_var, & idyn_timestep, & + is_iau_active, & + iau_wgt_dyn, & on_gpu) bind(c, name="solve_nh_run_wrapper") result(rc) import :: c_int, c_double, c_bool, c_ptr integer(c_int) :: rc ! Stores the return code @@ -351,6 +353,10 @@ function solve_nh_run_wrapper(rho_now, & integer(c_int), value, target :: idyn_timestep + logical(c_int), value, target :: is_iau_active + + real(c_double), value, target :: iau_wgt_dyn + logical(c_int), value :: on_gpu end function solve_nh_run_wrapper @@ -500,13 +506,12 @@ function solve_nh_init_wrapper(c_lin_e, & rayleigh_type, & rayleigh_coeff, & divdamp_order, & - is_iau_active, & - iau_wgt_dyn, & divdamp_type, & divdamp_trans_start, & divdamp_trans_end, & l_vert_nested, & ldeepatmo, & + iau_init, & rhotheta_offctr, & veladv_offctr, & nudge_max_coeff, & @@ -814,10 +819,6 @@ function solve_nh_init_wrapper(c_lin_e, & integer(c_int), value, target :: divdamp_order - logical(c_int), value, target :: is_iau_active - - real(c_double), value, target :: iau_wgt_dyn - integer(c_int), value, target :: divdamp_type real(c_double), value, target :: divdamp_trans_start @@ -828,6 +829,8 @@ function solve_nh_init_wrapper(c_lin_e, & logical(c_int), value, target :: ldeepatmo + logical(c_int), value, target :: iau_init + real(c_double), value, target :: rhotheta_offctr real(c_double), value, target :: veladv_offctr @@ -904,6 +907,8 @@ subroutine solve_nh_run(rho_now, & divdamp_fac_o2, & ndyn_substeps_var, & idyn_timestep, & + is_iau_active, & + iau_wgt_dyn, & rc) use, intrinsic :: iso_c_binding @@ -991,6 +996,10 @@ subroutine solve_nh_run(rho_now, & integer(c_int), value, target :: idyn_timestep + logical(c_int), value, target :: is_iau_active + + real(c_double), value, target :: iau_wgt_dyn + logical(c_int) :: on_gpu integer(c_int) :: rho_now_size_0 @@ -1421,6 +1430,8 @@ subroutine solve_nh_run(rho_now, & divdamp_fac_o2=divdamp_fac_o2, & ndyn_substeps_var=ndyn_substeps_var, & idyn_timestep=idyn_timestep, & + is_iau_active=is_iau_active, & + iau_wgt_dyn=iau_wgt_dyn, & on_gpu=on_gpu) !$acc end host_data !$acc end host_data @@ -1514,13 +1525,12 @@ subroutine solve_nh_init(c_lin_e, & rayleigh_type, & rayleigh_coeff, & divdamp_order, & - is_iau_active, & - iau_wgt_dyn, & divdamp_type, & divdamp_trans_start, & divdamp_trans_end, & l_vert_nested, & ldeepatmo, & + iau_init, & rhotheta_offctr, & veladv_offctr, & nudge_max_coeff, & @@ -1647,10 +1657,6 @@ subroutine solve_nh_init(c_lin_e, & integer(c_int), value, target :: divdamp_order - logical(c_int), value, target :: is_iau_active - - real(c_double), value, target :: iau_wgt_dyn - integer(c_int), value, target :: divdamp_type real(c_double), value, target :: divdamp_trans_start @@ -1661,6 +1667,8 @@ subroutine solve_nh_init(c_lin_e, & logical(c_int), value, target :: ldeepatmo + logical(c_int), value, target :: iau_init + real(c_double), value, target :: rhotheta_offctr real(c_double), value, target :: veladv_offctr @@ -2233,13 +2241,12 @@ subroutine solve_nh_init(c_lin_e, & rayleigh_type=rayleigh_type, & rayleigh_coeff=rayleigh_coeff, & divdamp_order=divdamp_order, & - is_iau_active=is_iau_active, & - iau_wgt_dyn=iau_wgt_dyn, & divdamp_type=divdamp_type, & divdamp_trans_start=divdamp_trans_start, & divdamp_trans_end=divdamp_trans_end, & l_vert_nested=l_vert_nested, & ldeepatmo=ldeepatmo, & + iau_init=iau_init, & rhotheta_offctr=rhotheta_offctr, & veladv_offctr=veladv_offctr, & nudge_max_coeff=nudge_max_coeff, & diff --git a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.h b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.h index e12a36d50e..9846ab7eb5 100644 --- a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.h +++ b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.h @@ -36,7 +36,8 @@ extern int solve_nh_run_wrapper( int vn_traj_size_0, int vn_traj_size_1, double dtime, double *max_vcfl_size1_array, int max_vcfl_size1_array_size_0, int lprep_adv, int at_initial_timestep, double divdamp_fac_o2, - int ndyn_substeps_var, int idyn_timestep, int on_gpu); + int ndyn_substeps_var, int idyn_timestep, int is_iau_active, + double iau_wgt_dyn, int on_gpu); extern int solve_nh_init_wrapper( double *c_lin_e, int c_lin_e_size_0, int c_lin_e_size_1, double *c_intp, int c_intp_size_0, int c_intp_size_1, double *e_flx_avg, @@ -88,10 +89,10 @@ extern int solve_nh_init_wrapper( int coeff2_dwdz_size_1, double *coeff_gradekin, int coeff_gradekin_size_0, int coeff_gradekin_size_1, int *c_owner_mask, int c_owner_mask_size_0, int itime_scheme, int iadv_rhotheta, int igradp_method, int rayleigh_type, - double rayleigh_coeff, int divdamp_order, int is_iau_active, - double iau_wgt_dyn, int divdamp_type, double divdamp_trans_start, - double divdamp_trans_end, int l_vert_nested, int ldeepatmo, - double rhotheta_offctr, double veladv_offctr, double nudge_max_coeff, - double divdamp_fac, double divdamp_fac2, double divdamp_fac3, - double divdamp_fac4, double divdamp_z, double divdamp_z2, double divdamp_z3, - double divdamp_z4, int nflat_gradp, int backend, int on_gpu); \ No newline at end of file + double rayleigh_coeff, int divdamp_order, int divdamp_type, + double divdamp_trans_start, double divdamp_trans_end, int l_vert_nested, + int ldeepatmo, int iau_init, double rhotheta_offctr, double veladv_offctr, + double nudge_max_coeff, double divdamp_fac, double divdamp_fac2, + double divdamp_fac3, double divdamp_fac4, double divdamp_z, + double divdamp_z2, double divdamp_z3, double divdamp_z4, int nflat_gradp, + int backend, int on_gpu); \ No newline at end of file diff --git a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.py b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.py index e847aa080a..31cf654086 100644 --- a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.py +++ b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.py @@ -137,6 +137,8 @@ def solve_nh_run_wrapper( divdamp_fac_o2, ndyn_substeps_var, idyn_timestep, + is_iau_active, + iau_wgt_dyn, on_gpu, ): with runtime_config.HOOK_BINDINGS_FUNCTION["solve_nh_run"]: @@ -566,6 +568,8 @@ def solve_nh_run_wrapper( divdamp_fac_o2=divdamp_fac_o2, ndyn_substeps_var=ndyn_substeps_var, idyn_timestep=idyn_timestep, + is_iau_active=is_iau_active, + iau_wgt_dyn=iau_wgt_dyn, ) if __debug__: @@ -1315,13 +1319,12 @@ def solve_nh_init_wrapper( rayleigh_type, rayleigh_coeff, divdamp_order, - is_iau_active, - iau_wgt_dyn, divdamp_type, divdamp_trans_start, divdamp_trans_end, l_vert_nested, ldeepatmo, + iau_init, rhotheta_offctr, veladv_offctr, nudge_max_coeff, @@ -1825,13 +1828,12 @@ def solve_nh_init_wrapper( rayleigh_type=rayleigh_type, rayleigh_coeff=rayleigh_coeff, divdamp_order=divdamp_order, - is_iau_active=is_iau_active, - iau_wgt_dyn=iau_wgt_dyn, divdamp_type=divdamp_type, divdamp_trans_start=divdamp_trans_start, divdamp_trans_end=divdamp_trans_end, l_vert_nested=l_vert_nested, ldeepatmo=ldeepatmo, + iau_init=iau_init, rhotheta_offctr=rhotheta_offctr, veladv_offctr=veladv_offctr, nudge_max_coeff=nudge_max_coeff, diff --git a/tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py b/tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py index 717aedb692..259d4a30bf 100644 --- a/tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py +++ b/tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py @@ -44,13 +44,12 @@ def solve_nh_init( rayleigh_type = constants.RayleighType.KLEMP rayleigh_coeff = 0.05 divdamp_order = dycore_states.DivergenceDampingOrder.COMBINED - is_iau_active = False - iau_wgt_dyn = 1.0 divdamp_type = 3 divdamp_trans_start = 12500.0 divdamp_trans_end = 17500.0 l_vert_nested = False ldeepatmo = False + iau_init = False rhotheta_offctr = -0.1 veladv_offctr = 0.25 max_nudging_coefficient = 0.375 @@ -221,13 +220,12 @@ def solve_nh_init( rayleigh_type=rayleigh_type, rayleigh_coeff=rayleigh_coeff, divdamp_order=divdamp_order, - is_iau_active=is_iau_active, - iau_wgt_dyn=iau_wgt_dyn, divdamp_type=divdamp_type, divdamp_trans_start=divdamp_trans_start, divdamp_trans_end=divdamp_trans_end, l_vert_nested=l_vert_nested, ldeepatmo=ldeepatmo, + iau_init=iau_init, rhotheta_offctr=rhotheta_offctr, veladv_offctr=veladv_offctr, nudge_max_coeff=max_nudging_coefficient, @@ -293,13 +291,12 @@ def test_dycore_wrapper_granule_inputs( rayleigh_type = constants.RayleighType.KLEMP rayleigh_coeff = 0.05 divdamp_order = dycore_states.DivergenceDampingOrder.COMBINED - is_iau_active = False - iau_wgt_dyn = 1.0 divdamp_type = 3 divdamp_trans_start = 12500.0 divdamp_trans_end = 17500.0 l_vert_nested = False ldeepatmo = False + iau_init = False rhotheta_offctr = -0.1 veladv_offctr = 0.25 max_nudging_coefficient = 0.375 @@ -654,13 +651,12 @@ def test_dycore_wrapper_granule_inputs( rayleigh_type=rayleigh_type, rayleigh_coeff=rayleigh_coeff, divdamp_order=divdamp_order, - is_iau_active=is_iau_active, - iau_wgt_dyn=iau_wgt_dyn, divdamp_type=divdamp_type, divdamp_trans_start=divdamp_trans_start, divdamp_trans_end=divdamp_trans_end, l_vert_nested=l_vert_nested, ldeepatmo=ldeepatmo, + iau_init=iau_init, rhotheta_offctr=rhotheta_offctr, veladv_offctr=veladv_offctr, nudge_max_coeff=max_nudging_coefficient, @@ -774,6 +770,8 @@ def test_dycore_wrapper_granule_inputs( divdamp_fac_o2=second_order_divdamp_factor, ndyn_substeps_var=ndyn_substeps, idyn_timestep=substep, + is_iau_active=False, + iau_wgt_dyn=0.0, ) # Check input arguments to SolveNonhydro.time_step @@ -968,6 +966,8 @@ def test_granule_solve_nonhydro_single_step_regional( divdamp_fac_o2=second_order_divdamp_factor, # This is a scalar ndyn_substeps_var=ndyn_substeps, idyn_timestep=substep, + is_iau_active=False, + iau_wgt_dyn=0.0, ) # Comparison asserts should now use py2fgen.as_array @@ -1156,6 +1156,8 @@ def test_granule_solve_nonhydro_multi_step_regional( divdamp_fac_o2=second_order_divdamp_factor, ndyn_substeps_var=ndyn_substeps, idyn_timestep=i_substep, + is_iau_active=False, + iau_wgt_dyn=0.0, ) w_new, w_now = w_now, w_new