Add diamond to gridmanager

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

@@ -471,6 +471,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,

model/common/README.md

@@ -7,3 +7,45 @@ 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

model/common/src/icon4py/model/common/grid/grid_manager.py

@@ -37,6 +37,8 @@ def __init__(self, *args, **kwargs):
     C2EDim,
     C2VDim,
     CellDim,
+    E2C2EDim,
+    E2C2EODim,
     E2C2VDim,
     E2CDim,
     E2VDim,
@@ -80,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)
@@ -363,7 +366,9 @@ 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(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)
@@ -396,6 +401,8 @@ def _from_grid_dataset(self, dataset: Dataset) -> IconGrid:
                     C2E2CODim: c2e2c0,
                     E2C2VDim: e2c2v,
                     V2E2VDim: v2e2v,
+                    E2C2EDim: e2c2e,
+                    E2C2EODim: e2c2e0,
                 }
             )
             .with_start_end_indices(CellDim, start_indices[CellDim], end_indices[CellDim])
@@ -405,13 +412,102 @@ def _from_grid_dataset(self, dataset: Dataset) -> IconGrid:
 
         return icon_grid
 
-    def _construct_diamond_array(self, c2v: np.ndarray, e2c: np.ndarray):
-        dummy_c2v = np.append(
-            c2v,
-            GridFile.INVALID_INDEX * np.ones((1, c2v.shape[1]), dtype=np.int32),
-            axis=0,
-        )
+    @staticmethod
+    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.
+
+                     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
+        """
+        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 i in range(sh[0]):
+            far_indices[i, :] = flat[i, ~np.in1d(flat[i, :], e2v[i, :])][: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
+        /    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 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])
-        return np.apply_along_axis(np.unique, 1, flattened)
+
+        diamond_sides = 4
+        e2c2e = GridFile.INVALID_INDEX * np.ones((sh[0], diamond_sides), dtype=np.int32)
+        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
+
+
+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

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 = {

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

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,10 @@
     processor_props,
     ranked_data_path,
 )
-from icon4py.model.common.test_utils.datatest_utils import TEST_DATA_ROOT
 
+from .utils import MCH_GRID_FILE
 
-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")
-
-
-@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