Make iau runtime

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,

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,

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,

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