Skip to content

Commit

Permalink
Add diamond to gridmanager (#352)
Browse files Browse the repository at this point in the history
constructs the E2C2E and E2C2EO connectivities in grid_manager.py
fixes the index ordering of E2C2V connectivity such that it matches the one in ICON.
  • Loading branch information
halungge authored Feb 7, 2024
1 parent a5c8e18 commit 9081e6d
Showing 8 changed files with 696 additions and 178 deletions.
Original file line number Diff line number Diff line change
@@ -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,
42 changes: 42 additions & 0 deletions model/common/README.md
Original file line number Diff line number Diff line change
@@ -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
114 changes: 105 additions & 9 deletions model/common/src/icon4py/model/common/grid/grid_manager.py
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions model/common/src/icon4py/model/common/grid/horizontal.py
Original file line number Diff line number Diff line change
@@ -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 = {
12 changes: 12 additions & 0 deletions model/common/tests/grid_tests/__init__.py
Original file line number Diff line number Diff line change
@@ -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
39 changes: 4 additions & 35 deletions model/common/tests/grid_tests/conftest.py
Original file line number Diff line number Diff line change
@@ -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
Loading

0 comments on commit 9081e6d

Please sign in to comment.