Better torus support for mean cell area etc.

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

@@ -119,9 +119,9 @@ def __exit__(self, exc_type, exc_val, exc_tb):
     def __call__(self, backend: gtx_typing.Backend | None, keep_skip_values: bool):
         if not self._reader:
             self.open()
+        self._geometry = self._read_geometry_fields(backend)
         self._grid = self._construct_grid(backend=backend, with_skip_values=keep_skip_values)
         self._coordinates = self._read_coordinates(backend)
-        self._geometry = self._read_geometry_fields(backend)
         self.close()
 
     def _read_coordinates(self, backend: gtx_typing.Backend | None) -> CoordinateDict:
@@ -185,6 +185,16 @@ def _read_geometry_fields(self, backend: gtx_typing.Backend | None):
                 self._reader.variable(gridfile.GeometryName.DUAL_AREA),
                 allocator=backend,
             ),
+            gridfile.GeometryName.EDGE_LENGTH.value: gtx.as_field(
+                (dims.EdgeDim,),
+                self._reader.variable(gridfile.GeometryName.EDGE_LENGTH),
+                allocator=backend,
+            ),
+            gridfile.GeometryName.DUAL_EDGE_LENGTH.value: gtx.as_field(
+                (dims.EdgeDim,),
+                self._reader.variable(gridfile.GeometryName.DUAL_EDGE_LENGTH),
+                allocator=backend,
+            ),
             gridfile.GeometryName.EDGE_CELL_DISTANCE.value: gtx.as_field(
                 (dims.EdgeDim, dims.E2CDim),
                 self._reader.variable(gridfile.GeometryName.EDGE_CELL_DISTANCE, transpose=True),
@@ -335,11 +345,55 @@ def _construct_grid(
         grid_level = self._reader.attribute(gridfile.MandatoryPropertyName.LEVEL)
         if geometry_type := self._reader.try_attribute(gridfile.MPIMPropertyName.GEOMETRY):
             geometry_type = base.GeometryType(geometry_type)
-        global_params = icon.GlobalGridParams(
+        sphere_radius = self._reader.try_attribute(gridfile.MPIMPropertyName.SPHERE_RADIUS)
+        domain_length = self._reader.try_attribute(gridfile.MPIMPropertyName.DOMAIN_LENGTH)
+        domain_height = self._reader.try_attribute(gridfile.MPIMPropertyName.DOMAIN_HEIGHT)
+
+        mean_edge_length = self._reader.try_attribute(gridfile.MPIMPropertyName.MEAN_EDGE_LENGTH)
+        mean_dual_edge_length = self._reader.try_attribute(
+            gridfile.MPIMPropertyName.MEAN_DUAL_EDGE_LENGTH
+        )
+        mean_cell_area = self._reader.try_attribute(gridfile.MPIMPropertyName.MEAN_CELL_AREA)
+        mean_dual_cell_area = self._reader.try_attribute(
+            gridfile.MPIMPropertyName.MEAN_DUAL_CELL_AREA
+        )
+
+        # TODO(msimberg): EDGE_LENGTH requires GridGeometry (not just
+        # GeometryDict). GridGeometry requires the IconGrid.  IconGrid requires
+        # GlobalGridParams.  GlobalGridParams requires EDGE_LENGTH. EDGE_LENGTH
+        # requires GridGeometry.
+        # - Option 1: Only use MEAN_EDGE_LENGTH from grid file.
+        # - Option 2: Fix circular dependency.
+        # - Option 3: Use EDGE_LENGTH from grid file to make it available in
+        #   GeometryDict. This is implemented below. Is the field always there?
+        edge_lengths = self.geometry[gridfile.GeometryName.EDGE_LENGTH.value].ndarray
+        dual_edge_lengths = self.geometry[gridfile.GeometryName.DUAL_EDGE_LENGTH.value].ndarray
+        cell_areas = self.geometry[gridfile.GeometryName.CELL_AREA.value].ndarray
+        dual_cell_areas = mean_dual_cell_area = xp.mean(
+            self.geometry[gridfile.GeometryName.DUAL_AREA.value].ndarray
+        )
+
+        global_params = icon.GlobalGridParams.from_fields(
+            backend=backend,
             grid_shape=icon.GridShape(
                 geometry_type=geometry_type,
                 subdivision=icon.GridSubdivision(root=grid_root, level=grid_level),
-            )
+            ),
+            radius=sphere_radius,
+            domain_length=domain_length,
+            domain_height=domain_height,
+            # TODO(msimberg): The above is given for the grid, not derived from
+            # anything. The below can be derived in the worst case, so keep
+            # separate?
+            num_cells=num_cells,
+            mean_edge_length=mean_edge_length,
+            mean_dual_edge_length=mean_dual_edge_length,
+            mean_cell_area=mean_cell_area,
+            mean_dual_cell_area=mean_dual_cell_area,
+            edge_lengths=edge_lengths,
+            dual_edge_lengths=dual_edge_lengths,
+            cell_areas=cell_areas,
+            dual_cell_areas=dual_cell_areas,
         )
         grid_size = base.HorizontalGridSize(
             num_vertices=num_vertices, num_edges=num_edges, num_cells=num_cells

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

@@ -179,6 +179,8 @@ class GeometryName(FieldName):
     CELL_AREA = "cell_area"
     # TODO(halungge): compute from coordinates
     DUAL_AREA = "dual_area"
+    EDGE_LENGTH = "edge_length"
+    DUAL_EDGE_LENGTH = "dual_edge_length"
     CELL_NORMAL_ORIENTATION = "orientation_of_normal"
     TANGENT_ORIENTATION = "edge_system_orientation"
     EDGE_ORIENTATION_ON_VERTEX = "edge_orientation"

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

@@ -6,11 +6,10 @@
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
 import dataclasses
-import functools
 import logging
 import math
 from collections.abc import Mapping
-from typing import Final
+from typing import Final, TypeVar
 
 import gt4py.next as gtx
 from gt4py.next import allocators as gtx_allocators
@@ -76,60 +75,106 @@ def __init__(
         self.subdivision = subdivision
 
 
-@dataclasses.dataclass
+T = TypeVar("T")
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
 class GlobalGridParams:
+    # TODO(msimberg): Which of these belong here and which belong in e.g.
+    # GridConfig, HorizontalGridConfig, etc.?
     grid_shape: Final[GridShape | None] = None
     radius: float = constants.EARTH_RADIUS
-    num_cells: int
-    mean_cell_area: float
-
-    def __init__(
-        self,
-        *,
-        grid_shape: GridShape | None = None,
-        radius: float = constants.EARTH_RADIUS,
-        num_cells: int | None = None,
+    domain_length: float | None = None
+    domain_height: float | None = None
+    global_num_cells: int | None = None
+    num_cells: int | None = None
+    mean_edge_length: float | None = None
+    mean_dual_edge_length: float | None = None
+    mean_cell_area: float | None = None
+    mean_dual_cell_area: float | None = None
+    characteristic_length: float | None = None
+
+    @classmethod
+    def from_fields(
+        cls: type[T],
+        backend: gtx.typing.Backend | None,
+        mean_edge_length: float | None = None,
+        edge_lengths: data_alloc.NDArray | None = None,
+        mean_dual_edge_length: float | None = None,
+        dual_edge_lengths: data_alloc.NDArray | None = None,
         mean_cell_area: float | None = None,
-    ) -> None:
-        self.grid_shape = grid_shape
-        self.radius = radius
-
-        if num_cells is not None:
-            self.num_cells = num_cells
+        cell_areas: data_alloc.NDArray | None = None,
+        mean_dual_cell_area: float | None = None,
+        dual_cell_areas: data_alloc.NDArray | None = None,
+        **kwargs,
+    ) -> T:
+        xp = data_alloc.import_array_ns(backend)
+
+        def init_mean(value: float | None, data: data_alloc.NDArray | None) -> float | None:
+            if value is not None:
+                return value
+            if data is not None:
+                return xp.mean(data)
+            return None
+
+        mean_edge_length = init_mean(mean_edge_length, edge_lengths)
+        mean_dual_edge_length = init_mean(mean_dual_edge_length, dual_edge_lengths)
+        mean_cell_area = init_mean(mean_cell_area, cell_areas)
+        mean_dual_cell_area = init_mean(mean_dual_cell_area, dual_cell_areas)
+
+        return cls(
+            mean_edge_length=mean_edge_length,
+            mean_dual_edge_length=mean_dual_edge_length,
+            mean_cell_area=mean_cell_area,
+            mean_dual_cell_area=mean_dual_cell_area,
+            **kwargs,
+        )
 
-        if mean_cell_area is not None:
-            self.mean_cell_area = mean_cell_area
+    def __post_init__(self) -> None:
+        if self.geometry_type is not None:
+            match self.geometry_type:
+                case base.GeometryType.ICOSAHEDRON:
+                    object.__setattr__(self, "domain_length", None)
+                    object.__setattr__(self, "domain_height", None)
+                    if self.radius is None:
+                        object.__setattr__(self, "radius", constants.EARTH_RADIUS)
+                case base.GeometryType.TORUS:
+                    object.__setattr__(self, "radius", None)
+                case _:
+                    ...
+
+        if self.global_num_cells is None and self.geometry_type is base.GeometryType.ICOSAHEDRON:
+            object.__setattr__(
+                self,
+                "global_num_cells",
+                compute_icosahedron_num_cells(self.grid_shape.subdivision),
+            )
+
+        if self.num_cells is None and self.global_num_cells is not None:
+            object.__setattr__(self, "num_cells", self.global_num_cells)
+
+        if (
+            self.mean_cell_area is None
+            and self.radius is not None
+            and self.global_num_cells is not None
+            and self.geometry_type is base.GeometryType.ICOSAHEDRON
+        ):
+            object.__setattr__(
+                self,
+                "mean_cell_area",
+                compute_mean_cell_area_for_sphere(self.radius, self.global_num_cells),
+            )
+
+        if self.characteristic_length is None and self.mean_cell_area is not None:
+            object.__setattr__(self, "characteristic_length", math.sqrt(self.mean_cell_area))
 
     @property
     def geometry_type(self) -> base.GeometryType | None:
         return self.grid_shape.geometry_type if self.grid_shape else None
 
-    @functools.cached_property
-    def num_cells(self) -> int:
-        match self.geometry_type:
-            case base.GeometryType.ICOSAHEDRON:
-                assert self.grid_shape.subdivision is not None
-                return compute_icosahedron_num_cells(self.grid_shape.subdivision)
-            case base.GeometryType.TORUS:
-                raise NotImplementedError("TODO : lookup torus cell number computation")
-            case _:
-                raise ValueError(f"Unknown geometry type {self.geometry_type}")
-
-    @functools.cached_property
-    def characteristic_length(self) -> float:
-        return math.sqrt(self.mean_cell_area)
-
-    @functools.cached_property
-    def mean_cell_area(self) -> float:
-        match self.geometry_type:
-            case base.GeometryType.ICOSAHEDRON:
-                return compute_mean_cell_area_for_sphere(self.radius, self.num_cells)
-            case base.GeometryType.TORUS:
-                raise NotImplementedError(
-                    f"mean_cell_area not implemented for {self.geometry_type}"
-                )
-            case _:
-                raise NotImplementedError(f"Unknown geometry type {self.geometry_type}")
+    @property
+    def subdivision(self) -> GridSubdivision | None:
+        return self.grid_shape.subdivision if self.grid_shape else None
 
 
 def compute_icosahedron_num_cells(subdivision: GridSubdivision) -> int: