Cleanup nh constants (#339)

Calculate NhConstants inside NonHydrostaticParams and removes NhConstants from argument to the timestep.

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

@@ -137,7 +137,6 @@
 from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_68 import (
     mo_solve_nonhydro_stencil_68,
 )
-from icon4py.model.atmosphere.dycore.state_utils.nh_constants import NHConstants
 from icon4py.model.atmosphere.dycore.state_utils.states import (
     DiagnosticStateNonHydro,
     InterpolationState,
@@ -271,7 +270,19 @@ def __init__(self, config: NonHydrostaticConfig):
         self.grav_o_cpd: Final[float] = constants.GRAV / constants.CPD
 
         # start level for 3D divergence damping terms
-        self.kstart_dd3d: int = 0
+        self.kstart_dd3d: Final[int] = 0
+
+        #: Weighting coefficients for velocity advection if tendency averaging is used
+        #: The off-centering specified here turned out to be beneficial to numerical
+        #: stability in extreme situations
+        self.wgt_nnow_vel: Final[float] = 0.5 - config.veladv_offctr
+        self.wgt_nnew_vel: Final[float] = 0.5 + config.veladv_offctr
+
+        #: Weighting coefficients for rho and theta at interface levels in the corrector step
+        #: This empirically determined weighting minimizes the vertical wind off-centering
+        #: needed for numerical stability of vertical sound wave propagation
+        self.wgt_nnew_rth: Final[float] = 0.5 + config.rhotheta_offctr
+        self.wgt_nnow_rth: Final[float] = 1.0 - self.wgt_nnew_rth
 
         # start level for moist physics processes (specified by htop_moist_proc)
         self.kstart_moist: int = 1
@@ -410,7 +421,6 @@ def time_step(
         prognostic_state_ls: list[PrognosticState],
         prep_adv: PrepAdvection,
         z_fields: ZFields,  # TODO (Magdalena): move local fields to SolveNonHydro
-        nh_constants: NHConstants,
         divdamp_fac_o2: float,
         dtime: float,
         idyn_timestep: float,
@@ -424,6 +434,7 @@ def time_step(
         log.info(
             f"running timestep: dtime = {dtime}, dyn_timestep = {idyn_timestep}, init = {l_init}, recompute = {l_recompute}, prep_adv = {lprep_adv} "
         )
+
         start_cell_lb = self.grid.get_start_index(
             CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim)
         )
@@ -471,7 +482,6 @@ def time_step(
             nnew=nnew,
             nnow=nnow,
             lclean_mflx=lclean_mflx,
-            nh_constants=nh_constants,
             lprep_adv=lprep_adv,
         )
         log.info(
@@ -529,6 +539,7 @@ def time_step(
             offset_provider={},
         )
 
+    # flake8: noqa: C901
     def run_predictor_step(
         self,
         diagnostic_state_nh: DiagnosticStateNonHydro,
@@ -616,13 +627,11 @@ def run_predictor_step(
         start_cell_lb_plus2 = self.grid.get_start_index(
             CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2
         )
-        end_cell_local_minus1 = self.grid.get_end_index(
-            CellDim, HorizontalMarkerIndex.local(CellDim) - 1
-        )
+
+        end_cell_halo = self.grid.get_end_index(CellDim, HorizontalMarkerIndex.halo(CellDim))
         start_cell_nudging = self.grid.get_start_index(
             CellDim, HorizontalMarkerIndex.nudging(CellDim)
         )
-
         end_cell_local = self.grid.get_end_index(CellDim, HorizontalMarkerIndex.local(CellDim))
 
         #  Precompute Rayleigh damping factor
@@ -652,7 +661,7 @@ def run_predictor_step(
             exner_pr=diagnostic_state_nh.exner_pr,
             z_exner_ex_pr=self.z_exner_ex_pr,
             horizontal_start=start_cell_lb_plus2,
-            horizontal_end=end_cell_local_minus1,
+            horizontal_end=end_cell_halo,
             k_field=self.k_field,
             nlev=self.grid.num_levels,
             vertical_start=0,
@@ -671,7 +680,7 @@ def run_predictor_step(
                 k_field=self.k_field,
                 nlev=self.grid.num_levels,
                 horizontal_start=start_cell_lb_plus2,
-                horizontal_end=end_cell_local_minus1,
+                horizontal_end=end_cell_halo,
                 vertical_start=max(1, self.vertical_params.nflatlev),
                 vertical_end=self.grid.num_levels + 1,
                 offset_provider={"Koff": KDim},
@@ -699,7 +708,7 @@ def run_predictor_step(
             k_field=self.k_field,
             nlev=self.grid.num_levels,
             horizontal_start=start_cell_lb_plus2,
-            horizontal_end=end_cell_local_minus1,
+            horizontal_end=end_cell_halo,
             vertical_start=0,
             vertical_end=self.grid.num_levels,
             offset_provider={"Koff": KDim},
@@ -715,7 +724,7 @@ def run_predictor_step(
             k_field=self.k_field,
             nlev=self.grid.num_levels,
             horizontal_start=start_cell_lb_plus2,
-            horizontal_end=end_cell_local_minus1,
+            horizontal_end=end_cell_halo,
             vertical_start=0,
             vertical_end=self.grid.num_levels + 1,
             offset_provider={"Koff": KDim},
@@ -730,7 +739,7 @@ def run_predictor_step(
                 z_rth_pr_2=self.z_rth_pr_2,
                 z_dexner_dz_c_2=self.z_dexner_dz_c_2,
                 horizontal_start=start_cell_lb_plus2,
-                horizontal_end=end_cell_local_minus1,
+                horizontal_end=end_cell_halo,
                 vertical_start=self.vertical_params.nflat_gradp,
                 vertical_end=self.grid.num_levels,
                 offset_provider={"Koff": KDim},
@@ -791,7 +800,7 @@ def run_predictor_step(
                 geofac_grg_x=self.interpolation_state.geofac_grg_x,
                 geofac_grg_y=self.interpolation_state.geofac_grg_y,
                 horizontal_start=start_cell_lb_plus2,
-                horizontal_end=end_cell_local_minus1,
+                horizontal_end=end_cell_halo,
                 vertical_start=0,
                 vertical_end=self.grid.num_levels,  # UBOUND(p_ccpr,2)
                 offset_provider={
@@ -1092,7 +1101,7 @@ def run_predictor_step(
             nflatlev_startindex_plus1=int32(self.vertical_params.nflatlev + 1),
             nlev=self.grid.num_levels,
             horizontal_start=start_cell_lb_plus2,
-            horizontal_end=end_cell_local_minus1,
+            horizontal_end=end_cell_halo,
             vertical_start=0,
             vertical_end=self.grid.num_levels + 1,
             offset_provider={
@@ -1336,7 +1345,6 @@ def run_corrector_step(
         diagnostic_state_nh: DiagnosticStateNonHydro,
         prognostic_state: list[PrognosticState],
         z_fields: ZFields,
-        nh_constants: NHConstants,
         divdamp_fac_o2: float,
         prep_adv: PrepAdvection,
         dtime: float,
@@ -1439,8 +1447,8 @@ def run_corrector_step(
             theta_v_ic=diagnostic_state_nh.theta_v_ic,
             z_th_ddz_exner_c=self.z_th_ddz_exner_c,
             dtime=dtime,
-            wgt_nnow_rth=nh_constants.wgt_nnow_rth,
-            wgt_nnew_rth=nh_constants.wgt_nnew_rth,
+            wgt_nnow_rth=self.params.wgt_nnow_rth,
+            wgt_nnew_rth=self.params.wgt_nnew_rth,
             horizontal_start=start_cell_lb_plus2,
             horizontal_end=end_cell_local,
             vertical_start=1,
@@ -1475,8 +1483,8 @@ def run_corrector_step(
                 z_gradh_exner=z_fields.z_gradh_exner,
                 vn_nnew=prognostic_state[nnew].vn,
                 dtime=dtime,
-                wgt_nnow_vel=nh_constants.wgt_nnow_vel,
-                wgt_nnew_vel=nh_constants.wgt_nnew_vel,
+                wgt_nnow_vel=self.params.wgt_nnow_vel,
+                wgt_nnew_vel=self.params.wgt_nnew_vel,
                 cpd=constants.CPD,
                 horizontal_start=start_edge_nudging_plus1,
                 horizontal_end=end_edge_local,
@@ -1652,8 +1660,8 @@ def run_corrector_step(
                 cvd=constants.CVD,
                 dtime=dtime,
                 cpd=constants.CPD,
-                wgt_nnow_vel=nh_constants.wgt_nnow_vel,
-                wgt_nnew_vel=nh_constants.wgt_nnew_vel,
+                wgt_nnow_vel=self.params.wgt_nnow_vel,
+                wgt_nnew_vel=self.params.wgt_nnew_vel,
                 nlev=self.grid.num_levels,
                 horizontal_start=start_cell_nudging,
                 horizontal_end=end_cell_local,

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

@@ -20,7 +20,6 @@
     NonHydrostaticParams,
     SolveNonhydro,
 )
-from icon4py.model.atmosphere.dycore.state_utils.nh_constants import NHConstants
 from icon4py.model.atmosphere.dycore.state_utils.states import (
     DiagnosticStateNonHydro,
     PrepAdvection,
@@ -174,13 +173,6 @@ def test_run_solve_nonhydro_single_step(
         z_vt_ie=_allocate(EdgeDim, KDim, grid=icon_grid),
     )
 
-    nh_constants = NHConstants(
-        wgt_nnow_rth=sp.wgt_nnow_rth(),
-        wgt_nnew_rth=sp.wgt_nnew_rth(),
-        wgt_nnow_vel=sp.wgt_nnow_vel(),
-        wgt_nnew_vel=sp.wgt_nnew_vel(),
-    )
-
     interpolation_state = interpolation_savepoint.construct_interpolation_state_for_nonhydro()
     metric_state_nonhydro = metrics_savepoint.construct_nh_metric_state(icon_grid.num_levels)
 
@@ -211,7 +203,6 @@ def test_run_solve_nonhydro_single_step(
         prognostic_state_ls=prognostic_state_ls,
         prep_adv=prep_adv,
         z_fields=z_fields,
-        nh_constants=nh_constants,
         divdamp_fac_o2=0.032,
         dtime=dtime,
         idyn_timestep=dyn_timestep,

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

@@ -20,7 +20,6 @@
     NonHydrostaticParams,
     SolveNonhydro,
 )
-from icon4py.model.atmosphere.dycore.state_utils.nh_constants import NHConstants
 from icon4py.model.atmosphere.dycore.state_utils.states import (
     DiagnosticStateNonHydro,
     PrepAdvection,
@@ -575,7 +574,6 @@ def test_nonhydro_corrector_step(
         z_vt_ie=sp_v.z_vt_ie(),
     )
 
-    nh_constants = create_nh_constants(sp)
     divdamp_fac_o2 = sp.divdamp_fac_o2()
 
     interpolation_state = construct_interpolation_state_for_nonhydro(interpolation_savepoint)
@@ -610,7 +608,6 @@ def test_nonhydro_corrector_step(
         nnew=nnew,
         nnow=nnow,
         lclean_mflx=clean_mflx,
-        nh_constants=nh_constants,
         lprep_adv=lprep_adv,
     )
 
@@ -725,8 +722,6 @@ def test_run_solve_nonhydro_single_step(
 
     z_fields = allocate_z_fields(icon_grid)
 
-    nh_constants = create_nh_constants(sp)
-
     interpolation_state = construct_interpolation_state_for_nonhydro(interpolation_savepoint)
     metric_state_nonhydro = construct_nh_metric_state(metrics_savepoint, icon_grid.num_levels)
 
@@ -754,7 +749,6 @@ def test_run_solve_nonhydro_single_step(
         prognostic_state_ls=prognostic_state_ls,
         prep_adv=prep_adv,
         z_fields=z_fields,
-        nh_constants=nh_constants,
         divdamp_fac_o2=initial_divdamp_fac,
         dtime=dtime,
         idyn_timestep=dyn_timestep,
@@ -846,8 +840,6 @@ def test_run_solve_nonhydro_multi_step(
 
     z_fields = allocate_z_fields(icon_grid)
 
-    nh_constants = create_nh_constants(sp)
-
     interpolation_state = construct_interpolation_state_for_nonhydro(interpolation_savepoint)
     metric_state_nonhydro = construct_nh_metric_state(metrics_savepoint, icon_grid.num_levels)
 
@@ -873,7 +865,6 @@ def test_run_solve_nonhydro_multi_step(
             prognostic_state_ls=prognostic_state_ls,
             prep_adv=prep_adv,
             z_fields=z_fields,
-            nh_constants=nh_constants,
             divdamp_fac_o2=sp.divdamp_fac_o2(),
             dtime=dtime,
             idyn_timestep=dyn_timestep,
@@ -961,13 +952,15 @@ def test_run_solve_nonhydro_multi_step(
     )
 
 
-def create_nh_constants(sp):
-    return NHConstants(
-        wgt_nnow_rth=sp.wgt_nnow_rth(),
-        wgt_nnew_rth=sp.wgt_nnew_rth(),
-        wgt_nnow_vel=sp.wgt_nnow_vel(),
-        wgt_nnew_vel=sp.wgt_nnew_vel(),
-    )
+@pytest.mark.datatest
+def test_non_hydrostatic_params(savepoint_nonhydro_init):
+    config = NonHydrostaticConfig()
+    params = NonHydrostaticParams(config)
+
+    assert params.wgt_nnew_vel == savepoint_nonhydro_init.wgt_nnew_vel()
+    assert params.wgt_nnow_vel == savepoint_nonhydro_init.wgt_nnow_vel()
+    assert params.wgt_nnew_rth == savepoint_nonhydro_init.wgt_nnew_rth()
+    assert params.wgt_nnow_rth == savepoint_nonhydro_init.wgt_nnow_rth()
 
 
 def create_prognostic_states(sp):

model/driver/pyproject.toml

@@ -90,7 +90,7 @@ line_length = 100  # It should be the same as in `tool.black.line-length` above
 lines_after_imports = 2
 multi_line_output = 3
 profile = 'black'
-sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'FIRSTPARTY', 'TESTS', 'LOCALFOLDER']
+sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'FIRSTPARTY', 'LOCALFOLDER']
 skip_gitignore = true
 skip_glob = ['*.venv/**', '_local/**']
 use_parentheses = true

model/driver/src/icon4py/model/driver/dycore_driver.py

@@ -24,7 +24,6 @@
     NonHydrostaticParams,
     SolveNonhydro,
 )
-from icon4py.model.atmosphere.dycore.state_utils.nh_constants import NHConstants
 from icon4py.model.atmosphere.dycore.state_utils.states import (
     DiagnosticStateNonHydro,
     PrepAdvection,
@@ -149,7 +148,6 @@ def time_integration(
         # below is a long list of arguments for dycore time_step that many can be moved to initialization of SolveNonhydro)
         prep_adv: PrepAdvection,
         z_fields: ZFields,  # local constants in solve_nh
-        nh_constants: NHConstants,
         inital_divdamp_fac_o2: float,
         do_prep_adv: bool,
     ):
@@ -193,7 +191,6 @@ def time_integration(
                 prognostic_state_list,
                 prep_adv,
                 z_fields,
-                nh_constants,
                 inital_divdamp_fac_o2,
                 do_prep_adv,
             )
@@ -214,7 +211,6 @@ def _integrate_one_time_step(
         prognostic_state_list: list[PrognosticState],
         prep_adv: PrepAdvection,
         z_fields: ZFields,
-        nh_constants: NHConstants,
         inital_divdamp_fac_o2: float,
         do_prep_adv: bool,
     ):
@@ -223,7 +219,6 @@ def _integrate_one_time_step(
             prognostic_state_list,
             prep_adv,
             z_fields,
-            nh_constants,
             inital_divdamp_fac_o2,
             do_prep_adv,
         )
@@ -243,7 +238,6 @@ def _do_dyn_substepping(
         prognostic_state_list: list[PrognosticState],
         prep_adv: PrepAdvection,
         z_fields: ZFields,
-        nh_constants: NHConstants,
         inital_divdamp_fac_o2: float,
         do_prep_adv: bool,
     ):
@@ -260,7 +254,6 @@ def _do_dyn_substepping(
                 prognostic_state_list,
                 prep_adv=prep_adv,
                 z_fields=z_fields,
-                nh_constants=nh_constants,
                 divdamp_fac_o2=inital_divdamp_fac_o2,
                 dtime=self._substep_timestep,
                 idyn_timestep=dyn_substep,
@@ -359,7 +352,6 @@ def initialize(file_path: Path, props: ProcessProperties):
         diffusion_diagnostic_state,
         solve_nonhydro_diagnostic_state,
         z_fields,
-        nh_constants,
         prep_adv,
         inital_divdamp_fac_o2,
         prognostic_state_now,
@@ -378,7 +370,6 @@ def initialize(file_path: Path, props: ProcessProperties):
         solve_nonhydro_diagnostic_state,
         prognostic_state_list,
         z_fields,
-        nh_constants,
         prep_adv,
         inital_divdamp_fac_o2,
     )