Add diamond to gridmanager

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,

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,
@@ -363,7 +365,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 +400,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 +411,50 @@ 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):
+        n_edges = e2v.shape[0]
+        e2c2v = -55 * np.ones((n_edges, 4), 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])
+        # 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
+
+    @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])
-        return np.apply_along_axis(np.unique, 1, flattened)
+        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
+
+
+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/src/icon4py/model/common/grid/icon.py

@@ -19,6 +19,7 @@
     C2E2CDim,
     C2E2CODim,
     C2EDim,
+    C2VDim,
     CECDim,
     CEDim,
     CellDim,
@@ -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),

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