From aac1f687218a27e586d8009425dcaef14d3dbf87 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Fri, 5 Jan 2024 10:11:09 +0100
Subject: [PATCH 01/23] function skeleton for diamond edges

---
 .../common/src/icon4py/model/common/grid/grid_manager.py  | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py
index 1aeaf87bdc..24fa00e456 100644
--- a/model/common/src/icon4py/model/common/grid/grid_manager.py
+++ b/model/common/src/icon4py/model/common/grid/grid_manager.py
@@ -363,7 +363,8 @@ def _from_grid_dataset(self, dataset: Dataset) -> IconGrid:
         c2v = self._get_index_field(reader, GridFile.OffsetName.C2V)
         e2v = self._get_index_field(reader, GridFile.OffsetName.E2V)
 
-        e2c2v = self._construct_diamond_array(c2v, e2c)
+        e2c2v = self._construct_diamond_vertices(c2v, e2c)
+        e2c2e = self._construct_diamond_edges(c2e, e2c)
 
         v2c = self._get_index_field(reader, GridFile.OffsetName.V2C)
         v2e = self._get_index_field(reader, GridFile.OffsetName.V2E)
@@ -405,7 +406,7 @@ def _from_grid_dataset(self, dataset: Dataset) -> IconGrid:
 
         return icon_grid
 
-    def _construct_diamond_array(self, c2v: np.ndarray, e2c: np.ndarray):
+    def _construct_diamond_vertices(self, c2v: np.ndarray, e2c: np.ndarray):
         dummy_c2v = np.append(
             c2v,
             GridFile.INVALID_INDEX * np.ones((1, c2v.shape[1]), dtype=np.int32),
@@ -415,3 +416,6 @@ def _construct_diamond_array(self, c2v: np.ndarray, e2c: np.ndarray):
         sh = expanded.shape
         flattened = expanded.reshape(sh[0], sh[1] * sh[2])
         return np.apply_along_axis(np.unique, 1, flattened)
+
+    def _construct_diamond_edges(self, c2e, e2c):
+        pass

From ea57d190de67db57cf5a6384bc89c8f0ec17a176 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Fri, 5 Jan 2024 14:02:19 +0100
Subject: [PATCH 02/23] wip(1)

---
 model/common/src/icon4py/model/common/grid/grid_manager.py | 3 ++-
 model/common/tests/grid_tests/test_grid_manager.py         | 7 ++++---
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py
index 24fa00e456..8b02692c96 100644
--- a/model/common/src/icon4py/model/common/grid/grid_manager.py
+++ b/model/common/src/icon4py/model/common/grid/grid_manager.py
@@ -415,7 +415,8 @@ def _construct_diamond_vertices(self, c2v: np.ndarray, e2c: np.ndarray):
         expanded = dummy_c2v[e2c[:, :], :]
         sh = expanded.shape
         flattened = expanded.reshape(sh[0], sh[1] * sh[2])
-        return np.apply_along_axis(np.unique, 1, flattened)
+        result = np.apply_along_axis(np.unique, 1, flattened)
+        return result
 
     def _construct_diamond_edges(self, c2e, e2c):
         pass
diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py
index fc27101868..5a60322680 100644
--- a/model/common/tests/grid_tests/test_grid_manager.py
+++ b/model/common/tests/grid_tests/test_grid_manager.py
@@ -384,7 +384,7 @@ def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile):
 # e2c2e (e2c2eo) - diamond: exists in serial, simple_mesh
 @pytest.mark.datatest
 @pytest.mark.with_netcdf
-@pytest.mark.skip("does not directly exist in the grid file, needs to be constructed")
+# @pytest.mark.skip("does not directly exist in the grid file, needs to be constructed")
 # TODO (Magdalena) construct from adjacent_cell_of_edge and then edge_of_cell
 def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile):
     caplog.set_level(logging.DEBUG)
@@ -396,7 +396,7 @@ def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile):
     assert np.allclose(grid.get_offset_provider("E2C2E").table, serialized_e2c2e)
 
 
-@pytest.mark.xfail
+# @pytest.mark.xfail
 @pytest.mark.datatest
 @pytest.mark.with_netcdf
 def test_gridmanager_eval_e2c2v(caplog, grid_savepoint, r04b09_dsl_gridfile):
@@ -404,9 +404,10 @@ def test_gridmanager_eval_e2c2v(caplog, grid_savepoint, r04b09_dsl_gridfile):
     grid = init_grid_manager(r04b09_dsl_gridfile).get_grid()
     # the "far" (adjacent to edge normal ) is not there. why?
     # despite that: ordering is different
+    serialized_ref = grid_savepoint.e2c2v()[: grid.num_edges, :]
     assert np.allclose(
         grid.get_offset_provider("E2C2V").table,
-        grid_savepoint.e2c2v()[0 : grid.num_edges, :],
+        serialized_ref,
     )
 
 

From dd165c4ffdf25ed7ffe1d42f4bb1bd1a7058f104 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Fri, 5 Jan 2024 15:19:53 +0100
Subject: [PATCH 03/23] keep order when unifying indices

---
 .../src/icon4py/model/common/grid/grid_manager.py   | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py
index 8b02692c96..20652bf92a 100644
--- a/model/common/src/icon4py/model/common/grid/grid_manager.py
+++ b/model/common/src/icon4py/model/common/grid/grid_manager.py
@@ -415,8 +415,13 @@ def _construct_diamond_vertices(self, c2v: np.ndarray, e2c: np.ndarray):
         expanded = dummy_c2v[e2c[:, :], :]
         sh = expanded.shape
         flattened = expanded.reshape(sh[0], sh[1] * sh[2])
-        result = np.apply_along_axis(np.unique, 1, flattened)
-        return result
+        return np.apply_along_axis(_unique_keep_order, 1, flattened)
 
-    def _construct_diamond_edges(self, c2e, e2c):
-        pass
+    def _construct_diamond_edges(self, e2c: np.ndarray, c2e: np.ndarray):
+        expanded = c2e[e2c[:, :], :]
+        sh = expanded.shape
+
+
+def _unique_keep_order(ar: np.ndarray):
+    uniq, index = np.unique(ar, return_index=True)
+    return uniq[np.argsort(index)]

From 1f4cf3ae9f43e7277e510018a55b44278726b214 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Mon, 8 Jan 2024 21:09:30 +0100
Subject: [PATCH 04/23] add e2ce2 and fix e2c2v

---
 .../icon4py/model/common/grid/grid_manager.py | 38 +++++++++-----
 .../tests/grid_tests/test_grid_manager.py     | 49 ++++++++++++-------
 2 files changed, 57 insertions(+), 30 deletions(-)

diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py
index 20652bf92a..d95c7fe9f6 100644
--- a/model/common/src/icon4py/model/common/grid/grid_manager.py
+++ b/model/common/src/icon4py/model/common/grid/grid_manager.py
@@ -45,6 +45,7 @@ def __init__(self, *args, **kwargs):
     V2E2VDim,
     V2EDim,
     VertexDim,
+    E2C2EDim,
 )
 from icon4py.model.common.grid.base import GridConfig, VerticalGridSize
 from icon4py.model.common.grid.horizontal import HorizontalGridSize
@@ -363,8 +364,8 @@ def _from_grid_dataset(self, dataset: Dataset) -> IconGrid:
         c2v = self._get_index_field(reader, GridFile.OffsetName.C2V)
         e2v = self._get_index_field(reader, GridFile.OffsetName.E2V)
 
-        e2c2v = self._construct_diamond_vertices(c2v, e2c)
-        e2c2e = self._construct_diamond_edges(c2e, e2c)
+        e2c2v = self._construct_diamond_vertices(e2v, c2v, e2c)
+        e2c2e = self._construct_diamond_edges(e2c, c2e)
 
         v2c = self._get_index_field(reader, GridFile.OffsetName.V2C)
         v2e = self._get_index_field(reader, GridFile.OffsetName.V2E)
@@ -397,6 +398,7 @@ def _from_grid_dataset(self, dataset: Dataset) -> IconGrid:
                     C2E2CODim: c2e2c0,
                     E2C2VDim: e2c2v,
                     V2E2VDim: v2e2v,
+                    E2C2EDim: e2c2e,
                 }
             )
             .with_start_end_indices(CellDim, start_indices[CellDim], end_indices[CellDim])
@@ -406,7 +408,11 @@ def _from_grid_dataset(self, dataset: Dataset) -> IconGrid:
 
         return icon_grid
 
-    def _construct_diamond_vertices(self, c2v: np.ndarray, e2c: np.ndarray):
+    def _construct_diamond_vertices(self, e2v: np.ndarray, c2v: np.ndarray, e2c: np.ndarray):
+        n_edges = e2v.shape[0]
+        e2c2v = -55 * np.ones((n_edges, 4), dtype=np.int32)
+        e2c2v[:, 0:2] = e2v
+        # enlarge this table to contain an entry at the end for INVALID_VALUES
         dummy_c2v = np.append(
             c2v,
             GridFile.INVALID_INDEX * np.ones((1, c2v.shape[1]), dtype=np.int32),
@@ -414,14 +420,24 @@ def _construct_diamond_vertices(self, c2v: np.ndarray, e2c: np.ndarray):
         )
         expanded = dummy_c2v[e2c[:, :], :]
         sh = expanded.shape
-        flattened = expanded.reshape(sh[0], sh[1] * sh[2])
-        return np.apply_along_axis(_unique_keep_order, 1, flattened)
+        flat = expanded.reshape(sh[0], sh[1] * sh[2])
+        # TODO (magdalena) vectorize speed this up!
+        for x in range(sh[0]):
+            e2c2v[x, 2:] = flat[x, ~np.in1d(flat[x, :], e2c2v[x, :])][:2]
+
+        return e2c2v
 
     def _construct_diamond_edges(self, e2c: np.ndarray, c2e: np.ndarray):
-        expanded = c2e[e2c[:, :], :]
+        dummy_c2e = np.append(
+            c2e,
+            GridFile.INVALID_INDEX * np.ones((1, c2e.shape[1]), dtype=np.int32),
+            axis=0,
+        )
+        expanded = dummy_c2e[e2c, :]
         sh = expanded.shape
-
-
-def _unique_keep_order(ar: np.ndarray):
-    uniq, index = np.unique(ar, return_index=True)
-    return uniq[np.argsort(index)]
+        flattened = expanded.reshape(sh[0], sh[1] * sh[2])
+        e2c2e = GridFile.INVALID_INDEX * np.ones((sh[0], 4), dtype=np.int32)
+        for x in range(sh[0]):
+            var = flattened[x, (~np.in1d(flattened[x, :], np.asarray([x, GridFile.INVALID_INDEX])))]
+            e2c2e[x, : var.shape[0]] = var
+        return e2c2e
diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py
index 5a60322680..31e4af47e8 100644
--- a/model/common/tests/grid_tests/test_grid_manager.py
+++ b/model/common/tests/grid_tests/test_grid_manager.py
@@ -335,7 +335,11 @@ def test_gridmanager_eval_e2v(caplog, grid_savepoint, r04b09_dsl_gridfile):
 
 
 def has_invalid_index(ar: np.ndarray):
-    return np.any(np.where(ar == GridFile.INVALID_INDEX))
+    return np.any(invalid_index(ar))
+
+
+def invalid_index(ar):
+    return np.where(ar == GridFile.INVALID_INDEX)
 
 
 # e2c : exists in serial, simple, grid
@@ -384,31 +388,34 @@ def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile):
 # e2c2e (e2c2eo) - diamond: exists in serial, simple_mesh
 @pytest.mark.datatest
 @pytest.mark.with_netcdf
-# @pytest.mark.skip("does not directly exist in the grid file, needs to be constructed")
-# TODO (Magdalena) construct from adjacent_cell_of_edge and then edge_of_cell
 def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile):
     caplog.set_level(logging.DEBUG)
-    gm, num_cells, num_edges, num_vertex = init_grid_manager(r04b09_dsl_gridfile)
-    serialized_e2c2e = grid_savepoint.e2c2e()[0:num_cells, :]
+    grid = init_grid_manager(r04b09_dsl_gridfile).get_grid()
+    serialized_e2c2e = grid_savepoint.e2c2e()[0 : grid.num_edges, :]
     assert has_invalid_index(serialized_e2c2e)
-    grid = gm.get_grid()
-    assert has_invalid_index(grid.get_offset_provider("E2C2E").table)
-    assert np.allclose(grid.get_offset_provider("E2C2E").table, serialized_e2c2e)
+
+    e2c2e_table = grid.get_offset_provider("E2C2E").table
+    assert has_invalid_index(e2c2e_table)
+    # ICON calculates diamond edges only from rl_start = 2 (lateral_boundary(EdgeDim) + 1 for
+    # boundaries all values are INVALID even though the half diamond exists (see mo_model_domimp_setup.f90 ll 163ff.)
+    assert_unless_invalid(e2c2e_table, serialized_e2c2e)
+
+
+def assert_unless_invalid(table, serialized_ref):
+    invalid_positions = invalid_index(table)
+    np.allclose(table[invalid_positions], serialized_ref[invalid_positions])
 
 
-# @pytest.mark.xfail
 @pytest.mark.datatest
 @pytest.mark.with_netcdf
 def test_gridmanager_eval_e2c2v(caplog, grid_savepoint, r04b09_dsl_gridfile):
     caplog.set_level(logging.DEBUG)
     grid = init_grid_manager(r04b09_dsl_gridfile).get_grid()
-    # the "far" (adjacent to edge normal ) is not there. why?
-    # despite that: ordering is different
+    # the "far" (adjacent to edge normal ) is not always there, because ICON only calculates those starting from
+    #   (lateral_boundary(EdgeDim) + 1) to end(EdgeDim)  (see mo_intp_coeffs.f90) and only for owned cells
     serialized_ref = grid_savepoint.e2c2v()[: grid.num_edges, :]
-    assert np.allclose(
-        grid.get_offset_provider("E2C2V").table,
-        serialized_ref,
-    )
+    table = grid.get_offset_provider("E2C2V").table
+    assert_unless_invalid(table, serialized_ref)
 
 
 @pytest.mark.datatest
@@ -435,6 +442,12 @@ def test_grid_manager_getsize(simple_grid_gridfile, dim, size, caplog):
     assert size == gm.get_size(dim)
 
 
+def assert_up_to_order(table, diamond_table):
+    assert table.shape == diamond_table.shape
+    for n in range(table.shape[0]):
+        assert np.all(np.in1d(table[n, :], diamond_table[n, :]))
+
+
 @pytest.mark.with_netcdf
 def test_grid_manager_diamond_offset(simple_grid_gridfile):
     simple_grid = SimpleGrid()
@@ -445,10 +458,8 @@ def test_grid_manager_diamond_offset(simple_grid_gridfile):
     )
     gm()
     icon_grid = gm.get_grid()
-    assert np.allclose(
-        icon_grid.get_offset_provider("E2C2V").table,
-        np.sort(simple_grid.diamond_table, 1),
-    )
+    table = icon_grid.get_offset_provider("E2C2V").table
+    assert_up_to_order(table, simple_grid.diamond_table)
 
 
 @pytest.mark.with_netcdf

From 2bdeb68008ba997bf93b9ae663512b145a8aa120 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Tue, 9 Jan 2024 09:11:58 +0100
Subject: [PATCH 05/23] refactor diamond calculations

---
 .../icon4py/model/common/grid/grid_manager.py | 44 ++++++++++++-------
 1 file changed, 29 insertions(+), 15 deletions(-)

diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py
index d95c7fe9f6..7c580fb0e8 100644
--- a/model/common/src/icon4py/model/common/grid/grid_manager.py
+++ b/model/common/src/icon4py/model/common/grid/grid_manager.py
@@ -37,6 +37,7 @@ def __init__(self, *args, **kwargs):
     C2EDim,
     C2VDim,
     CellDim,
+    E2C2EDim,
     E2C2VDim,
     E2CDim,
     E2VDim,
@@ -45,7 +46,6 @@ def __init__(self, *args, **kwargs):
     V2E2VDim,
     V2EDim,
     VertexDim,
-    E2C2EDim,
 )
 from icon4py.model.common.grid.base import GridConfig, VerticalGridSize
 from icon4py.model.common.grid.horizontal import HorizontalGridSize
@@ -408,31 +408,24 @@ def _from_grid_dataset(self, dataset: Dataset) -> IconGrid:
 
         return icon_grid
 
-    def _construct_diamond_vertices(self, e2v: np.ndarray, c2v: np.ndarray, e2c: np.ndarray):
+    @staticmethod
+    def _construct_diamond_vertices(e2v: np.ndarray, c2v: np.ndarray, e2c: np.ndarray):
         n_edges = e2v.shape[0]
         e2c2v = -55 * np.ones((n_edges, 4), dtype=np.int32)
         e2c2v[:, 0:2] = e2v
-        # enlarge this table to contain an entry at the end for INVALID_VALUES
-        dummy_c2v = np.append(
-            c2v,
-            GridFile.INVALID_INDEX * np.ones((1, c2v.shape[1]), dtype=np.int32),
-            axis=0,
-        )
+        dummy_c2v = _patch_with_dummy_lastline(c2v)
         expanded = dummy_c2v[e2c[:, :], :]
         sh = expanded.shape
         flat = expanded.reshape(sh[0], sh[1] * sh[2])
-        # TODO (magdalena) vectorize speed this up!
+        # TODO (magdalena) vectorize speed this up?
         for x in range(sh[0]):
             e2c2v[x, 2:] = flat[x, ~np.in1d(flat[x, :], e2c2v[x, :])][:2]
 
         return e2c2v
 
-    def _construct_diamond_edges(self, e2c: np.ndarray, c2e: np.ndarray):
-        dummy_c2e = np.append(
-            c2e,
-            GridFile.INVALID_INDEX * np.ones((1, c2e.shape[1]), dtype=np.int32),
-            axis=0,
-        )
+    @staticmethod
+    def _construct_diamond_edges(e2c: np.ndarray, c2e: np.ndarray):
+        dummy_c2e = _patch_with_dummy_lastline(c2e)
         expanded = dummy_c2e[e2c, :]
         sh = expanded.shape
         flattened = expanded.reshape(sh[0], sh[1] * sh[2])
@@ -441,3 +434,24 @@ def _construct_diamond_edges(self, e2c: np.ndarray, c2e: np.ndarray):
             var = flattened[x, (~np.in1d(flattened[x, :], np.asarray([x, GridFile.INVALID_INDEX])))]
             e2c2e[x, : var.shape[0]] = var
         return e2c2e
+
+
+def _patch_with_dummy_lastline(ar):
+    """
+    Patch an array for easy access with an another offset containing invalid indices (-1).
+
+    Enlarges this table to contain a fake last line to account for numpy wrap around when
+    encountering a -1 = GridFile.INVALID_INDEX value
+
+    Args:
+        ar: np.ndarray connectivity array to be patched
+
+    Returns: same array with an additional line containing only GridFile.INVALID_INDEX
+
+    """
+    patched_ar = np.append(
+        ar,
+        GridFile.INVALID_INDEX * np.ones((1, ar.shape[1]), dtype=np.int32),
+        axis=0,
+    )
+    return patched_ar

From d1069f5c3973e563512c2fdaf6247aeadeaf24d0 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Tue, 9 Jan 2024 10:47:24 +0100
Subject: [PATCH 06/23] add E2C2EODim connectivity

---
 model/common/src/icon4py/model/common/grid/grid_manager.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py
index 7c580fb0e8..7d2416e4bc 100644
--- a/model/common/src/icon4py/model/common/grid/grid_manager.py
+++ b/model/common/src/icon4py/model/common/grid/grid_manager.py
@@ -38,6 +38,7 @@ def __init__(self, *args, **kwargs):
     C2VDim,
     CellDim,
     E2C2EDim,
+    E2C2EODim,
     E2C2VDim,
     E2CDim,
     E2VDim,
@@ -366,6 +367,7 @@ def _from_grid_dataset(self, dataset: Dataset) -> IconGrid:
 
         e2c2v = self._construct_diamond_vertices(e2v, c2v, e2c)
         e2c2e = self._construct_diamond_edges(e2c, c2e)
+        e2c2e0 = np.column_stack((e2c2e, np.asarray(range(e2c2e.shape[0]))))
 
         v2c = self._get_index_field(reader, GridFile.OffsetName.V2C)
         v2e = self._get_index_field(reader, GridFile.OffsetName.V2E)
@@ -399,6 +401,7 @@ def _from_grid_dataset(self, dataset: Dataset) -> IconGrid:
                     E2C2VDim: e2c2v,
                     V2E2VDim: v2e2v,
                     E2C2EDim: e2c2e,
+                    E2C2EODim: e2c2e0,
                 }
             )
             .with_start_end_indices(CellDim, start_indices[CellDim], end_indices[CellDim])

From b15c3cfd47fab96c008cfd6185c732dd234d55dc Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Tue, 9 Jan 2024 15:22:48 +0100
Subject: [PATCH 07/23] add C2V in offset provider mapping update README.md in
 model/common

---
 model/common/README.md                        | 35 +++++++++++++++++++
 .../src/icon4py/model/common/grid/icon.py     |  2 ++
 2 files changed, 37 insertions(+)

diff --git a/model/common/README.md b/model/common/README.md
index aa6a42598a..3db0ffd3d2 100644
--- a/model/common/README.md
+++ b/model/common/README.md
@@ -7,3 +7,38 @@ Utilities shared by several ICON4Py components.
 ## Installation instructions
 
 Check the `README.md` at the root of the `model` folder for installation instructions.
+
+
+## Contents
+### decomposition
+Contains infrastructure for parallel implementation of `icon4py/model`.
+`icon4py` uses [GHEX](https://github.com/ghex-org/GHEX) for halo exchanges. In order to run in parallel
+optional dependencies `mpi4py` and `ghex` need to be installed, which can be done through
+
+```bash
+pip install -r requirements-dev-opt.txt
+```
+
+
+### grid
+Contains basic infrastructure regarding the (unstructured) grid used in `icon4py`. There are
+two implementations of the general grid a small simple grid with periodic boundaries in
+[simple.py](src/icon4py/model/common/grid/simple.py) used for testing and the
+ICON grid [icon.py](src/icon4py/model/common/grid/icon.py) both implement the same protocl.
+The ICON grid can be initialized from an ICON grid file via the [grid_manager.py](src/icon4py/model/common/grid/grid_manager.py) 
+(THIS is still EXPERIMENTAL!!) or from serialized data.
+The `grid_manager.py` needs netcdf as an optional dependency, which can be installed with
+```bash
+pip install -r requirements-dev-opt.txt
+```
+
+### interpolation
+Contains interpolation stencils and port of interpolation fields in ICON.
+
+### math
+math utilities.
+### states
+contains type for the ICON prognostic state used by several packages.
+
+### test_utils
+Utilities used in tests made available here for usage in other packages
diff --git a/model/common/src/icon4py/model/common/grid/icon.py b/model/common/src/icon4py/model/common/grid/icon.py
index 5672ea9886..6b9ae1ac94 100644
--- a/model/common/src/icon4py/model/common/grid/icon.py
+++ b/model/common/src/icon4py/model/common/grid/icon.py
@@ -34,6 +34,7 @@
     V2CDim,
     V2EDim,
     VertexDim,
+    C2VDim,
 )
 from icon4py.model.common.grid.base import BaseGrid
 from icon4py.model.common.utils import builder
@@ -55,6 +56,7 @@ def __init__(self):
             "E2C2V": (self._get_offset_provider, E2C2VDim, EdgeDim, VertexDim),
             "V2E": (self._get_offset_provider, V2EDim, VertexDim, EdgeDim),
             "V2C": (self._get_offset_provider, V2CDim, VertexDim, CellDim),
+            "C2V": (self._get_offset_provider, C2VDim, CellDim, VertexDim),
             "E2ECV": (self._get_offset_provider_for_sparse_fields, E2C2VDim, EdgeDim, ECVDim),
             "C2CEC": (self._get_offset_provider_for_sparse_fields, C2E2CDim, CellDim, CECDim),
             "C2CE": (self._get_offset_provider_for_sparse_fields, C2EDim, CellDim, CEDim),

From 534042ecbe0a72878ed9e64edfe027fe76706c2e Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Fri, 12 Jan 2024 10:07:14 +0100
Subject: [PATCH 08/23] add tests for R02B04_R.nc

---
 model/common/tests/grid_tests/__init__.py     |   0
 model/common/tests/grid_tests/conftest.py     |  42 +----
 .../tests/grid_tests/test_grid_manager.py     | 162 ++++++++++++++----
 model/common/tests/grid_tests/utils.py        |  36 ++++
 4 files changed, 166 insertions(+), 74 deletions(-)
 create mode 100644 model/common/tests/grid_tests/__init__.py
 create mode 100644 model/common/tests/grid_tests/utils.py

diff --git a/model/common/tests/grid_tests/__init__.py b/model/common/tests/grid_tests/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/model/common/tests/grid_tests/conftest.py b/model/common/tests/grid_tests/conftest.py
index 639e8e1990..91f499c8ba 100644
--- a/model/common/tests/grid_tests/conftest.py
+++ b/model/common/tests/grid_tests/conftest.py
@@ -12,7 +12,6 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 import pytest
 
-from icon4py.model.common.test_utils.data_handling import download_and_extract
 from icon4py.model.common.test_utils.datatest_fixtures import (  # noqa: F401
     damping_height,
     data_provider,
@@ -26,40 +25,11 @@
     processor_props,
     ranked_data_path,
 )
-from icon4py.model.common.test_utils.datatest_utils import TEST_DATA_ROOT
-
-
-grids_path = TEST_DATA_ROOT.joinpath("grids")
-r04b09_dsl_grid_path = grids_path.joinpath("mch_ch_r04b09_dsl")
-r04b09_dsl_data_file = r04b09_dsl_grid_path.joinpath("mch_ch_r04b09_dsl_grids_v1.tar.gz").name
-r02b04_global_grid_path = grids_path.joinpath("r02b04_global")
-r02b04_global_data_file = r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_G.tar.gz").name
-
-
-mch_ch_r04b09_dsl_grid_uri = "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download"
-r02b04_global_grid_uri = "https://polybox.ethz.ch/index.php/s/0EM8O8U53GKGsst/download"
-
-
-@pytest.fixture()
-def r04b09_dsl_gridfile(get_grid_files):
-    return r04b09_dsl_grid_path.joinpath("grid.nc")
-
+from .utils import (
+    MCH_GRID_FILE,
+)
 
-@pytest.fixture(scope="session")
-def get_grid_files(pytestconfig):
-    """
-    Get the grid files used for testing.
 
-    Session scoped fixture which is a prerequisite of all the other fixtures in this file.
-    """
-    if not pytestconfig.getoption("datatest"):
-        pytest.skip("not running datatest marked tests")
-    download_and_extract(
-        mch_ch_r04b09_dsl_grid_uri, r04b09_dsl_grid_path, r04b09_dsl_grid_path, r04b09_dsl_data_file
-    )
-    download_and_extract(
-        r02b04_global_grid_uri,
-        r02b04_global_grid_path,
-        r02b04_global_grid_path,
-        r02b04_global_data_file,
-    )
+@pytest.fixture
+def grid_file():
+    return MCH_GRID_FILE
diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py
index 31e4af47e8..4500e22e79 100644
--- a/model/common/tests/grid_tests/test_grid_manager.py
+++ b/model/common/tests/grid_tests/test_grid_manager.py
@@ -19,6 +19,7 @@
 import numpy as np
 import pytest
 
+from icon4py.model.common.test_utils.datatest_utils import REGIONAL_EXPERIMENT, GLOBAL_EXPERIMENT
 
 if typing.TYPE_CHECKING:
     import netCDF4
@@ -52,6 +53,7 @@
 from icon4py.model.common.grid.simple import SimpleGrid
 from icon4py.model.common.grid.vertical import VerticalGridSize
 
+from .utils import MCH_GRID_FILE, R02B04_GLOBAL, resolve_file_from_gridfile_name
 
 SIMPLE_GRID_NC = "simple_grid.nc"
 
@@ -232,9 +234,14 @@ def test_gridparser_dimension(simple_grid_gridfile):
 
 @pytest.mark.datatest
 @pytest.mark.with_netcdf
-def test_gridfile_vertex_cell_edge_dimensions(grid_savepoint, r04b09_dsl_gridfile):
-    data = netCDF4.Dataset(r04b09_dsl_gridfile, "r")
-    grid_file = GridFile(data)
+@pytest.mark.parametrize(
+    "grid_file, experiment",
+    [(MCH_GRID_FILE, REGIONAL_EXPERIMENT), (R02B04_GLOBAL, GLOBAL_EXPERIMENT)],
+)
+def test_gridfile_vertex_cell_edge_dimensions(grid_savepoint, grid_file):
+    file = resolve_file_from_gridfile_name(grid_file)
+    dataset = netCDF4.Dataset(file, "r")
+    grid_file = GridFile(dataset)
 
     assert grid_file.dimension(GridFile.DimensionName.CELL_NAME) == grid_savepoint.num(CellDim)
     assert grid_file.dimension(GridFile.DimensionName.EDGE_NAME) == grid_savepoint.num(EdgeDim)
@@ -261,25 +268,36 @@ def test_grid_parser_index_fields(simple_grid_gridfile, caplog):
 # v2e: exists in serial, simple, grid
 @pytest.mark.datatest
 @pytest.mark.with_netcdf
-def test_gridmanager_eval_v2e(caplog, grid_savepoint, r04b09_dsl_gridfile):
+@pytest.mark.parametrize(
+    "grid_file, experiment",
+    [(MCH_GRID_FILE, REGIONAL_EXPERIMENT), (R02B04_GLOBAL, GLOBAL_EXPERIMENT)],
+)
+def test_gridmanager_eval_v2e(caplog, grid_savepoint, grid_file):
     caplog.set_level(logging.DEBUG)
-    grid = init_grid_manager(r04b09_dsl_gridfile).get_grid()
+    file = resolve_file_from_gridfile_name(grid_file)
+    grid = init_grid_manager(file).get_grid()
     seralized_v2e = grid_savepoint.v2e()[0 : grid.num_vertices, :]
     # there are vertices at the boundary of a local domain or at a pentagon point that have less than
     # 6 neighbors hence there are "Missing values" in the grid file
     # they get substituted by the "last valid index" in preprocessing step in icon.
     assert not has_invalid_index(seralized_v2e)
-    assert has_invalid_index(grid.get_offset_provider("V2E").table)
+    v2e_table = grid.get_offset_provider("V2E").table
+    assert has_invalid_index(v2e_table)
     reset_invalid_index(seralized_v2e)
-    assert np.allclose(grid.get_offset_provider("V2E").table, seralized_v2e)
+    assert np.allclose(v2e_table, seralized_v2e)
 
 
 # v2c: exists in serial, simple, grid
 @pytest.mark.datatest
 @pytest.mark.with_netcdf
-def test_gridmanager_eval_v2c(caplog, grid_savepoint, r04b09_dsl_gridfile):
+@pytest.mark.parametrize(
+    "grid_file, experiment",
+    [(MCH_GRID_FILE, REGIONAL_EXPERIMENT), (R02B04_GLOBAL, GLOBAL_EXPERIMENT)],
+)
+def test_gridmanager_eval_v2c(caplog, grid_savepoint, grid_file):
     caplog.set_level(logging.DEBUG)
-    grid = init_grid_manager(r04b09_dsl_gridfile).get_grid()
+    file = resolve_file_from_gridfile_name(grid_file)
+    grid = init_grid_manager(file).get_grid()
     serialized_v2c = grid_savepoint.v2c()[0 : grid.num_vertices, :]
     # there are vertices that have less than 6 neighboring cells: either pentagon points or
     # vertices at the boundary of the domain for a limited area mode
@@ -322,9 +340,14 @@ def reset_invalid_index(index_array: np.ndarray):
 # e2v: exists in serial, simple, grid
 @pytest.mark.datatest
 @pytest.mark.with_netcdf
-def test_gridmanager_eval_e2v(caplog, grid_savepoint, r04b09_dsl_gridfile):
+@pytest.mark.parametrize(
+    "grid_file, experiment",
+    [(MCH_GRID_FILE, REGIONAL_EXPERIMENT), (R02B04_GLOBAL, GLOBAL_EXPERIMENT)],
+)
+def test_gridmanager_eval_e2v(caplog, grid_savepoint, grid_file):
     caplog.set_level(logging.DEBUG)
-    grid = init_grid_manager(r04b09_dsl_gridfile).get_grid()
+    file = resolve_file_from_gridfile_name(grid_file)
+    grid = init_grid_manager(file).get_grid()
 
     serialized_e2v = grid_savepoint.e2v()[0 : grid.num_edges, :]
     # all vertices in the system have to neighboring edges, there no edges that point nowhere
@@ -342,27 +365,62 @@ def invalid_index(ar):
     return np.where(ar == GridFile.INVALID_INDEX)
 
 
+def _is_local(grid_file: str):
+    return grid_file == MCH_GRID_FILE
+
+
+def assert_invalid_indices(e2c_table: np.ndarray, grid_file: str):
+    """
+    Assert invalid indices for E2C connectivity.
+
+    Local grids: there are edges at the boundary that have only one
+    neighboring cell, there are "missing values" in the grid file
+    and for E2C they do not get substituted in the ICON preprocessing.
+
+    Global grids have no "missing values" indices since all edges always have 2 neighboring cells.
+
+    Args:
+
+        e2c_table: E2C connectivity
+        grid_file: name of grid file used
+
+
+    """
+    if _is_local(grid_file):
+        assert has_invalid_index(e2c_table)
+    else:
+        assert not has_invalid_index(e2c_table)
+
+
 # e2c : exists in serial, simple, grid
 @pytest.mark.datatest
 @pytest.mark.with_netcdf
-def test_gridmanager_eval_e2c(caplog, grid_savepoint, r04b09_dsl_gridfile):
+@pytest.mark.parametrize(
+    "grid_file, experiment",
+    [(MCH_GRID_FILE, REGIONAL_EXPERIMENT), (R02B04_GLOBAL, GLOBAL_EXPERIMENT)],
+)
+def test_gridmanager_eval_e2c(caplog, grid_savepoint, grid_file):
     caplog.set_level(logging.DEBUG)
-    grid = init_grid_manager(r04b09_dsl_gridfile).get_grid()
+    file = resolve_file_from_gridfile_name(grid_file)
+    grid = init_grid_manager(file).get_grid()
     serialized_e2c = grid_savepoint.e2c()[0 : grid.num_edges, :]
-    # there are edges at the boundary that have only one
-    # neighboring cell, there are "missing values" in the grid file
-    # and here they do not get substituted in the ICON preprocessing
-    assert has_invalid_index(serialized_e2c)
-    assert has_invalid_index(grid.get_offset_provider("E2C").table)
-    assert np.allclose(grid.get_offset_provider("E2C").table, serialized_e2c)
+    e2c_table = grid.get_offset_provider("E2C").table
+    assert_invalid_indices(serialized_e2c, grid_file)
+    assert_invalid_indices(e2c_table, grid_file)
+    assert np.allclose(e2c_table, serialized_e2c)
 
 
 # c2e: serial, simple, grid
 @pytest.mark.datatest
 @pytest.mark.with_netcdf
-def test_gridmanager_eval_c2e(caplog, grid_savepoint, r04b09_dsl_gridfile):
+@pytest.mark.parametrize(
+    "grid_file, experiment",
+    [(MCH_GRID_FILE, REGIONAL_EXPERIMENT), (R02B04_GLOBAL, GLOBAL_EXPERIMENT)],
+)
+def test_gridmanager_eval_c2e(caplog, grid_savepoint, grid_file):
     caplog.set_level(logging.DEBUG)
-    grid = init_grid_manager(r04b09_dsl_gridfile).get_grid()
+    file = resolve_file_from_gridfile_name(grid_file)
+    grid = init_grid_manager(file).get_grid()
 
     serialized_c2e = grid_savepoint.c2e()[0 : grid.num_cells, :]
     # no cells with less than 3 neighboring edges exist, otherwise the cell is not there in the
@@ -376,9 +434,14 @@ def test_gridmanager_eval_c2e(caplog, grid_savepoint, r04b09_dsl_gridfile):
 # c2e2c: exists in  serial, simple_mesh, grid
 @pytest.mark.datatest
 @pytest.mark.with_netcdf
-def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile):
+@pytest.mark.parametrize(
+    "grid_file, experiment",
+    [(MCH_GRID_FILE, REGIONAL_EXPERIMENT), (R02B04_GLOBAL, GLOBAL_EXPERIMENT)],
+)
+def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, grid_file):
     caplog.set_level(logging.DEBUG)
-    grid = init_grid_manager(r04b09_dsl_gridfile).get_grid()
+    file = resolve_file_from_gridfile_name(grid_file)
+    grid = init_grid_manager(file).get_grid()
     assert np.allclose(
         grid.get_offset_provider("C2E2C").table,
         grid_savepoint.c2e2c()[0 : grid.num_cells, :],
@@ -388,14 +451,20 @@ def test_gridmanager_eval_c2e2c(caplog, grid_savepoint, r04b09_dsl_gridfile):
 # e2c2e (e2c2eo) - diamond: exists in serial, simple_mesh
 @pytest.mark.datatest
 @pytest.mark.with_netcdf
-def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, r04b09_dsl_gridfile):
+@pytest.mark.parametrize(
+    "grid_file, experiment",
+    [(MCH_GRID_FILE, REGIONAL_EXPERIMENT), (R02B04_GLOBAL, GLOBAL_EXPERIMENT)],
+)
+def test_gridmanager_eval_e2c2e(caplog, grid_savepoint, grid_file):
     caplog.set_level(logging.DEBUG)
-    grid = init_grid_manager(r04b09_dsl_gridfile).get_grid()
-    serialized_e2c2e = grid_savepoint.e2c2e()[0 : grid.num_edges, :]
-    assert has_invalid_index(serialized_e2c2e)
+    file = resolve_file_from_gridfile_name(grid_file)
+    grid = init_grid_manager(file).get_grid()
+    serialized_e2c2e = grid_savepoint.e2c2e()[0 : grid.num_cells, :]
+    assert_invalid_indices(serialized_e2c2e, grid_file)
 
     e2c2e_table = grid.get_offset_provider("E2C2E").table
-    assert has_invalid_index(e2c2e_table)
+
+    assert_invalid_indices(e2c2e_table, grid_file)
     # ICON calculates diamond edges only from rl_start = 2 (lateral_boundary(EdgeDim) + 1 for
     # boundaries all values are INVALID even though the half diamond exists (see mo_model_domimp_setup.f90 ll 163ff.)
     assert_unless_invalid(e2c2e_table, serialized_e2c2e)
@@ -408,9 +477,14 @@ def assert_unless_invalid(table, serialized_ref):
 
 @pytest.mark.datatest
 @pytest.mark.with_netcdf
-def test_gridmanager_eval_e2c2v(caplog, grid_savepoint, r04b09_dsl_gridfile):
+@pytest.mark.parametrize(
+    "grid_file, experiment",
+    [(MCH_GRID_FILE, REGIONAL_EXPERIMENT), (R02B04_GLOBAL, GLOBAL_EXPERIMENT)],
+)
+def test_gridmanager_eval_e2c2v(caplog, grid_savepoint, grid_file):
     caplog.set_level(logging.DEBUG)
-    grid = init_grid_manager(r04b09_dsl_gridfile).get_grid()
+    file = resolve_file_from_gridfile_name(grid_file)
+    grid = init_grid_manager(file).get_grid()
     # the "far" (adjacent to edge normal ) is not always there, because ICON only calculates those starting from
     #   (lateral_boundary(EdgeDim) + 1) to end(EdgeDim)  (see mo_intp_coeffs.f90) and only for owned cells
     serialized_ref = grid_savepoint.e2c2v()[: grid.num_edges, :]
@@ -420,15 +494,20 @@ def test_gridmanager_eval_e2c2v(caplog, grid_savepoint, r04b09_dsl_gridfile):
 
 @pytest.mark.datatest
 @pytest.mark.with_netcdf
-def test_gridmanager_eval_c2v(caplog, grid_savepoint, r04b09_dsl_gridfile):
+@pytest.mark.parametrize(
+    "grid_file, experiment",
+    [(MCH_GRID_FILE, REGIONAL_EXPERIMENT), (R02B04_GLOBAL, GLOBAL_EXPERIMENT)],
+)
+def test_gridmanager_eval_c2v(caplog, grid_savepoint, grid_file):
     caplog.set_level(logging.DEBUG)
-    grid = init_grid_manager(r04b09_dsl_gridfile).get_grid()
+    file = resolve_file_from_gridfile_name(grid_file)
+    grid = init_grid_manager(file).get_grid()
     c2v = grid.get_offset_provider("C2V").table
     assert np.allclose(c2v, grid_savepoint.c2v()[0 : grid.num_cells, :])
 
 
-def init_grid_manager(fname):
-    grid_manager = GridManager(ToGt4PyTransformation(), fname, VerticalGridSize(65))
+def init_grid_manager(fname, num_levels=65):
+    grid_manager = GridManager(ToGt4PyTransformation(), fname, VerticalGridSize(num_levels))
     grid_manager()
     return grid_manager
 
@@ -482,6 +561,7 @@ def test_gt4py_transform_offset_by_1_where_valid(size):
     assert np.allclose(expected, offset)
 
 
+# TODO (magdalena) speed this up...
 @pytest.mark.datatest
 @pytest.mark.with_netcdf
 @pytest.mark.parametrize(
@@ -521,8 +601,9 @@ def test_gt4py_transform_offset_by_1_where_valid(size):
         (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 0),
     ],
 )
-def test_get_start_index(r04b09_dsl_gridfile, icon_grid, dim, marker, index):
-    grid_from_manager = init_grid_manager(r04b09_dsl_gridfile).get_grid()
+def test_get_start_index(icon_grid, dim, marker, index):
+    file = resolve_file_from_gridfile_name(MCH_GRID_FILE)
+    grid_from_manager = init_grid_manager(file).get_grid()
     assert grid_from_manager.get_start_index(dim, marker) == index
     assert grid_from_manager.get_start_index(dim, marker) == icon_grid.get_start_index(dim, marker)
 
@@ -570,7 +651,12 @@ def test_get_start_index(r04b09_dsl_gridfile, icon_grid, dim, marker, index):
         (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 428),
     ],
 )
-def test_get_end_index(r04b09_dsl_gridfile, icon_grid, dim, marker, index):
-    grid_from_manager = init_grid_manager(r04b09_dsl_gridfile).get_grid()
+@pytest.mark.parametrize(
+    "grid_file, experiment, num_levels", [(MCH_GRID_FILE, REGIONAL_EXPERIMENT, 65)]
+)
+def test_get_end_index_for_local_grid(grid_file, num_levels, grid_savepoint, dim, marker, index):
+    file = resolve_file_from_gridfile_name(grid_file)
+    icon_grid = grid_savepoint.construct_icon_grid()
+    grid_from_manager = init_grid_manager(file, num_levels=num_levels).get_grid()
     assert grid_from_manager.get_end_index(dim, marker) == index
     assert grid_from_manager.get_end_index(dim, marker) == icon_grid.get_end_index(dim, marker)
diff --git a/model/common/tests/grid_tests/utils.py b/model/common/tests/grid_tests/utils.py
new file mode 100644
index 0000000000..f100133328
--- /dev/null
+++ b/model/common/tests/grid_tests/utils.py
@@ -0,0 +1,36 @@
+from icon4py.model.common.test_utils.data_handling import download_and_extract
+from icon4py.model.common.test_utils.datatest_utils import TEST_DATA_ROOT
+
+MCH_GRID_FILE = "mch_ch_r04b09_dsl"
+R02B04_GLOBAL = "r02b04_global"
+
+grids_path = TEST_DATA_ROOT.joinpath("grids")
+r04b09_dsl_grid_path = grids_path.joinpath(MCH_GRID_FILE)
+r04b09_dsl_data_file = r04b09_dsl_grid_path.joinpath("mch_ch_r04b09_dsl_grids_v1.tar.gz").name
+r02b04_global_grid_path = grids_path.joinpath(R02B04_GLOBAL)
+r02b04_global_data_file = r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_R.tar.gz").name
+
+
+mch_ch_r04b09_dsl_grid_uri = "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download"
+r02b04_global_grid_uri = "https://polybox.ethz.ch/index.php/s/AKAO6ImQdIatnkB/download"
+
+
+def resolve_file_from_gridfile_name(name: str):
+    if name == MCH_GRID_FILE:
+        if not r04b09_dsl_grid_path.joinpath("grid.nc").exists():
+            download_and_extract(
+                mch_ch_r04b09_dsl_grid_uri,
+                r04b09_dsl_grid_path,
+                r04b09_dsl_grid_path,
+                r04b09_dsl_data_file,
+            )
+        return r04b09_dsl_grid_path.joinpath("grid.nc")
+    elif name == R02B04_GLOBAL:
+        if not r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_R.nc").exists():
+            download_and_extract(
+                r02b04_global_grid_uri,
+                r02b04_global_grid_path,
+                r02b04_global_grid_path,
+                r02b04_global_data_file,
+            )
+        return r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_R.nc")

From 2341106cd458d422b8d32e63e1f8f3cb1a901bb2 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Wed, 17 Jan 2024 08:38:14 +0100
Subject: [PATCH 09/23] add test for domain bounds for global grid file

---
 .../icon4py/model/common/grid/horizontal.py   |  2 +
 .../tests/grid_tests/test_grid_manager.py     | 41 +++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/model/common/src/icon4py/model/common/grid/horizontal.py b/model/common/src/icon4py/model/common/grid/horizontal.py
index 6eb62ea959..e606656d2e 100644
--- a/model/common/src/icon4py/model/common/grid/horizontal.py
+++ b/model/common/src/icon4py/model/common/grid/horizontal.py
@@ -102,6 +102,7 @@ class HorizontalMarkerIndex:
     _nudging = {
         dimension.CellDim: _NUDGING_CELLS,
         dimension.EdgeDim: _NUDGING_EDGES,
+        # TODO [magdalena] there is no nudging for vertices?
         dimension.VertexDim: _NUDGING_VERTICES,
     }
     _end = {
@@ -131,6 +132,7 @@ def halo(cls, dim: Dimension) -> int:
     @classmethod
     def nudging(cls, dim: Dimension) -> int:
         """Indicate the nudging zone."""
+
         return cls._nudging[dim]
 
     @classmethod
diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py
index 4500e22e79..0e717a4244 100644
--- a/model/common/tests/grid_tests/test_grid_manager.py
+++ b/model/common/tests/grid_tests/test_grid_manager.py
@@ -660,3 +660,44 @@ def test_get_end_index_for_local_grid(grid_file, num_levels, grid_savepoint, dim
     grid_from_manager = init_grid_manager(file, num_levels=num_levels).get_grid()
     assert grid_from_manager.get_end_index(dim, marker) == index
     assert grid_from_manager.get_end_index(dim, marker) == icon_grid.get_end_index(dim, marker)
+
+
+@pytest.mark.datatest
+@pytest.mark.with_netcdf
+@pytest.mark.parametrize(
+    "dim, marker,start_index,end_index",
+    [
+        (CellDim, HorizontalMarkerIndex.interior(CellDim), 0, 0),
+        (CellDim, HorizontalMarkerIndex.local(CellDim), 0, 20480),
+        (CellDim, HorizontalMarkerIndex.nudging(CellDim), 0, 0),
+        (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim), 0, 0),
+        (CellDim, HorizontalMarkerIndex.end(CellDim), 20480, 20480),
+        (CellDim, HorizontalMarkerIndex.halo(CellDim), 20480, 20480),
+        (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 0, 0),
+        (EdgeDim, HorizontalMarkerIndex.local(EdgeDim), 0, 30720),
+        (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 0, 0),
+        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim), 0, 0),
+        (EdgeDim, HorizontalMarkerIndex.end(EdgeDim), 30720, 30720),
+        (EdgeDim, HorizontalMarkerIndex.halo(EdgeDim), 30720, 30720),
+        (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 0, 0),
+        (VertexDim, HorizontalMarkerIndex.local(VertexDim), 0, 10242),
+        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim), 0, 0),
+        (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10242, 10242),
+        (VertexDim, HorizontalMarkerIndex.halo(VertexDim), 10242, 10242),
+    ],
+)
+@pytest.mark.parametrize(
+    "grid_file, experiment, num_levels", [(R02B04_GLOBAL, GLOBAL_EXPERIMENT, 80)]
+)
+def test_get_start_end_index_for_global_grid(
+    grid_file, num_levels, grid_savepoint, dim, marker, start_index, end_index
+):
+    file = resolve_file_from_gridfile_name(grid_file)
+    from_savepoint = grid_savepoint.construct_icon_grid()
+    from_grid_file = init_grid_manager(file, num_levels=num_levels).get_grid()
+    assert from_grid_file.get_start_index(dim, marker) == start_index
+    assert from_grid_file.get_start_index(dim, marker) == from_savepoint.get_start_index(
+        dim, marker
+    )
+    assert from_grid_file.get_end_index(dim, marker) == end_index
+    assert from_grid_file.get_end_index(dim, marker) == from_savepoint.get_end_index(dim, marker)

From e0050134cedc764bc3a0ce2c18f0235115782597 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Wed, 17 Jan 2024 10:11:57 +0100
Subject: [PATCH 10/23] use cached gridfile in tests

---
 .../model/atmosphere/diffusion/diffusion.py      |  1 +
 .../common/tests/grid_tests/test_grid_manager.py | 16 ++++++++++------
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py
index 4b4e027d42..1aaec7c817 100644
--- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py
+++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py
@@ -466,6 +466,7 @@ def _get_start_index_for_w_diffusion() -> int32:
             offset_provider={"Koff": KDim},
         )
 
+        # TODO (magdalena) port to gt4py?
         self.diff_multfac_n2w = init_nabla2_factor_in_upper_damping_zone(
             k_size=self.grid.num_levels,
             nshift=0,
diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py
index 0e717a4244..5edde5b8a0 100644
--- a/model/common/tests/grid_tests/test_grid_manager.py
+++ b/model/common/tests/grid_tests/test_grid_manager.py
@@ -12,6 +12,7 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 from __future__ import annotations
 
+import functools
 import logging
 import typing
 from uuid import uuid4
@@ -506,8 +507,9 @@ def test_gridmanager_eval_c2v(caplog, grid_savepoint, grid_file):
     assert np.allclose(c2v, grid_savepoint.c2v()[0 : grid.num_cells, :])
 
 
-def init_grid_manager(fname, num_levels=65):
-    grid_manager = GridManager(ToGt4PyTransformation(), fname, VerticalGridSize(num_levels))
+@functools.cache
+def init_grid_manager(fname, num_levels=65, transformation=ToGt4PyTransformation()):
+    grid_manager = GridManager(transformation, fname, VerticalGridSize(num_levels))
     grid_manager()
     return grid_manager
 
@@ -516,7 +518,9 @@ def init_grid_manager(fname, num_levels=65):
 @pytest.mark.with_netcdf
 def test_grid_manager_getsize(simple_grid_gridfile, dim, size, caplog):
     caplog.set_level(logging.DEBUG)
-    gm = GridManager(IndexTransformation(), simple_grid_gridfile, VerticalGridSize(num_lev=80))
+    gm = init_grid_manager(
+        simple_grid_gridfile, num_levels=10, transformation=IndexTransformation()
+    )
     gm()
     assert size == gm.get_size(dim)
 
@@ -530,10 +534,10 @@ def assert_up_to_order(table, diamond_table):
 @pytest.mark.with_netcdf
 def test_grid_manager_diamond_offset(simple_grid_gridfile):
     simple_grid = SimpleGrid()
-    gm = GridManager(
-        IndexTransformation(),
+    gm = init_grid_manager(
         simple_grid_gridfile,
-        VerticalGridSize(num_lev=simple_grid.num_levels),
+        num_levels=simple_grid.num_levels,
+        transformation=IndexTransformation(),
     )
     gm()
     icon_grid = gm.get_grid()

From 603e33dd9d7aba6c4e930d8c08098132946bce1e Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Thu, 18 Jan 2024 09:40:59 +0100
Subject: [PATCH 11/23] pre-commit fixes

---
 model/common/README.md                        | 13 ++++++--
 .../icon4py/model/common/grid/horizontal.py   |  1 -
 .../src/icon4py/model/common/grid/icon.py     |  2 +-
 model/common/tests/grid_tests/__init__.py     | 12 ++++++++
 model/common/tests/grid_tests/conftest.py     |  6 ++--
 .../tests/grid_tests/test_grid_manager.py     | 30 ++++++++++---------
 .../common/tests/grid_tests/test_icon_grid.py |  6 ++--
 model/common/tests/grid_tests/utils.py        | 14 +++++++++
 8 files changed, 59 insertions(+), 25 deletions(-)

diff --git a/model/common/README.md b/model/common/README.md
index 3db0ffd3d2..1a81a07a85 100644
--- a/model/common/README.md
+++ b/model/common/README.md
@@ -8,9 +8,10 @@ Utilities shared by several ICON4Py components.
 
 Check the `README.md` at the root of the `model` folder for installation instructions.
 
-
 ## Contents
+
 ### decomposition
+
 Contains infrastructure for parallel implementation of `icon4py/model`.
 `icon4py` uses [GHEX](https://github.com/ghex-org/GHEX) for halo exchanges. In order to run in parallel
 optional dependencies `mpi4py` and `ghex` need to be installed, which can be done through
@@ -19,26 +20,32 @@ optional dependencies `mpi4py` and `ghex` need to be installed, which can be don
 pip install -r requirements-dev-opt.txt
 ```
 
-
 ### grid
+
 Contains basic infrastructure regarding the (unstructured) grid used in `icon4py`. There are
 two implementations of the general grid a small simple grid with periodic boundaries in
 [simple.py](src/icon4py/model/common/grid/simple.py) used for testing and the
 ICON grid [icon.py](src/icon4py/model/common/grid/icon.py) both implement the same protocl.
-The ICON grid can be initialized from an ICON grid file via the [grid_manager.py](src/icon4py/model/common/grid/grid_manager.py) 
+The ICON grid can be initialized from an ICON grid file via the [grid_manager.py](src/icon4py/model/common/grid/grid_manager.py)
 (THIS is still EXPERIMENTAL!!) or from serialized data.
 The `grid_manager.py` needs netcdf as an optional dependency, which can be installed with
+
 ```bash
 pip install -r requirements-dev-opt.txt
 ```
 
 ### interpolation
+
 Contains interpolation stencils and port of interpolation fields in ICON.
 
 ### math
+
 math utilities.
+
 ### states
+
 contains type for the ICON prognostic state used by several packages.
 
 ### test_utils
+
 Utilities used in tests made available here for usage in other packages
diff --git a/model/common/src/icon4py/model/common/grid/horizontal.py b/model/common/src/icon4py/model/common/grid/horizontal.py
index e606656d2e..c1db28f286 100644
--- a/model/common/src/icon4py/model/common/grid/horizontal.py
+++ b/model/common/src/icon4py/model/common/grid/horizontal.py
@@ -132,7 +132,6 @@ def halo(cls, dim: Dimension) -> int:
     @classmethod
     def nudging(cls, dim: Dimension) -> int:
         """Indicate the nudging zone."""
-
         return cls._nudging[dim]
 
     @classmethod
diff --git a/model/common/src/icon4py/model/common/grid/icon.py b/model/common/src/icon4py/model/common/grid/icon.py
index 6b9ae1ac94..c99ad6c93a 100644
--- a/model/common/src/icon4py/model/common/grid/icon.py
+++ b/model/common/src/icon4py/model/common/grid/icon.py
@@ -19,6 +19,7 @@
     C2E2CDim,
     C2E2CODim,
     C2EDim,
+    C2VDim,
     CECDim,
     CEDim,
     CellDim,
@@ -34,7 +35,6 @@
     V2CDim,
     V2EDim,
     VertexDim,
-    C2VDim,
 )
 from icon4py.model.common.grid.base import BaseGrid
 from icon4py.model.common.utils import builder
diff --git a/model/common/tests/grid_tests/__init__.py b/model/common/tests/grid_tests/__init__.py
index e69de29bb2..15dfdb0098 100644
--- a/model/common/tests/grid_tests/__init__.py
+++ b/model/common/tests/grid_tests/__init__.py
@@ -0,0 +1,12 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
diff --git a/model/common/tests/grid_tests/conftest.py b/model/common/tests/grid_tests/conftest.py
index 91f499c8ba..ecfc2301d0 100644
--- a/model/common/tests/grid_tests/conftest.py
+++ b/model/common/tests/grid_tests/conftest.py
@@ -20,14 +20,12 @@
     download_ser_data,
     experiment,
     grid_savepoint,
-    icon_grid,
     interpolation_savepoint,
     processor_props,
     ranked_data_path,
 )
-from .utils import (
-    MCH_GRID_FILE,
-)
+
+from .utils import MCH_GRID_FILE
 
 
 @pytest.fixture
diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py
index 5edde5b8a0..eb3293f180 100644
--- a/model/common/tests/grid_tests/test_grid_manager.py
+++ b/model/common/tests/grid_tests/test_grid_manager.py
@@ -20,7 +20,8 @@
 import numpy as np
 import pytest
 
-from icon4py.model.common.test_utils.datatest_utils import REGIONAL_EXPERIMENT, GLOBAL_EXPERIMENT
+from icon4py.model.common.test_utils.datatest_utils import GLOBAL_EXPERIMENT, REGIONAL_EXPERIMENT
+
 
 if typing.TYPE_CHECKING:
     import netCDF4
@@ -56,6 +57,7 @@
 
 from .utils import MCH_GRID_FILE, R02B04_GLOBAL, resolve_file_from_gridfile_name
 
+
 SIMPLE_GRID_NC = "simple_grid.nc"
 
 
@@ -381,11 +383,9 @@ def assert_invalid_indices(e2c_table: np.ndarray, grid_file: str):
     Global grids have no "missing values" indices since all edges always have 2 neighboring cells.
 
     Args:
-
         e2c_table: E2C connectivity
         grid_file: name of grid file used
 
-
     """
     if _is_local(grid_file):
         assert has_invalid_index(e2c_table)
@@ -565,11 +565,10 @@ def test_gt4py_transform_offset_by_1_where_valid(size):
     assert np.allclose(expected, offset)
 
 
-# TODO (magdalena) speed this up...
 @pytest.mark.datatest
 @pytest.mark.with_netcdf
 @pytest.mark.parametrize(
-    "dim, marker, index",
+    "dim, marker, start_index",
     [
         (CellDim, HorizontalMarkerIndex.interior(CellDim), 4104),
         (CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, 0),
@@ -605,11 +604,14 @@ def test_gt4py_transform_offset_by_1_where_valid(size):
         (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 0),
     ],
 )
-def test_get_start_index(icon_grid, dim, marker, index):
+def test_get_start_index_for_local_grid(grid_savepoint, dim, marker, start_index):
     file = resolve_file_from_gridfile_name(MCH_GRID_FILE)
-    grid_from_manager = init_grid_manager(file).get_grid()
-    assert grid_from_manager.get_start_index(dim, marker) == index
-    assert grid_from_manager.get_start_index(dim, marker) == icon_grid.get_start_index(dim, marker)
+    from_savepoint = grid_savepoint.construct_icon_grid()
+    from_grid_file = init_grid_manager(file).get_grid()
+    assert from_grid_file.get_start_index(dim, marker) == start_index
+    assert from_grid_file.get_start_index(dim, marker) == from_savepoint.get_start_index(
+        dim, marker
+    )
 
 
 @pytest.mark.datatest
@@ -660,10 +662,10 @@ def test_get_start_index(icon_grid, dim, marker, index):
 )
 def test_get_end_index_for_local_grid(grid_file, num_levels, grid_savepoint, dim, marker, index):
     file = resolve_file_from_gridfile_name(grid_file)
-    icon_grid = grid_savepoint.construct_icon_grid()
-    grid_from_manager = init_grid_manager(file, num_levels=num_levels).get_grid()
-    assert grid_from_manager.get_end_index(dim, marker) == index
-    assert grid_from_manager.get_end_index(dim, marker) == icon_grid.get_end_index(dim, marker)
+    from_savepoint = grid_savepoint.construct_icon_grid()
+    from_grid_file = init_grid_manager(file, num_levels=num_levels).get_grid()
+    assert from_grid_file.get_end_index(dim, marker) == index
+    assert from_grid_file.get_end_index(dim, marker) == from_savepoint.get_end_index(dim, marker)
 
 
 @pytest.mark.datatest
@@ -696,8 +698,8 @@ def test_get_end_index_for_local_grid(grid_file, num_levels, grid_savepoint, dim
 def test_get_start_end_index_for_global_grid(
     grid_file, num_levels, grid_savepoint, dim, marker, start_index, end_index
 ):
-    file = resolve_file_from_gridfile_name(grid_file)
     from_savepoint = grid_savepoint.construct_icon_grid()
+    file = resolve_file_from_gridfile_name(grid_file)
     from_grid_file = init_grid_manager(file, num_levels=num_levels).get_grid()
     assert from_grid_file.get_start_index(dim, marker) == start_index
     assert from_grid_file.get_start_index(dim, marker) == from_savepoint.get_start_index(
diff --git a/model/common/tests/grid_tests/test_icon_grid.py b/model/common/tests/grid_tests/test_icon_grid.py
index 8546a44968..496266e5e9 100644
--- a/model/common/tests/grid_tests/test_icon_grid.py
+++ b/model/common/tests/grid_tests/test_icon_grid.py
@@ -60,7 +60,8 @@
         (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 428),
     ],
 )
-def test_horizontal_end_index(icon_grid, dim, marker, index):
+def test_horizontal_end_index(grid_savepoint, dim, marker, index):
+    icon_grid = grid_savepoint.construct_icon_grid()
     assert index == icon_grid.get_end_index(dim, marker)
 
 
@@ -106,7 +107,8 @@ def test_horizontal_end_index(icon_grid, dim, marker, index):
         (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 0),
     ],
 )
-def test_horizontal_start_index(icon_grid, dim, marker, index):
+def test_horizontal_start_index(grid_savepoint, dim, marker, index):
+    icon_grid = grid_savepoint.construct_icon_grid()
     assert index == icon_grid.get_start_index(dim, marker)
 
 
diff --git a/model/common/tests/grid_tests/utils.py b/model/common/tests/grid_tests/utils.py
index f100133328..423e7b229f 100644
--- a/model/common/tests/grid_tests/utils.py
+++ b/model/common/tests/grid_tests/utils.py
@@ -1,6 +1,20 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
 from icon4py.model.common.test_utils.data_handling import download_and_extract
 from icon4py.model.common.test_utils.datatest_utils import TEST_DATA_ROOT
 
+
 MCH_GRID_FILE = "mch_ch_r04b09_dsl"
 R02B04_GLOBAL = "r02b04_global"
 

From 3d73048ba31cdbc09e5c1a97ad7caafb6a1105c7 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Thu, 18 Jan 2024 16:14:17 +0100
Subject: [PATCH 12/23] revert changes in icon_grid

---
 model/common/tests/grid_tests/conftest.py       | 1 +
 model/common/tests/grid_tests/test_icon_grid.py | 6 ++----
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/model/common/tests/grid_tests/conftest.py b/model/common/tests/grid_tests/conftest.py
index ecfc2301d0..5df1882b54 100644
--- a/model/common/tests/grid_tests/conftest.py
+++ b/model/common/tests/grid_tests/conftest.py
@@ -20,6 +20,7 @@
     download_ser_data,
     experiment,
     grid_savepoint,
+    icon_grid,
     interpolation_savepoint,
     processor_props,
     ranked_data_path,
diff --git a/model/common/tests/grid_tests/test_icon_grid.py b/model/common/tests/grid_tests/test_icon_grid.py
index 496266e5e9..8546a44968 100644
--- a/model/common/tests/grid_tests/test_icon_grid.py
+++ b/model/common/tests/grid_tests/test_icon_grid.py
@@ -60,8 +60,7 @@
         (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 428),
     ],
 )
-def test_horizontal_end_index(grid_savepoint, dim, marker, index):
-    icon_grid = grid_savepoint.construct_icon_grid()
+def test_horizontal_end_index(icon_grid, dim, marker, index):
     assert index == icon_grid.get_end_index(dim, marker)
 
 
@@ -107,8 +106,7 @@ def test_horizontal_end_index(grid_savepoint, dim, marker, index):
         (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 0),
     ],
 )
-def test_horizontal_start_index(grid_savepoint, dim, marker, index):
-    icon_grid = grid_savepoint.construct_icon_grid()
+def test_horizontal_start_index(icon_grid, dim, marker, index):
     assert index == icon_grid.get_start_index(dim, marker)
 
 

From 11dd011b87cc97c5e8cb0cce989966fccbc53988 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Mon, 22 Jan 2024 17:40:32 +0100
Subject: [PATCH 13/23] add docstrings

---
 .../icon4py/model/common/grid/grid_manager.py | 64 +++++++++++++++++--
 1 file changed, 60 insertions(+), 4 deletions(-)

diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py
index 7d2416e4bc..df9cf8ef71 100644
--- a/model/common/src/icon4py/model/common/grid/grid_manager.py
+++ b/model/common/src/icon4py/model/common/grid/grid_manager.py
@@ -82,7 +82,8 @@ class PropertyName(GridFileName):
     class OffsetName(GridFileName):
         """Names for connectivities used in the grid file."""
 
-        # e2c2e/e2c2eO: diamond edges (including origin) not present in grid file-> calculate?
+        # e2c2e/e2c2eO: diamond edges (including origin) not present in grid file-> construct
+        #               from e2c and c2e
         # e2c2v: diamond vertices: not present in grid file -> constructed from e2c and c2v
 
         #: name of C2E2C connectivity in grid file: dims(nv=3, cell)
@@ -412,9 +413,39 @@ def _from_grid_dataset(self, dataset: Dataset) -> IconGrid:
         return icon_grid
 
     @staticmethod
-    def _construct_diamond_vertices(e2v: np.ndarray, c2v: np.ndarray, e2c: np.ndarray):
+    def _construct_diamond_vertices(
+        e2v: np.ndarray, c2v: np.ndarray, e2c: np.ndarray
+    ) -> np.ndarray:
+        """
+        Construct the connectivity table for the vertices of a diamond in the ICON triangular grid.
+
+        Starting from the e2v and c2v connectivity the connectivity table for e2c2v is built up.
+
+                     v0
+                    / \
+                  /    \
+                 /      \
+                /        \
+               v1---e0---v3
+                \       /
+                 \     /
+                  \   /
+                   \ /
+                    v2
+        For example for this diamond: e0 -> (v0, v1, v2, v3)
+        Ordering is the same as ICON uses.
+
+        Args:
+            e2v: np.ndarray containing the connectivity table for edge-to-vertex
+            c2v: np.ndarray containing the connectivity table for cell-to-vertex
+            e2c: np.ndarray containing the connectivity table for edge-to-cell
+
+        Returns: np.ndarray containing the connectivity table for edge-to-vertex on the diamond
+        """
+
         n_edges = e2v.shape[0]
-        e2c2v = -55 * np.ones((n_edges, 4), dtype=np.int32)
+        diamond_size = 4
+        e2c2v = -55 * np.ones((n_edges, diamond_size), dtype=np.int32)
         e2c2v[:, 0:2] = e2v
         dummy_c2v = _patch_with_dummy_lastline(c2v)
         expanded = dummy_c2v[e2c[:, :], :]
@@ -427,7 +458,32 @@ def _construct_diamond_vertices(e2v: np.ndarray, c2v: np.ndarray, e2c: np.ndarra
         return e2c2v
 
     @staticmethod
-    def _construct_diamond_edges(e2c: np.ndarray, c2e: np.ndarray):
+    def _construct_diamond_edges(e2c: np.ndarray, c2e: np.ndarray) -> np.ndarray:
+        """
+        Construct the connectivity table for the edges of a diamond in the ICON triangular grid.
+
+        Starting from the e2c and c2e connectivity the connectivity table for e2c2e is built up.
+
+
+            / \
+          /    \
+         e2    e1
+        /    c0  \
+        ----e0----
+        \   c1   /
+         e3    e4
+          \   /
+           \ /
+
+        For example, for this diamond for e0 -> (e1, e2, e3, e4)
+
+
+        Args:
+            e2c: np.ndarray containing the connectivity table for edge-to-cell
+            c2e: np.ndarray containing the connectivity table for cell-to-edge
+
+        Returns: np.ndarray containing the connectivity table for edge-to-edge on the diamond
+        """
         dummy_c2e = _patch_with_dummy_lastline(c2e)
         expanded = dummy_c2e[e2c, :]
         sh = expanded.shape

From 1c0749ab6f963e2160dd5d6dc687bf537bc795d4 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Mon, 22 Jan 2024 17:57:41 +0100
Subject: [PATCH 14/23] review fixes in utils.py

---
 model/common/tests/grid_tests/utils.py | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/model/common/tests/grid_tests/utils.py b/model/common/tests/grid_tests/utils.py
index 423e7b229f..7e96648e8a 100644
--- a/model/common/tests/grid_tests/utils.py
+++ b/model/common/tests/grid_tests/utils.py
@@ -10,6 +10,7 @@
 # distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
+from pathlib import Path
 
 from icon4py.model.common.test_utils.data_handling import download_and_extract
 from icon4py.model.common.test_utils.datatest_utils import TEST_DATA_ROOT
@@ -21,30 +22,35 @@
 grids_path = TEST_DATA_ROOT.joinpath("grids")
 r04b09_dsl_grid_path = grids_path.joinpath(MCH_GRID_FILE)
 r04b09_dsl_data_file = r04b09_dsl_grid_path.joinpath("mch_ch_r04b09_dsl_grids_v1.tar.gz").name
+
 r02b04_global_grid_path = grids_path.joinpath(R02B04_GLOBAL)
 r02b04_global_data_file = r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_R.tar.gz").name
 
 
-mch_ch_r04b09_dsl_grid_uri = "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download"
-r02b04_global_grid_uri = "https://polybox.ethz.ch/index.php/s/AKAO6ImQdIatnkB/download"
+MC_CH_R04B09_DSL_GRID_URI = "https://polybox.ethz.ch/index.php/s/hD232znfEPBh4Oh/download"
+R02B04_GLOBAL_GRID_URI = "https://polybox.ethz.ch/index.php/s/AKAO6ImQdIatnkB/download"
 
 
-def resolve_file_from_gridfile_name(name: str):
+def resolve_file_from_gridfile_name(name: str) -> Path:
     if name == MCH_GRID_FILE:
-        if not r04b09_dsl_grid_path.joinpath("grid.nc").exists():
+        gridfile = r04b09_dsl_grid_path.joinpath("grid.nc")
+        if not gridfile.exists():
             download_and_extract(
-                mch_ch_r04b09_dsl_grid_uri,
+                MC_CH_R04B09_DSL_GRID_URI,
                 r04b09_dsl_grid_path,
                 r04b09_dsl_grid_path,
                 r04b09_dsl_data_file,
             )
-        return r04b09_dsl_grid_path.joinpath("grid.nc")
+        return gridfile
     elif name == R02B04_GLOBAL:
-        if not r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_R.nc").exists():
+        gridfile = r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_R.nc")
+        if not gridfile.exists():
             download_and_extract(
-                r02b04_global_grid_uri,
+                R02B04_GLOBAL_GRID_URI,
                 r02b04_global_grid_path,
                 r02b04_global_grid_path,
                 r02b04_global_data_file,
             )
-        return r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_R.nc")
+        return gridfile
+    else:
+        raise ValueError(f"invalid name: use one of {R02B04_GLOBAL, MCH_GRID_FILE}")

From e1fde4cebce08d2592f4e726afa1a0b28d780d3f Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Tue, 23 Jan 2024 17:46:18 +0100
Subject: [PATCH 15/23] rework tests

---
 .../icon4py/model/common/grid/grid_manager.py | 23 ++++++++-----------
 1 file changed, 10 insertions(+), 13 deletions(-)

diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py
index df9cf8ef71..ca9e89ec36 100644
--- a/model/common/src/icon4py/model/common/grid/grid_manager.py
+++ b/model/common/src/icon4py/model/common/grid/grid_manager.py
@@ -416,7 +416,7 @@ def _from_grid_dataset(self, dataset: Dataset) -> IconGrid:
     def _construct_diamond_vertices(
         e2v: np.ndarray, c2v: np.ndarray, e2c: np.ndarray
     ) -> np.ndarray:
-        """
+        r"""
         Construct the connectivity table for the vertices of a diamond in the ICON triangular grid.
 
         Starting from the e2v and c2v connectivity the connectivity table for e2c2v is built up.
@@ -442,29 +442,23 @@ def _construct_diamond_vertices(
 
         Returns: np.ndarray containing the connectivity table for edge-to-vertex on the diamond
         """
-
-        n_edges = e2v.shape[0]
-        diamond_size = 4
-        e2c2v = -55 * np.ones((n_edges, diamond_size), dtype=np.int32)
-        e2c2v[:, 0:2] = e2v
         dummy_c2v = _patch_with_dummy_lastline(c2v)
         expanded = dummy_c2v[e2c[:, :], :]
         sh = expanded.shape
         flat = expanded.reshape(sh[0], sh[1] * sh[2])
+        far_indices = np.zeros_like(e2v)
         # TODO (magdalena) vectorize speed this up?
         for x in range(sh[0]):
-            e2c2v[x, 2:] = flat[x, ~np.in1d(flat[x, :], e2c2v[x, :])][:2]
-
-        return e2c2v
+            far_indices[x, :] = flat[x, ~np.in1d(flat[x, :], e2v[x, :])][:2]
+        return np.hstack((e2v, far_indices))
 
     @staticmethod
     def _construct_diamond_edges(e2c: np.ndarray, c2e: np.ndarray) -> np.ndarray:
-        """
+        r"""
         Construct the connectivity table for the edges of a diamond in the ICON triangular grid.
 
         Starting from the e2c and c2e connectivity the connectivity table for e2c2e is built up.
 
-
             / \
           /    \
          e2    e1
@@ -482,13 +476,16 @@ def _construct_diamond_edges(e2c: np.ndarray, c2e: np.ndarray) -> np.ndarray:
             e2c: np.ndarray containing the connectivity table for edge-to-cell
             c2e: np.ndarray containing the connectivity table for cell-to-edge
 
-        Returns: np.ndarray containing the connectivity table for edge-to-edge on the diamond
+        Returns: np.ndarray containing the connectivity table for central edge-to- boundary edges
+                 on the diamond
         """
         dummy_c2e = _patch_with_dummy_lastline(c2e)
         expanded = dummy_c2e[e2c, :]
         sh = expanded.shape
         flattened = expanded.reshape(sh[0], sh[1] * sh[2])
-        e2c2e = GridFile.INVALID_INDEX * np.ones((sh[0], 4), dtype=np.int32)
+
+        diamond_sides = 4
+        e2c2e = GridFile.INVALID_INDEX * np.ones((sh[0], diamond_sides), dtype=np.int32)
         for x in range(sh[0]):
             var = flattened[x, (~np.in1d(flattened[x, :], np.asarray([x, GridFile.INVALID_INDEX])))]
             e2c2e[x, : var.shape[0]] = var

From b99e820dd2fe6e2173d41b5939c4966df7435a95 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Wed, 24 Jan 2024 09:11:30 +0100
Subject: [PATCH 16/23] extract constants and remove serialbox comparison from
 test_grid_manager.py::test_get_start_end_index_for_global_grid

---
 .../tests/grid_tests/test_grid_manager.py     | 65 +++++++++++++------
 1 file changed, 46 insertions(+), 19 deletions(-)

diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py
index eb3293f180..a8b43530a8 100644
--- a/model/common/tests/grid_tests/test_grid_manager.py
+++ b/model/common/tests/grid_tests/test_grid_manager.py
@@ -59,6 +59,11 @@
 
 
 SIMPLE_GRID_NC = "simple_grid.nc"
+R02B04_GLOBLA_NUM_VERTICES = 10242
+
+R02B04_GLOBAL_NUM_EDGES = 30720
+
+R02B04_GLOBAL_NUM_CELLS = 20480
 
 
 @pytest.fixture
@@ -668,42 +673,64 @@ def test_get_end_index_for_local_grid(grid_file, num_levels, grid_savepoint, dim
     assert from_grid_file.get_end_index(dim, marker) == from_savepoint.get_end_index(dim, marker)
 
 
-@pytest.mark.datatest
 @pytest.mark.with_netcdf
 @pytest.mark.parametrize(
     "dim, marker,start_index,end_index",
     [
         (CellDim, HorizontalMarkerIndex.interior(CellDim), 0, 0),
-        (CellDim, HorizontalMarkerIndex.local(CellDim), 0, 20480),
+        (CellDim, HorizontalMarkerIndex.local(CellDim), 0, R02B04_GLOBAL_NUM_CELLS),
         (CellDim, HorizontalMarkerIndex.nudging(CellDim), 0, 0),
         (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim), 0, 0),
-        (CellDim, HorizontalMarkerIndex.end(CellDim), 20480, 20480),
-        (CellDim, HorizontalMarkerIndex.halo(CellDim), 20480, 20480),
+        (
+            CellDim,
+            HorizontalMarkerIndex.end(CellDim),
+            R02B04_GLOBAL_NUM_CELLS,
+            R02B04_GLOBAL_NUM_CELLS,
+        ),
+        (
+            CellDim,
+            HorizontalMarkerIndex.halo(CellDim),
+            R02B04_GLOBAL_NUM_CELLS,
+            R02B04_GLOBAL_NUM_CELLS,
+        ),
         (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 0, 0),
-        (EdgeDim, HorizontalMarkerIndex.local(EdgeDim), 0, 30720),
+        (EdgeDim, HorizontalMarkerIndex.local(EdgeDim), 0, R02B04_GLOBAL_NUM_EDGES),
         (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 0, 0),
         (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim), 0, 0),
-        (EdgeDim, HorizontalMarkerIndex.end(EdgeDim), 30720, 30720),
-        (EdgeDim, HorizontalMarkerIndex.halo(EdgeDim), 30720, 30720),
+        (
+            EdgeDim,
+            HorizontalMarkerIndex.end(EdgeDim),
+            R02B04_GLOBAL_NUM_EDGES,
+            R02B04_GLOBAL_NUM_EDGES,
+        ),
+        (
+            EdgeDim,
+            HorizontalMarkerIndex.halo(EdgeDim),
+            R02B04_GLOBAL_NUM_EDGES,
+            R02B04_GLOBAL_NUM_EDGES,
+        ),
         (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 0, 0),
-        (VertexDim, HorizontalMarkerIndex.local(VertexDim), 0, 10242),
+        (VertexDim, HorizontalMarkerIndex.local(VertexDim), 0, R02B04_GLOBLA_NUM_VERTICES),
         (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim), 0, 0),
-        (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10242, 10242),
-        (VertexDim, HorizontalMarkerIndex.halo(VertexDim), 10242, 10242),
+        (
+            VertexDim,
+            HorizontalMarkerIndex.end(VertexDim),
+            R02B04_GLOBLA_NUM_VERTICES,
+            R02B04_GLOBLA_NUM_VERTICES,
+        ),
+        (
+            VertexDim,
+            HorizontalMarkerIndex.halo(VertexDim),
+            R02B04_GLOBLA_NUM_VERTICES,
+            R02B04_GLOBLA_NUM_VERTICES,
+        ),
     ],
 )
-@pytest.mark.parametrize(
-    "grid_file, experiment, num_levels", [(R02B04_GLOBAL, GLOBAL_EXPERIMENT, 80)]
-)
+@pytest.mark.parametrize("grid_file, num_levels", [(R02B04_GLOBAL, 80)])
 def test_get_start_end_index_for_global_grid(
-    grid_file, num_levels, grid_savepoint, dim, marker, start_index, end_index
+    grid_file, num_levels, dim, marker, start_index, end_index
 ):
-    from_savepoint = grid_savepoint.construct_icon_grid()
     file = resolve_file_from_gridfile_name(grid_file)
     from_grid_file = init_grid_manager(file, num_levels=num_levels).get_grid()
     assert from_grid_file.get_start_index(dim, marker) == start_index
-    assert from_grid_file.get_start_index(dim, marker) == from_savepoint.get_start_index(
-        dim, marker
-    )
     assert from_grid_file.get_end_index(dim, marker) == end_index
-    assert from_grid_file.get_end_index(dim, marker) == from_savepoint.get_end_index(dim, marker)

From f087703938de267d540feba65b1ddb366c0ef997 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Thu, 25 Jan 2024 13:28:31 +0100
Subject: [PATCH 17/23] extract constants (WIP)

---
 .../tests/grid_tests/test_grid_manager.py     | 191 +++++++++---------
 1 file changed, 91 insertions(+), 100 deletions(-)

diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py
index a8b43530a8..0829105274 100644
--- a/model/common/tests/grid_tests/test_grid_manager.py
+++ b/model/common/tests/grid_tests/test_grid_manager.py
@@ -22,6 +22,13 @@
 
 from icon4py.model.common.test_utils.datatest_utils import GLOBAL_EXPERIMENT, REGIONAL_EXPERIMENT
 
+MCH_CH_R04B09_CELL_2ND_BOUNDARY = 850
+
+MCH_CH_RO4B09_CELL_INTERIOR = 4104
+
+MCH_CH_R04B09_LOCAL_NUM_EDGES = 31558
+
+MCH_CH_RO4B09_LOCAL_NUM_CELLS = 20896
 
 if typing.TYPE_CHECKING:
     import netCDF4
@@ -59,7 +66,7 @@
 
 
 SIMPLE_GRID_NC = "simple_grid.nc"
-R02B04_GLOBLA_NUM_VERTICES = 10242
+R02B04_GLOBAL_NUM_VERTICES = 10242
 
 R02B04_GLOBAL_NUM_EDGES = 30720
 
@@ -570,107 +577,91 @@ def test_gt4py_transform_offset_by_1_where_valid(size):
     assert np.allclose(expected, offset)
 
 
-@pytest.mark.datatest
 @pytest.mark.with_netcdf
 @pytest.mark.parametrize(
-    "dim, marker, start_index",
+    "dim, marker, start_index, end_index",
     [
-        (CellDim, HorizontalMarkerIndex.interior(CellDim), 4104),
-        (CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, 0),
-        (CellDim, HorizontalMarkerIndex.local(CellDim) - 1, 20896),
-        (CellDim, HorizontalMarkerIndex.halo(CellDim), 20896),
-        (CellDim, HorizontalMarkerIndex.nudging(CellDim), 3316),
-        (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 2511),
-        (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 1688),
-        (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 850),
-        (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 0),
-        (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 6176),
-        (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2, 31558),
-        (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 1, 31558),
-        (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, 5387),
-        (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 4989),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4184),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 3777),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 2954),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2538),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 1700),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1278),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 428),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 0),
-        (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 2071),
-        (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, 10663),
-        (VertexDim, HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663),
-        (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663),
-        (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10663),
-        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 1673),
-        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1266),
-        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 850),
-        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 428),
-        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 0),
-    ],
-)
-def test_get_start_index_for_local_grid(grid_savepoint, dim, marker, start_index):
-    file = resolve_file_from_gridfile_name(MCH_GRID_FILE)
-    from_savepoint = grid_savepoint.construct_icon_grid()
-    from_grid_file = init_grid_manager(file).get_grid()
-    assert from_grid_file.get_start_index(dim, marker) == start_index
-    assert from_grid_file.get_start_index(dim, marker) == from_savepoint.get_start_index(
-        dim, marker
-    )
-
-
-@pytest.mark.datatest
-@pytest.mark.with_netcdf
-@pytest.mark.parametrize(
-    "dim, marker, index",
-    [
-        (CellDim, HorizontalMarkerIndex.interior(CellDim), 20896),
-        (CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, 850),
-        (CellDim, HorizontalMarkerIndex.local(CellDim) - 2, 20896),
-        (CellDim, HorizontalMarkerIndex.local(CellDim) - 1, 20896),
-        (CellDim, HorizontalMarkerIndex.local(CellDim), 20896),
-        (CellDim, HorizontalMarkerIndex.nudging(CellDim), 4104),
-        (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 3316),
-        (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 2511),
-        (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 1688),
-        (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 850),
-        (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 31558),
-        (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 2, 31558),
-        (EdgeDim, HorizontalMarkerIndex.local(EdgeDim) - 1, 31558),
-        (EdgeDim, HorizontalMarkerIndex.local(EdgeDim), 31558),
-        (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, 6176),
-        (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 5387),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4989),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 4184),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 3777),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2954),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 2538),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1700),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 1278),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 428),
-        (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 10663),
-        (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 2, 10663),
-        (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, 10663),
-        (VertexDim, HorizontalMarkerIndex.local(VertexDim), 10663),
-        (VertexDim, HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663),
-        (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663),
-        (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10663),
-        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 2071),
-        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1673),
-        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 1266),
-        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 850),
-        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 428),
+        (
+            CellDim,
+            HorizontalMarkerIndex.interior(CellDim),
+            MCH_CH_RO4B09_CELL_INTERIOR,
+            MCH_CH_RO4B09_LOCAL_NUM_CELLS,
+        ),
+        (CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, 0, MCH_CH_R04B09_CELL_2ND_BOUNDARY),
+        (
+            CellDim,
+            HorizontalMarkerIndex.local(CellDim) - 2,
+            MCH_CH_RO4B09_LOCAL_NUM_CELLS,
+            MCH_CH_RO4B09_LOCAL_NUM_CELLS,
+        ),
+        (
+            CellDim,
+            HorizontalMarkerIndex.local(CellDim) - 1,
+            MCH_CH_RO4B09_LOCAL_NUM_CELLS,
+            MCH_CH_RO4B09_LOCAL_NUM_CELLS,
+        ),
+        (
+            CellDim,
+            HorizontalMarkerIndex.local(CellDim),
+            MCH_CH_RO4B09_LOCAL_NUM_CELLS,
+            MCH_CH_RO4B09_LOCAL_NUM_CELLS,
+        ),
+        (CellDim, HorizontalMarkerIndex.nudging(CellDim), 3316, MCH_CH_RO4B09_CELL_INTERIOR),
+        (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 2511, 3316),
+        (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 1688, 2511),
+        (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 850, 1688),
+        (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 0, 850),
+        (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 6176, MCH_CH_R04B09_LOCAL_NUM_EDGES),
+        (
+            EdgeDim,
+            HorizontalMarkerIndex.local(EdgeDim) - 2,
+            MCH_CH_R04B09_LOCAL_NUM_EDGES,
+            MCH_CH_R04B09_LOCAL_NUM_EDGES,
+        ),
+        (
+            EdgeDim,
+            HorizontalMarkerIndex.local(EdgeDim) - 1,
+            MCH_CH_R04B09_LOCAL_NUM_EDGES,
+            MCH_CH_R04B09_LOCAL_NUM_EDGES,
+        ),
+        (
+            EdgeDim,
+            HorizontalMarkerIndex.local(EdgeDim),
+            MCH_CH_R04B09_LOCAL_NUM_EDGES,
+            MCH_CH_R04B09_LOCAL_NUM_EDGES,
+        ),
+        (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 4989, 5387),
+        (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, 5387, 6176),
+        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4184, 4989),
+        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 3777, 4184),
+        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 2954, 3777),
+        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2538, 2954),
+        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 1700, 2538),
+        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1278, 1700),
+        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 428, 1278),
+        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 0, 428),
+        (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 2071, 10663),
+        (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 2, 10663, 10663),
+        (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, 10663, 10663),
+        (VertexDim, HorizontalMarkerIndex.local(VertexDim), 10663, 10663),
+        (VertexDim, HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663, 10663),
+        (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663, 10663),
+        (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10663, 10663),
+        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 1673, 2071),
+        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1266, 1673),
+        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 850, 1266),
+        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 428, 850),
+        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 0, 428),
     ],
 )
-@pytest.mark.parametrize(
-    "grid_file, experiment, num_levels", [(MCH_GRID_FILE, REGIONAL_EXPERIMENT, 65)]
-)
-def test_get_end_index_for_local_grid(grid_file, num_levels, grid_savepoint, dim, marker, index):
+@pytest.mark.parametrize("grid_file, num_levels", [(MCH_GRID_FILE, 65)])
+def test_get_start_end_index_for_local_grid(
+    grid_file, num_levels, dim, marker, start_index, end_index
+):
     file = resolve_file_from_gridfile_name(grid_file)
-    from_savepoint = grid_savepoint.construct_icon_grid()
     from_grid_file = init_grid_manager(file, num_levels=num_levels).get_grid()
-    assert from_grid_file.get_end_index(dim, marker) == index
-    assert from_grid_file.get_end_index(dim, marker) == from_savepoint.get_end_index(dim, marker)
+    assert from_grid_file.get_start_index(dim, marker) == start_index
+    assert from_grid_file.get_end_index(dim, marker) == end_index
 
 
 @pytest.mark.with_netcdf
@@ -710,19 +701,19 @@ def test_get_end_index_for_local_grid(grid_file, num_levels, grid_savepoint, dim
             R02B04_GLOBAL_NUM_EDGES,
         ),
         (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 0, 0),
-        (VertexDim, HorizontalMarkerIndex.local(VertexDim), 0, R02B04_GLOBLA_NUM_VERTICES),
+        (VertexDim, HorizontalMarkerIndex.local(VertexDim), 0, R02B04_GLOBAL_NUM_VERTICES),
         (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim), 0, 0),
         (
             VertexDim,
             HorizontalMarkerIndex.end(VertexDim),
-            R02B04_GLOBLA_NUM_VERTICES,
-            R02B04_GLOBLA_NUM_VERTICES,
+            R02B04_GLOBAL_NUM_VERTICES,
+            R02B04_GLOBAL_NUM_VERTICES,
         ),
         (
             VertexDim,
             HorizontalMarkerIndex.halo(VertexDim),
-            R02B04_GLOBLA_NUM_VERTICES,
-            R02B04_GLOBLA_NUM_VERTICES,
+            R02B04_GLOBAL_NUM_VERTICES,
+            R02B04_GLOBAL_NUM_VERTICES,
         ),
     ],
 )

From f5833e671647c6f4a5f6cf141d41b5f53ffa8846 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Thu, 1 Feb 2024 11:22:10 +0100
Subject: [PATCH 18/23] extract constants for domain bounds (WIP)

---
 model/common/tests/grid_tests/test_grid_manager.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py
index 0829105274..74d86611e7 100644
--- a/model/common/tests/grid_tests/test_grid_manager.py
+++ b/model/common/tests/grid_tests/test_grid_manager.py
@@ -30,6 +30,16 @@
 
 MCH_CH_RO4B09_LOCAL_NUM_CELLS = 20896
 
+MCH_CH_R04B09_CELL_DOMAINS = {
+    "2ND_BOUNDARY_LINE": 428,
+    "3D_BOUNDARY_LINE": 850,
+    "4TH_BOUNDARY_LINE": 1266,
+    "INTERIOR": 2071,
+    "HALO": 10663,
+    "2ND_HALO_LINE": 10663,
+    "LOCAL": 0,
+}
+
 if typing.TYPE_CHECKING:
     import netCDF4
 

From 48f13972830cdf6ee310707428dd7d56bf9d26e9 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Mon, 5 Feb 2024 10:26:47 +0100
Subject: [PATCH 19/23] extract domain starts for local grid into constants

---
 .../tests/grid_tests/test_grid_manager.py     | 267 ++++++++++++++----
 1 file changed, 217 insertions(+), 50 deletions(-)

diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py
index 74d86611e7..d94cd22fb6 100644
--- a/model/common/tests/grid_tests/test_grid_manager.py
+++ b/model/common/tests/grid_tests/test_grid_manager.py
@@ -22,24 +22,6 @@
 
 from icon4py.model.common.test_utils.datatest_utils import GLOBAL_EXPERIMENT, REGIONAL_EXPERIMENT
 
-MCH_CH_R04B09_CELL_2ND_BOUNDARY = 850
-
-MCH_CH_RO4B09_CELL_INTERIOR = 4104
-
-MCH_CH_R04B09_LOCAL_NUM_EDGES = 31558
-
-MCH_CH_RO4B09_LOCAL_NUM_CELLS = 20896
-
-MCH_CH_R04B09_CELL_DOMAINS = {
-    "2ND_BOUNDARY_LINE": 428,
-    "3D_BOUNDARY_LINE": 850,
-    "4TH_BOUNDARY_LINE": 1266,
-    "INTERIOR": 2071,
-    "HALO": 10663,
-    "2ND_HALO_LINE": 10663,
-    "LOCAL": 0,
-}
-
 if typing.TYPE_CHECKING:
     import netCDF4
 
@@ -76,12 +58,52 @@
 
 
 SIMPLE_GRID_NC = "simple_grid.nc"
-R02B04_GLOBAL_NUM_VERTICES = 10242
 
+R02B04_GLOBAL_NUM_VERTICES = 10242
 R02B04_GLOBAL_NUM_EDGES = 30720
-
 R02B04_GLOBAL_NUM_CELLS = 20480
 
+MCH_CH_04B09_NUM_VERTICES = 10663
+MCH_CH_R04B09_LOCAL_NUM_EDGES = 31558
+MCH_CH_RO4B09_LOCAL_NUM_CELLS = 20896
+
+
+MCH_CH_R04B09_CELL_DOMAINS = {
+    "2ND_BOUNDARY_LINE": 850,
+    "3D_BOUNDARY_LINE": 1688,
+    "4TH_BOUNDARY_LINE": 2511,
+    "NUDGING": 3316,
+    "INTERIOR": 4104,
+    "HALO": 20896,
+    "LOCAL": 0,
+}
+
+MCH_CH_R04B09_VERTEX_DOMAINS = {
+    "2ND_BOUNDARY_LINE": 428,
+    "3D_BOUNDARY_LINE": 850,
+    "4TH_BOUNDARY_LINE": 1266,
+    "5TH_BOUNDARY_LINE": 1673,
+    "INTERIOR": 2071,
+    "HALO": 10663,
+    "LOCAL": 0,
+}
+
+MCH_CH_R04B09_EDGE_DOMAINS = {
+    "2ND_BOUNDARY_LINE": 428,
+    "3D_BOUNDARY_LINE": 1278,
+    "4TH_BOUNDARY_LINE": 1700,
+    "5TH_BOUNDARY_LINE": 2538,
+    "6TH_BOUNDARY_LINE": 2954,
+    "7TH_BOUNDARY_LINE": 3777,
+    "8TH_BOUNDARY_LINE": 4184,
+    "NUDGING": 4989,
+    "2ND_NUDGING": 5387,
+    "INTERIOR": 6176,
+    "HALO": 31558,
+    "LOCAL": 0,
+    "END": 31558,
+}
+
 
 @pytest.fixture
 def simple_grid_gridfile(tmp_path):
@@ -594,10 +616,15 @@ def test_gt4py_transform_offset_by_1_where_valid(size):
         (
             CellDim,
             HorizontalMarkerIndex.interior(CellDim),
-            MCH_CH_RO4B09_CELL_INTERIOR,
+            MCH_CH_R04B09_CELL_DOMAINS["INTERIOR"],
             MCH_CH_RO4B09_LOCAL_NUM_CELLS,
         ),
-        (CellDim, HorizontalMarkerIndex.interior(CellDim) + 1, 0, MCH_CH_R04B09_CELL_2ND_BOUNDARY),
+        (
+            CellDim,
+            HorizontalMarkerIndex.interior(CellDim) + 1,
+            0,
+            MCH_CH_R04B09_CELL_DOMAINS["2ND_BOUNDARY_LINE"],
+        ),
         (
             CellDim,
             HorizontalMarkerIndex.local(CellDim) - 2,
@@ -616,12 +643,42 @@ def test_gt4py_transform_offset_by_1_where_valid(size):
             MCH_CH_RO4B09_LOCAL_NUM_CELLS,
             MCH_CH_RO4B09_LOCAL_NUM_CELLS,
         ),
-        (CellDim, HorizontalMarkerIndex.nudging(CellDim), 3316, MCH_CH_RO4B09_CELL_INTERIOR),
-        (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 3, 2511, 3316),
-        (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2, 1688, 2511),
-        (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, 850, 1688),
-        (CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 0, 0, 850),
-        (EdgeDim, HorizontalMarkerIndex.interior(EdgeDim), 6176, MCH_CH_R04B09_LOCAL_NUM_EDGES),
+        (
+            CellDim,
+            HorizontalMarkerIndex.nudging(CellDim),
+            MCH_CH_R04B09_CELL_DOMAINS["NUDGING"],
+            MCH_CH_R04B09_CELL_DOMAINS["INTERIOR"],
+        ),
+        (
+            CellDim,
+            HorizontalMarkerIndex.lateral_boundary(CellDim) + 3,
+            MCH_CH_R04B09_CELL_DOMAINS["4TH_BOUNDARY_LINE"],
+            MCH_CH_R04B09_CELL_DOMAINS["NUDGING"],
+        ),
+        (
+            CellDim,
+            HorizontalMarkerIndex.lateral_boundary(CellDim) + 2,
+            MCH_CH_R04B09_CELL_DOMAINS["3D_BOUNDARY_LINE"],
+            MCH_CH_R04B09_CELL_DOMAINS["4TH_BOUNDARY_LINE"],
+        ),
+        (
+            CellDim,
+            HorizontalMarkerIndex.lateral_boundary(CellDim) + 1,
+            MCH_CH_R04B09_CELL_DOMAINS["2ND_BOUNDARY_LINE"],
+            MCH_CH_R04B09_CELL_DOMAINS["3D_BOUNDARY_LINE"],
+        ),
+        (
+            CellDim,
+            HorizontalMarkerIndex.lateral_boundary(CellDim) + 0,
+            0,
+            MCH_CH_R04B09_CELL_DOMAINS["2ND_BOUNDARY_LINE"],
+        ),
+        (
+            EdgeDim,
+            HorizontalMarkerIndex.interior(EdgeDim),
+            MCH_CH_R04B09_EDGE_DOMAINS["INTERIOR"],
+            MCH_CH_R04B09_LOCAL_NUM_EDGES,
+        ),
         (
             EdgeDim,
             HorizontalMarkerIndex.local(EdgeDim) - 2,
@@ -640,28 +697,138 @@ def test_gt4py_transform_offset_by_1_where_valid(size):
             MCH_CH_R04B09_LOCAL_NUM_EDGES,
             MCH_CH_R04B09_LOCAL_NUM_EDGES,
         ),
-        (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim), 4989, 5387),
-        (EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1, 5387, 6176),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7, 4184, 4989),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6, 3777, 4184),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5, 2954, 3777),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4, 2538, 2954),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3, 1700, 2538),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2, 1278, 1700),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1, 428, 1278),
-        (EdgeDim, HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0, 0, 428),
-        (VertexDim, HorizontalMarkerIndex.interior(VertexDim), 2071, 10663),
-        (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 2, 10663, 10663),
-        (VertexDim, HorizontalMarkerIndex.local(VertexDim) - 1, 10663, 10663),
-        (VertexDim, HorizontalMarkerIndex.local(VertexDim), 10663, 10663),
-        (VertexDim, HorizontalMarkerIndex.nudging(VertexDim) + 1, 10663, 10663),
-        (VertexDim, HorizontalMarkerIndex.nudging(VertexDim), 10663, 10663),
-        (VertexDim, HorizontalMarkerIndex.end(VertexDim), 10663, 10663),
-        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4, 1673, 2071),
-        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3, 1266, 1673),
-        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2, 850, 1266),
-        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1, 428, 850),
-        (VertexDim, HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0, 0, 428),
+        (
+            EdgeDim,
+            HorizontalMarkerIndex.nudging(EdgeDim),
+            MCH_CH_R04B09_EDGE_DOMAINS["NUDGING"],
+            MCH_CH_R04B09_EDGE_DOMAINS["2ND_NUDGING"],
+        ),
+        (
+            EdgeDim,
+            HorizontalMarkerIndex.nudging(EdgeDim) + 1,
+            MCH_CH_R04B09_EDGE_DOMAINS["2ND_NUDGING"],
+            MCH_CH_R04B09_EDGE_DOMAINS["INTERIOR"],
+        ),
+        (
+            EdgeDim,
+            HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 7,
+            MCH_CH_R04B09_EDGE_DOMAINS["8TH_BOUNDARY_LINE"],
+            MCH_CH_R04B09_EDGE_DOMAINS["NUDGING"],
+        ),
+        (
+            EdgeDim,
+            HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 6,
+            MCH_CH_R04B09_EDGE_DOMAINS["7TH_BOUNDARY_LINE"],
+            MCH_CH_R04B09_EDGE_DOMAINS["8TH_BOUNDARY_LINE"],
+        ),
+        (
+            EdgeDim,
+            HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 5,
+            MCH_CH_R04B09_EDGE_DOMAINS["6TH_BOUNDARY_LINE"],
+            MCH_CH_R04B09_EDGE_DOMAINS["7TH_BOUNDARY_LINE"],
+        ),
+        (
+            EdgeDim,
+            HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 4,
+            MCH_CH_R04B09_EDGE_DOMAINS["5TH_BOUNDARY_LINE"],
+            MCH_CH_R04B09_EDGE_DOMAINS["6TH_BOUNDARY_LINE"],
+        ),
+        (
+            EdgeDim,
+            HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 3,
+            MCH_CH_R04B09_EDGE_DOMAINS["4TH_BOUNDARY_LINE"],
+            MCH_CH_R04B09_EDGE_DOMAINS["5TH_BOUNDARY_LINE"],
+        ),
+        (
+            EdgeDim,
+            HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2,
+            MCH_CH_R04B09_EDGE_DOMAINS["3D_BOUNDARY_LINE"],
+            MCH_CH_R04B09_EDGE_DOMAINS["4TH_BOUNDARY_LINE"],
+        ),
+        (
+            EdgeDim,
+            HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1,
+            MCH_CH_R04B09_EDGE_DOMAINS["2ND_BOUNDARY_LINE"],
+            MCH_CH_R04B09_EDGE_DOMAINS["3D_BOUNDARY_LINE"],
+        ),
+        (
+            EdgeDim,
+            HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 0,
+            0,
+            MCH_CH_R04B09_EDGE_DOMAINS["2ND_BOUNDARY_LINE"],
+        ),
+        (
+            VertexDim,
+            HorizontalMarkerIndex.interior(VertexDim),
+            MCH_CH_R04B09_VERTEX_DOMAINS["INTERIOR"],
+            MCH_CH_04B09_NUM_VERTICES,
+        ),
+        (
+            VertexDim,
+            HorizontalMarkerIndex.local(VertexDim) - 2,
+            MCH_CH_04B09_NUM_VERTICES,
+            MCH_CH_04B09_NUM_VERTICES,
+        ),
+        (
+            VertexDim,
+            HorizontalMarkerIndex.local(VertexDim) - 1,
+            MCH_CH_04B09_NUM_VERTICES,
+            MCH_CH_04B09_NUM_VERTICES,
+        ),
+        (
+            VertexDim,
+            HorizontalMarkerIndex.local(VertexDim),
+            MCH_CH_04B09_NUM_VERTICES,
+            MCH_CH_04B09_NUM_VERTICES,
+        ),
+        (
+            VertexDim,
+            HorizontalMarkerIndex.nudging(VertexDim) + 1,
+            MCH_CH_04B09_NUM_VERTICES,
+            MCH_CH_04B09_NUM_VERTICES,
+        ),
+        (
+            VertexDim,
+            HorizontalMarkerIndex.nudging(VertexDim),
+            MCH_CH_04B09_NUM_VERTICES,
+            MCH_CH_04B09_NUM_VERTICES,
+        ),
+        (
+            VertexDim,
+            HorizontalMarkerIndex.end(VertexDim),
+            MCH_CH_04B09_NUM_VERTICES,
+            MCH_CH_04B09_NUM_VERTICES,
+        ),
+        (
+            VertexDim,
+            HorizontalMarkerIndex.lateral_boundary(VertexDim) + 4,
+            MCH_CH_R04B09_VERTEX_DOMAINS["5TH_BOUNDARY_LINE"],
+            MCH_CH_R04B09_VERTEX_DOMAINS["INTERIOR"],
+        ),
+        (
+            VertexDim,
+            HorizontalMarkerIndex.lateral_boundary(VertexDim) + 3,
+            MCH_CH_R04B09_VERTEX_DOMAINS["4TH_BOUNDARY_LINE"],
+            MCH_CH_R04B09_VERTEX_DOMAINS["5TH_BOUNDARY_LINE"],
+        ),
+        (
+            VertexDim,
+            HorizontalMarkerIndex.lateral_boundary(VertexDim) + 2,
+            MCH_CH_R04B09_VERTEX_DOMAINS["3D_BOUNDARY_LINE"],
+            MCH_CH_R04B09_VERTEX_DOMAINS["4TH_BOUNDARY_LINE"],
+        ),
+        (
+            VertexDim,
+            HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1,
+            MCH_CH_R04B09_VERTEX_DOMAINS["2ND_BOUNDARY_LINE"],
+            MCH_CH_R04B09_VERTEX_DOMAINS["3D_BOUNDARY_LINE"],
+        ),
+        (
+            VertexDim,
+            HorizontalMarkerIndex.lateral_boundary(VertexDim) + 0,
+            0,
+            MCH_CH_R04B09_VERTEX_DOMAINS["2ND_BOUNDARY_LINE"],
+        ),
     ],
 )
 @pytest.mark.parametrize("grid_file, num_levels", [(MCH_GRID_FILE, 65)])

From a6425722813932690961b77f3324d67150efb899 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Mon, 5 Feb 2024 10:29:56 +0100
Subject: [PATCH 20/23] pre-commit

---
 model/common/tests/grid_tests/test_grid_manager.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py
index d94cd22fb6..2c0661b673 100644
--- a/model/common/tests/grid_tests/test_grid_manager.py
+++ b/model/common/tests/grid_tests/test_grid_manager.py
@@ -22,6 +22,7 @@
 
 from icon4py.model.common.test_utils.datatest_utils import GLOBAL_EXPERIMENT, REGIONAL_EXPERIMENT
 
+
 if typing.TYPE_CHECKING:
     import netCDF4
 

From 74c215f1773e33ead1fd9fed1a098f810dc97f51 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Mon, 5 Feb 2024 12:14:38 +0100
Subject: [PATCH 21/23] (FIXES) ad hoc fixes - add default argument for on_gpu
 - run test_smagorinsky.py on roundtrip backend

---
 .../src/icon4py/model/common/test_utils/serialbox_utils.py    | 2 +-
 model/common/tests/math_tests/test_smagorinsky.py             | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py
index af8a1c22b3..85d0233608 100644
--- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py
+++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py
@@ -314,7 +314,7 @@ def _get_decomp_fields(self, dim: Dimension):
         mask = self.owner_mask(dim)[0 : self.num(dim)]
         return dim, global_index, mask
 
-    def construct_icon_grid(self, on_gpu: bool) -> IconGrid:
+    def construct_icon_grid(self, on_gpu: bool = False) -> IconGrid:
         cell_starts = self.cells_start_index()
         cell_ends = self.cells_end_index()
         vertex_starts = self.vertex_start_index()
diff --git a/model/common/tests/math_tests/test_smagorinsky.py b/model/common/tests/math_tests/test_smagorinsky.py
index 7413392932..1b7860e7b1 100644
--- a/model/common/tests/math_tests/test_smagorinsky.py
+++ b/model/common/tests/math_tests/test_smagorinsky.py
@@ -12,6 +12,7 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 import numpy as np
+from gt4py.next.program_processors.runners.roundtrip import backend as roundtrip
 
 from icon4py.model.common.dimension import KDim
 from icon4py.model.common.grid.simple import SimpleGrid
@@ -28,7 +29,8 @@ def test_init_enh_smag_fac():
     z = (0.1, 0.2, 0.3, 0.4)
 
     enhanced_smag_fac_np = enhanced_smagorinski_factor_numpy(fac, z, a_vec.asnumpy())
-    en_smag_fac_for_zero_nshift(
+    # TODO (magdalena) fails with embedded backend, because broadcast(0,0, (KDim,)) returns a scalar
+    en_smag_fac_for_zero_nshift.with_backend(roundtrip)(
         a_vec,
         *fac,
         *z,

From a9f0d19808f851a40f2266f144630a4e83421e48 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Mon, 5 Feb 2024 13:16:45 +0100
Subject: [PATCH 22/23] rename loop variables

---
 .../src/icon4py/model/common/grid/grid_manager.py      | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py
index ca9e89ec36..62e09a2828 100644
--- a/model/common/src/icon4py/model/common/grid/grid_manager.py
+++ b/model/common/src/icon4py/model/common/grid/grid_manager.py
@@ -448,8 +448,8 @@ def _construct_diamond_vertices(
         flat = expanded.reshape(sh[0], sh[1] * sh[2])
         far_indices = np.zeros_like(e2v)
         # TODO (magdalena) vectorize speed this up?
-        for x in range(sh[0]):
-            far_indices[x, :] = flat[x, ~np.in1d(flat[x, :], e2v[x, :])][:2]
+        for i in range(sh[0]):
+            far_indices[i, :] = flat[i, ~np.in1d(flat[i, :], e2v[i, :])][:2]
         return np.hstack((e2v, far_indices))
 
     @staticmethod
@@ -486,9 +486,9 @@ def _construct_diamond_edges(e2c: np.ndarray, c2e: np.ndarray) -> np.ndarray:
 
         diamond_sides = 4
         e2c2e = GridFile.INVALID_INDEX * np.ones((sh[0], diamond_sides), dtype=np.int32)
-        for x in range(sh[0]):
-            var = flattened[x, (~np.in1d(flattened[x, :], np.asarray([x, GridFile.INVALID_INDEX])))]
-            e2c2e[x, : var.shape[0]] = var
+        for i in range(sh[0]):
+            var = flattened[i, (~np.in1d(flattened[i, :], np.asarray([i, GridFile.INVALID_INDEX])))]
+            e2c2e[i, : var.shape[0]] = var
         return e2c2e
 
 

From 82400598d527138ee95dab50744550f7b19ca5c7 Mon Sep 17 00:00:00 2001
From: Magdalena Luz <luzm@ethz.ch>
Date: Tue, 6 Feb 2024 13:05:36 +0100
Subject: [PATCH 23/23] (UNDO) remove default value from construct_icon grid

---
 .../src/icon4py/model/common/test_utils/serialbox_utils.py      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py
index 06565ea249..fa7c6062d3 100644
--- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py
+++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py
@@ -314,7 +314,7 @@ def _get_decomp_fields(self, dim: Dimension):
         mask = self.owner_mask(dim)[0 : self.num(dim)]
         return dim, global_index, mask
 
-    def construct_icon_grid(self, on_gpu: bool = False) -> IconGrid:
+    def construct_icon_grid(self, on_gpu: bool) -> IconGrid:
         cell_starts = self.cells_start_index()
         cell_ends = self.cells_end_index()
         vertex_starts = self.vertex_start_index()