Renaming symbols

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

@@ -151,14 +151,14 @@ def test_nonhydro_predictor_step(
         backend=backend,
     )
     nlev = icon_grid.num_levels
-    at_first_substep = (jstep_init == 0)
+    at_first_substep = jstep_init == 0
 
-    prognostic_state_swp = utils.create_prognostic_states(sp)
+    prognostic_states = utils.create_prognostic_states(sp)
     _ = solve_nonhydro.update_time_levels(diagnostic_state_nh, at_first_substep=at_first_substep)
 
     solve_nonhydro.run_predictor_step(
         diagnostic_state_nh=diagnostic_state_nh,
-        prognostic_state_swp=prognostic_state_swp,
+        prognostic_states=prognostic_states,
         z_fields=solve_nonhydro.intermediate_fields,
         dtime=dtime,
         l_recompute=recompute,
@@ -303,7 +303,7 @@ def test_nonhydro_predictor_step(
         sp_exit.z_hydro_corr().asnumpy()[edge_start_nudging_level_2:, nlev - 1],
         atol=1e-20,
     )
-    prognostic_state_nnew = prognostic_state_swp.next
+    prognostic_state_nnew = prognostic_states.next
     vn_new_reference = sp_exit.vn_new().asnumpy()
 
     # stencils 24
@@ -559,15 +559,17 @@ def test_nonhydro_corrector_step(
         owner_mask=grid_savepoint.c_owner_mask(),
         backend=backend,
     )
-    at_first_substep = (jstep_init == 0)
-    at_last_substep=jstep_init == (ndyn_substeps - 1)
+    at_first_substep = jstep_init == 0
+    at_last_substep = jstep_init == (ndyn_substeps - 1)
 
-    prognostic_state_swp = utils.create_prognostic_states(sp)
-    corrector_tl = solve_nonhydro.update_time_levels(diagnostic_state_nh, at_first_substep=at_first_substep)
+    prognostic_states = utils.create_prognostic_states(sp)
+    corrector_tl = solve_nonhydro.update_time_levels(
+        diagnostic_state_nh, at_first_substep=at_first_substep
+    )
 
     solve_nonhydro.run_corrector_step(
         diagnostic_state_nh=diagnostic_state_nh,
-        prognostic_state_swp=prognostic_state_swp,
+        prognostic_states=prognostic_states,
         z_fields=z_fields,
         prep_adv=prep_adv,
         divdamp_fac_o2=divdamp_fac_o2,
@@ -602,29 +604,29 @@ def test_nonhydro_corrector_step(
 
     # stencil 23,26, 27, 4th_order_divdamp
     assert helpers.dallclose(
-        prognostic_state_swp.next.vn.asnumpy(),
+        prognostic_states.next.vn.asnumpy(),
         savepoint_nonhydro_exit.vn_new().asnumpy(),
         rtol=1e-9,  # TODO (magdalena) was 1e-10 for local experiment only
     )
 
     assert helpers.dallclose(
-        prognostic_state_swp.next.exner.asnumpy(),
+        prognostic_states.next.exner.asnumpy(),
         savepoint_nonhydro_exit.exner_new().asnumpy(),
     )
 
     assert helpers.dallclose(
-        prognostic_state_swp.next.rho.asnumpy(),
+        prognostic_states.next.rho.asnumpy(),
         savepoint_nonhydro_exit.rho_new().asnumpy(),
     )
 
     assert helpers.dallclose(
-        prognostic_state_swp.next.w.asnumpy(),
+        prognostic_states.next.w.asnumpy(),
         savepoint_nonhydro_exit.w_new().asnumpy(),
         atol=8e-14,
     )
 
     assert helpers.dallclose(
-        prognostic_state_swp.next.theta_v.asnumpy(),
+        prognostic_states.next.theta_v.asnumpy(),
         savepoint_nonhydro_exit.theta_v_new().asnumpy(),
     )
     # stencil 31
@@ -748,12 +750,12 @@ def test_run_solve_nonhydro_single_step(
         backend=backend,
     )
 
-    prognostic_state_swp = utils.create_prognostic_states(sp)
+    prognostic_states = utils.create_prognostic_states(sp)
 
     initial_divdamp_fac = sp.divdamp_fac_o2()
     solve_nonhydro.time_step(
         diagnostic_state_nh=diagnostic_state_nh,
-        prognostic_state_swp=prognostic_state_swp,
+        prognostic_states=prognostic_states,
         prep_adv=prep_adv,
         divdamp_fac_o2=initial_divdamp_fac,
         dtime=dtime,
@@ -764,7 +766,7 @@ def test_run_solve_nonhydro_single_step(
         at_first_substep=jstep_init == 0,
         at_last_substep=jstep_init == (ndyn_substeps - 1),
     )
-    prognostic_state_nnew = prognostic_state_swp.next
+    prognostic_state_nnew = prognostic_states.next
     assert helpers.dallclose(
         prognostic_state_nnew.theta_v.asnumpy(),
         sp_step_exit.theta_v_new().asnumpy(),
@@ -852,7 +854,7 @@ def test_run_solve_nonhydro_multi_step(
     linit = sp.get_metadata("linit").get("linit")
 
     diagnostic_state_nh = utils.construct_diagnostics(sp)
-    prognostic_state_swp = utils.create_prognostic_states(sp)
+    prognostic_states = utils.create_prognostic_states(sp)
 
     interpolation_state = utils.construct_interpolation_state(interpolation_savepoint)
     metric_state_nonhydro = utils.construct_metric_state(metrics_savepoint, icon_grid.num_levels)
@@ -879,7 +881,7 @@ def test_run_solve_nonhydro_multi_step(
 
         solve_nonhydro.time_step(
             diagnostic_state_nh=diagnostic_state_nh,
-            prognostic_state_swp=prognostic_state_swp,
+            prognostic_states=prognostic_states,
             prep_adv=prep_adv,
             divdamp_fac_o2=sp.divdamp_fac_o2(),
             dtime=dtime,
@@ -895,7 +897,7 @@ def test_run_solve_nonhydro_multi_step(
         recompute = False
         clean_mflx = False
         if not at_last_substep:
-            prognostic_state_swp.swap()
+            prognostic_states.swap()
 
     cell_start_lb_plus2 = icon_grid.start_index(
         h_grid.domain(dims.CellDim)(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3)
@@ -939,28 +941,28 @@ def test_run_solve_nonhydro_multi_step(
     )
 
     assert helpers.dallclose(
-        prognostic_state_swp.next.theta_v.asnumpy(),
+        prognostic_states.next.theta_v.asnumpy(),
         sp_step_exit.theta_v_new().asnumpy(),
     )
 
     assert helpers.dallclose(
-        prognostic_state_swp.next.rho.asnumpy(),
+        prognostic_states.next.rho.asnumpy(),
         savepoint_nonhydro_exit.rho_new().asnumpy(),
     )
 
     assert helpers.dallclose(
-        prognostic_state_swp.next.exner.asnumpy(),
+        prognostic_states.next.exner.asnumpy(),
         sp_step_exit.exner_new().asnumpy(),
     )
 
     assert helpers.dallclose(
-        prognostic_state_swp.next.w.asnumpy(),
+        prognostic_states.next.w.asnumpy(),
         savepoint_nonhydro_exit.w_new().asnumpy(),
         atol=8e-14,
     )
 
     assert helpers.dallclose(
-        prognostic_state_swp.next.vn.asnumpy(),
+        prognostic_states.next.vn.asnumpy(),
         savepoint_nonhydro_exit.vn_new().asnumpy(),
         atol=5e-13,
     )

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

@@ -159,5 +159,5 @@ def create_prognostic_states(sp) -> common_utils.NextStepPair[prognostics.Progno
         rho=sp.rho_new(),
         exner=sp.exner_new(),
     )
-    prognostic_state_swp = common_utils.NextStepPair(prognostic_state_nnow, prognostic_state_nnew)
-    return prognostic_state_swp
+    prognostic_states = common_utils.NextStepPair(prognostic_state_nnow, prognostic_state_nnew)
+    return prognostic_states

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

@@ -112,7 +112,7 @@ def time_integration(
         diffusion_diagnostic_state: diffusion_states.DiffusionDiagnosticState,
         solve_nonhydro_diagnostic_state: dycore_states.DiagnosticStateNonHydro,
         # TODO (Chia Rui): expand the PrognosticState to include indices of now and next, now it is always assumed that now = 0, next = 1 at the beginning
-        prognostic_state_swp: common_utils.NextStepPair[prognostics.PrognosticState],
+        prognostic_states: common_utils.NextStepPair[prognostics.PrognosticState],
         # below is a long list of arguments for dycore time_step that many can be moved to initialization of SolveNonhydro)
         prep_adv: dycore_states.PrepAdvection,
         inital_divdamp_fac_o2: float,
@@ -139,7 +139,7 @@ def time_integration(
             log.info("running initial step to diffuse fields before timeloop starts")
             self.diffusion.initial_run(
                 diffusion_diagnostic_state,
-                prognostic_state_swp.current,
+                prognostic_states.current,
                 self.dtime_in_seconds,
             )
         log.info(
@@ -149,10 +149,10 @@ def time_integration(
         for time_step in range(self._n_time_steps):
             log.info(f"simulation date : {self._simulation_date} run timestep : {time_step}")
             log.info(
-                f" MAX VN: {prognostic_state_swp.current.vn.asnumpy().max():.15e} , MAX W: {prognostic_state_swp.current.w.asnumpy().max():.15e}"
+                f" MAX VN: {prognostic_states.current.vn.asnumpy().max():.15e} , MAX W: {prognostic_states.current.w.asnumpy().max():.15e}"
             )
             log.info(
-                f" MAX RHO: {prognostic_state_swp.current.rho.asnumpy().max():.15e} , MAX THETA_V: {prognostic_state_swp.current.theta_v.asnumpy().max():.15e}"
+                f" MAX RHO: {prognostic_states.current.rho.asnumpy().max():.15e} , MAX THETA_V: {prognostic_states.current.theta_v.asnumpy().max():.15e}"
             )
             # TODO (Chia Rui): check with Anurag about printing of max and min of variables.
 
@@ -164,7 +164,7 @@ def time_integration(
             self._integrate_one_time_step(
                 diffusion_diagnostic_state,
                 solve_nonhydro_diagnostic_state,
-                prognostic_state_swp,
+                prognostic_states,
                 prep_adv,
                 initial_divdamp_fac_o2,
                 do_prep_adv,
@@ -183,7 +183,7 @@ def _integrate_one_time_step(
         self,
         diffusion_diagnostic_state: diffusion_states.DiffusionDiagnosticState,
         solve_nonhydro_diagnostic_state: dycore_states.DiagnosticStateNonHydro,
-        prognostic_state_swp: common_utils.NextStepPair[prognostics.PrognosticState],
+        prognostic_states: common_utils.NextStepPair[prognostics.PrognosticState],
         prep_adv: dycore_states.PrepAdvection,
         inital_divdamp_fac_o2: float,
         do_prep_adv: bool,
@@ -192,7 +192,7 @@ def _integrate_one_time_step(
 
         self._do_dyn_substepping(
             solve_nonhydro_diagnostic_state,
-            prognostic_state_swp,
+            prognostic_states,
             prep_adv,
             initial_divdamp_fac_o2,
             do_prep_adv,
@@ -201,18 +201,18 @@ def _integrate_one_time_step(
         if self.diffusion.config.apply_to_horizontal_wind:
             self.diffusion.run(
                 diffusion_diagnostic_state,
-                prognostic_state_swp.next,
+                prognostic_states.next,
                 self.dtime_in_seconds,
             )
 
-        prognostic_state_swp.swap_buffers()
+        prognostic_states.swap_buffers()
 
         # TODO (Chia Rui): add tracer advection here
 
     def _do_dyn_substepping(
         self,
         solve_nonhydro_diagnostic_state: dycore_states.DiagnosticStateNonHydro,
-        prognostic_state_swp: common_utils.NextStepPair[prognostics.PrognosticState],
+        prognostic_states: common_utils.NextStepPair[prognostics.PrognosticState],
         prep_adv: dycore_states.PrepAdvection,
         inital_divdamp_fac_o2: float,
         do_prep_adv: bool,
@@ -228,7 +228,7 @@ def _do_dyn_substepping(
             )
             self.solve_nonhydro.time_step(
                 solve_nonhydro_diagnostic_state,
-                prognostic_state_swp,
+                prognostic_states,
                 prep_adv=prep_adv,
                 divdamp_fac_o2=initial_divdamp_fac_o2,
                 dtime=self._substep_timestep,
@@ -244,7 +244,7 @@ def _do_dyn_substepping(
             do_clean_mflx = False
 
             if not self._is_last_substep(dyn_substep):
-                prognostic_state_swp.swap_buffers()
+                prognostic_states.swap_buffers()
 
             self._is_first_step_in_simulation = False
 
@@ -259,14 +259,14 @@ class DriverStates(NamedTuple):
         prep_advection_prognostic: Fields collecting data for advection during the solve nonhydro timestep.
         solve_nonhydro_diagnostic: Initial state for solve_nonhydro diagnostic variables.
         diffusion_diagnostic: Initial state for diffusion diagnostic variables.
-        prognostic_swp: Initial state for prognostic variables (double buffered).
+        prognostics: Initial state for prognostic variables (double buffered).
         diagnostic: Initial state for global diagnostic variables.
     """
 
     prep_advection_prognostic: solve_nh_states.PrepAdvection
     solve_nonhydro_diagnostic: solve_nh_states.DiagnosticStateNonHydro
     diffusion_diagnostic: diffusion_states.DiffusionDiagnosticState
-    prognostic_swp: common_utils.NextStepPair[prognostics.PrognosticState]
+    prognostics: common_utils.NextStepPair[prognostics.PrognosticState]
     diagnostic: diagnostics.DiagnosticState
 
 
@@ -405,7 +405,7 @@ def initialize(
         rank=props.rank,
         experiment_type=experiment_type,
     )
-    prognostics_swp = common_utils.NextStepPair(prognostic_state_now, prognostic_state_next)
+    prognosticss = common_utils.NextStepPair(prognostic_state_now, prognostic_state_next)
 
     timeloop = TimeLoop(
         run_config=config.run_config,
@@ -419,7 +419,7 @@ def initialize(
             prep_advection_prognostic=prep_adv,
             solve_nonhydro_diagnostic=solve_nonhydro_diagnostic_state,
             diffusion_diagnostic=diffusion_diagnostic_state,
-            prognostic_swp=prognostics_swp,
+            prognostics=prognosticss,
             diagnostic=diagnostic_state,
         ),
         DriverParams(divdamp_fac_o2=initial_divdamp_fac_o2),
@@ -522,7 +522,7 @@ def icon4py_driver(
     time_loop.time_integration(
         ds.diffusion_diagnostic,
         ds.solve_nonhydro_diagnostic,
-        ds.prognostic_swp,
+        ds.prognostics,
         ds.prep_advection_prognostic,
         dp.divdamp_fac_o2,
         do_prep_adv=False,

model/driver/tests/driver_tests/test_timeloop.py

@@ -306,12 +306,12 @@ def test_run_timeloop_single_step(
         exner=sp.exner_new(),
     )
 
-    prognostic_state_swp = common_utils.NextStepPair(prognostic_state, prognostic_state_new)
+    prognostic_states = common_utils.NextStepPair(prognostic_state, prognostic_state_new)
 
     timeloop.time_integration(
         diffusion_diagnostic_state,
         nonhydro_diagnostic_state,
-        prognostic_state_swp,
+        prognostic_states,
         prep_adv,
         sp.divdamp_fac_o2(),
         do_prep_adv,
@@ -324,29 +324,29 @@ def test_run_timeloop_single_step(
     w_sp = timeloop_diffusion_savepoint_exit.w()
 
     assert helpers.dallclose(
-        prognostic_state_swp.current.vn.asnumpy(),
+        prognostic_states.current.vn.asnumpy(),
         vn_sp.asnumpy(),
         atol=6e-12,
     )
 
     assert helpers.dallclose(
-        prognostic_state_swp.current.w.asnumpy(),
+        prognostic_states.current.w.asnumpy(),
         w_sp.asnumpy(),
         atol=8e-14,
     )
 
     assert helpers.dallclose(
-        prognostic_state_swp.current.exner.asnumpy(),
+        prognostic_states.current.exner.asnumpy(),
         exner_sp.asnumpy(),
     )
 
     assert helpers.dallclose(
-        prognostic_state_swp.current.theta_v.asnumpy(),
+        prognostic_states.current.theta_v.asnumpy(),
         theta_sp.asnumpy(),
         atol=4e-12,
     )
 
     assert helpers.dallclose(
-        prognostic_state_swp.current.rho.asnumpy(),
+        prognostic_states.current.rho.asnumpy(),
         rho_sp.asnumpy(),
     )

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

@@ -432,7 +432,7 @@ def solve_nh_run(
         rho=rho_new,
         exner=exner_new,
     )
-    prognostic_state_swp = common_utils.NextStepPair(prognostic_state_nnow, prognostic_state_nnew)
+    prognostic_states = common_utils.NextStepPair(prognostic_state_nnow, prognostic_state_nnew)
 
     # adjust for Fortran indexes
     nnow = nnow - 1
@@ -441,7 +441,7 @@ def solve_nh_run(
 
     dycore_wrapper_state["granule"].time_step(
         diagnostic_state_nh=diagnostic_state_nh,
-        prognostic_state_swp=prognostic_state_swp,
+        prognostic_states=prognostic_states,
         prep_adv=prep_adv,
         divdamp_fac_o2=divdamp_fac_o2,
         dtime=dtime,