Init vertical grid

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py

@@ -330,7 +330,7 @@ def __init__(
         self.grid: Optional[icon_grid.IconGrid] = None
         self.config: Optional[DiffusionConfig] = None
         self.params: Optional[DiffusionParams] = None
-        self.vertical_params: Optional[v_grid.VerticalModelParams] = None
+        self.vertical_params: Optional[v_grid.VerticalGridParams] = None
         self.interpolation_state: diffusion_states.DiffusionInterpolationState = None
         self.metric_state: diffusion_states.DiffusionMetricState = None
         self.diff_multfac_w: Optional[float] = None
@@ -348,7 +348,7 @@ def init(
         grid: icon_grid.IconGrid,
         config: DiffusionConfig,
         params: DiffusionParams,
-        vertical_params: v_grid.VerticalModelParams,
+        vertical_params: v_grid.VerticalGridParams,
         metric_state: diffusion_states.DiffusionMetricState,
         interpolation_state: diffusion_states.DiffusionInterpolationState,
         edge_params: h_grid.EdgeParams,
@@ -404,7 +404,7 @@ def _get_start_index_for_w_diffusion() -> gtx.int32:
             config.substep_as_float,
             *params.smagorinski_factor,
             *params.smagorinski_height,
-            self.vertical_params.physical_heights,
+            self.vertical_params.inteface_physical_height,
             self.diff_multfac_vn,
             self.smag_limit,
             self.enh_smag_fac,
@@ -415,8 +415,8 @@ def _get_start_index_for_w_diffusion() -> gtx.int32:
         self.diff_multfac_n2w = diffusion_utils.init_nabla2_factor_in_upper_damping_zone(
             k_size=self.grid.num_levels,
             nshift=0,
-            physical_heights=self.vertical_params.physical_heights,
-            nrdmax=self.vertical_params.index_of_damping_layer,
+            physical_heights=self.vertical_params.inteface_physical_height,
+            nrdmax=self.vertical_params.end_index_of_damping_layer,
         )
         self._horizontal_start_index_w_diffusion = _get_start_index_for_w_diffusion()
         self._initialized = True
@@ -734,7 +734,7 @@ def _do_diffusion_step(
             k=self.vertical_index,
             cell=self.horizontal_cell_index,
             nrdmax=gtx.int32(
-                self.vertical_params.index_of_damping_layer + 1
+                self.vertical_params.end_index_of_damping_layer + 1
             ),  # +1 since Fortran includes boundaries
             interior_idx=gtx.int32(cell_start_interior),
             halo_idx=gtx.int32(cell_end_local),

model/atmosphere/diffusion/tests/diffusion_tests/conftest.py

@@ -19,16 +19,23 @@
     decomposition_info,
     download_ser_data,
     experiment,
+    flat_height,
     grid_savepoint,
+    htop_moist_proc,
     icon_grid,
     interpolation_savepoint,
     linit,
+    lowest_layer_thickness,
+    maximal_layer_thickness,
     metrics_savepoint,
+    model_top_height,
     ndyn_substeps,
     processor_props,
     ranked_data_path,
     step_date_exit,
     step_date_init,
+    stretch_factor,
+    top_height_limit_for_maximal_layer_thickness,
 )
 
 

model/atmosphere/diffusion/tests/diffusion_tests/mpi_tests/test_parallel_diffusion.py

@@ -17,7 +17,7 @@
 from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams
 from icon4py.model.common.decomposition import definitions
 from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim
-from icon4py.model.common.grid.vertical import VerticalModelParams
+from icon4py.model.common.grid.vertical import VerticalGridConfig, VerticalGridParams
 from icon4py.model.common.test_utils.datatest_utils import REGIONAL_EXPERIMENT
 from icon4py.model.common.test_utils.parallel_helpers import (  # noqa: F401  # import fixtures from test_utils package
     check_comm_size,
@@ -50,6 +50,9 @@ def test_parallel_diffusion(
     grid_savepoint,
     metrics_savepoint,
     interpolation_savepoint,
+    lowest_layer_thickness,
+    model_top_height,
+    stretch_factor,
     damping_height,
 ):
     check_comm_size(processor_props)
@@ -73,6 +76,13 @@ def test_parallel_diffusion(
     cell_geometry = grid_savepoint.construct_cell_geometry()
     edge_geometry = grid_savepoint.construct_edge_geometry()
     interpolation_state = construct_interpolation_state(interpolation_savepoint)
+    vertical_config = VerticalGridConfig(
+        icon_grid.num_levels,
+        lowest_layer_thickness=lowest_layer_thickness,
+        model_top_height=model_top_height,
+        stretch_factor=stretch_factor,
+        rayleigh_damping_height=damping_height,
+    )
     config = construct_config(experiment, ndyn_substeps=ndyn_substeps)
     diffusion_params = DiffusionParams(config)
     dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime")
@@ -87,7 +97,9 @@ def test_parallel_diffusion(
         grid=icon_grid,
         config=config,
         params=diffusion_params,
-        vertical_params=VerticalModelParams(grid_savepoint.vct_a(), damping_height),
+        vertical_params=VerticalGridParams(
+            vertical_config, grid_savepoint.vct_a(), grid_savepoint.vct_b()
+        ),
         metric_state=metric_state,
         interpolation_state=interpolation_state,
         edge_params=edge_geometry,

model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py

@@ -17,7 +17,7 @@
 from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams
 from icon4py.model.atmosphere.diffusion.diffusion_utils import scale_k
 from icon4py.model.common.grid.horizontal import CellParams, EdgeParams
-from icon4py.model.common.grid.vertical import VerticalModelParams
+from icon4py.model.common.grid.vertical import VerticalGridConfig, VerticalGridParams
 from icon4py.model.common.settings import backend
 from icon4py.model.common.test_utils.datatest_utils import GLOBAL_EXPERIMENT, REGIONAL_EXPERIMENT
 from icon4py.model.common.test_utils.helpers import dallclose
@@ -96,6 +96,15 @@ def test_smagorinski_factor_diffusion_type_5(experiment):
     assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor)))
 
 
+def create_vertical_params(vertical_config, grid_savepoint):
+    return VerticalGridParams(
+        vertical_config=vertical_config,
+        vct_a=grid_savepoint.vct_a(),
+        vct_b=grid_savepoint.vct_b(),
+        _min_index_flat_horizontal_grad_pressure=grid_savepoint.nflat_gradp(),
+    )
+
+
 @pytest.mark.datatest
 def test_diffusion_init(
     diffusion_savepoint_init,
@@ -105,18 +114,23 @@ def test_diffusion_init(
     icon_grid,
     experiment,
     step_date_init,
+    lowest_layer_thickness,
+    model_top_height,
+    stretch_factor,
     damping_height,
     ndyn_substeps,
 ):
     config = construct_config(experiment, ndyn_substeps=ndyn_substeps)
     additional_parameters = DiffusionParams(config)
 
-    vertical_params = VerticalModelParams(
-        vct_a=grid_savepoint.vct_a(),
+    vertical_config = VerticalGridConfig(
+        icon_grid.num_levels,
+        lowest_layer_thickness=lowest_layer_thickness,
+        model_top_height=model_top_height,
+        stretch_factor=stretch_factor,
         rayleigh_damping_height=damping_height,
-        nflatlev=grid_savepoint.nflatlev(),
-        nflat_gradp=grid_savepoint.nflat_gradp(),
     )
+    vertical_params = create_vertical_params(vertical_config, grid_savepoint)
 
     meta = diffusion_savepoint_init.get_metadata("linit", "date")
 
@@ -199,12 +213,12 @@ def _verify_init_values_against_savepoint(
 
 @pytest.mark.datatest
 @pytest.mark.parametrize(
-    "experiment,step_date_init,damping_height",
+    "experiment,step_date_init",
     [
-        (REGIONAL_EXPERIMENT, "2021-06-20T12:00:10.000", 12500.0),
-        (REGIONAL_EXPERIMENT, "2021-06-20T12:00:20.000", 12500.0),
-        (GLOBAL_EXPERIMENT, "2000-01-01T00:00:02.000", 50000.0),
-        (GLOBAL_EXPERIMENT, "2000-01-01T00:00:04.000", 50000.0),
+        (REGIONAL_EXPERIMENT, "2021-06-20T12:00:10.000"),
+        (REGIONAL_EXPERIMENT, "2021-06-20T12:00:20.000"),
+        (GLOBAL_EXPERIMENT, "2000-01-01T00:00:02.000"),
+        (GLOBAL_EXPERIMENT, "2000-01-01T00:00:04.000"),
     ],
 )
 @pytest.mark.parametrize("ndyn_substeps", (2,))
@@ -215,17 +229,22 @@ def test_verify_diffusion_init_against_savepoint(
     interpolation_savepoint,
     metrics_savepoint,
     diffusion_savepoint_init,
+    lowest_layer_thickness,
+    model_top_height,
+    stretch_factor,
     damping_height,
     ndyn_substeps,
 ):
     config = construct_config(experiment, ndyn_substeps=ndyn_substeps)
     additional_parameters = DiffusionParams(config)
-    vertical_params = VerticalModelParams(
-        vct_a=grid_savepoint.vct_a(),
+    vertical_config = VerticalGridConfig(
+        icon_grid.num_levels,
+        lowest_layer_thickness=lowest_layer_thickness,
+        model_top_height=model_top_height,
+        stretch_factor=stretch_factor,
         rayleigh_damping_height=damping_height,
-        nflatlev=grid_savepoint.nflatlev(),
-        nflat_gradp=grid_savepoint.nflat_gradp(),
     )
+    vertical_params = create_vertical_params(vertical_config, grid_savepoint)
     interpolation_state = construct_interpolation_state(interpolation_savepoint)
     metric_state = construct_metric_state(metrics_savepoint)
     edge_params = grid_savepoint.construct_edge_geometry()
@@ -248,10 +267,10 @@ def test_verify_diffusion_init_against_savepoint(
 
 @pytest.mark.datatest
 @pytest.mark.parametrize(
-    "experiment, step_date_init, step_date_exit, damping_height",
+    "experiment, step_date_init, step_date_exit",
     [
-        (REGIONAL_EXPERIMENT, "2021-06-20T12:00:10.000", "2021-06-20T12:00:10.000", 12500.0),
-        (GLOBAL_EXPERIMENT, "2000-01-01T00:00:02.000", "2000-01-01T00:00:02.000", 50000.0),
+        (REGIONAL_EXPERIMENT, "2021-06-20T12:00:10.000", "2021-06-20T12:00:10.000"),
+        (GLOBAL_EXPERIMENT, "2000-01-01T00:00:02.000", "2000-01-01T00:00:02.000"),
     ],
 )
 @pytest.mark.parametrize("ndyn_substeps", (2,))
@@ -263,6 +282,9 @@ def test_run_diffusion_single_step(
     grid_savepoint,
     icon_grid,
     experiment,
+    lowest_layer_thickness,
+    model_top_height,
+    stretch_factor,
     damping_height,
     ndyn_substeps,
 ):
@@ -273,12 +295,14 @@ def test_run_diffusion_single_step(
     metric_state = construct_metric_state(metrics_savepoint)
     diagnostic_state = construct_diagnostics(diffusion_savepoint_init)
     prognostic_state = diffusion_savepoint_init.construct_prognostics()
-    vertical_params = VerticalModelParams(
-        vct_a=grid_savepoint.vct_a(),
+    vertical_config = VerticalGridConfig(
+        icon_grid.num_levels,
+        lowest_layer_thickness=lowest_layer_thickness,
+        model_top_height=model_top_height,
+        stretch_factor=stretch_factor,
         rayleigh_damping_height=damping_height,
-        nflatlev=grid_savepoint.nflatlev(),
-        nflat_gradp=grid_savepoint.nflat_gradp(),
     )
+    vertical_params = create_vertical_params(vertical_config, grid_savepoint)
     config = construct_config(experiment, ndyn_substeps)
     additional_parameters = DiffusionParams(config)
 
@@ -307,11 +331,12 @@ def test_run_diffusion_single_step(
 
 
 @pytest.mark.datatest
-@pytest.mark.parametrize(
-    "linit, experiment, damping_height", [(True, REGIONAL_EXPERIMENT, 12500.0)]
-)
+@pytest.mark.parametrize("linit, experiment", [(True, REGIONAL_EXPERIMENT)])
 def test_run_diffusion_initial_step(
     experiment,
+    lowest_layer_thickness,
+    model_top_height,
+    stretch_factor,
     damping_height,
     diffusion_savepoint_init,
     diffusion_savepoint_exit,
@@ -327,8 +352,14 @@ def test_run_diffusion_initial_step(
     metric_state = construct_metric_state(metrics_savepoint)
     diagnostic_state = construct_diagnostics(diffusion_savepoint_init)
     prognostic_state = diffusion_savepoint_init.construct_prognostics()
-    vct_a = grid_savepoint.vct_a()
-    vertical_params = VerticalModelParams(vct_a=vct_a, rayleigh_damping_height=damping_height)
+    vertical_config = VerticalGridConfig(
+        icon_grid.num_levels,
+        lowest_layer_thickness=lowest_layer_thickness,
+        model_top_height=model_top_height,
+        stretch_factor=stretch_factor,
+        rayleigh_damping_height=damping_height,
+    )
+    vertical_params = create_vertical_params(vertical_config, grid_savepoint)
     config = construct_config(experiment, ndyn_substeps=2)
     additional_parameters = DiffusionParams(config)