Better torus support for mean cell area etc.

ci/default.yml

@@ -42,7 +42,7 @@ test_tools_datatests_aarch64:
     - if: $COMPONENT == 'common' && $LEVEL == 'integration'
       variables:
         NUM_PROCESSES: 1
-        SLURM_TIMELIMIT: '00:30:00'
+        SLURM_TIMELIMIT: '00:45:00'
     - if: $BACKEND == 'dace_gpu'
       variables:
         NUM_PROCESSES: 8

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

@@ -140,6 +140,8 @@ def __init__(
         input_fields_provider = factory.PrecomputedFieldProvider(
             {
                 # TODO(halungge): rescaled by grid_length_rescale_factor (mo_grid_tools.f90)
+                attrs.EDGE_LENGTH: extra_fields[gridfile.GeometryName.EDGE_LENGTH],
+                attrs.DUAL_EDGE_LENGTH: extra_fields[gridfile.GeometryName.DUAL_EDGE_LENGTH],
                 attrs.EDGE_CELL_DISTANCE: extra_fields[gridfile.GeometryName.EDGE_CELL_DISTANCE],
                 attrs.EDGE_VERTEX_DISTANCE: extra_fields[
                     gridfile.GeometryName.EDGE_VERTEX_DISTANCE
@@ -177,50 +179,12 @@ def __init__(
         self._register_computed_fields()
 
     def _register_computed_fields(self):
-        edge_length_provider = factory.ProgramFieldProvider(
-            func=stencils.compute_edge_length,
-            domain={
-                dims.EdgeDim: (
-                    self._edge_domain(h_grid.Zone.LOCAL),
-                    self._edge_domain(h_grid.Zone.LOCAL),
-                )
-            },
-            fields={
-                "length": attrs.EDGE_LENGTH,
-            },
-            deps={
-                "vertex_lat": attrs.VERTEX_LAT,
-                "vertex_lon": attrs.VERTEX_LON,
-            },
-            params={"radius": self._grid.global_properties.radius},
-        )
-        self.register_provider(edge_length_provider)
         meta = attrs.metadata_for_inverse(attrs.attrs[attrs.EDGE_LENGTH])
         name = meta["standard_name"]
         self._attrs.update({name: meta})
         inverse_edge_length = self._inverse_field_provider(attrs.EDGE_LENGTH)
         self.register_provider(inverse_edge_length)
 
-        dual_length_provider = factory.ProgramFieldProvider(
-            func=stencils.compute_cell_center_arc_distance,
-            domain={
-                dims.EdgeDim: (
-                    self._edge_domain(h_grid.Zone.LOCAL),
-                    self._edge_domain(h_grid.Zone.LOCAL),
-                )
-            },
-            fields={
-                "dual_edge_length": attrs.DUAL_EDGE_LENGTH,
-            },
-            deps={
-                "edge_neighbor_0_lat": "latitude_of_edge_cell_neighbor_0",
-                "edge_neighbor_0_lon": "longitude_of_edge_cell_neighbor_0",
-                "edge_neighbor_1_lat": "latitude_of_edge_cell_neighbor_1",
-                "edge_neighbor_1_lon": "longitude_of_edge_cell_neighbor_1",
-            },
-            params={"radius": self._grid.global_properties.radius},
-        )
-        self.register_provider(dual_length_provider)
         inverse_dual_edge_length = self._inverse_field_provider(attrs.DUAL_EDGE_LENGTH)
         self.register_provider(inverse_dual_edge_length)
 

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,44 @@ 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
+        )
+
+        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,
+            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,104 @@ def __init__(
         self.subdivision = subdivision
 
 
-@dataclasses.dataclass
+_T = TypeVar("_T")
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
 class GlobalGridParams:
     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:

model/common/tests/common/grid/unit_tests/test_grid_manager.py

@@ -361,7 +361,7 @@ def test_grid_manager_grid_level_and_root(
         global_num_cells
         == utils.run_grid_manager(
             grid_descriptor, keep_skip_values=True, backend=backend
-        ).grid.global_properties.num_cells
+        ).grid.global_properties.global_num_cells
     )