diff --git a/model/common/src/icon4py/model/common/grid/geometry.py b/model/common/src/icon4py/model/common/grid/geometry.py index df18139e89..a11ae1f6fc 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,253 @@ 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) + + 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/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/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()