Make iau runtime

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py

@@ -154,8 +154,6 @@ 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,
@@ -286,12 +284,6 @@ 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
-
         self._validate()
 
     def _validate(self):
@@ -458,10 +450,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._grid.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 +482,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 +491,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._grid.iau_init else [False],
             },
             horizontal_sizes={
                 "horizontal_start": gtx.int32(self._start_edge_nudging_level_2),
@@ -574,13 +564,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._grid.iau_init else [False],
             },
             horizontal_sizes={
                 "start_cell_index_nudging": self._start_cell_nudging,
@@ -609,14 +598,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._grid.iau_init else [False],
             },
             horizontal_sizes={
                 "start_cell_index_nudging": self._start_cell_nudging,
@@ -1008,6 +996,8 @@ def time_step(
         lprep_adv: bool,
         at_first_substep: bool,
         at_last_substep: bool,
+        is_iau_active: bool,
+        iau_wgt_dyn: float,
     ):
         """
         Update prognostic variables (prognostic_states.next) after the dynamical process over one substep.
@@ -1042,6 +1032,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 +1047,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 +1073,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 +1143,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 +1197,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 +1252,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 +1328,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 +1378,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,

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,
     )

model/atmosphere/dycore/tests/dycore/mpi_tests/test_parallel_solve_nonhydro.py

@@ -157,6 +157,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=False,
+        iau_wgt_dyn=0.0,
     )
     _log.info(f"rank={processor_props.rank}/{processor_props.comm_size}: dycore step run ")
 

model/common/src/icon4py/model/common/grid/base.py

@@ -51,6 +51,7 @@ class GridConfig:
     # TODO(halungge): Decouple the vertical from horizontal grid.
     vertical_size: int
     limited_area: bool = True
+    iau_init: bool = False
     distributed: bool = False
     n_shift_total: int = 0
     length_rescale_factor: float = 1.0
@@ -152,6 +153,10 @@ def num_levels(self) -> int:
     def limited_area(self) -> bool:
         return self.config.limited_area
 
+    @property
+    def iau_init(self) -> bool:
+        return self.config.iau_init
+
     def get_connectivity(self, offset: str | gtx.FieldOffset) -> gtx_common.NeighborTable:
         """Get the connectivity by its name."""
         if isinstance(offset, gtx.FieldOffset):

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,
         )

model/driver/src/icon4py/model/driver/icon4py_driver.py

@@ -299,6 +299,8 @@ def _do_dyn_substepping(
                 lprep_adv=do_prep_adv,
                 at_first_substep=self._is_first_substep(dyn_substep),
                 at_last_substep=self._is_last_substep(dyn_substep),
+                is_iau_active=False,
+                iau_wgt_dyn=0.0,
             )
 
             if not self._is_last_substep(dyn_substep):

model/standalone_driver/src/icon4py/model/standalone_driver/standalone_driver.py

@@ -271,6 +271,8 @@ def _do_dyn_substepping(
                 lprep_adv=do_prep_adv,
                 at_first_substep=self._is_first_substep(dyn_substep),
                 at_last_substep=self._is_last_substep(dyn_substep),
+                is_iau_active=False,
+                iau_wgt_dyn=0.0,
             )
             timer_solve_nh.capture()
 

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,
         )

tools/src/icon4py/tools/py2fgen/wrappers/common.py

@@ -194,6 +194,7 @@ def construct_icon_grid(
     num_edges: int,
     vertical_size: int,
     limited_area: bool,
+    iau_init: bool,
     distributed: bool,
     allocator: gtx_typing.Allocator | None,
 ) -> icon.IconGrid:
@@ -239,6 +240,7 @@ def construct_icon_grid(
         ),
         vertical_size=vertical_size,
         limited_area=limited_area,
+        iau_init=iau_init,
         distributed=distributed,
         keep_skip_values=False,
     )

tools/src/icon4py/tools/py2fgen/wrappers/dycore_wrapper.py

@@ -105,8 +105,6 @@ 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,
@@ -160,8 +158,6 @@ 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,
@@ -331,6 +327,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.")
@@ -418,6 +416,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

tools/src/icon4py/tools/py2fgen/wrappers/grid_wrapper.py

@@ -115,6 +115,7 @@ def grid_init(
     num_edges: gtx.int32,
     vertical_size: gtx.int32,
     limited_area: bool,
+    iau_init: bool,
     backend: gtx.int32,
 ) -> None:
     on_gpu = c2e.array_ns != np  # TODO(havogt): expose `on_gpu` from py2fgen
@@ -168,6 +169,7 @@ def grid_init(
         num_edges=num_edges,
         vertical_size=vertical_size,
         limited_area=limited_area,
+        iau_init=iau_init,
         distributed=not processor_props.is_single_rank(),
         allocator=allocator,
     )

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, &