Rename and enhance documentation related to the velocity tendencies i…

…n the diagnostic state

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

@@ -41,8 +41,8 @@ class DiagnosticStateNonHydro:
     mass_fl_e: fa.EdgeKField[float]
     ddt_vn_phy: fa.EdgeKField[float]
     grf_tend_vn: fa.EdgeKField[float]
-    ddt_vn_apc_pc: common_utils.Pair[fa.EdgeKField[float]]
-    ddt_w_adv_pc: common_utils.Pair[fa.CellKField[float]]
+    ddt_vn_apc_pc: common_utils.TimeStepPair[fa.EdgeKField[float]]
+    ddt_w_adv_pc: common_utils.TimeStepPair[fa.CellKField[float]]
 
     # Analysis increments
     rho_incr: Optional[fa.EdgeKField[float]]  # moist density increment [kg/m^3]

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

@@ -798,22 +798,36 @@ def _determine_local_domains(self):
         )
         self._end_vertex_halo = self._grid.end_index(vertex_domain(h_grid.Zone.HALO))
 
-    def update_time_levels(
+    def update_time_levels_for_velocity_tendencies(
         self, diagnostic_state_nh: dycore_states.DiagnosticStateNonHydro, at_first_substep: bool
     ) -> Literal[0, 1]:
-        """Set time levels of ddt_adv fields for call to velocity_tendencies."""
+        """
+        Set time levels of ddt_adv fields for call to velocity_tendencies.
+
+        When using `TimeSteppingScheme.MOST_EFFICIENT` (itime_scheme=4 in ICON Fortran),
+        `ddt_w_adv_pc.current` is not computed at the beginning of each substep.
+        Instead, the value computed during the previous substep is reused
+        for efficiency (except, of course, in the very first substep of the
+        initial time step). Additionally, in this scheme the predictor and
+        corrector outputs are kept in separated elements of the pair (.current for
+        the predictor and .next for the corrector) and interpoolated at the end
+        of the corrector step to get the final output.
+
+        If a different time stepping scheme is used, the first element of the pair
+        is used for both predictor and corrector outputs.
+
+        Args:
+            diagnostic_state_nh: Diagnostic fields calculated in the dynamical core (SolveNonHydro)
+            at_first_substep: Flag indicating if this is the first substep of the time step.
+
+        Returns:
+            The index of the pair element to be used for the corrector output.
+        """
+
         if self._config.itime_scheme == TimeSteppingScheme.MOST_EFFICIENT:
-            # When using TimeSteppingScheme.MOST_EFFICIENT (itime_scheme=4 in ICON Fortran),
-            # `ddt_w_adv_pc.first` is not computed at the beginning of each substep.
-            # Instead, the value computed during the previous substep is reused
-            # for efficiency (except, of course, in the very first substep of the
-            # initial time step). In this case, one of the pair's items (.first)
-            # is used for the predictor output, while the other item (.second) is
-            # used for the corrector.
             if not at_first_substep:
                 diagnostic_state_nh.ddt_w_adv_pc.swap()
-                # `diagnostic_state_nh.ddt_vn_apc_pc`` does not need to be swapped
-                # because is always recomputed in the predictor step
+                diagnostic_state_nh.ddt_vn_apc_pc.swap()
             return 1
         else:
             # Use only the first item of the pairs for both predictor and corrector outputs
@@ -853,7 +867,9 @@ def time_step(
                 offset_provider={},
             )
 
-        corrector_tl = self.update_time_levels(diagnostic_state_nh, at_first_substep)
+        corrector_tl = self.update_time_levels_for_velocity_tendencies(
+            diagnostic_state_nh, at_first_substep
+        )
 
         self.run_predictor_step(
             diagnostic_state_nh=diagnostic_state_nh,
@@ -1304,7 +1320,7 @@ def run_predictor_step(
 
         self._add_temporal_tendencies_to_vn(
             vn_nnow=prognostic_states.current.vn,
-            ddt_vn_apc_ntl1=diagnostic_state_nh.ddt_vn_apc_pc.first,
+            ddt_vn_apc_ntl1=diagnostic_state_nh.ddt_vn_apc_pc.current,
             ddt_vn_phy=diagnostic_state_nh.ddt_vn_phy,
             z_theta_v_e=z_fields.z_theta_v_e,
             z_gradh_exner=z_fields.z_gradh_exner,
@@ -1440,7 +1456,7 @@ def run_predictor_step(
         self._stencils_43_44_45_45b(
             z_w_expl=z_fields.z_w_expl,
             w_nnow=prognostic_states.current.w,
-            ddt_w_adv_ntl1=diagnostic_state_nh.ddt_w_adv_pc.first,
+            ddt_w_adv_ntl1=diagnostic_state_nh.ddt_w_adv_pc.current,
             z_th_ddz_exner_c=self.z_th_ddz_exner_c,
             z_contr_w_fl_l=z_fields.z_contr_w_fl_l,
             rho_ic=diagnostic_state_nh.rho_ic,
@@ -1750,7 +1766,7 @@ def run_corrector_step(
             log.debug(f"corrector: start stencil 23")
             self._add_temporal_tendencies_to_vn_by_interpolating_between_time_levels(
                 vn_nnow=prognostic_states.current.vn,
-                ddt_vn_apc_ntl1=diagnostic_state_nh.ddt_vn_apc_pc.first,
+                ddt_vn_apc_ntl1=diagnostic_state_nh.ddt_vn_apc_pc.current,
                 ddt_vn_apc_ntl2=diagnostic_state_nh.ddt_vn_apc_pc[time_level],
                 ddt_vn_phy=diagnostic_state_nh.ddt_vn_phy,
                 z_theta_v_e=z_fields.z_theta_v_e,
@@ -1921,7 +1937,7 @@ def run_corrector_step(
             self._stencils_42_44_45_45b(
                 z_w_expl=z_fields.z_w_expl,
                 w_nnow=prognostic_states.current.w,
-                ddt_w_adv_ntl1=diagnostic_state_nh.ddt_w_adv_pc.first,
+                ddt_w_adv_ntl1=diagnostic_state_nh.ddt_w_adv_pc.current,
                 ddt_w_adv_ntl2=diagnostic_state_nh.ddt_w_adv_pc[time_level],
                 z_th_ddz_exner_c=self.z_th_ddz_exner_c,
                 z_contr_w_fl_l=z_fields.z_contr_w_fl_l,
@@ -1956,7 +1972,7 @@ def run_corrector_step(
             self._stencils_43_44_45_45b(
                 z_w_expl=z_fields.z_w_expl,
                 w_nnow=prognostic_states.current.w,
-                ddt_w_adv_ntl1=diagnostic_state_nh.ddt_w_adv_pc.first,
+                ddt_w_adv_ntl1=diagnostic_state_nh.ddt_w_adv_pc.current,
                 z_th_ddz_exner_c=self.z_th_ddz_exner_c,
                 z_contr_w_fl_l=z_fields.z_contr_w_fl_l,
                 rho_ic=diagnostic_state_nh.rho_ic,

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

@@ -386,7 +386,7 @@ def run_predictor_step(
                 local_z_w_con_c=self.z_w_con_c,
                 coeff1_dwdz=self.metric_state.coeff1_dwdz,
                 coeff2_dwdz=self.metric_state.coeff2_dwdz,
-                ddt_w_adv=diagnostic_state.ddt_w_adv_pc.first,
+                ddt_w_adv=diagnostic_state.ddt_w_adv_pc.current,
                 horizontal_start=self._start_cell_nudging,
                 horizontal_end=self._end_cell_local,
                 vertical_start=1,
@@ -403,7 +403,7 @@ def run_predictor_step(
                 area=cell_areas,
                 geofac_n2s=self.interpolation_state.geofac_n2s,
                 w=prognostic_state.w,
-                ddt_w_adv=diagnostic_state.ddt_w_adv_pc.first,
+                ddt_w_adv=diagnostic_state.ddt_w_adv_pc.current,
                 scalfac_exdiff=scalfac_exdiff,
                 cfl_w_limit=cfl_w_limit,
                 dtime=dtime,
@@ -429,7 +429,7 @@ def run_predictor_step(
             z_w_con_c_full=self.z_w_con_c_full,
             vn_ie=diagnostic_state.vn_ie,
             ddqz_z_full_e=self.metric_state.ddqz_z_full_e,
-            ddt_vn_apc=diagnostic_state.ddt_vn_apc_pc.first,
+            ddt_vn_apc=diagnostic_state.ddt_vn_apc_pc.current,
             horizontal_start=self._start_edge_nudging_level_2,
             horizontal_end=self._end_edge_local,
             vertical_start=0,
@@ -448,7 +448,7 @@ def run_predictor_step(
             zeta=self.zeta,
             geofac_grdiv=self.interpolation_state.geofac_grdiv,
             vn=prognostic_state.vn,
-            ddt_vn_apc=diagnostic_state.ddt_vn_apc_pc.first,
+            ddt_vn_apc=diagnostic_state.ddt_vn_apc_pc.current,
             cfl_w_limit=cfl_w_limit,
             scalfac_exdiff=scalfac_exdiff,
             dtime=dtime,

model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py

@@ -154,7 +154,9 @@ def test_nonhydro_predictor_step(
     at_first_substep = jstep_init == 0
 
     prognostic_states = utils.create_prognostic_states(sp)
-    _ = solve_nonhydro.update_time_levels(diagnostic_state_nh, at_first_substep=at_first_substep)
+    _ = solve_nonhydro.update_time_levels_for_velocity_tendencies(
+        diagnostic_state_nh, at_first_substep=at_first_substep
+    )
 
     solve_nonhydro.run_predictor_step(
         diagnostic_state_nh=diagnostic_state_nh,
@@ -563,7 +565,7 @@ def test_nonhydro_corrector_step(
     at_last_substep = jstep_init == (ndyn_substeps - 1)
 
     prognostic_states = utils.create_prognostic_states(sp)
-    corrector_tl = solve_nonhydro.update_time_levels(
+    corrector_tl = solve_nonhydro.update_time_levels_for_velocity_tendencies(
         diagnostic_state_nh, at_first_substep=at_first_substep
     )
 

model/atmosphere/dycore/tests/dycore_tests/test_velocity_advection.py

@@ -295,14 +295,14 @@ def test_velocity_predictor_step(
     )
     # stencil 16
     assert helpers.dallclose(
-        diagnostic_state.ddt_w_adv_pc.first.asnumpy()[start_cell_nudging:, :],
+        diagnostic_state.ddt_w_adv_pc.current.asnumpy()[start_cell_nudging:, :],
         icon_result_ddt_w_adv_pc[start_cell_nudging:, :],
         atol=5.0e-16,
         rtol=1.0e-10,
     )
     # stencil 19
     assert helpers.dallclose(
-        diagnostic_state.ddt_vn_apc_pc.first.asnumpy(),
+        diagnostic_state.ddt_vn_apc_pc.current.asnumpy(),
         icon_result_ddt_vn_apc_pc,
         atol=1.0e-15,
     )
@@ -432,13 +432,13 @@ def test_velocity_corrector_step(
     )
     # stencil 16
     assert helpers.dallclose(
-        diagnostic_state.ddt_w_adv_pc.second.asnumpy()[start_cell_nudging:, :],
+        diagnostic_state.ddt_w_adv_pc.next.asnumpy()[start_cell_nudging:, :],
         icon_result_ddt_w_adv_pc[start_cell_nudging:, :],
         atol=5.0e-16,
     )
     # stencil 19
     assert helpers.dallclose(
-        diagnostic_state.ddt_vn_apc_pc.second.asnumpy(),
+        diagnostic_state.ddt_vn_apc_pc.next.asnumpy(),
         icon_result_ddt_vn_apc_pc,
         atol=5.0e-16,
     )