diff --git a/model/common/src/icon4py/model/common/grid/geometry.py b/model/common/src/icon4py/model/common/grid/geometry.py index df18139e89..79b59d602c 100644 --- a/model/common/src/icon4py/model/common/grid/geometry.py +++ b/model/common/src/icon4py/model/common/grid/geometry.py @@ -12,6 +12,7 @@ import gt4py.next.typing as gtx_typing from gt4py import next as gtx +from typing_extensions import assert_never import icon4py.model.common.math.helpers as math_helpers from icon4py.model.common import ( @@ -41,7 +42,7 @@ class GridGeometry(factory.FieldSource): """ Factory for the ICON grid geometry fields. - Computes geometry fields from the grid geographical coordinates fo cells, egdes, vertices. + Computes geometry fields from the grid geographical coordinates fo cells, edges, vertices. Computations are triggered upon first request. Can be queried for geometry fields and metadata @@ -80,33 +81,107 @@ def __init__( grid: icon.IconGrid, decomposition_info: definitions.DecompositionInfo, backend: gtx_typing.Backend | None, - coordinates: gm.CoordinateDict, - extra_fields: gm.GeometryDict, metadata: dict[str, model.FieldMetaData], - ): + ) -> None: """ Args: grid: IconGrid the grid topology decomposition_info: data structure containing owner masks for field dimensions backend: backend used for memory allocation and computation - coordinates: dictionary containing geographical coordinates for grid cells, edges and vertices, - extra_fields: fields that are not computed but directly read off the grid file, - currently only the edge_system_orientation cell_area. Should eventually disappear. metadata: a dictionary of FieldMetaData for all fields computed in GridGeometry. """ self._providers = {} self._backend = backend + self._xp = data_alloc.import_array_ns(backend) self._allocator = gtx.constructors.zeros.partial(allocator=backend) self._grid = grid self._decomposition_info = decomposition_info self._attrs = metadata - self._geometry_type: base.GeometryType | None = grid.global_properties.geometry_type + self._geometry_type: base.GeometryType = grid.global_properties.geometry_type self._edge_domain = h_grid.domain(dims.EdgeDim) log.info( f"initialized geometry for backend = '{self._backend_name()}' and grid = '{self._grid}'" ) + @staticmethod + def with_geometry_type( + grid: icon.IconGrid, + decomposition_info: definitions.DecompositionInfo, + backend: gtx_typing.Backend | None, + coordinates: gm.CoordinateDict, + extra_fields: gm.GeometryDict, + metadata: dict[str, model.FieldMetaData], + ) -> "GridGeometry": + match grid.global_properties.geometry_type: + case base.GeometryType.ICOSAHEDRON: + constructor = IcosahedronGridGeometry + case base.GeometryType.TORUS: + constructor = TorusGridGeometry + case _: + assert_never(grid.global_properties.geometry_type) + + return constructor( + grid, + decomposition_info, + backend, + coordinates, + extra_fields, + metadata, + ) + + def _inverse_field_provider(self, field_name: str) -> factory.FieldProvider: + meta = attrs.metadata_for_inverse(attrs.attrs[field_name]) + name = meta["standard_name"] + self._attrs.update({name: meta}) + provider = factory.ProgramFieldProvider( + func=math_helpers.compute_inverse_on_edges, + deps={"f": field_name}, + fields={"f_inverse": name}, + domain={ + dims.EdgeDim: ( + self._edge_domain(h_grid.Zone.LOCAL), + self._edge_domain(h_grid.Zone.LOCAL), + ) + }, + ) + return provider + + def __repr__(self) -> str: + geometry_name = self._geometry_type._name_ if self._geometry_type else "" + return ( + f"{self.__class__.__name__} for geometry_type={geometry_name} (grid={self._grid.id!r})" + ) + + @property + def metadata(self) -> dict[str, model.FieldMetaData]: + return self._attrs + + @property + def backend(self) -> gtx_typing.Backend: + return self._backend + + @property + def grid(self) -> icon.IconGrid: + return self._grid + + @property + def vertical_grid(self) -> None: + return None + + +class IcosahedronGridGeometry(GridGeometry): + def __init__( + self, + grid: icon.IconGrid, + decomposition_info: definitions.DecompositionInfo, + backend: gtx_typing.Backend | None, + coordinates: gm.CoordinateDict, + extra_fields: gm.GeometryDict, + metadata: dict[str, model.FieldMetaData], + ) -> None: + super().__init__(grid, decomposition_info, backend, metadata) + ( edge_orientation0_lat, edge_orientation0_lon, @@ -174,9 +249,11 @@ def __init__( } ) self.register_provider(input_fields_provider) - self._register_computed_fields() + self._register_computed_fields(input_fields_provider) - def _register_computed_fields(self) -> None: + def _register_computed_fields( + self, input_fields_provider: factory.PrecomputedFieldProvider + ) -> None: meta = attrs.metadata_for_inverse(attrs.attrs[attrs.EDGE_LENGTH]) name = meta["standard_name"] self._attrs.update({name: meta}) @@ -501,50 +578,255 @@ def _register_computed_fields(self) -> None: ) self.register_provider(cartesian_cell_centers) - def _inverse_field_provider(self, field_name: str) -> factory.FieldProvider: - meta = attrs.metadata_for_inverse(attrs.attrs[field_name]) + +class TorusGridGeometry(GridGeometry): + def __init__( + self, + grid: icon.IconGrid, + decomposition_info: definitions.DecompositionInfo, + backend: gtx_typing.Backend | None, + coordinates: gm.CoordinateDict, + extra_fields: gm.GeometryDict, + metadata: dict[str, model.FieldMetaData], + ) -> None: + super().__init__(grid, decomposition_info, backend, metadata) + + log.info( + f"initialized geometry for backend = '{self._backend_name()}' and grid = '{self._grid}'" + ) + + coordinates_ = { + attrs.CELL_LAT: coordinates[dims.CellDim]["lat"], + attrs.CELL_LON: coordinates[dims.CellDim]["lon"], + attrs.CELL_CENTER_X: coordinates[dims.CellDim]["x"], + attrs.CELL_CENTER_Y: coordinates[dims.CellDim]["y"], + attrs.CELL_CENTER_Z: coordinates[dims.CellDim]["z"], + attrs.VERTEX_LAT: coordinates[dims.VertexDim]["lat"], + attrs.VERTEX_LON: coordinates[dims.VertexDim]["lon"], + attrs.VERTEX_X: coordinates[dims.VertexDim]["x"], + attrs.VERTEX_Y: coordinates[dims.VertexDim]["y"], + attrs.VERTEX_Z: coordinates[dims.VertexDim]["z"], + attrs.EDGE_LON: coordinates[dims.EdgeDim]["lon"], + attrs.EDGE_LAT: coordinates[dims.EdgeDim]["lat"], + attrs.EDGE_CENTER_X: coordinates[dims.EdgeDim]["x"], + attrs.EDGE_CENTER_Y: coordinates[dims.EdgeDim]["y"], + attrs.EDGE_CENTER_Z: coordinates[dims.EdgeDim]["z"], + } + coodinate_provider = factory.PrecomputedFieldProvider(coordinates_) + self.register_provider(coodinate_provider) + + input_fields_provider = factory.PrecomputedFieldProvider( + { + 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 + ], + attrs.CELL_AREA: extra_fields[gridfile.GeometryName.CELL_AREA], + attrs.DUAL_AREA: extra_fields[gridfile.GeometryName.DUAL_AREA], + attrs.TANGENT_ORIENTATION: extra_fields[gridfile.GeometryName.TANGENT_ORIENTATION], + "edge_owner_mask": gtx.as_field( + (dims.EdgeDim,), + decomposition_info.owner_mask(dims.EdgeDim), + dtype=bool, + allocator=self._backend, + ), + attrs.CELL_NORMAL_ORIENTATION: extra_fields[ + gridfile.GeometryName.CELL_NORMAL_ORIENTATION + ], + attrs.VERTEX_EDGE_ORIENTATION: extra_fields[ + gridfile.GeometryName.EDGE_ORIENTATION_ON_VERTEX + ], + "vertex_owner_mask": gtx.as_field( + (dims.VertexDim,), + decomposition_info.owner_mask(dims.VertexDim), + allocator=self._backend, + dtype=bool, + ), + "cell_owner_mask": gtx.as_field( + (dims.CellDim,), + decomposition_info.owner_mask(dims.CellDim), + allocator=self._backend, + dtype=bool, + ), + } + ) + self.register_provider(input_fields_provider) + self._register_computed_fields(input_fields_provider) + + def _register_computed_fields( + self, input_fields_provider: factory.PrecomputedFieldProvider + ) -> None: + meta = attrs.metadata_for_inverse(attrs.attrs[attrs.EDGE_LENGTH]) name = meta["standard_name"] self._attrs.update({name: meta}) - provider = factory.ProgramFieldProvider( - func=math_helpers.compute_inverse_on_edges, - deps={"f": field_name}, - fields={"f_inverse": name}, + inverse_edge_length = self._inverse_field_provider(attrs.EDGE_LENGTH) + self.register_provider(inverse_edge_length) + + inverse_dual_edge_length = self._inverse_field_provider(attrs.DUAL_EDGE_LENGTH) + self.register_provider(inverse_dual_edge_length) + + # TODO(msimberg): Is this needed for torus? Even the serialized data + # doesn't seem to make sense. Earth radius is used even for the torus. + vertex_vertex_distance = factory.ProgramFieldProvider( + func=stencils.compute_distance_of_far_edges_in_diamond_torus, domain={ dims.EdgeDim: ( - self._edge_domain(h_grid.Zone.LOCAL), + self._edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2), self._edge_domain(h_grid.Zone.LOCAL), ) }, + fields={"far_vertex_distance": attrs.VERTEX_VERTEX_LENGTH}, + deps={ + "vertex_x": attrs.VERTEX_X, + "vertex_y": attrs.VERTEX_Y, + }, + params={ + "domain_length": self._grid.global_properties.domain_length, + "domain_height": self._grid.global_properties.domain_height, + }, ) - return provider - - def __repr__(self) -> str: - geometry_name = self._geometry_type._name_ if self._geometry_type else "" - return ( - f"{self.__class__.__name__} for geometry_type={geometry_name} (grid={self._grid.id!r})" + self.register_provider(vertex_vertex_distance) + inverse_far_edge_distance_provider = self._inverse_field_provider( + attrs.VERTEX_VERTEX_LENGTH ) + self.register_provider(inverse_far_edge_distance_provider) - @property - def metadata(self) -> dict[str, model.FieldMetaData]: - return self._attrs - - @property - def backend(self) -> gtx_typing.Backend | None: - return self._backend + edge_areas = factory.ProgramFieldProvider( + func=stencils.compute_edge_area, + deps={ + "owner_mask": "edge_owner_mask", + "primal_edge_length": attrs.EDGE_LENGTH, + "dual_edge_length": attrs.DUAL_EDGE_LENGTH, + }, + fields={"area": attrs.EDGE_AREA}, + domain={ + dims.EdgeDim: ( + self._edge_domain(h_grid.Zone.LOCAL), + self._edge_domain(h_grid.Zone.END), + ) + }, + ) + self.register_provider(edge_areas) + coriolis_params = factory.PrecomputedFieldProvider( + { + "coriolis_parameter": gtx.as_field( + (dims.EdgeDim,), + self._xp.zeros( + self._grid.start_index(self._edge_domain(h_grid.Zone.END)) + - self._grid.start_index(self._edge_domain(h_grid.Zone.LOCAL)) + ), + dtype=ta.wpfloat, + allocator=self._backend, + ) + } + ) + self.register_provider(coriolis_params) - @property - def grid(self) -> icon.IconGrid: - return self._grid + # normals and tangents + tangent_normal_coordinates = factory.ProgramFieldProvider( + func=stencils.compute_cartesian_coordinates_of_edge_tangent_and_normal_torus, + deps={ + "vertex_x": attrs.VERTEX_X, + "vertex_y": attrs.VERTEX_Y, + "edge_x": attrs.EDGE_CENTER_X, + "edge_y": attrs.EDGE_CENTER_Y, + "edge_orientation": attrs.TANGENT_ORIENTATION, + }, + fields={ + "tangent_x": attrs.EDGE_TANGENT_X, + "tangent_y": attrs.EDGE_TANGENT_Y, + "tangent_z": attrs.EDGE_TANGENT_Z, + "tangent_u": attrs.EDGE_DUAL_U, + "tangent_v": attrs.EDGE_DUAL_V, + "normal_x": attrs.EDGE_NORMAL_X, + "normal_y": attrs.EDGE_NORMAL_Y, + "normal_z": attrs.EDGE_NORMAL_Z, + "normal_u": attrs.EDGE_NORMAL_U, + "normal_v": attrs.EDGE_NORMAL_V, + }, + domain={ + dims.EdgeDim: ( + self._edge_domain(h_grid.Zone.LOCAL), + self._edge_domain(h_grid.Zone.END), + ) + }, + params={ + "domain_length": self._grid.global_properties.domain_length, + "domain_height": self._grid.global_properties.domain_height, + }, + ) + self.register_provider(tangent_normal_coordinates) - @property - def vertical_grid(self) -> None: - return None + # 3. primal_normal_vert, primal_normal_cell + normal_vert_wrapper = SparseFieldProviderWrapper( + tangent_normal_coordinates, + target_dims=attrs.attrs[attrs.EDGE_NORMAL_VERTEX_U]["dims"], + fields=(attrs.EDGE_NORMAL_VERTEX_U, attrs.EDGE_NORMAL_VERTEX_V), + pairs=( + ( + attrs.EDGE_NORMAL_X, + attrs.EDGE_NORMAL_X, + attrs.EDGE_NORMAL_X, + attrs.EDGE_NORMAL_X, + ), + ( + attrs.EDGE_NORMAL_Y, + attrs.EDGE_NORMAL_Y, + attrs.EDGE_NORMAL_Y, + attrs.EDGE_NORMAL_Y, + ), + ), + ) + self.register_provider(normal_vert_wrapper) + normal_cell_wrapper = SparseFieldProviderWrapper( + tangent_normal_coordinates, + target_dims=attrs.attrs[attrs.EDGE_NORMAL_CELL_U]["dims"], + fields=(attrs.EDGE_NORMAL_CELL_U, attrs.EDGE_NORMAL_CELL_V), + pairs=( + (attrs.EDGE_NORMAL_X, attrs.EDGE_NORMAL_X), + (attrs.EDGE_NORMAL_Y, attrs.EDGE_NORMAL_Y), + ), + ) + self.register_provider(normal_cell_wrapper) + # 3. dual normals: the dual normals are the edge tangents + tangent_vert_wrapper = SparseFieldProviderWrapper( + tangent_normal_coordinates, + target_dims=attrs.attrs[attrs.EDGE_TANGENT_VERTEX_U]["dims"], + fields=(attrs.EDGE_TANGENT_VERTEX_U, attrs.EDGE_TANGENT_VERTEX_V), + pairs=( + ( + attrs.EDGE_TANGENT_X, + attrs.EDGE_TANGENT_X, + attrs.EDGE_TANGENT_X, + attrs.EDGE_TANGENT_X, + ), + ( + attrs.EDGE_TANGENT_Y, + attrs.EDGE_TANGENT_Y, + attrs.EDGE_TANGENT_Y, + attrs.EDGE_TANGENT_Y, + ), + ), + ) + self.register_provider(tangent_vert_wrapper) + tangent_cell_wrapper = SparseFieldProviderWrapper( + tangent_normal_coordinates, + target_dims=attrs.attrs[attrs.EDGE_TANGENT_CELL_U]["dims"], + fields=(attrs.EDGE_TANGENT_CELL_U, attrs.EDGE_TANGENT_CELL_V), + pairs=( + (attrs.EDGE_TANGENT_X, attrs.EDGE_TANGENT_X), + (attrs.EDGE_TANGENT_Y, attrs.EDGE_TANGENT_Y), + ), + ) + self.register_provider(tangent_cell_wrapper) class SparseFieldProviderWrapper(factory.FieldProvider): def __init__( self, - field_provider: factory.ProgramFieldProvider, + field_provider: factory.FieldProvider, target_dims: Sequence[gtx.Dimension], fields: Sequence[str], pairs: Sequence[tuple[str, ...]], diff --git a/model/common/src/icon4py/model/common/grid/geometry_attributes.py b/model/common/src/icon4py/model/common/grid/geometry_attributes.py index ca5d1587b0..047e25267e 100644 --- a/model/common/src/icon4py/model/common/grid/geometry_attributes.py +++ b/model/common/src/icon4py/model/common/grid/geometry_attributes.py @@ -50,9 +50,6 @@ EDGE_LENGTH: Final[str] = "edge_length" DUAL_EDGE_LENGTH: Final[str] = "length_of_dual_edge" -EDGE_TANGENT_X: Final[str] = "x_component_of_edge_tangential_unit_vector" -EDGE_TANGENT_Y: Final[str] = "y_component_of_edge_tangential_unit_vector" -EDGE_TANGENT_Z: Final[str] = "z_component_of_edge_tangential_unit_vector" EDGE_TANGENT_VERTEX_U: Final[str] = "eastward_component_of_edge_tangent_on_vertex" EDGE_TANGENT_VERTEX_V: Final[str] = "northward_component_of_edge_tangent_on_vertex" EDGE_TANGENT_CELL_U: Final[str] = "eastward_component_of_edge_tangent_on_cell" @@ -63,6 +60,9 @@ EDGE_NORMAL_Z: Final[str] = "z_component_of_edge_normal_unit_vector" EDGE_NORMAL_U: Final[str] = "eastward_component_of_edge_normal" EDGE_NORMAL_V: Final[str] = "northward_component_of_edge_normal" +EDGE_TANGENT_X: Final[str] = "x_component_of_edge_tangential_unit_vector" +EDGE_TANGENT_Y: Final[str] = "y_component_of_edge_tangential_unit_vector" +EDGE_TANGENT_Z: Final[str] = "z_component_of_edge_tangential_unit_vector" EDGE_DUAL_U: Final[str] = "eastward_component_of_edge_tangent" EDGE_DUAL_V: Final[str] = "northward_component_of_edge_tangent" EDGE_NORMAL_VERTEX_U: Final[str] = "eastward_component_of_edge_normal_on_vertex" diff --git a/model/common/src/icon4py/model/common/grid/geometry_stencils.py b/model/common/src/icon4py/model/common/grid/geometry_stencils.py index 90addf46df..92ad02701a 100644 --- a/model/common/src/icon4py/model/common/grid/geometry_stencils.py +++ b/model/common/src/icon4py/model/common/grid/geometry_stencils.py @@ -17,6 +17,8 @@ from icon4py.model.common.math.helpers import ( arc_length_on_edges, cross_product_on_edges, + diff_on_edges_torus, + distance_on_edges_torus, geographical_to_cartesian_on_edges, geographical_to_cartesian_on_vertices, normalize_cartesian_vector_on_edges, @@ -56,6 +58,49 @@ def cartesian_coordinates_of_edge_tangent( return normalize_cartesian_vector_on_edges(x, y, z) +@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) +def cartesian_coordinates_of_edge_tangent_torus( + vertex_x: fa.VertexField[ta.wpfloat], + vertex_y: fa.VertexField[ta.wpfloat], + edge_orientation: fa.EdgeField[ta.wpfloat], + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, +) -> tuple[ + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], +]: + """ + Compute normalized cartesian vector tangential to an edge on a torus grid. + + That is: computes the distance between the two vertices adjacent to the edge: + t = d(v1, v2) + + Args: + vertex_x: x coordinates of vertices + vertex_y: y coordinates of vertices + edge_orientation: encoding of the edge orientation: (-1, +1) depending on whether the + edge is directed from first to second neighbor of vice versa. + Returns: + x: x coordinate of normalized tangent vector + y: y coordinate of normalized tangent vector + z: z coordinate of normalized tangent vector + """ + xdiff, ydiff = diff_on_edges_torus( + vertex_x(E2V[0]), + vertex_x(E2V[1]), + vertex_y(E2V[0]), + vertex_y(E2V[1]), + domain_length, + domain_height, + ) + x = edge_orientation * xdiff + y = edge_orientation * ydiff + z = 0.0 * x # TODO(msimberg): zeros + + return normalize_cartesian_vector_on_edges(x, y, z) + + @gtx.field_operator def cartesian_coordinates_of_edge_normal( edge_lat: fa.EdgeField[ta.wpfloat], @@ -93,6 +138,30 @@ def cartesian_coordinates_of_edge_normal( return normalize_cartesian_vector_on_edges(x, y, z) +@gtx.field_operator +def cartesian_coordinates_of_edge_normal_torus( + edge_tangent_x: fa.EdgeField[ta.wpfloat], + edge_tangent_y: fa.EdgeField[ta.wpfloat], +) -> tuple[ + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], +]: + """ + Compute the normal to the edge tangent vector on a torus grid. + + Args: + edge_tangent_x: x coordinate of the tangent + edge_tangent_y: y coordinate of the tangent + Returns: + edge_normal_x: x coordinate of the normal + edge_normal_y: y coordinate of the normal + edge_normal_z: y coordinate of the normal + """ + z = 0.0 * edge_tangent_x # TODO(msimberg): zeros + return normalize_cartesian_vector_on_edges(-edge_tangent_y, edge_tangent_x, z) + + @gtx.field_operator def cartesian_coordinates_edge_tangent_and_normal( vertex_lat: fa.VertexField[ta.wpfloat], @@ -123,6 +192,59 @@ def cartesian_coordinates_edge_tangent_and_normal( return tangent_x, tangent_y, tangent_z, normal_x, normal_y, normal_z +@gtx.field_operator +def cartesian_coordinates_edge_tangent_and_normal_torus( + vertex_x: fa.VertexField[ta.wpfloat], + vertex_y: fa.VertexField[ta.wpfloat], + edge_x: fa.EdgeField[ta.wpfloat], + edge_y: fa.EdgeField[ta.wpfloat], + edge_orientation: fa.EdgeField[ta.wpfloat], + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, +) -> tuple[ + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], +]: + """Compute normalized cartesian vectors of edge tangent and edge normal.""" + tangent_x, tangent_y, tangent_z = cartesian_coordinates_of_edge_tangent_torus( + vertex_x, + vertex_y, + edge_orientation, + domain_length, + domain_height, + ) + tangent_u = tangent_x + tangent_v = tangent_y + + normal_x, normal_y, normal_z = cartesian_coordinates_of_edge_normal_torus( + tangent_x, + tangent_y, + ) + normal_u = normal_x + normal_v = normal_y + + return ( + tangent_x, + tangent_y, + tangent_z, + tangent_u, + tangent_v, + normal_x, + normal_y, + normal_z, + normal_u, + normal_v, + ) + + @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_cartesian_coordinates_of_edge_tangent_and_normal( vertex_lat: fa.VertexField[ta.wpfloat], @@ -150,6 +272,52 @@ def compute_cartesian_coordinates_of_edge_tangent_and_normal( ) +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_cartesian_coordinates_of_edge_tangent_and_normal_torus( + vertex_x: fa.VertexField[ta.wpfloat], + vertex_y: fa.VertexField[ta.wpfloat], + edge_x: fa.EdgeField[ta.wpfloat], + edge_y: fa.EdgeField[ta.wpfloat], + edge_orientation: fa.EdgeField[ta.wpfloat], + tangent_x: fa.EdgeField[ta.wpfloat], + tangent_y: fa.EdgeField[ta.wpfloat], + tangent_z: fa.EdgeField[ta.wpfloat], + tangent_u: fa.EdgeField[ta.wpfloat], + tangent_v: fa.EdgeField[ta.wpfloat], + normal_x: fa.EdgeField[ta.wpfloat], + normal_y: fa.EdgeField[ta.wpfloat], + normal_z: fa.EdgeField[ta.wpfloat], + normal_u: fa.EdgeField[ta.wpfloat], + normal_v: fa.EdgeField[ta.wpfloat], + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, +): + cartesian_coordinates_edge_tangent_and_normal_torus( + vertex_x, + vertex_y, + edge_x, + edge_y, + edge_orientation, + domain_length, + domain_height, + out=( + tangent_x, + tangent_y, + tangent_z, + tangent_u, + tangent_v, + normal_x, + normal_y, + normal_z, + normal_u, + normal_v, + ), + domain={dims.EdgeDim: (horizontal_start, horizontal_end)}, + ) + + @gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) def zonal_and_meridional_component_of_edge_field_at_vertex( vertex_lat: fa.VertexField[ta.wpfloat], @@ -410,6 +578,38 @@ def arc_distance_of_far_edges_in_diamond( return far_vertex_vertex_length +@gtx.field_operator +def distance_of_far_edges_in_diamond_torus( + vertex_x: fa.VertexField[ta.wpfloat], + vertex_y: fa.VertexField[ta.wpfloat], + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, +) -> fa.EdgeField[ta.wpfloat]: + """ + Compute the distance between the "far" vertices of an edge on a torus grid. + + See arc_distance_of_far_edges_in_diamond for details. + + Args: + vertex_x: x coordinate of vertices + vertex_y: y coordinate of vertices + domain_length: length of the domain + domain_height: height of the domain + + Returns: + distance between the "far" vertices in the diamond. + + """ + return distance_on_edges_torus( + vertex_x(E2C2V[2]), + vertex_x(E2C2V[3]), + vertex_y(E2C2V[2]), + vertex_y(E2C2V[3]), + domain_length, + domain_height, + ) + + @gtx.field_operator def edge_length( vertex_lat: fa.VertexField[ta.wpfloat], @@ -502,6 +702,26 @@ def compute_arc_distance_of_far_edges_in_diamond( ) +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_distance_of_far_edges_in_diamond_torus( + vertex_x: fa.VertexField[ta.wpfloat], + vertex_y: fa.VertexField[ta.wpfloat], + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, + far_vertex_distance: fa.EdgeField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, +): + distance_of_far_edges_in_diamond_torus( + vertex_x, + vertex_y, + domain_length, + domain_height, + out=far_vertex_distance, + domain={dims.EdgeDim: (horizontal_start, horizontal_end)}, + ) + + @gtx.field_operator def edge_area( owner_mask: fa.EdgeField[bool], diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py index cdb39bf3b9..4d2a9866ed 100644 --- a/model/common/src/icon4py/model/common/grid/grid_manager.py +++ b/model/common/src/icon4py/model/common/grid/grid_manager.py @@ -63,7 +63,10 @@ def __call__(self, array: data_alloc.NDArray): ) -CoordinateDict: TypeAlias = dict[gtx.Dimension, dict[Literal["lat", "lon"], gtx.Field]] +# TODO(msimberg): x, y, z added temporarily for torus grids, but should they be in a separate dict? +CoordinateDict: TypeAlias = dict[ + gtx.Dimension, dict[Literal["lat", "lon", "x", "y", "z"], gtx.Field] +] # TODO (halungge): use a TypeDict for that GeometryDict: TypeAlias = dict[gridfile.GeometryName, gtx.Field] @@ -120,25 +123,35 @@ def __exit__(self, exc_type, exc_val, exc_tb): def __call__(self, allocator: gtx_typing.FieldBufferAllocationUtil, keep_skip_values: bool): if not self._reader: self.open() - self._geometry = self._read_geometry_fields(allocator) - self._grid = self._construct_grid(allocator=allocator, with_skip_values=keep_skip_values) - self._coordinates = self._read_coordinates(allocator) + + if geometry_type := self._reader.try_attribute(gridfile.MPIMPropertyName.GEOMETRY): + geometry_type = base.GeometryType(geometry_type) + + self._geometry = self._read_geometry_fields(allocator, geometry_type) + self._grid = self._construct_grid( + allocator=allocator, with_skip_values=keep_skip_values, geometry_type=geometry_type + ) + self._coordinates = self._read_coordinates(allocator, geometry_type) self.close() - def _read_coordinates(self, backend: gtx_typing.Backend | None) -> CoordinateDict: - return { + def _read_coordinates( + self, + allocator: gtx_typing.FieldBufferAllocationUtil | None, + geometry_type: base.GeometryType | None, + ) -> CoordinateDict: + coordinates = { dims.CellDim: { "lat": gtx.as_field( (dims.CellDim,), self._reader.variable(gridfile.CoordinateName.CELL_LATITUDE), dtype=ta.wpfloat, - allocator=backend, + allocator=allocator, ), "lon": gtx.as_field( (dims.CellDim,), self._reader.variable(gridfile.CoordinateName.CELL_LONGITUDE), dtype=ta.wpfloat, - allocator=backend, + allocator=allocator, ), }, dims.EdgeDim: { @@ -146,32 +159,94 @@ def _read_coordinates(self, backend: gtx_typing.Backend | None) -> CoordinateDic (dims.EdgeDim,), self._reader.variable(gridfile.CoordinateName.EDGE_LATITUDE), dtype=ta.wpfloat, - allocator=backend, + allocator=allocator, ), "lon": gtx.as_field( (dims.EdgeDim,), self._reader.variable(gridfile.CoordinateName.EDGE_LONGITUDE), dtype=ta.wpfloat, - allocator=backend, + allocator=allocator, ), }, dims.VertexDim: { "lat": gtx.as_field( (dims.VertexDim,), self._reader.variable(gridfile.CoordinateName.VERTEX_LATITUDE), - allocator=backend, + allocator=allocator, dtype=ta.wpfloat, ), "lon": gtx.as_field( (dims.VertexDim,), self._reader.variable(gridfile.CoordinateName.VERTEX_LONGITUDE), - allocator=backend, + allocator=allocator, dtype=ta.wpfloat, ), }, } - def _read_geometry_fields(self, allocator: gtx_typing.FieldBufferAllocationUtil): + if geometry_type == base.GeometryType.TORUS: + coordinates[dims.CellDim]["x"] = gtx.as_field( + (dims.CellDim,), + self._reader.variable(gridfile.CoordinateName.CELL_X), + dtype=ta.wpfloat, + allocator=allocator, + ) + coordinates[dims.CellDim]["y"] = gtx.as_field( + (dims.CellDim,), + self._reader.variable(gridfile.CoordinateName.CELL_Y), + dtype=ta.wpfloat, + allocator=allocator, + ) + coordinates[dims.CellDim]["z"] = gtx.as_field( + (dims.CellDim,), + self._reader.variable(gridfile.CoordinateName.CELL_Z), + dtype=ta.wpfloat, + allocator=allocator, + ) + coordinates[dims.EdgeDim]["x"] = gtx.as_field( + (dims.EdgeDim,), + self._reader.variable(gridfile.CoordinateName.EDGE_X), + dtype=ta.wpfloat, + allocator=allocator, + ) + coordinates[dims.EdgeDim]["y"] = gtx.as_field( + (dims.EdgeDim,), + self._reader.variable(gridfile.CoordinateName.EDGE_Y), + dtype=ta.wpfloat, + allocator=allocator, + ) + coordinates[dims.EdgeDim]["z"] = gtx.as_field( + (dims.EdgeDim,), + self._reader.variable(gridfile.CoordinateName.EDGE_Z), + dtype=ta.wpfloat, + allocator=allocator, + ) + coordinates[dims.VertexDim]["x"] = gtx.as_field( + (dims.VertexDim,), + self._reader.variable(gridfile.CoordinateName.VERTEX_X), + dtype=ta.wpfloat, + allocator=allocator, + ) + coordinates[dims.VertexDim]["y"] = gtx.as_field( + (dims.VertexDim,), + self._reader.variable(gridfile.CoordinateName.VERTEX_Y), + dtype=ta.wpfloat, + allocator=allocator, + ) + coordinates[dims.VertexDim]["z"] = gtx.as_field( + (dims.VertexDim,), + self._reader.variable(gridfile.CoordinateName.VERTEX_Z), + dtype=ta.wpfloat, + allocator=allocator, + ) + + return coordinates + + def _read_geometry_fields( + self, + allocator: gtx_typing.FieldBufferAllocationUtil, + geometry_type: base.GeometryType | None, + ) -> GeometryDict: return { # TODO(halungge): still needs to ported, values from "our" grid files contains (wrong) values: # based on bug in generator fixed with this [PR40](https://gitlab.dkrz.de/dwd-sw/dwd_icon_tools/-/merge_requests/40) . @@ -242,7 +317,7 @@ def _read_grid_refinement_fields( Args: decomposition_info: Optional decomposition information, if not provided the grid is assumed to be a single node run. - backend: Optional backend to use for reading the fields, if not provided the default backend is used. + allocator: Optional allocator to use for reading the fields, if not provided the default backend is used. Returns: dict[gtx.Dimension, gtx.Field]: A dictionary containing the refinement control fields for each dimension. """ @@ -274,7 +349,10 @@ def coordinates(self) -> CoordinateDict: return self._coordinates def _construct_grid( - self, allocator: gtx_typing.FieldBufferAllocationUtil, with_skip_values: bool + self, + allocator: gtx_typing.FieldBufferAllocationUtil | None, + with_skip_values: bool, + geometry_type: base.GeometryType | None, ) -> icon.IconGrid: """Construct the grid topology from the icon grid file. diff --git a/model/common/src/icon4py/model/common/grid/gridfile.py b/model/common/src/icon4py/model/common/grid/gridfile.py index 0eaeee8c2b..5873089d67 100644 --- a/model/common/src/icon4py/model/common/grid/gridfile.py +++ b/model/common/src/icon4py/model/common/grid/gridfile.py @@ -187,6 +187,12 @@ class GeometryName(FieldName): # TODO(halungge): compute from coordinates EDGE_CELL_DISTANCE = "edge_cell_distance" EDGE_VERTEX_DISTANCE = "edge_vert_distance" + EDGE_NORMAL_X = "edge_primal_normal_cartesian_x" + EDGE_NORMAL_Y = "edge_primal_normal_cartesian_y" + EDGE_NORMAL_Z = "edge_primal_normal_cartesian_z" + EDGE_TANGENT_X = "edge_dual_normal_cartesian_x" + EDGE_TANGENT_Y = "edge_dual_normal_cartesian_y" + EDGE_TANGENT_Z = "edge_dual_normal_cartesian_z" class CoordinateName(FieldName): @@ -202,6 +208,16 @@ class CoordinateName(FieldName): VERTEX_LONGITUDE = "vlon" VERTEX_LATITUDE = "vlat" + CELL_X = "cell_circumcenter_cartesian_x" + CELL_Y = "cell_circumcenter_cartesian_y" + CELL_Z = "cell_circumcenter_cartesian_z" + EDGE_X = "edge_middle_cartesian_x" + EDGE_Y = "edge_middle_cartesian_y" + EDGE_Z = "edge_middle_cartesian_z" + VERTEX_X = "cartesian_x_vertices" + VERTEX_Y = "cartesian_y_vertices" + VERTEX_Z = "cartesian_z_vertices" + class GridRefinementName(FieldName): """Names of arrays in grid file defining the grid control, definition of boundaries layers, start and end indices of horizontal zones.""" diff --git a/model/common/src/icon4py/model/common/grid/icon.py b/model/common/src/icon4py/model/common/grid/icon.py index 8382b8ea72..40ac367f4b 100644 --- a/model/common/src/icon4py/model/common/grid/icon.py +++ b/model/common/src/icon4py/model/common/grid/icon.py @@ -166,8 +166,8 @@ def __post_init__(self) -> 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 + def geometry_type(self) -> base.GeometryType: + return self.grid_shape.geometry_type @property def subdivision(self) -> GridSubdivision | None: diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py index cb154ea4ed..480cb9d048 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py @@ -15,6 +15,7 @@ from icon4py.model.common import constants, dimension as dims from icon4py.model.common.decomposition import definitions from icon4py.model.common.grid import ( + base, geometry, geometry_attributes as geometry_attrs, grid_refinement as refinement, @@ -54,7 +55,9 @@ def __init__( self._attrs = metadata self._providers: dict[str, factory.FieldProvider] = {} self._geometry = geometry_source + geometry_type = self._grid.global_properties.geometry_type characteristic_length = self._grid.global_properties.characteristic_length + mean_dual_edge_length = self._grid.global_properties.mean_dual_edge_length # TODO @halungge: Dummy config dict - to be replaced by real configuration self._config = { "divavg_cntrwgt": 0.5, @@ -66,13 +69,13 @@ def __init__( "rbf_kernel_edge": rbf.DEFAULT_RBF_KERNEL[rbf.RBFDimension.EDGE], "rbf_kernel_vertex": rbf.DEFAULT_RBF_KERNEL[rbf.RBFDimension.VERTEX], "rbf_scale_cell": rbf.compute_default_rbf_scale( - characteristic_length, rbf.RBFDimension.CELL + geometry_type, characteristic_length, mean_dual_edge_length, rbf.RBFDimension.CELL ), "rbf_scale_edge": rbf.compute_default_rbf_scale( - characteristic_length, rbf.RBFDimension.EDGE + geometry_type, characteristic_length, mean_dual_edge_length, rbf.RBFDimension.EDGE ), "rbf_scale_vertex": rbf.compute_default_rbf_scale( - characteristic_length, rbf.RBFDimension.VERTEX + geometry_type, characteristic_length, mean_dual_edge_length, rbf.RBFDimension.VERTEX ), } log.info( @@ -182,31 +185,59 @@ def _register_computed_fields(self) -> None: self.register_provider(geofac_grdiv) - cell_average_weight = factory.NumpyDataProvider( - func=functools.partial( - interpolation_fields.compute_mass_conserving_bilinear_cell_average_weight, - array_ns=self._xp, - ), - fields=(attrs.C_BLN_AVG,), - domain=(dims.CellDim, dims.C2E2CODim), - deps={ - "lat": geometry_attrs.CELL_LAT, - "lon": geometry_attrs.CELL_LON, - "cell_areas": geometry_attrs.CELL_AREA, - "cell_owner_mask": "cell_owner_mask", - }, - connectivities={"c2e2c0": dims.C2E2CODim}, - params={ - "horizontal_start": self.grid.start_index( - cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) - ), - "horizontal_start_level_3": self.grid.start_index( - cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) - ), - "divavg_cntrwgt": self._config["divavg_cntrwgt"], - }, - ) - self.register_provider(cell_average_weight) + match self.grid.global_properties.geometry_type: + case base.GeometryType.ICOSAHEDRON: + cell_average_weight = factory.NumpyDataProvider( + func=functools.partial( + interpolation_fields.compute_mass_conserving_bilinear_cell_average_weight, + array_ns=self._xp, + ), + fields=(attrs.C_BLN_AVG,), + domain=(dims.CellDim, dims.C2E2CODim), + deps={ + "lat": geometry_attrs.CELL_LAT, + "lon": geometry_attrs.CELL_LON, + "cell_areas": geometry_attrs.CELL_AREA, + "cell_owner_mask": "cell_owner_mask", + }, + connectivities={"c2e2c0": dims.C2E2CODim}, + params={ + "horizontal_start": self.grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ), + "horizontal_start_level_3": self.grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) + ), + "divavg_cntrwgt": self._config["divavg_cntrwgt"], + }, + ) + self.register_provider(cell_average_weight) + case base.GeometryType.TORUS: + cell_average_weight = factory.NumpyDataProvider( + func=functools.partial( + interpolation_fields.compute_mass_conserving_bilinear_cell_average_weight_torus, + array_ns=self._xp, + ), + fields=(attrs.C_BLN_AVG,), + domain=(dims.CellDim, dims.C2E2CODim), + deps={ + "cell_areas": geometry_attrs.CELL_AREA, + "cell_owner_mask": "cell_owner_mask", + }, + connectivities={"c2e2c0": dims.C2E2CODim}, + params={ + "horizontal_start": self.grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ), + "horizontal_start_level_3": self.grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) + ), + "divavg_cntrwgt": self._config["divavg_cntrwgt"], + }, + ) + self.register_provider(cell_average_weight) + case _: + raise ValueError("TODO") c_lin_e = factory.NumpyDataProvider( func=functools.partial(interpolation_fields.compute_c_lin_e, array_ns=self._xp), @@ -274,47 +305,83 @@ def _register_computed_fields(self) -> None: ) self.register_provider(e_flx_avg) - e_bln_c_s = factory.NumpyDataProvider( - func=functools.partial(interpolation_fields.compute_e_bln_c_s, array_ns=self._xp), - fields=(attrs.E_BLN_C_S,), - domain=(dims.CellDim, dims.C2EDim), - deps={ - "cells_lat": geometry_attrs.CELL_LAT, - "cells_lon": geometry_attrs.CELL_LON, - "edges_lat": geometry_attrs.EDGE_LAT, - "edges_lon": geometry_attrs.EDGE_LON, - }, - connectivities={"c2e": dims.C2EDim}, - params={"weighting_factor": self._config["weighting_factor"]}, - ) - self.register_provider(e_bln_c_s) + match self.grid.global_properties.geometry_type: + case base.GeometryType.ICOSAHEDRON: + e_bln_c_s = factory.NumpyDataProvider( + func=functools.partial( + interpolation_fields.compute_e_bln_c_s, array_ns=self._xp + ), + fields=(attrs.E_BLN_C_S,), + domain=(dims.CellDim, dims.C2EDim), + deps={ + "cells_lat": geometry_attrs.CELL_LAT, + "cells_lon": geometry_attrs.CELL_LON, + "edges_lat": geometry_attrs.EDGE_LAT, + "edges_lon": geometry_attrs.EDGE_LON, + }, + connectivities={"c2e": dims.C2EDim}, + params={"weighting_factor": self._config["weighting_factor"]}, + ) + self.register_provider(e_bln_c_s) + case base.GeometryType.TORUS: + e_bln_c_s = factory.NumpyDataProvider( + func=functools.partial( + interpolation_fields.compute_e_bln_c_s_torus, array_ns=self._xp + ), + fields=(attrs.E_BLN_C_S,), + domain=(dims.CellDim, dims.C2EDim), + deps={}, + connectivities={"c2e": dims.C2EDim}, + params={}, + ) + self.register_provider(e_bln_c_s) + case _: + raise ValueError("TODO") - pos_on_tplane_e_x_y = factory.NumpyDataProvider( - func=functools.partial( - interpolation_fields.compute_pos_on_tplane_e_x_y, array_ns=self._xp - ), - fields=(attrs.POS_ON_TPLANE_E_X, attrs.POS_ON_TPLANE_E_Y), - domain=(dims.EdgeDim, dims.E2CDim), - deps={ - "primal_normal_v1": geometry_attrs.EDGE_NORMAL_U, - "primal_normal_v2": geometry_attrs.EDGE_NORMAL_V, - "dual_normal_v1": geometry_attrs.EDGE_DUAL_U, - "dual_normal_v2": geometry_attrs.EDGE_DUAL_V, - "cells_lon": geometry_attrs.CELL_LON, - "cells_lat": geometry_attrs.CELL_LAT, - "edges_lon": geometry_attrs.EDGE_LON, - "edges_lat": geometry_attrs.EDGE_LAT, - "owner_mask": "edge_owner_mask", - }, - connectivities={"e2c": dims.E2CDim}, - params={ - "grid_sphere_radius": constants.EARTH_RADIUS, - "horizontal_start": self.grid.start_index( - edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) - ), - }, - ) - self.register_provider(pos_on_tplane_e_x_y) + match self.grid.global_properties.geometry_type: + case base.GeometryType.ICOSAHEDRON: + pos_on_tplane_e_x_y = factory.NumpyDataProvider( + func=functools.partial( + interpolation_fields.compute_pos_on_tplane_e_x_y, array_ns=self._xp + ), + fields=(attrs.POS_ON_TPLANE_E_X, attrs.POS_ON_TPLANE_E_Y), + domain=(dims.EdgeDim, dims.E2CDim), + deps={ + "primal_normal_v1": geometry_attrs.EDGE_NORMAL_U, + "primal_normal_v2": geometry_attrs.EDGE_NORMAL_V, + "dual_normal_v1": geometry_attrs.EDGE_DUAL_U, + "dual_normal_v2": geometry_attrs.EDGE_DUAL_V, + "cells_lon": geometry_attrs.CELL_LON, + "cells_lat": geometry_attrs.CELL_LAT, + "edges_lon": geometry_attrs.EDGE_LON, + "edges_lat": geometry_attrs.EDGE_LAT, + "owner_mask": "edge_owner_mask", + }, + connectivities={"e2c": dims.E2CDim}, + params={ + "grid_sphere_radius": constants.EARTH_RADIUS, + "horizontal_start": self.grid.start_index( + edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ), + }, + ) + self.register_provider(pos_on_tplane_e_x_y) + case base.GeometryType.TORUS: + pos_on_tplane_e_x_y = factory.NumpyDataProvider( + func=functools.partial( + interpolation_fields.compute_pos_on_tplane_e_x_y_torus, array_ns=self._xp + ), + fields=(attrs.POS_ON_TPLANE_E_X, attrs.POS_ON_TPLANE_E_Y), + domain=(dims.EdgeDim, dims.E2CDim), + deps={ + "dual_edge_length": geometry_attrs.DUAL_EDGE_LENGTH, + }, + connectivities={"e2c": dims.E2CDim}, + params={}, + ) + self.register_provider(pos_on_tplane_e_x_y) + case _: + raise ValueError("TODO") cells_aw_verts = factory.NumpyDataProvider( func=functools.partial(interpolation_fields.compute_cells_aw_verts, array_ns=self._xp), @@ -359,10 +426,17 @@ def _register_computed_fields(self) -> None: connectivities={"rbf_offset": dims.C2E2C2EDim}, params={ "rbf_kernel": self._config["rbf_kernel_cell"].value, + "geometry_type": self._grid.global_properties.geometry_type.value, "scale_factor": self._config["rbf_scale_cell"], "horizontal_start": self.grid.start_index( cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) ), + "domain_length": self._grid.global_properties.domain_length + if self._grid.global_properties.domain_length + else -1.0, + "domain_height": self._grid.global_properties.domain_height + if self._grid.global_properties.domain_height + else -1.0, }, ) self.register_provider(rbf_vec_coeff_c) @@ -386,10 +460,17 @@ def _register_computed_fields(self) -> None: connectivities={"rbf_offset": dims.E2C2EDim}, params={ "rbf_kernel": self._config["rbf_kernel_edge"].value, + "geometry_type": self._grid.global_properties.geometry_type.value, "scale_factor": self._config["rbf_scale_edge"], "horizontal_start": self.grid.start_index( edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) ), + "domain_length": self._grid.global_properties.domain_length + if self._grid.global_properties.domain_length + else -1.0, + "domain_height": self._grid.global_properties.domain_height + if self._grid.global_properties.domain_height + else -1.0, }, ) self.register_provider(rbf_vec_coeff_e) @@ -414,10 +495,17 @@ def _register_computed_fields(self) -> None: connectivities={"rbf_offset": dims.V2EDim}, params={ "rbf_kernel": self._config["rbf_kernel_vertex"].value, + "geometry_type": self._grid.global_properties.geometry_type.value, "scale_factor": self._config["rbf_scale_vertex"], "horizontal_start": self.grid.start_index( vertex_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) ), + "domain_length": self._grid.global_properties.domain_length + if self._grid.global_properties.domain_length + else -1.0, + "domain_height": self._grid.global_properties.domain_height + if self._grid.global_properties.domain_height + else -1.0, }, ) self.register_provider(rbf_vec_coeff_v) diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py index c39ed10956..982e95fe09 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py @@ -622,6 +622,33 @@ def _enforce_mass_conservation( return c_bln_avg +def _compute_uniform_c_bln_avg( + c2e2c: data_alloc.NDArray, + divavg_cntrwgt: ta.wpfloat, + horizontal_start: gtx.int32, + array_ns: ModuleType = np, +) -> data_alloc.NDArray: + """ + Compute bilinear cell average weight for a torus grid. + + Args: + divavg_cntrwgt: + c2e2c: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, C2E2CDim], gtx.int32] + horizontal_start: + + Returns: + c_bln_avg: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], ta.wpfloat] + """ + local_weight = divavg_cntrwgt + neighbor_weight = (1.0 - divavg_cntrwgt) / 3.0 + c_bln_avg = array_ns.full( + (c2e2c.shape[0], c2e2c.shape[1] + 1), + [local_weight, neighbor_weight, neighbor_weight, neighbor_weight], + ) + + return c_bln_avg + + def compute_mass_conserving_bilinear_cell_average_weight( c2e2c0: data_alloc.NDArray, lat: data_alloc.NDArray, @@ -647,6 +674,31 @@ def compute_mass_conserving_bilinear_cell_average_weight( ) +def compute_mass_conserving_bilinear_cell_average_weight_torus( + c2e2c0: data_alloc.NDArray, + cell_areas: data_alloc.NDArray, + cell_owner_mask: data_alloc.NDArray, + divavg_cntrwgt: ta.wpfloat, + horizontal_start: gtx.int32, + horizontal_start_level_3: gtx.int32, + array_ns: ModuleType = np, +) -> data_alloc.NDArray: + c_bln_avg = _compute_uniform_c_bln_avg( + c2e2c0[:, 1:], divavg_cntrwgt, horizontal_start, array_ns + ) + # TODO(msimberg): Exact result for torus without the following. 1e-16 error + # with the the following. Is it needed? + return _force_mass_conservation_to_c_bln_avg( + c2e2c0, + c_bln_avg, + cell_areas, + cell_owner_mask, + divavg_cntrwgt, + horizontal_start_level_3, + array_ns, + ) + + def _create_inverse_neighbor_index( source_offset, inverse_offset, array_ns: ModuleType ) -> data_alloc.NDArray: @@ -956,6 +1008,22 @@ def compute_e_bln_c_s( return e_bln_c_s +def compute_e_bln_c_s_torus( + c2e: data_alloc.NDArray, + array_ns: ModuleType = np, +) -> data_alloc.NDArray: + """ + Compute e_bln_c_s. + + Args: + c2e: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], gtx.int32] + + Returns: + e_bln_c_s: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], ta.wpfloat] + """ + return array_ns.full_like(c2e, 1.0 / 3.0, dtype=ta.wpfloat) + + def compute_pos_on_tplane_e_x_y( grid_sphere_radius: ta.wpfloat, primal_normal_v1: data_alloc.NDArray, @@ -1047,3 +1115,44 @@ def compute_pos_on_tplane_e_x_y( ) return pos_on_tplane_e[:, :, 0], pos_on_tplane_e[:, :, 1] + + +def compute_pos_on_tplane_e_x_y_torus( + dual_edge_length: data_alloc.NDArray, + e2c: data_alloc.NDArray, + array_ns: ModuleType = np, +) -> data_alloc.NDArray: + """ + Compute pos_on_tplane_e_x_y. + get geographical coordinates of edge midpoint + get line and block indices of neighbour cells + get geographical coordinates of first cell center + projection first cell center into local \\lambda-\\Phi-system + get geographical coordinates of second cell center + projection second cell center into local \\lambda-\\Phi-system + + Args: + dual_edge_length: numpy_array, representing a gtx.Field[gtx.Dims[EdgeDim], ta.wpfloat] + e2c: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], gtx.int32] + + Returns: + pos_on_tplane_e_x: \\ numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], ta.wpfloat] + pos_on_tplane_e_y: // + """ + # The implementation makes the simplifying assumptions that: + # - The torus grid consists of equilateral triangles, which means that the + # neighboring cell centers must always be at 0.5 * dual_edge_length from + # the edge center (the edge lies symmetrically perpendicular to the primal + # edge). + # - The neighboring cell centers are exactly along the primal normal/dual + # tangent direction, which means the x component in the local coordinate + # system is always zero, and the y component is always 0.5 * + # dual_edge_length. + # - The first neighbor cell is in the opposite direction of the primal + # normal and the second neighbor is in the direction of the primal normal. + half_dual_edge_length = 0.5 * dual_edge_length[0] + num_edges = e2c.shape[0] + return ( + array_ns.full([num_edges, 2], [-half_dual_edge_length, half_dual_edge_length]), + array_ns.zeros([num_edges, 2]), + ) diff --git a/model/common/src/icon4py/model/common/interpolation/rbf_interpolation.py b/model/common/src/icon4py/model/common/interpolation/rbf_interpolation.py index 38b9cf0912..9a00aaadcc 100644 --- a/model/common/src/icon4py/model/common/interpolation/rbf_interpolation.py +++ b/model/common/src/icon4py/model/common/interpolation/rbf_interpolation.py @@ -47,26 +47,38 @@ class InterpolationKernel(enum.Enum): } -def compute_default_rbf_scale(mean_characteristic_length: ta.wpfloat, dim: RBFDimension): - """Compute the default RBF scale factor. This assumes that the Gaussian - kernel is used for vertices and cells, and that the inverse multiquadratic - kernel is used for edges.""" - - threshold = 2.5 if dim == RBFDimension.CELL else 2.0 - c1 = 0.4 if dim == RBFDimension.EDGE else 1.8 - if dim == RBFDimension.CELL: - c2 = 3.75 - c3 = 0.9 - elif dim == RBFDimension.VERTEX: - c2 = 3.0 - c3 = 0.96 - else: - c2 = 2.0 - c3 = 0.325 - - resol = mean_characteristic_length / 1000.0 - scale = 0.5 / (1.0 + c1 * math.log(threshold / resol) ** c2) if resol < threshold else 0.5 - return scale * (resol / 0.125) ** c3 if resol <= 0.125 else scale +def compute_default_rbf_scale( + geometry_type: base_grid.GeometryType, + mean_characteristic_length: ta.wpfloat, + mean_dual_edge_length: ta.wpfloat, + dim: RBFDimension, +): + """Compute the default RBF scale factor. This assumes that the default + interpolation kernels are used for each dimension""" + + match geometry_type: + case base_grid.GeometryType.ICOSAHEDRON: + threshold = 2.5 if dim == RBFDimension.CELL else 2.0 + c1 = 0.4 if dim == RBFDimension.EDGE else 1.8 + if dim == RBFDimension.CELL: + c2 = 3.75 + c3 = 0.9 + elif dim == RBFDimension.VERTEX: + c2 = 3.0 + c3 = 0.96 + else: + c2 = 2.0 + c3 = 0.325 + + resol = mean_characteristic_length / 1000.0 + scale = ( + 0.5 / (1.0 + c1 * math.log(threshold / resol) ** c2) if resol < threshold else 0.5 + ) + return scale * (resol / 0.125) ** c3 if resol <= 0.125 else scale + case base_grid.GeometryType.TORUS: + return mean_dual_edge_length + case _: + raise ValueError(f"Unsupported geometry type: {geometry_type}") def construct_rbf_matrix_offsets_tables_for_cells( @@ -105,7 +117,15 @@ def _dot_product( return array_ns.matmul(v1, v2_tilde) -def _arc_length_pairwise(v: data_alloc.NDArray, array_ns: ModuleType = np) -> data_alloc.NDArray: +# TODO(msimberg): Rename again. This is arc length on a unit sphere for +# icosahedral grids and real distance for torus grids. +def _distance_pairwise( + geometry_type: base_grid.GeometryType, + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, + v: data_alloc.NDArray, + array_ns: ModuleType = np, +) -> data_alloc.NDArray: """ Compute the pairwise arc lengths between points in each row of v. @@ -115,23 +135,41 @@ def _arc_length_pairwise(v: data_alloc.NDArray, array_ns: ModuleType = np) -> da dimension of the points. array_ns: numpy or cupy module to use for computations. """ - # For pairs of points p1 and p2 compute: - # arccos(dot(p1, p2) / (norm(p1) * norm(p2))) noqa: ERA001 - # Compute all pairs of dot products - arc_lengths = _dot_product(v, v, array_ns=array_ns) - # Use the dot product of the diagonals to get the norm of each point - norms = array_ns.sqrt(array_ns.diagonal(arc_lengths, axis1=1, axis2=2)) - # Divide the dot products by the broadcasted norms - array_ns.divide(arc_lengths, norms[:, :, array_ns.newaxis], out=arc_lengths) - array_ns.divide(arc_lengths, norms[:, array_ns.newaxis, :], out=arc_lengths) - # Ensure all points are within [-1.0, 1.0] (may be outside due to numerical - # inaccuracies) - array_ns.clip(arc_lengths, -1.0, 1.0, out=arc_lengths) - return array_ns.arccos(arc_lengths) - - -def _arc_length_vector_matrix( - v1: data_alloc.NDArray, v2: data_alloc.NDArray, array_ns: ModuleType = np + match geometry_type: + case base_grid.GeometryType.ICOSAHEDRON: + # For pairs of points p1 and p2 compute: + # arccos(dot(p1, p2) / (norm(p1) * norm(p2))) noqa: ERA001 + # Compute all pairs of dot products + arc_lengths = _dot_product(v, v, array_ns=array_ns) + # Use the dot product of the diagonals to get the norm of each point + norms = array_ns.sqrt(array_ns.diagonal(arc_lengths, axis1=1, axis2=2)) + # Divide the dot products by the broadcasted norms + array_ns.divide(arc_lengths, norms[:, :, array_ns.newaxis], out=arc_lengths) + array_ns.divide(arc_lengths, norms[:, array_ns.newaxis, :], out=arc_lengths) + # Ensure all points are within [-1.0, 1.0] (may be outside due to numerical + # inaccuracies) + array_ns.clip(arc_lengths, -1.0, 1.0, out=arc_lengths) + return array_ns.arccos(arc_lengths) + case base_grid.GeometryType.TORUS: + # For pairs of points p1 and p2 compute: + # norm(p1 - p2), taking into account the periodic boundaries noqa: ERA001 + diff = array_ns.abs(v[:, :, array_ns.newaxis, :] - v[:, array_ns.newaxis, :, :]) + domain_size = array_ns.asarray([domain_length, domain_height, ta.wpfloat(0.0)]) + domain_size_expanded = domain_size[array_ns.newaxis, array_ns.newaxis, :] + inverted_diff = array_ns.subtract(domain_size_expanded, diff) + array_ns.minimum(diff, inverted_diff, out=diff) + return array_ns.linalg.norm(diff, axis=-1) + case _: + raise ValueError(f"Unsupported geometry type: {geometry_type}") + + +def _distance_vector_matrix( + geometry_type: base_grid.GeometryType, + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, + v1: data_alloc.NDArray, + v2: data_alloc.NDArray, + array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ Compute the arc lengths between each point in v1 and the points in v2 at the same row. @@ -144,19 +182,32 @@ def _arc_length_vector_matrix( of the points. array_ns: numpy or cupy module to use for computations. """ - # For pairs of points p1 and p2 compute: - # arccos(dot(p1, p2) / (norm(p1) * norm(p2))) noqa: ERA001 - # Compute all pairs of dot products - arc_lengths = _dot_product(v1, v2, array_ns=array_ns) - v1_norm = array_ns.linalg.norm(v1, axis=-1) - v2_norm = array_ns.linalg.norm(v2, axis=-1) - # Divide the dot products by the broadcasted norms - array_ns.divide(arc_lengths, v1_norm[:, :, array_ns.newaxis], out=arc_lengths) - array_ns.divide(arc_lengths, v2_norm[:, array_ns.newaxis, :], out=arc_lengths) - # Ensure all points are within [-1.0, 1.0] (may be outside due to numerical - # inaccuracies) - array_ns.clip(arc_lengths, -1.0, 1.0, out=arc_lengths) - return array_ns.squeeze(array_ns.arccos(arc_lengths), axis=1) + match geometry_type: + case base_grid.GeometryType.ICOSAHEDRON: + # For pairs of points p1 and p2 compute: + # arccos(dot(p1, p2) / (norm(p1) * norm(p2))) noqa: ERA001 + # Compute all pairs of dot products + arc_lengths = _dot_product(v1, v2, array_ns=array_ns) + v1_norm = array_ns.linalg.norm(v1, axis=-1) + v2_norm = array_ns.linalg.norm(v2, axis=-1) + # Divide the dot products by the broadcasted norms + array_ns.divide(arc_lengths, v1_norm[:, :, array_ns.newaxis], out=arc_lengths) + array_ns.divide(arc_lengths, v2_norm[:, array_ns.newaxis, :], out=arc_lengths) + # Ensure all points are within [-1.0, 1.0] (may be outside due to numerical + # inaccuracies) + array_ns.clip(arc_lengths, -1.0, 1.0, out=arc_lengths) + return array_ns.squeeze(array_ns.arccos(arc_lengths), axis=1) + case base_grid.GeometryType.TORUS: + # For pairs of points p1 and p2 compute: + # norm(p1 - p2) noqa: ERA001 + diff = np.abs(v1 - v2) + domain_size = array_ns.asarray([domain_length, domain_height, ta.wpfloat(0.0)]) + domain_size_expanded = domain_size[array_ns.newaxis, array_ns.newaxis, :] + inverted_diff = array_ns.subtract(domain_size_expanded, diff) + diff = array_ns.minimum(diff, inverted_diff, out=diff) + return array_ns.linalg.norm(diff, axis=-1) + case _: + raise ValueError(f"Unsupported geometry type: {geometry_type}") def _gaussian( @@ -191,22 +242,29 @@ def _kernel( def _cartesian_coordinates_from_zonal_and_meridional_components( + geometry_type: base_grid.GeometryType, lat: data_alloc.NDArray, lon: data_alloc.NDArray, u: data_alloc.NDArray, v: data_alloc.NDArray, array_ns: ModuleType = np, ) -> tuple[data_alloc.NDArray, data_alloc.NDArray, data_alloc.NDArray]: - cos_lat = array_ns.cos(lat) - sin_lat = array_ns.sin(lat) - cos_lon = array_ns.cos(lon) - sin_lon = array_ns.sin(lon) - - x = -u * sin_lon - v * sin_lat * cos_lon - y = u * cos_lon - v * sin_lat * sin_lon - z = cos_lat * v - - return x, y, z + match geometry_type: + case base_grid.GeometryType.ICOSAHEDRON: + cos_lat = array_ns.cos(lat) + sin_lat = array_ns.sin(lat) + cos_lon = array_ns.cos(lon) + sin_lon = array_ns.sin(lon) + + x = -u * sin_lon - v * sin_lat * cos_lon + y = u * cos_lon - v * sin_lat * sin_lon + z = cos_lat * v + + return x, y, z + case base_grid.GeometryType.TORUS: + return u, v, array_ns.zeros_like(u) + case _: + raise ValueError(f"Unsupported geometry type: {geometry_type}") def _compute_rbf_interpolation_coeffs( @@ -224,8 +282,11 @@ def _compute_rbf_interpolation_coeffs( uv: tuple[tuple[data_alloc.NDArray, data_alloc.NDArray], ...], rbf_offset: data_alloc.NDArray, rbf_kernel: InterpolationKernel, + geometry_type: base_grid.GeometryType, scale_factor: ta.wpfloat, horizontal_start: gtx.int32, + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, array_ns: ModuleType = np, ) -> tuple[data_alloc.NDArray, ...]: rbf_offset_shape_full = rbf_offset.shape @@ -271,8 +332,13 @@ def index_offset(f): axis=-1, ) assert element_center.shape == (rbf_offset.shape[0], 3) - vector_dist = _arc_length_vector_matrix( - element_center[:, array_ns.newaxis, :], edge_center, array_ns=array_ns + vector_dist = _distance_vector_matrix( + geometry_type, + domain_length, + domain_height, + element_center[:, array_ns.newaxis, :], + edge_center, + array_ns=array_ns, ) assert vector_dist.shape == rbf_offset.shape rbf_val = _kernel(rbf_kernel, vector_dist, scale_factor, array_ns=array_ns) @@ -287,6 +353,7 @@ def index_offset(f): assert 1 <= num_zonal_meridional_components <= 2 for i in range(num_zonal_meridional_components): z_nx_x, z_nx_y, z_nx_z = _cartesian_coordinates_from_zonal_and_meridional_components( + geometry_type, element_center_lat[horizontal_start:], element_center_lon[horizontal_start:], uv[i][0][horizontal_start:], @@ -311,7 +378,9 @@ def index_offset(f): ) # Distance between edge midpoints for RBF interpolation matrix - z_dist = _arc_length_pairwise(edge_center, array_ns=array_ns) + z_dist = _distance_pairwise( + geometry_type, domain_length, domain_height, edge_center, array_ns=array_ns + ) assert z_dist.shape == ( rbf_offset.shape[0], rbf_offset.shape[1], @@ -333,7 +402,7 @@ def index_offset(f): # https://github.com/cupy/cupy/pull/9116. rbf_vec_coeff_np = [ np.zeros(rbf_offset_shape_full, dtype=ta.wpfloat) - for j in range(num_zonal_meridional_components) + for _ in range(num_zonal_meridional_components) ] rbf_offset_np = data_alloc.as_numpy(rbf_offset) z_rbfmat_np = data_alloc.as_numpy(z_rbfmat) @@ -372,8 +441,11 @@ def compute_rbf_interpolation_coeffs_cell( rbf_offset: data_alloc.NDArray, # TODO(): Can't pass enum as "params" in NumpyFieldsProvider? rbf_kernel: int, + geometry_type: int, scale_factor: ta.wpfloat, horizontal_start: gtx.int32, + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, array_ns: ModuleType = np, ) -> tuple[data_alloc.NDArray]: zeros = array_ns.zeros(rbf_offset.shape[0], dtype=ta.wpfloat) @@ -394,8 +466,11 @@ def compute_rbf_interpolation_coeffs_cell( ((ones, zeros), (zeros, ones)), rbf_offset, InterpolationKernel(rbf_kernel), + base_grid.GeometryType(geometry_type), scale_factor, horizontal_start, + domain_length, + domain_height, array_ns=array_ns, ) @@ -413,13 +488,16 @@ def compute_rbf_interpolation_coeffs_edge( edge_dual_normal_v: data_alloc.NDArray, rbf_offset: data_alloc.NDArray, rbf_kernel: int, + geometry_type: int, scale_factor: ta.wpfloat, horizontal_start: gtx.int32, + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, array_ns: ModuleType = np, ) -> data_alloc.NDArray: return _compute_rbf_interpolation_coeffs( - edge_lat, - edge_lon, + edge_lat, # these should be ignored for torus + edge_lon, # these should be ignored for torus edge_center_x, edge_center_y, edge_center_z, @@ -432,8 +510,11 @@ def compute_rbf_interpolation_coeffs_edge( ((edge_dual_normal_u, edge_dual_normal_v),), rbf_offset, InterpolationKernel(rbf_kernel), + base_grid.GeometryType(geometry_type), scale_factor, horizontal_start, + domain_length, + domain_height, array_ns=array_ns, )[0] @@ -452,8 +533,11 @@ def compute_rbf_interpolation_coeffs_vertex( edge_normal_z: data_alloc.NDArray, rbf_offset: data_alloc.NDArray, rbf_kernel: int, + geometry_type: int, scale_factor: ta.wpfloat, horizontal_start: gtx.int32, + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, array_ns: ModuleType = np, ) -> tuple[data_alloc.NDArray, data_alloc.NDArray]: zeros = array_ns.zeros(rbf_offset.shape[0], dtype=ta.wpfloat) @@ -474,7 +558,10 @@ def compute_rbf_interpolation_coeffs_vertex( ((ones, zeros), (zeros, ones)), rbf_offset, InterpolationKernel(rbf_kernel), + base_grid.GeometryType(geometry_type), scale_factor, horizontal_start, + domain_length, + domain_height, array_ns=array_ns, ) diff --git a/model/common/src/icon4py/model/common/math/helpers.py b/model/common/src/icon4py/model/common/math/helpers.py index 8d47cf0a16..746bd6b419 100644 --- a/model/common/src/icon4py/model/common/math/helpers.py +++ b/model/common/src/icon4py/model/common/math/helpers.py @@ -7,7 +7,14 @@ # SPDX-License-Identifier: BSD-3-Clause from gt4py import next as gtx -from gt4py.next import arccos, cos, sin, sqrt, where +from gt4py.next import ( + abs, # noqa: A004 + arccos, + cos, + sin, + sqrt, + where, +) from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta from icon4py.model.common.dimension import E2C, E2V, Koff @@ -543,6 +550,79 @@ def arc_length_on_edges( return radius * arccos(dot_product_on_edges(x0, x1, y0, y1, z0, z1)) +@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) +def diff_on_edges_torus( + x0: fa.EdgeField[ta.wpfloat], + x1: fa.EdgeField[ta.wpfloat], + y0: fa.EdgeField[ta.wpfloat], + y1: fa.EdgeField[ta.wpfloat], + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, +) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: + """ + Compute the difference between two points on the torus. + + Inputs are cartesian coordinates of the points. Z is assumed zero and is + ignored. Distances are computed modulo the domain length and height. + + Args: + x0: x coordinate of point_0 + x1: x coordinate of point_1 + y0: y coordinate of point_0 + y1: y coordinate of point_1 + domain_length: length of the domain + domain_height: height of the domain + + Returns: + dx, dy + + """ + x1 = where( + abs(x1 - x0) <= 0.5 * domain_length, + x1, + where(x0 > x1, x1 + domain_length, x1 - domain_length), + ) + + y1 = where( + abs(y1 - y0) <= 0.5 * domain_height, + y1, + where(y0 > y1, y1 + domain_height, y1 - domain_height), + ) + + return (x1 - x0, y1 - y0) + + +@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) +def distance_on_edges_torus( + x0: fa.EdgeField[ta.wpfloat], + x1: fa.EdgeField[ta.wpfloat], + y0: fa.EdgeField[ta.wpfloat], + y1: fa.EdgeField[ta.wpfloat], + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, +) -> fa.EdgeField[ta.wpfloat]: + """ + Compute the distance between two points on the torus. + + Inputs are cartesian coordinates of the points. Z is assumed zero and is + ignored. Distances are computed modulo the domain length and height. + + Args: + x0: x coordinate of point_0 + x1: x coordinate of point_1 + y0: y coordinate of point_0 + y1: y coordinate of point_1 + domain_length: length of the domain + domain_height: height of the domain + + Returns: + distance + + """ + xdiff, ydiff = diff_on_edges_torus(x0, x1, y0, y1, domain_length, domain_height) + return sqrt(xdiff**2 + ydiff**2) + + @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def average_two_vertical_levels_downwards_on_edges( input_field: fa.EdgeKField[wpfloat], diff --git a/model/common/tests/common/grid/unit_tests/test_geometry.py b/model/common/tests/common/grid/unit_tests/test_geometry.py index 003c3e0cc1..7baca91a15 100644 --- a/model/common/tests/common/grid/unit_tests/test_geometry.py +++ b/model/common/tests/common/grid/unit_tests/test_geometry.py @@ -13,8 +13,9 @@ import numpy as np import pytest -from icon4py.model.common import dimension as dims +from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import ( + base, geometry, geometry_attributes as attrs, horizontal as h_grid, @@ -41,8 +42,11 @@ from icon4py.model.testing import serialbox as sb -def test_geometry_raises_for_unknown_field(backend: gtx_typing.Backend) -> None: - geometry = grid_utils.get_grid_geometry(backend, definitions.Experiments.EXCLAIM_APE) +@pytest.mark.datatest +def test_geometry_raises_for_unknown_field( + backend: gtx_typing.Backend, experiment: definitions.Experiment +) -> None: + geometry = grid_utils.get_grid_geometry(backend, experiment) with pytest.raises(ValueError) as e: geometry.get("foo") assert "'foo'" in e.value # type: ignore[operator] @@ -54,6 +58,7 @@ def test_geometry_raises_for_unknown_field(backend: gtx_typing.Backend) -> None: [ (definitions.Experiments.MCH_CH_R04B09, 1e-7), (definitions.Experiments.EXCLAIM_APE, 3e-12), + (definitions.Experiments.GAUSS3D, 1e-13), ], ) @pytest.mark.datatest @@ -87,6 +92,7 @@ def test_coriolis_parameter( [ (definitions.Experiments.MCH_CH_R04B09, 1e-9), (definitions.Experiments.EXCLAIM_APE, 1e-12), + (definitions.Experiments.GAUSS3D, 1e-13), ], ) @pytest.mark.datatest @@ -107,6 +113,7 @@ def test_compute_edge_length( [ (definitions.Experiments.MCH_CH_R04B09, 1e-9), (definitions.Experiments.EXCLAIM_APE, 1e-12), + (definitions.Experiments.GAUSS3D, 1e-13), ], ) @pytest.mark.datatest @@ -128,6 +135,7 @@ def test_compute_inverse_edge_length( [ (definitions.Experiments.MCH_CH_R04B09, 1e-7), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 1e-13), ], ) @pytest.mark.datatest @@ -149,6 +157,7 @@ def test_compute_dual_edge_length( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 1e-13), ], ) @pytest.mark.datatest @@ -175,6 +184,7 @@ def test_compute_inverse_dual_edge_length( [ (definitions.Experiments.MCH_CH_R04B09, 5e-10), (definitions.Experiments.EXCLAIM_APE, 1e-12), + (definitions.Experiments.GAUSS3D, 1e-14), ], ) @pytest.mark.datatest @@ -186,9 +196,14 @@ def test_compute_inverse_vertex_vertex_length( ) -> None: grid_geometry = grid_utils.get_grid_geometry(backend, experiment) - expected = grid_savepoint.inv_vert_vert_length() - result = grid_geometry.get(attrs.INVERSE_VERTEX_VERTEX_LENGTH) - assert test_utils.dallclose(result.asnumpy(), expected.asnumpy(), rtol=rtol) + expected = grid_savepoint.inv_vert_vert_length().asnumpy() + result = grid_geometry.get(attrs.INVERSE_VERTEX_VERTEX_LENGTH).asnumpy() + if grid_geometry.grid.geometry_type == base.GeometryType.TORUS: + # TODO(msimberg, jcanton): icon fortran multiplies sphere radius even for torus + # grids. Fix submitted upstream. The following can be removed when fixed + # serialized data is available. + result = result / constants.EARTH_RADIUS + assert test_utils.dallclose(result, expected, rtol=rtol) @pytest.mark.datatest @@ -338,6 +353,7 @@ def test_dual_normal_vert( assert test_utils.dallclose(dual_normal_vert_v.asnumpy(), dual_normal_vert_v_ref, atol=1e-12) +@pytest.mark.datatest def test_cartesian_centers_edge( backend: gtx_typing.Backend, experiment: definitions.Experiment ) -> None: @@ -349,12 +365,23 @@ def test_cartesian_centers_edge( assert x.ndarray.shape == (grid.num_edges,) assert y.ndarray.shape == (grid.num_edges,) assert z.ndarray.shape == (grid.num_edges,) - # those are coordinates on the unit sphere: hence norm should be 1 - norm = data_alloc.zero_field(grid, dims.EdgeDim, dtype=x.dtype, allocator=backend) - math_helpers.norm2_on_edges(x, z, y, out=norm, offset_provider={}) - assert test_utils.dallclose(norm.asnumpy(), 1.0) + match grid.geometry_type: + case base.GeometryType.ICOSAHEDRON: + # those are coordinates on the unit sphere: hence norm should be 1 + norm = data_alloc.zero_field(grid, dims.EdgeDim, dtype=x.dtype, allocator=backend) + math_helpers.norm2_on_edges(x, z, y, out=norm, offset_provider={}) + assert test_utils.dallclose(norm.asnumpy(), 1.0) + case base.GeometryType.TORUS: + # on a torus coordinates should be within the domain + assert all(x.asnumpy() >= 0.0) + assert all(x.asnumpy() <= grid.global_properties.domain_length) + assert all(y.asnumpy() >= 0.0) + assert all(y.asnumpy() <= grid.global_properties.domain_height) + assert all(z.asnumpy() == 0.0) + +@pytest.mark.datatest def test_cartesian_centers_cell( backend: gtx_typing.Backend, experiment: definitions.Experiment ) -> None: @@ -366,12 +393,23 @@ def test_cartesian_centers_cell( assert x.ndarray.shape == (grid.num_cells,) assert y.ndarray.shape == (grid.num_cells,) assert z.ndarray.shape == (grid.num_cells,) - # those are coordinates on the unit sphere: hence norm should be 1 - norm = data_alloc.zero_field(grid, dims.CellDim, dtype=x.dtype, allocator=backend) - math_helpers.norm2_on_cells(x, z, y, out=norm, offset_provider={}) - assert test_utils.dallclose(norm.asnumpy(), 1.0) + + match grid.geometry_type: + case base.GeometryType.ICOSAHEDRON: + # those are coordinates on the unit sphere: hence norm should be 1 + norm = data_alloc.zero_field(grid, dims.CellDim, dtype=x.dtype, allocator=backend) + math_helpers.norm2_on_cells(x, z, y, out=norm, offset_provider={}) + assert test_utils.dallclose(norm.asnumpy(), 1.0) + case base.GeometryType.TORUS: + # on a torus coordinates should be within the domain + assert all(x.asnumpy() >= 0.0) + assert all(x.asnumpy() <= grid.global_properties.domain_length) + assert all(y.asnumpy() >= 0.0) + assert all(y.asnumpy() <= grid.global_properties.domain_height) + assert all(z.asnumpy() == 0.0) +@pytest.mark.datatest def test_vertex(backend: gtx_typing.Backend, experiment: definitions.Experiment) -> None: grid_geometry = grid_utils.get_grid_geometry(backend, experiment) grid = grid_geometry.grid @@ -381,10 +419,20 @@ def test_vertex(backend: gtx_typing.Backend, experiment: definitions.Experiment) assert x.ndarray.shape == (grid.num_vertices,) assert y.ndarray.shape == (grid.num_vertices,) assert z.ndarray.shape == (grid.num_vertices,) - # those are coordinates on the unit sphere: hence norm should be 1 - norm = data_alloc.zero_field(grid, dims.VertexDim, dtype=x.dtype, allocator=backend) - math_helpers.norm2_on_vertices(x, z, y, out=norm, offset_provider={}) - assert test_utils.dallclose(norm.asnumpy(), 1.0) + + match grid.geometry_type: + case base.GeometryType.ICOSAHEDRON: + # those are coordinates on the unit sphere: hence norm should be 1 + norm = data_alloc.zero_field(grid, dims.VertexDim, dtype=x.dtype, allocator=backend) + math_helpers.norm2_on_vertices(x, z, y, out=norm, offset_provider={}) + assert test_utils.dallclose(norm.asnumpy(), 1.0) + case base.GeometryType.TORUS: + # on a torus coordinates should be within the domain + assert all(x.asnumpy() >= 0.0) + assert all(x.asnumpy() <= grid.global_properties.domain_length) + assert all(y.asnumpy() >= 0.0) + assert all(y.asnumpy() <= grid.global_properties.domain_height) + assert all(z.asnumpy() == 0.0) def test_sparse_fields_creator() -> None: diff --git a/model/common/tests/common/interpolation/unit_tests/test_interpolation_factory.py b/model/common/tests/common/interpolation/unit_tests/test_interpolation_factory.py index 30bb81024a..179a982334 100644 --- a/model/common/tests/common/interpolation/unit_tests/test_interpolation_factory.py +++ b/model/common/tests/common/interpolation/unit_tests/test_interpolation_factory.py @@ -101,6 +101,7 @@ def test_factory_raises_error_on_unknown_field( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -124,6 +125,7 @@ def test_get_c_lin_e( [ (definitions.Experiments.MCH_CH_R04B09, 1e-9), (definitions.Experiments.EXCLAIM_APE, 1e-12), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -149,6 +151,7 @@ def test_get_geofac_div( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -183,6 +186,7 @@ def assert_reordered(val: np.ndarray, ref: np.ndarray, **kwargs: Any) -> None: [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -209,6 +213,7 @@ def test_get_geofac_rot( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -262,6 +267,7 @@ def test_get_geofac_grg( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 1e-15), ], ) @pytest.mark.datatest @@ -301,6 +307,7 @@ def test_e_flx_avg( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -324,6 +331,7 @@ def test_e_bln_c_s( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 1e-13), ], ) @pytest.mark.datatest @@ -348,6 +356,7 @@ def test_pos_on_tplane_e_x_y( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -372,6 +381,7 @@ def test_cells_aw_verts( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -394,6 +404,7 @@ def test_nudgecoeffs( [ (definitions.Experiments.EXCLAIM_APE, 3e-9), (definitions.Experiments.MCH_CH_R04B09, 4e-2), + (definitions.Experiments.GAUSS3D, 1e-14), ], ) @pytest.mark.datatest @@ -427,6 +438,7 @@ def test_rbf_interpolation_coeffs_cell( [ (definitions.Experiments.EXCLAIM_APE, 8e-14), (definitions.Experiments.MCH_CH_R04B09, 2e-9), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -455,6 +467,7 @@ def test_rbf_interpolation_coeffs_edge( [ (definitions.Experiments.EXCLAIM_APE, 3e-10), (definitions.Experiments.MCH_CH_R04B09, 3e-3), + (definitions.Experiments.GAUSS3D, 1e-15), ], ) @pytest.mark.datatest diff --git a/model/common/tests/common/interpolation/unit_tests/test_interpolation_fields.py b/model/common/tests/common/interpolation/unit_tests/test_interpolation_fields.py index b623edc8f2..55ef66a918 100644 --- a/model/common/tests/common/interpolation/unit_tests/test_interpolation_fields.py +++ b/model/common/tests/common/interpolation/unit_tests/test_interpolation_fields.py @@ -20,6 +20,7 @@ compute_c_lin_e, compute_cells_aw_verts, compute_e_bln_c_s, + compute_e_bln_c_s_torus, compute_e_flx_avg, compute_geofac_div, compute_geofac_grdiv, @@ -27,7 +28,9 @@ compute_geofac_n2s, compute_geofac_rot, compute_mass_conserving_bilinear_cell_average_weight, + compute_mass_conserving_bilinear_cell_average_weight_torus, compute_pos_on_tplane_e_x_y, + compute_pos_on_tplane_e_x_y_torus, ) from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import definitions, serialbox as sb @@ -245,7 +248,11 @@ def test_compute_geofac_grdiv( @pytest.mark.datatest @pytest.mark.parametrize( "experiment, atol", - [(definitions.Experiments.MCH_CH_R04B09, 1e-10), (definitions.Experiments.EXCLAIM_APE, 1e-10)], + [ + (definitions.Experiments.MCH_CH_R04B09, 1e-10), + (definitions.Experiments.EXCLAIM_APE, 1e-10), + (definitions.Experiments.GAUSS3D, 1e-15), + ], ) def test_compute_c_bln_avg( grid_savepoint: sb.IconGridSavepoint, @@ -268,18 +275,30 @@ def test_compute_c_bln_avg( c2e2c0 = icon_grid.get_connectivity(dims.C2E2CO).ndarray - c_bln_avg = functools.partial( - compute_mass_conserving_bilinear_cell_average_weight, array_ns=xp - )( - c2e2c0, - lat, - lon, - cell_areas, - cell_owner_mask, - divavg_cntrwgt, - horizontal_start, - horizontal_start_p2, - ) + match icon_grid.global_properties.geometry_type: + case base_grid.GeometryType.ICOSAHEDRON: + c_bln_avg = compute_mass_conserving_bilinear_cell_average_weight( + c2e2c0, + lat, + lon, + cell_areas, + cell_owner_mask, + divavg_cntrwgt, + horizontal_start, + horizontal_start_p2, + array_ns=xp, + ) + case base_grid.GeometryType.TORUS: + c_bln_avg = compute_mass_conserving_bilinear_cell_average_weight_torus( + c2e2c0, + cell_areas, + cell_owner_mask, + divavg_cntrwgt, + horizontal_start, + horizontal_start_p2, + array_ns=xp, + ) + assert test_helpers.dallclose(data_alloc.as_numpy(c_bln_avg), c_bln_avg_ref, atol=atol) @@ -373,9 +392,14 @@ def test_compute_e_bln_c_s( edges_lon = grid_savepoint.edges_center_lon().ndarray xp = data_alloc.import_array_ns(backend) - e_bln_c_s = functools.partial(compute_e_bln_c_s, array_ns=xp)( - c2e, cells_lat, cells_lon, edges_lat, edges_lon, 0.0 - ) + # TODO(msimberg): pass geometry type to compute_pos_on_tplane_e_x_y? + match icon_grid.global_properties.geometry_type: + case base_grid.GeometryType.ICOSAHEDRON: + e_bln_c_s = compute_e_bln_c_s( + c2e, cells_lat, cells_lon, edges_lat, edges_lon, 0.0, array_ns=xp + ) + case base_grid.GeometryType.TORUS: + e_bln_c_s = compute_e_bln_c_s_torus(c2e, array_ns=xp) assert test_helpers.dallclose( data_alloc.as_numpy(e_bln_c_s), e_bln_c_s_ref.asnumpy(), atol=1e-6, rtol=1e-7 ) @@ -393,6 +417,7 @@ def test_compute_pos_on_tplane_e( pos_on_tplane_e_x_ref = interpolation_savepoint.pos_on_tplane_e_x().asnumpy() pos_on_tplane_e_y_ref = interpolation_savepoint.pos_on_tplane_e_y().asnumpy() sphere_radius = constants.EARTH_RADIUS + dual_edge_length = grid_savepoint.dual_edge_length().ndarray primal_normal_v1 = grid_savepoint.primal_normal_v1().ndarray primal_normal_v2 = grid_savepoint.primal_normal_v2().ndarray dual_normal_v1 = grid_savepoint.dual_normal_v1().ndarray @@ -404,20 +429,29 @@ def test_compute_pos_on_tplane_e( edges_lat = grid_savepoint.edges_center_lat().ndarray e2c = icon_grid.get_connectivity(dims.E2C).ndarray horizontal_start = icon_grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)) - pos_on_tplane_e_x, pos_on_tplane_e_y = compute_pos_on_tplane_e_x_y( - sphere_radius, - primal_normal_v1, - primal_normal_v2, - dual_normal_v1, - dual_normal_v2, - cells_lon, - cells_lat, - edges_lon, - edges_lat, - owner_mask, - e2c, - horizontal_start, - array_ns=xp, - ) + # TODO(msimberg): pass geometry type to compute_pos_on_tplane_e_x_y? + match icon_grid.global_properties.geometry_type: + case base_grid.GeometryType.ICOSAHEDRON: + pos_on_tplane_e_x, pos_on_tplane_e_y = compute_pos_on_tplane_e_x_y( + sphere_radius, + primal_normal_v1, + primal_normal_v2, + dual_normal_v1, + dual_normal_v2, + cells_lon, + cells_lat, + edges_lon, + edges_lat, + owner_mask, + e2c, + horizontal_start, + array_ns=xp, + ) + case base_grid.GeometryType.TORUS: + pos_on_tplane_e_x, pos_on_tplane_e_y = compute_pos_on_tplane_e_x_y_torus( + dual_edge_length, + e2c, + array_ns=xp, + ) assert test_helpers.dallclose(pos_on_tplane_e_x, pos_on_tplane_e_x_ref, atol=1e-6, rtol=1e-7) assert test_helpers.dallclose(pos_on_tplane_e_y, pos_on_tplane_e_y_ref, atol=1e-6, rtol=1e-7) diff --git a/model/common/tests/common/interpolation/unit_tests/test_rbf_interpolation.py b/model/common/tests/common/interpolation/unit_tests/test_rbf_interpolation.py index 9912cc3b94..4e172e82b2 100644 --- a/model/common/tests/common/interpolation/unit_tests/test_rbf_interpolation.py +++ b/model/common/tests/common/interpolation/unit_tests/test_rbf_interpolation.py @@ -145,6 +145,7 @@ def test_construct_rbf_matrix_offsets_tables_for_vertices( [ (definitions.Experiments.EXCLAIM_APE, 3e-9), (definitions.Experiments.MCH_CH_R04B09, 3e-2), + (definitions.Experiments.GAUSS3D, 1e-14), ], ) def test_rbf_interpolation_coeffs_cell( @@ -177,8 +178,16 @@ def test_rbf_interpolation_coeffs_cell( geometry.get(geometry_attrs.EDGE_NORMAL_Z).ndarray, rbf.construct_rbf_matrix_offsets_tables_for_cells(grid), rbf.DEFAULT_RBF_KERNEL[rbf_dim], - rbf.compute_default_rbf_scale(math.sqrt(grid_savepoint.mean_cell_area()), rbf_dim), + grid.global_properties.geometry_type.value, + rbf.compute_default_rbf_scale( + grid.global_properties.geometry_type, + grid.global_properties.characteristic_length, + grid.global_properties.mean_dual_edge_length, + rbf_dim, + ), horizontal_start, + grid.global_properties.domain_length, + grid.global_properties.domain_height, array_ns=data_alloc.import_array_ns(backend), ) @@ -214,6 +223,7 @@ def test_rbf_interpolation_coeffs_cell( [ (definitions.Experiments.EXCLAIM_APE, 3e-10), (definitions.Experiments.MCH_CH_R04B09, 3e-3), + (definitions.Experiments.GAUSS3D, 1e-15), ], ) def test_rbf_interpolation_coeffs_vertex( @@ -246,8 +256,16 @@ def test_rbf_interpolation_coeffs_vertex( geometry.get(geometry_attrs.EDGE_NORMAL_Z).ndarray, rbf.construct_rbf_matrix_offsets_tables_for_vertices(grid), rbf.DEFAULT_RBF_KERNEL[rbf_dim], - rbf.compute_default_rbf_scale(math.sqrt(grid_savepoint.mean_cell_area()), rbf_dim), + grid.global_properties.geometry_type.value, + rbf.compute_default_rbf_scale( + grid.global_properties.geometry_type, + grid.global_properties.characteristic_length, + grid.global_properties.mean_dual_edge_length, + rbf_dim, + ), horizontal_start, + grid.global_properties.domain_length, + grid.global_properties.domain_height, array_ns=data_alloc.import_array_ns(backend), ) @@ -283,6 +301,7 @@ def test_rbf_interpolation_coeffs_vertex( [ (definitions.Experiments.EXCLAIM_APE, 8e-14), (definitions.Experiments.MCH_CH_R04B09, 2e-9), + (definitions.Experiments.GAUSS3D, 0), ], ) def test_rbf_interpolation_coeffs_edge( @@ -317,8 +336,16 @@ def test_rbf_interpolation_coeffs_edge( # coefficients in savepoint. grid_savepoint.e2c2e(), rbf.DEFAULT_RBF_KERNEL[rbf_dim], - rbf.compute_default_rbf_scale(math.sqrt(grid_savepoint.mean_cell_area()), rbf_dim), + grid.global_properties.geometry_type.value, + rbf.compute_default_rbf_scale( + grid.global_properties.geometry_type, + grid.global_properties.characteristic_length, + grid.global_properties.mean_dual_edge_length, + rbf_dim, + ), horizontal_start, + grid.global_properties.domain_length, + grid.global_properties.domain_height, array_ns=data_alloc.import_array_ns(backend), ) diff --git a/model/testing/src/icon4py/model/testing/fixtures/benchmark.py b/model/testing/src/icon4py/model/testing/fixtures/benchmark.py index f230f43e82..658dc605ac 100644 --- a/model/testing/src/icon4py/model/testing/fixtures/benchmark.py +++ b/model/testing/src/icon4py/model/testing/fixtures/benchmark.py @@ -43,7 +43,7 @@ def geometry_field_source( generic_concrete_backend = model_options.customize_backend(None, backend_like) decomposition_info = grid_utils.construct_decomposition_info(mesh, allocator) - geometry_field_source = grid_geometry.GridGeometry( + geometry_field_source = grid_geometry.GridGeometry.with_geometry_type( grid=mesh, decomposition_info=decomposition_info, backend=generic_concrete_backend, diff --git a/model/testing/src/icon4py/model/testing/fixtures/datatest.py b/model/testing/src/icon4py/model/testing/fixtures/datatest.py index 123884f464..3058fc7210 100644 --- a/model/testing/src/icon4py/model/testing/fixtures/datatest.py +++ b/model/testing/src/icon4py/model/testing/fixtures/datatest.py @@ -88,7 +88,11 @@ def cpu_allocator() -> gtx_typing.FieldBufferAllocationUtil: @pytest.fixture( - params=[definitions.Experiments.MCH_CH_R04B09, definitions.Experiments.EXCLAIM_APE], + params=[ + definitions.Experiments.MCH_CH_R04B09, + definitions.Experiments.EXCLAIM_APE, + definitions.Experiments.GAUSS3D, + ], ids=lambda r: r.name, ) def experiment(request: pytest.FixtureRequest) -> definitions.Experiment: diff --git a/model/testing/src/icon4py/model/testing/grid_utils.py b/model/testing/src/icon4py/model/testing/grid_utils.py index 5d2dc9696d..4e5f920538 100644 --- a/model/testing/src/icon4py/model/testing/grid_utils.py +++ b/model/testing/src/icon4py/model/testing/grid_utils.py @@ -135,7 +135,7 @@ def _construct_grid_geometry() -> geometry.GridGeometry: ) grid = gm.grid decomposition_info = construct_decomposition_info(grid, backend) - geometry_source = geometry.GridGeometry( + return geometry.GridGeometry.with_geometry_type( grid, decomposition_info, backend, @@ -143,7 +143,6 @@ def _construct_grid_geometry() -> geometry.GridGeometry: gm.geometry_fields, geometry_attrs.attrs, ) - return geometry_source if not grid_geometries.get(register_name): grid_geometries[register_name] = _construct_grid_geometry()