diff --git a/model/common/src/icon4py/model/common/grid/geometry.py b/model/common/src/icon4py/model/common/grid/geometry.py index 831d457347..76ac473e9e 100644 --- a/model/common/src/icon4py/model/common/grid/geometry.py +++ b/model/common/src/icon4py/model/common/grid/geometry.py @@ -41,7 +41,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 @@ -84,7 +84,7 @@ def __init__( extra_fields: gm.GeometryDict, metadata: dict[str, model.FieldMetaData], exchange: decomposition.ExchangeRuntime = decomposition.single_node_default, - ): + ) -> None: """ Args: grid: IconGrid the grid topology @@ -98,45 +98,42 @@ def __init__( """ 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) self._exchange = exchange log.info( - f"initialized geometry for backend = '{self._backend_name()}' and grid = '{self._grid}'" + f"initializing geometry for backend = '{self._backend_name()}' and grid = '{self._grid}'" ) - ( - edge_orientation0_lat, - edge_orientation0_lon, - edge_orientation1_lat, - edge_orientation1_lon, - ) = create_auxiliary_coordinate_arrays_for_orientation( - self._grid, - coordinates[dims.CellDim]["lat"], - coordinates[dims.CellDim]["lon"], - coordinates[dims.EdgeDim]["lat"], - coordinates[dims.EdgeDim]["lon"], - self._backend, - ) + # Setup coordinates based on geometry type coordinates_ = { - attrs.CELL_LAT: coordinates[dims.CellDim]["lat"], attrs.CELL_LON: coordinates[dims.CellDim]["lon"], - attrs.VERTEX_LAT: coordinates[dims.VertexDim]["lat"], + attrs.CELL_LAT: coordinates[dims.CellDim]["lat"], attrs.EDGE_LON: coordinates[dims.EdgeDim]["lon"], attrs.EDGE_LAT: coordinates[dims.EdgeDim]["lat"], attrs.VERTEX_LON: coordinates[dims.VertexDim]["lon"], - "latitude_of_edge_cell_neighbor_0": edge_orientation0_lat, - "longitude_of_edge_cell_neighbor_0": edge_orientation0_lon, - "latitude_of_edge_cell_neighbor_1": edge_orientation1_lat, - "longitude_of_edge_cell_neighbor_1": edge_orientation1_lon, + attrs.VERTEX_LAT: coordinates[dims.VertexDim]["lat"], } + if self._geometry_type == base.GeometryType.TORUS: + coordinates_[attrs.CELL_CENTER_X] = coordinates[dims.CellDim]["x"] + coordinates_[attrs.CELL_CENTER_Y] = coordinates[dims.CellDim]["y"] + coordinates_[attrs.CELL_CENTER_Z] = coordinates[dims.CellDim]["z"] + coordinates_[attrs.EDGE_CENTER_X] = coordinates[dims.EdgeDim]["x"] + coordinates_[attrs.EDGE_CENTER_Y] = coordinates[dims.EdgeDim]["y"] + coordinates_[attrs.EDGE_CENTER_Z] = coordinates[dims.EdgeDim]["z"] + coordinates_[attrs.VERTEX_X] = coordinates[dims.VertexDim]["x"] + coordinates_[attrs.VERTEX_Y] = coordinates[dims.VertexDim]["y"] + coordinates_[attrs.VERTEX_Z] = coordinates[dims.VertexDim]["z"] + coordinate_provider = factory.PrecomputedFieldProvider(coordinates_) self.register_provider(coordinate_provider) + # Setup input fields input_fields_provider = factory.PrecomputedFieldProvider( { # TODO(halungge): rescaled by grid_length_rescale_factor (mo_grid_tools.f90) @@ -178,39 +175,120 @@ def __init__( self.register_provider(input_fields_provider) self._register_computed_fields() + 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), + ) + }, + do_exchange=False, + ) + return provider + def _register_computed_fields(self) -> None: + """Register all computed geometry fields.""" + # Common fields for both geometries meta = attrs.metadata_for_inverse(attrs.attrs[attrs.EDGE_LENGTH]) name = meta["standard_name"] self._attrs.update({name: meta}) + inverse_edge_length = self._inverse_field_provider(attrs.EDGE_LENGTH) self.register_provider(inverse_edge_length) 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_arc_distance_of_far_edges_in_diamond, - domain={ - dims.EdgeDim: ( - self._edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2), - self._edge_domain(h_grid.Zone.LOCAL), + match self._geometry_type: + case base.GeometryType.ICOSAHEDRON: + self._register_cartesian_coordinates_icosahedron() + + vertex_vertex_distance = factory.ProgramFieldProvider( + func=stencils.compute_arc_distance_of_far_edges_in_diamond, + domain={ + dims.EdgeDim: ( + 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_lat": attrs.VERTEX_LAT, + "vertex_lon": attrs.VERTEX_LON, + }, + params={"radius": self._grid.global_properties.radius}, + do_exchange=True, ) - }, - fields={"far_vertex_distance": attrs.VERTEX_VERTEX_LENGTH}, - deps={ - "vertex_lat": attrs.VERTEX_LAT, - "vertex_lon": attrs.VERTEX_LON, - }, - params={"radius": self._grid.global_properties.radius}, - do_exchange=True, - ) - self.register_provider(vertex_vertex_distance) + self.register_provider(vertex_vertex_distance) + + coriolis_param = factory.ProgramFieldProvider( + func=stencils.compute_coriolis_parameter_on_edges, + deps={"edge_center_lat": attrs.EDGE_LAT}, + params={"angular_velocity": constants.EARTH_ANGULAR_VELOCITY}, + fields={"coriolis_parameter": attrs.CORIOLIS_PARAMETER}, + domain={ + dims.EdgeDim: ( + self._edge_domain(h_grid.Zone.LOCAL), + self._edge_domain(h_grid.Zone.END), + ) + }, + do_exchange=False, + ) + self.register_provider(coriolis_param) + + self._register_normals_and_tangents_icosahedron() + + case base.GeometryType.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.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, + }, + do_exchange=True, + ) + self.register_provider(vertex_vertex_distance) + + coriolis_param = factory.PrecomputedFieldProvider( + { + # TODO(jcanton): this constant (0.0) should eventually + # come from the config + "coriolis_parameter": stencils.coriolis_parameter_on_edges_torus( + coriolis_coefficient=0.0, + num_edges=self._grid.num_edges, + backend=self._backend, + ) + } + ) + self.register_provider(coriolis_param) + + self._register_normals_and_tangents_torus() + # Inverse of vertex-vertex distance inverse_far_edge_distance_provider = self._inverse_field_provider( attrs.VERTEX_VERTEX_LENGTH ) self.register_provider(inverse_far_edge_distance_provider) + # Edge areas edge_areas = factory.ProgramFieldProvider( func=stencils.compute_edge_area, deps={ @@ -228,22 +306,9 @@ def _register_computed_fields(self) -> None: do_exchange=True, ) self.register_provider(edge_areas) - coriolis_params = factory.ProgramFieldProvider( - func=stencils.compute_coriolis_parameter_on_edges, - deps={"edge_center_lat": attrs.EDGE_LAT}, - params={"angular_velocity": constants.EARTH_ANGULAR_VELOCITY}, - fields={"coriolis_parameter": attrs.CORIOLIS_PARAMETER}, - domain={ - dims.EdgeDim: ( - self._edge_domain(h_grid.Zone.LOCAL), - self._edge_domain(h_grid.Zone.END), - ) - }, - do_exchange=False, - ) - self.register_provider(coriolis_params) - # normals: + def _register_normals_and_tangents_icosahedron(self) -> None: + """Register normals and tangents specific to icosahedron geometry.""" # 1. edges%primal_cart_normal (cartesian coordinates for primal_normal) tangent_normal_coordinates = factory.ProgramFieldProvider( func=stencils.compute_cartesian_coordinates_of_edge_tangent_and_normal, @@ -358,6 +423,7 @@ def _register_computed_fields(self) -> None: do_exchange=True, ) self.register_provider(normal_vert_wrapper) + normal_cell = factory.ProgramFieldProvider( func=stencils.compute_zonal_and_meridional_component_of_edge_field_at_cell_center, deps={ @@ -389,7 +455,8 @@ def _register_computed_fields(self) -> None: do_exchange=True, ) self.register_provider(normal_cell_wrapper) - # 3. dual normals: the dual normals are the edge tangents + + # dual normals: the dual normals are the edge tangents tangent_vert = factory.ProgramFieldProvider( func=stencils.compute_zonal_and_meridional_component_of_edge_field_at_vertex, deps={ @@ -428,6 +495,7 @@ def _register_computed_fields(self) -> None: do_exchange=True, ) self.register_provider(tangent_vert_wrapper) + tangent_cell = factory.ProgramFieldProvider( func=stencils.compute_zonal_and_meridional_component_of_edge_field_at_cell_center, deps={ @@ -459,6 +527,117 @@ def _register_computed_fields(self) -> None: do_exchange=True, ) self.register_provider(tangent_cell_wrapper) + + def _register_normals_and_tangents_torus(self) -> None: + """Register normals and tangents specific to torus geometry.""" + # cartesian coordinates for primal_normal + 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, + }, + do_exchange=False, + ) + self.register_provider(tangent_normal_coordinates) + + # 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, + ), + ), + do_exchange=True, + ) + 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), + ), + do_exchange=True, + ) + self.register_provider(normal_cell_wrapper) + + # 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, + ), + ), + do_exchange=True, + ) + 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), + ), + do_exchange=True, + ) + self.register_provider(tangent_cell_wrapper) + + def _register_cartesian_coordinates_icosahedron(self) -> None: + """Register Cartesian coordinate conversions for icosahedron geometry.""" cartesian_vertices = factory.EmbeddedFieldOperatorProvider( func=math_helpers.geographical_to_cartesian_on_vertices.with_backend(self.backend), domain={ @@ -520,24 +699,6 @@ 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]) - 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.END), - ) - }, - do_exchange=False, - ) - return provider - def __repr__(self) -> str: geometry_name = self._geometry_type._name_ if self._geometry_type else "" return ( @@ -549,7 +710,7 @@ def metadata(self) -> dict[str, model.FieldMetaData]: return self._attrs @property - def backend(self) -> gtx_typing.Backend | None: + def backend(self) -> gtx_typing.Backend: return self._backend @property @@ -564,7 +725,7 @@ def vertical_grid(self) -> None: class SparseFieldProviderWrapper(factory.FieldProvider, factory.NeedsExchange): 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..775eea4122 100644 --- a/model/common/src/icon4py/model/common/grid/geometry_stencils.py +++ b/model/common/src/icon4py/model/common/grid/geometry_stencils.py @@ -8,6 +8,7 @@ from types import ModuleType +import gt4py.next.typing as gtx_typing import numpy as np from gt4py import next as gtx from gt4py.next import sin, where @@ -17,12 +18,14 @@ 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, zonal_and_meridional_components_on_edges, ) -from icon4py.model.common.utils import data_allocation as alloc +from icon4py.model.common.utils import data_allocation as data_alloc @gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) @@ -56,6 +59,51 @@ 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. + + This is done by taking the difference 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): This should use something like numpy.zeros_like if and + # when that becomes available in gt4py. + + 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 +141,32 @@ 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): This should use something like numpy.zeros_like if and + # when that becomes available in gt4py. + 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 +197,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 +277,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], @@ -377,14 +550,14 @@ def arc_distance_of_far_edges_in_diamond( Neighboring edges of an edge span up a diamond with 4 edges (E2C2E) and 4 vertices (E2C2V). Here we compute the arc length between the two vertices in this diamond that are not directly connected to the edge: - between d(v1, v3) - v1-------v4 + between d(v2, v3) + v2-------v1 | /| | / | | e | | / | |/ | - v2 ------v3 + v0 ------v3 @@ -398,16 +571,47 @@ def arc_distance_of_far_edges_in_diamond( """ x, y, z = geographical_to_cartesian_on_vertices(vertex_lat, vertex_lon) - x2 = x(E2C2V[2]) - x3 = x(E2C2V[3]) - y2 = y(E2C2V[2]) - y3 = y(E2C2V[3]) - z2 = z(E2C2V[2]) - z3 = z(E2C2V[3]) - # (xi, yi, zi) are normalized by construction + return arc_length_on_edges( + x(E2C2V[2]), + x(E2C2V[3]), + y(E2C2V[2]), + y(E2C2V[3]), + z(E2C2V[2]), + z(E2C2V[3]), + radius, + ) - far_vertex_vertex_length = arc_length_on_edges(x2, x3, y2, y3, z2, z3, radius) - 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 @@ -502,6 +706,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], @@ -557,6 +781,29 @@ def coriolis_parameter_on_edges( return 2.0 * angular_velocity * sin(edge_center_lat) +def coriolis_parameter_on_edges_torus( + coriolis_coefficient: float, + num_edges: int, + backend: gtx_typing.Backend, +) -> fa.EdgeField[ta.wpfloat]: + """ + Create a coriolis parameter field on edges for a torus grid. + Args: + coriolis_coefficient: coriolis coefficient + + Returns: + coriolis parameter + """ + xp = data_alloc.import_array_ns(backend) + coriolis_parameter = gtx.as_field( + (dims.EdgeDim,), + coriolis_coefficient * xp.ones(num_edges), + dtype=ta.wpfloat, + allocator=backend, + ) + return coriolis_parameter + + @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_coriolis_parameter_on_edges( edge_center_lat: fa.EdgeField[ta.wpfloat], @@ -574,11 +821,11 @@ def compute_coriolis_parameter_on_edges( def compute_primal_cart_normal( - primal_cart_normal_x: alloc.NDArray, - primal_cart_normal_y: alloc.NDArray, - primal_cart_normal_z: alloc.NDArray, + primal_cart_normal_x: data_alloc.NDArray, + primal_cart_normal_y: data_alloc.NDArray, + primal_cart_normal_z: data_alloc.NDArray, array_ns: ModuleType = np, -) -> alloc.NDArray: +) -> data_alloc.NDArray: primal_cart_normal = array_ns.transpose( array_ns.stack( ( 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 503a587c27..9779614499 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,9 @@ def __call__(self, array: data_alloc.NDArray): ) -CoordinateDict: TypeAlias = dict[gtx.Dimension, dict[Literal["lat", "lon"], gtx.Field]] +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,13 +122,25 @@ 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() + + if geometry_type := self._reader.try_attribute(gridfile.MPIMPropertyName.GEOMETRY): + geometry_type = base.GeometryType(geometry_type) + else: + geometry_type = base.GeometryType.ICOSAHEDRON + 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) + 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, allocator: gtx_typing.FieldBufferAllocationUtil) -> CoordinateDict: - return { + def _read_coordinates( + self, + allocator: gtx_typing.FieldBufferAllocationUtil, + geometry_type: base.GeometryType, + ) -> CoordinateDict: + coordinates = { dims.CellDim: { "lat": gtx.as_field( (dims.CellDim,), @@ -171,7 +185,68 @@ def _read_coordinates(self, allocator: gtx_typing.FieldBufferAllocationUtil) -> }, } - 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, + ) -> 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, ) -> icon.IconGrid: """Construct the grid topology from the icon grid file. @@ -294,8 +372,6 @@ def _construct_grid( uuid_ = self._reader.attribute(gridfile.MandatoryPropertyName.GRID_UUID) grid_root = self._reader.attribute(gridfile.MandatoryPropertyName.ROOT) grid_level = self._reader.attribute(gridfile.MandatoryPropertyName.LEVEL) - if geometry_type := self._reader.try_attribute(gridfile.MPIMPropertyName.GEOMETRY): - geometry_type = base.GeometryType(geometry_type) sphere_radius = self._reader.try_attribute(gridfile.MPIMPropertyName.SPHERE_RADIUS) domain_length = self._reader.try_attribute(gridfile.MPIMPropertyName.DOMAIN_LENGTH) domain_height = self._reader.try_attribute(gridfile.MPIMPropertyName.DOMAIN_HEIGHT) 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..f8588c4260 100644 --- a/model/common/src/icon4py/model/common/grid/icon.py +++ b/model/common/src/icon4py/model/common/grid/icon.py @@ -14,7 +14,6 @@ import gt4py.next as gtx from gt4py.next import allocators as gtx_allocators -from typing_extensions import assert_never from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid @@ -54,23 +53,16 @@ def __init__( ) -> None: if geometry_type is None and subdivision is None: raise ValueError("Either geometry_type or subdivision must be provided") - - if geometry_type is None: - geometry_type = base.GeometryType.ICOSAHEDRON - match geometry_type: case base.GeometryType.ICOSAHEDRON: if subdivision is None: raise ValueError("Subdivision must be provided for icosahedron geometry type") - if subdivision.root < 1 or subdivision.level < 0: raise ValueError( f"For icosahedron geometry type, root must be >= 1 and level must be >= 0, got {subdivision.root=} and {subdivision.level=}" ) case base.GeometryType.TORUS: subdivision = None - case _: - assert_never(geometry_type) self.geometry_type = geometry_type self.subdivision = subdivision @@ -137,8 +129,6 @@ def __post_init__(self) -> None: object.__setattr__(self, "radius", constants.EARTH_RADIUS) case base.GeometryType.TORUS: object.__setattr__(self, "radius", None) - case _: - ... if self.global_num_cells is None and self.geometry_type is base.GeometryType.ICOSAHEDRON: object.__setattr__( 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 79cd0d7823..46f4d40219 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 as decomposition from icon4py.model.common.grid import ( + base, geometry, geometry_attributes as geometry_attrs, grid_refinement as refinement, @@ -56,10 +57,12 @@ def __init__( self._providers: dict[str, factory.FieldProvider] = {} self._geometry = geometry_source self._exchange = exchange + 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, + "divergence_averaging_central_cell_weight": 0.5, # divavg_cntrwgt in ICON "weighting_factor": 0.0, "max_nudging_coefficient": 0.375, "nudge_efold_width": 2.0, @@ -68,13 +71,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( @@ -154,8 +157,8 @@ def _register_computed_fields(self) -> None: geofac_n2s = factory.NumpyDataProvider( func=functools.partial( interpolation_fields.compute_geofac_n2s, - array_ns=self._xp, exchange=functools.partial(self._exchange.exchange_and_wait, dims.CellDim), + array_ns=self._xp, ), fields=(attrs.GEOFAC_N2S,), domain=(dims.CellDim, dims.C2E2CODim), @@ -195,38 +198,146 @@ 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, - exchange=functools.partial(self._exchange.exchange_and_wait, dims.CellDim), - ), - 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, + exchange=functools.partial(self._exchange.exchange_and_wait, dims.CellDim), + 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) + ), + "divergence_averaging_central_cell_weight": self._config[ + "divergence_averaging_central_cell_weight" + ], + }, + ) + self.register_provider(cell_average_weight) + + 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) + + pos_on_tplane_e_x_y = factory.NumpyDataProvider( + func=functools.partial( + interpolation_fields.compute_pos_on_tplane_e_x_y, + exchange=functools.partial(self._exchange.exchange_and_wait, dims.EdgeDim), + 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: + cell_average_weight = factory.NumpyDataProvider( + func=functools.partial( + interpolation_fields.compute_mass_conserving_bilinear_cell_average_weight_torus, + exchange=functools.partial(self._exchange.exchange_and_wait, dims.CellDim), + 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) + ), + "divergence_averaging_central_cell_weight": self._config[ + "divergence_averaging_central_cell_weight" + ], + }, + ) + self.register_provider(cell_average_weight) + + 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) + + pos_on_tplane_e_x_y = factory.NumpyDataProvider( + func=functools.partial( + interpolation_fields.compute_pos_on_tplane_e_x_y_torus, + exchange=functools.partial(self._exchange.exchange_and_wait, dims.EdgeDim), + 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) c_lin_e = factory.NumpyDataProvider( func=functools.partial( interpolation_fields.compute_c_lin_e, - array_ns=self._xp, exchange=functools.partial(self._exchange.exchange_and_wait, dims.EdgeDim), + array_ns=self._xp, ), fields=(attrs.C_LIN_E,), domain=(dims.EdgeDim, dims.E2CDim), @@ -246,8 +357,8 @@ def _register_computed_fields(self) -> None: geofac_grg = factory.NumpyDataProvider( func=functools.partial( interpolation_fields.compute_geofac_grg, - array_ns=self._xp, exchange=functools.partial(self._exchange.exchange_and_wait, dims.CellDim), + array_ns=self._xp, ), fields=(attrs.GEOFAC_GRG_X, attrs.GEOFAC_GRG_Y), domain=(dims.CellDim, dims.C2E2CODim), @@ -270,8 +381,8 @@ def _register_computed_fields(self) -> None: e_flx_avg = factory.NumpyDataProvider( func=functools.partial( interpolation_fields.compute_e_flx_avg, - array_ns=self._xp, exchange=functools.partial(self._exchange.exchange_and_wait, dims.EdgeDim), + array_ns=self._xp, ), fields=(attrs.E_FLX_AVG,), domain=(dims.EdgeDim, dims.E2C2EODim), @@ -300,55 +411,11 @@ 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) - - pos_on_tplane_e_x_y = factory.NumpyDataProvider( - func=functools.partial( - interpolation_fields.compute_pos_on_tplane_e_x_y, - array_ns=self._xp, - exchange=functools.partial(self._exchange.exchange_and_wait, dims.EdgeDim), - ), - 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) - cells_aw_verts = factory.NumpyDataProvider( func=functools.partial( interpolation_fields.compute_cells_aw_verts, - array_ns=self._xp, exchange=functools.partial(self._exchange.exchange_and_wait, dims.VertexDim), + array_ns=self._xp, ), fields=(attrs.CELL_AW_VERTS,), domain=(dims.VertexDim, dims.V2CDim), @@ -374,8 +441,8 @@ def _register_computed_fields(self) -> None: rbf_vec_coeff_c = factory.NumpyDataProvider( func=functools.partial( rbf.compute_rbf_interpolation_coeffs_cell, - array_ns=self._xp, exchange=functools.partial(self._exchange.exchange_and_wait, dims.CellDim), + array_ns=self._xp, ), fields=(attrs.RBF_VEC_COEFF_C1, attrs.RBF_VEC_COEFF_C2), domain=(dims.CellDim, dims.C2E2C2EDim), @@ -395,10 +462,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) @@ -406,8 +480,8 @@ def _register_computed_fields(self) -> None: rbf_vec_coeff_e = factory.NumpyDataProvider( func=functools.partial( rbf.compute_rbf_interpolation_coeffs_edge, - array_ns=self._xp, exchange=functools.partial(self._exchange.exchange_and_wait, dims.EdgeDim), + array_ns=self._xp, ), fields=(attrs.RBF_VEC_COEFF_E,), domain=(dims.EdgeDim, dims.E2C2EDim), @@ -426,10 +500,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) @@ -437,8 +518,8 @@ def _register_computed_fields(self) -> None: rbf_vec_coeff_v = factory.NumpyDataProvider( func=functools.partial( rbf.compute_rbf_interpolation_coeffs_vertex, - array_ns=self._xp, exchange=functools.partial(self._exchange.exchange_and_wait, dims.VertexDim), + array_ns=self._xp, ), fields=(attrs.RBF_VEC_COEFF_V1, attrs.RBF_VEC_COEFF_V2), domain=(dims.VertexDim, dims.V2EDim), @@ -458,10 +539,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 58e7064744..a4e30d4659 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py @@ -390,7 +390,7 @@ def _compute_c_bln_avg( c2e2c: data_alloc.NDArray, lat: data_alloc.NDArray, lon: data_alloc.NDArray, - divavg_cntrwgt: ta.wpfloat, + divergence_averaging_central_cell_weight: ta.wpfloat, horizontal_start: gtx.int32, array_ns: ModuleType = np, ) -> data_alloc.NDArray: @@ -398,7 +398,7 @@ def _compute_c_bln_avg( Compute bilinear cell average weight. Args: - divavg_cntrwgt: + divergence_averaging_central_cell_weight: owner_mask: numpy array, representing a gtx.Field[gtx.Dims[CellDim], bool] c2e2c: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, C2E2CDim], gtx.int32] lat: \\ numpy array, representing a gtx.Field[gtx.Dims[CellDim], ta.wpfloat] @@ -421,11 +421,11 @@ def _compute_c_bln_avg( xtemp, lat[horizontal_start:], lon[horizontal_start:], - divavg_cntrwgt, + divergence_averaging_central_cell_weight, array_ns=array_ns, ) c_bln_avg = array_ns.zeros((c2e2c.shape[0], c2e2c.shape[1] + 1)) - c_bln_avg[horizontal_start:, 0] = divavg_cntrwgt + c_bln_avg[horizontal_start:, 0] = divergence_averaging_central_cell_weight c_bln_avg[horizontal_start:, 1] = wgt[0] c_bln_avg[horizontal_start:, 2] = wgt[1] c_bln_avg[horizontal_start:, 3] = wgt[2] @@ -437,7 +437,7 @@ def _force_mass_conservation_to_c_bln_avg( c_bln_avg: data_alloc.NDArray, cell_areas: data_alloc.NDArray, cell_owner_mask: data_alloc.NDArray, - divavg_cntrwgt: ta.wpfloat, + divergence_averaging_central_cell_weight: ta.wpfloat, horizontal_start: gtx.int32, exchange: Callable[[data_alloc.NDArray], None], array_ns: ModuleType = np, @@ -458,7 +458,7 @@ def _force_mass_conservation_to_c_bln_avg( c_bln_avg: input field: bilinear cell weight average cell_areas: area of cells cell_owner_mask: - divavg_cntrwgt: configured central weight + divergence_averaging_central_cell_weight: configured central weight horizontal_start: niter: max number of iterations @@ -498,12 +498,12 @@ def _apply_correction( c_bln_avg: data_alloc.NDArray, residual: data_alloc.NDArray, c2e2c0: data_alloc.NDArray, - divavg_cntrwgt: float, + divergence_averaging_central_cell_weight: float, horizontal_start: gtx.int32, ) -> data_alloc.NDArray: """Apply correction to local weigths based on the computed residuals.""" - maxwgt_loc = divavg_cntrwgt + 0.003 - minwgt_loc = divavg_cntrwgt - 0.003 + maxwgt_loc = divergence_averaging_central_cell_weight + 0.003 + minwgt_loc = divergence_averaging_central_cell_weight - 0.003 relax_coeff = 0.46 c_bln_avg[horizontal_start:, :] = ( c_bln_avg[horizontal_start:, :] - relax_coeff * residual[c2e2c0][horizontal_start:, :] @@ -569,7 +569,7 @@ def _enforce_mass_conservation( c_bln_avg=c_bln_avg, residual=residual, c2e2c0=c2e2c0, - divavg_cntrwgt=divavg_cntrwgt, + divergence_averaging_central_cell_weight=divergence_averaging_central_cell_weight, horizontal_start=horizontal_start, ) @@ -578,13 +578,43 @@ def _enforce_mass_conservation( return c_bln_avg +def _compute_uniform_c_bln_avg( + c2e2c: data_alloc.NDArray, + divergence_averaging_central_cell_weight: ta.wpfloat, + horizontal_start: gtx.int32, + array_ns: ModuleType = np, +) -> data_alloc.NDArray: + """ + Compute bilinear cell average weight for a torus grid. + + Args: + c2e2c + divergence_averaging_central_cell_weight: weight for local / center contribution + horizontal_start: start index of the horizontal domain + + Returns: + c_bln_avg + """ + local_weight = divergence_averaging_central_cell_weight + neighbor_weight = (1.0 - divergence_averaging_central_cell_weight) / 3.0 + + weights = array_ns.asarray([local_weight, neighbor_weight, neighbor_weight, neighbor_weight]) + + c_bln_avg = array_ns.tile( + weights, + (c2e2c.shape[0], 1), + ) + + return c_bln_avg + + def compute_mass_conserving_bilinear_cell_average_weight( c2e2c0: data_alloc.NDArray, lat: data_alloc.NDArray, lon: data_alloc.NDArray, cell_areas: data_alloc.NDArray, cell_owner_mask: data_alloc.NDArray, - divavg_cntrwgt: ta.wpfloat, + divergence_averaging_central_cell_weight: ta.wpfloat, horizontal_start: gtx.int32, horizontal_start_level_3: gtx.int32, exchange: Callable[[data_alloc.NDArray], None], @@ -594,7 +624,7 @@ def compute_mass_conserving_bilinear_cell_average_weight( c2e2c0[:, 1:], lat, lon, - divavg_cntrwgt, + divergence_averaging_central_cell_weight, horizontal_start, array_ns=array_ns, ) @@ -606,10 +636,41 @@ def compute_mass_conserving_bilinear_cell_average_weight( c_bln_avg, cell_areas, cell_owner_mask, - divavg_cntrwgt, + divergence_averaging_central_cell_weight, horizontal_start_level_3, + exchange=exchange, array_ns=array_ns, + ) + + +def compute_mass_conserving_bilinear_cell_average_weight_torus( + c2e2c0: data_alloc.NDArray, + cell_areas: data_alloc.NDArray, + cell_owner_mask: data_alloc.NDArray, + divergence_averaging_central_cell_weight: ta.wpfloat, + horizontal_start: gtx.int32, + horizontal_start_level_3: gtx.int32, + exchange: Callable[[data_alloc.NDArray], None], + array_ns: ModuleType = np, +) -> data_alloc.NDArray: + c_bln_avg = _compute_uniform_c_bln_avg( + c2e2c0[:, 1:], + divergence_averaging_central_cell_weight, + horizontal_start, + array_ns=array_ns, + ) + exchange(c_bln_avg) + # 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, + divergence_averaging_central_cell_weight, + horizontal_start_level_3, exchange=exchange, + array_ns=array_ns, ) @@ -929,6 +990,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: connectivity from cell to its neighboring edges + + Returns: + e_bln_c_s + """ + 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, @@ -1021,5 +1098,54 @@ def compute_pos_on_tplane_e_x_y( ), pos_on_tplane_e_y[llb:, 1], ) + + exchange(pos_on_tplane_e_x, pos_on_tplane_e_y) + return pos_on_tplane_e_x, pos_on_tplane_e_y + + +def compute_pos_on_tplane_e_x_y_torus( + dual_edge_length: data_alloc.NDArray, + e2c: data_alloc.NDArray, + exchange: Callable[[data_alloc.NDArray], None], + 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 + e2c + exchange: halo exchange callback + + Returns: + pos_on_tplane_e_x + 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] + + pos_on_tplane_e_x = array_ns.empty((num_edges, 2), dtype=dual_edge_length.dtype) + pos_on_tplane_e_x[:, 0] = -half_dual_edge_length + pos_on_tplane_e_x[:, 1] = half_dual_edge_length + + pos_on_tplane_e_y = array_ns.zeros((num_edges, 2), dtype=dual_edge_length.dtype) + exchange(pos_on_tplane_e_x, pos_on_tplane_e_y) return pos_on_tplane_e_x, pos_on_tplane_e_y 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 9af69817a8..1ed8b566cd 100644 --- a/model/common/src/icon4py/model/common/interpolation/rbf_interpolation.py +++ b/model/common/src/icon4py/model/common/interpolation/rbf_interpolation.py @@ -14,6 +14,7 @@ import gt4py.next as gtx import numpy as np import scipy.linalg as sla +from gt4py.next import astype from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base as base_grid @@ -48,26 +49,37 @@ class InterpolationKernel(enum.Enum): } -def compute_default_rbf_scale(mean_characteristic_length: ta.wpfloat, dim: RBFDimension): +def compute_default_rbf_scale( + geometry_type: base_grid.GeometryType, + mean_characteristic_length: ta.wpfloat, + mean_dual_edge_length: ta.wpfloat, + dim: RBFDimension, +) -> ta.wpfloat: """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 + 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 astype(scale * (resol / 0.125) ** c3 if resol <= 0.125 else scale, ta.wpfloat) + case base_grid.GeometryType.TORUS: + return mean_dual_edge_length def construct_rbf_matrix_offsets_tables_for_cells( @@ -106,9 +118,17 @@ 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: +def _compute_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. + Compute the distance between points in each row of v. + For the icosahedron geometry this is the arc lengths (in radians), for the + torus it is the Euclidean distance taking into account periodic boundaries. Args: v: 3D array of shape (n, m, 3) where n is the number of elements, @@ -116,26 +136,45 @@ 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) + + +def _compute_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. + Compute the distance between each point in v1 and the points in v2 at the + same row. + For the icosahedron geometry this is the arc lengths (in radians), for the + torus it is the Euclidean distance taking into account periodic boundaries. Args: v1: 2D array of shape (n, 3) where n is the number of elements and 3 is @@ -145,19 +184,30 @@ 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) def _gaussian( @@ -192,22 +242,27 @@ 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) + 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 + 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 + return x, y, z + case base_grid.GeometryType.TORUS: + return u, v, array_ns.zeros_like(u) def _compute_rbf_interpolation_coeffs( @@ -225,8 +280,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, exchange: Callable[[data_alloc.NDArray], None], array_ns: ModuleType = np, ) -> tuple[data_alloc.NDArray, ...]: @@ -273,8 +331,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 = _compute_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) @@ -289,6 +352,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:], @@ -313,7 +377,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 = _compute_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], @@ -335,7 +401,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) @@ -374,8 +440,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, exchange: Callable[[data_alloc.NDArray], None], array_ns: ModuleType = np, ) -> tuple[data_alloc.NDArray]: @@ -397,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, exchange=exchange, array_ns=array_ns, ) @@ -417,8 +489,11 @@ 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, exchange: Callable[[data_alloc.NDArray], None], array_ns: ModuleType = np, ) -> data_alloc.NDArray: @@ -437,8 +512,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, exchange=exchange, array_ns=array_ns, )[0] @@ -458,8 +536,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, exchange: Callable[[data_alloc.NDArray], None], array_ns: ModuleType = np, ) -> tuple[data_alloc.NDArray, data_alloc.NDArray]: @@ -481,8 +562,11 @@ 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, exchange=exchange, 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/fixtures.py b/model/common/tests/common/fixtures.py index 0de81237a7..91dba29f83 100644 --- a/model/common/tests/common/fixtures.py +++ b/model/common/tests/common/fixtures.py @@ -16,7 +16,7 @@ from icon4py.model.common.interpolation import interpolation_attributes, interpolation_factory from icon4py.model.common.metrics import metrics_attributes, metrics_factory from icon4py.model.testing import serialbox -from icon4py.model.testing.definitions import metrics_config +from icon4py.model.testing.definitions import construct_metrics_config from icon4py.model.testing.fixtures.datatest import ( backend, backend_like, @@ -147,7 +147,7 @@ def metrics_factory_from_savepoint( exner_expol, vwind_offctr, rayleigh_type, - ) = metrics_config(experiment) + ) = construct_metrics_config(experiment) vertical_config = vertical.VerticalGridConfig( geometry_source.grid.num_levels, lowest_layer_thickness=lowest_layer_thickness, 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 ce82fb4bbe..1f306f3e86 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 @@ -141,6 +146,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 @@ -152,9 +158,16 @@ 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. + # https://gitlab.dkrz.de/icon-libraries/libiconmath/-/merge_requests/82 + # https://gitlab.dkrz.de/icon/icon-nwp/-/merge_requests/1916 + result = result / constants.EARTH_RADIUS + assert test_utils.dallclose(result, expected, rtol=rtol) @pytest.mark.datatest @@ -304,6 +317,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: @@ -315,12 +329,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: @@ -332,12 +357,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 @@ -347,10 +383,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/grid/unit_tests/test_grid_manager.py b/model/common/tests/common/grid/unit_tests/test_grid_manager.py index df7e54d44e..3ecb8cae98 100644 --- a/model/common/tests/common/grid/unit_tests/test_grid_manager.py +++ b/model/common/tests/common/grid/unit_tests/test_grid_manager.py @@ -76,7 +76,13 @@ def test_grid_manager_eval_v2e( # they get substituted by the "last valid index" in preprocessing step in icon. assert not has_invalid_index(seralized_v2e) v2e_table = grid.get_connectivity("V2E").asnumpy() - assert has_invalid_index(v2e_table) + # Torus grids have no pentagon points and no boundaries hence no invalid + # indexes (while REGIONAL and GLOBAL grids can have) + assert ( + not has_invalid_index(v2e_table) + if experiment.grid.kind == definitions.GridKind.TORUS + else has_invalid_index(v2e_table) + ) _reset_invalid_index(seralized_v2e) assert np.allclose(v2e_table, seralized_v2e) @@ -116,9 +122,14 @@ def test_grid_manager_eval_v2c( # hence in the grid file there are "missing values" # they get substituted by the "last valid index" in preprocessing step in icon. assert not has_invalid_index(serialized_v2c) - assert has_invalid_index(v2c_table) + # Torus grids have no pentagon points and no boundaries hence no invalid + # indexes (while REGIONAL and GLOBAL grids can have) + assert ( + not has_invalid_index(v2c_table) + if experiment.grid.kind == definitions.GridKind.TORUS + else has_invalid_index(v2c_table) + ) _reset_invalid_index(serialized_v2c) - assert np.allclose(v2c_table, serialized_v2c) diff --git a/model/common/tests/common/grid/unit_tests/test_icon.py b/model/common/tests/common/grid/unit_tests/test_icon.py index 84656a9faf..a1dabf1e25 100644 --- a/model/common/tests/common/grid/unit_tests/test_icon.py +++ b/model/common/tests/common/grid/unit_tests/test_icon.py @@ -545,13 +545,19 @@ def test_global_grid_params_from_grid_manager( grid = utils.run_grid_manager(grid_descriptor, keep_skip_values=True, backend=backend).grid params = grid.global_properties assert params is not None - assert pytest.approx(params.geometry_type) == geometry_type - assert pytest.approx(params.subdivision) == subdivision + assert params.geometry_type == geometry_type + match geometry_type: + case base.GeometryType.ICOSAHEDRON: + assert params.subdivision == subdivision + case base.GeometryType.TORUS: + # get the value for torus' subdivision without hardcoding it here + # (it's actually not relevant to check this) + assert params.subdivision == icon.GridShape(base.GeometryType.TORUS).subdivision assert pytest.approx(params.radius) == radius assert pytest.approx(params.domain_length) == domain_length assert pytest.approx(params.domain_height) == domain_height - assert pytest.approx(params.global_num_cells) == global_num_cells - assert pytest.approx(params.num_cells) == num_cells + assert params.global_num_cells == global_num_cells + assert params.num_cells == num_cells assert pytest.approx(params.mean_edge_length) == mean_edge_length assert pytest.approx(params.mean_dual_edge_length) == mean_dual_edge_length assert pytest.approx(params.mean_cell_area) == mean_cell_area diff --git a/model/common/tests/common/grid/unit_tests/test_vertical.py b/model/common/tests/common/grid/unit_tests/test_vertical.py index 048bd6c44b..11f3a42e48 100644 --- a/model/common/tests/common/grid/unit_tests/test_vertical.py +++ b/model/common/tests/common/grid/unit_tests/test_vertical.py @@ -96,7 +96,8 @@ def test_damping_layer_calculation_from_icon_input( ) assert nrdmax == vertical_grid.end_index_of_damping_layer a_array = a.ndarray - assert a_array[nrdmax] > damping_height + damping_height = min(damping_height, a_array[0]) + assert a_array[nrdmax] >= damping_height assert a_array[nrdmax + 1] < damping_height assert vertical_grid.index(v_grid.Domain(dims.KDim, v_grid.Zone.DAMPING)) == nrdmax 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 345d971849..30e8c3dfca 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 @@ -187,7 +187,8 @@ def test_get_geofac_grg( assert field_x.shape == (grid.num_cells, 4) field_y = factory.get(attrs.GEOFAC_GRG_Y).asnumpy() assert field_y.shape == (grid.num_cells, 4) - assert test_helpers.dallclose(field_ref[0].asnumpy(), field_x, rtol=1e-11, atol=1e-16) + # less than 1.1e-16 does not pass on mac for mch_ch_r04b09_dsl (but still passes on CI) + assert test_helpers.dallclose(field_ref[0].asnumpy(), field_x, rtol=1e-11, atol=1.1e-16) assert test_helpers.dallclose(field_ref[1].asnumpy(), field_y, rtol=1e-11, atol=1e-16) @@ -228,6 +229,7 @@ def test_e_flx_avg( [ (definitions.Experiments.MCH_CH_R04B09, 1e-10), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -297,6 +299,7 @@ def test_nudgecoeffs( [ (definitions.Experiments.EXCLAIM_APE, 3.1e-9), (definitions.Experiments.MCH_CH_R04B09, 4e-2), + (definitions.Experiments.GAUSS3D, 1e-14), ], ) @pytest.mark.datatest @@ -330,6 +333,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 @@ -357,6 +361,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 9710c8090b..a175a3a5a1 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 @@ -19,6 +19,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, @@ -26,7 +27,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 @@ -260,7 +263,7 @@ def test_compute_c_bln_avg( xp = data_alloc.import_array_ns(backend) cell_areas = grid_savepoint.cell_areas().ndarray # both experiment use the default value - divavg_cntrwgt = 0.5 + divergence_averaging_central_cell_weight = 0.5 c_bln_avg_ref = interpolation_savepoint.c_bln_avg().asnumpy() horizontal_start = icon_grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)) horizontal_start_p2 = icon_grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3)) @@ -271,20 +274,32 @@ 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, - exchange=utils.dummy_exchange, - )( - 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, + divergence_averaging_central_cell_weight, + horizontal_start, + horizontal_start_p2, + exchange=utils.dummy_exchange, + 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, + divergence_averaging_central_cell_weight, + horizontal_start, + horizontal_start_p2, + exchange=utils.dummy_exchange, + array_ns=xp, + ) + assert test_helpers.dallclose(data_alloc.as_numpy(c_bln_avg), c_bln_avg_ref, rtol=1e-11) @@ -380,9 +395,13 @@ 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 - ) + 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(), rtol=1e-10 ) @@ -400,6 +419,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 @@ -411,21 +431,31 @@ 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, - utils.dummy_exchange, - array_ns=xp, - ) + + 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, + exchange=utils.dummy_exchange, + 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, + exchange=utils.dummy_exchange, + array_ns=xp, + ) assert test_helpers.dallclose(pos_on_tplane_e_x, pos_on_tplane_e_x_ref, atol=1e-8, rtol=1e-9) assert test_helpers.dallclose(pos_on_tplane_e_y, pos_on_tplane_e_y_ref, atol=1e-8, rtol=1e-9) 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 be00a14c03..7fdce1f3ab 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 @@ -147,6 +147,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( @@ -165,6 +166,12 @@ def test_rbf_interpolation_coeffs_cell( ) assert horizontal_start < grid.num_cells + geometry_type = ( + grid.global_properties.geometry_type + if grid.global_properties.geometry_type + else pytest.fail("geometry_type cannot be None") + ) + rbf_vec_coeff_c1, rbf_vec_coeff_c2 = rbf.compute_rbf_interpolation_coeffs_cell( # type: ignore[misc] # function returns two vars geometry.get(geometry_attrs.CELL_LAT).ndarray, geometry.get(geometry_attrs.CELL_LON).ndarray, @@ -179,8 +186,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), + geometry_type.value, + rbf.compute_default_rbf_scale( + 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, exchange=utils.dummy_exchange, array_ns=data_alloc.import_array_ns(backend), ) @@ -217,6 +232,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( @@ -235,6 +251,12 @@ def test_rbf_interpolation_coeffs_vertex( ) assert horizontal_start < grid.num_vertices + geometry_type = ( + grid.global_properties.geometry_type + if grid.global_properties.geometry_type + else pytest.fail("geometry_type cannot be None") + ) + rbf_vec_coeff_v1, rbf_vec_coeff_v2 = rbf.compute_rbf_interpolation_coeffs_vertex( geometry.get(geometry_attrs.VERTEX_LAT).ndarray, geometry.get(geometry_attrs.VERTEX_LON).ndarray, @@ -249,8 +271,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), + geometry_type.value, + rbf.compute_default_rbf_scale( + 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, exchange=utils.dummy_exchange, array_ns=data_alloc.import_array_ns(backend), ) @@ -287,6 +317,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( @@ -305,6 +336,12 @@ def test_rbf_interpolation_coeffs_edge( ) assert horizontal_start < grid.num_edges + geometry_type = ( + grid.global_properties.geometry_type + if grid.global_properties.geometry_type + else pytest.fail("geometry_type cannot be None") + ) + rbf_vec_coeff_e = rbf.compute_rbf_interpolation_coeffs_edge( geometry.get(geometry_attrs.EDGE_LAT).ndarray, geometry.get(geometry_attrs.EDGE_LON).ndarray, @@ -321,8 +358,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), + geometry_type.value, + rbf.compute_default_rbf_scale( + 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, exchange=utils.dummy_exchange, array_ns=data_alloc.import_array_ns(backend), ) diff --git a/model/common/tests/common/metrics/unit_tests/test_metric_fields.py b/model/common/tests/common/metrics/unit_tests/test_metric_fields.py index 33466f444f..33277759f8 100644 --- a/model/common/tests/common/metrics/unit_tests/test_metric_fields.py +++ b/model/common/tests/common/metrics/unit_tests/test_metric_fields.py @@ -18,6 +18,7 @@ from icon4py.model.common.metrics import metric_fields as mf from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import definitions, test_utils as testing_helpers +from icon4py.model.testing.definitions import construct_metrics_config from icon4py.model.testing.fixtures.datatest import ( backend, data_provider, @@ -137,7 +138,7 @@ def test_compute_scaling_factor_for_3d_divdamp( @pytest.mark.datatest def test_compute_rayleigh_w( icon_grid: base_grid.Grid, - experiment: definitions.Experiments, + experiment: definitions.Experiment, metrics_savepoint: sb.MetricSavepoint, grid_savepoint: sb.IconGridSavepoint, backend: gtx_typing.Backend, @@ -147,9 +148,16 @@ def test_compute_rayleigh_w( rayleigh_w_full = data_alloc.zero_field( icon_grid, dims.KDim, extend={dims.KDim: 1}, allocator=backend ) - rayleigh_type = 2 - rayleigh_coeff = 0.1 if experiment == definitions.Experiments.EXCLAIM_APE else 5.0 - damping_height = 50000.0 if experiment == definitions.Experiments.EXCLAIM_APE else 12500.0 + ( + _lowest_layer_thickness, + _model_top_height, + _stretch_factor, + damping_height, + rayleigh_coeff, + _exner_expol, + _vwind_offctr, + rayleigh_type, + ) = construct_metrics_config(experiment) mf.compute_rayleigh_w.with_backend(backend=backend)( rayleigh_w=rayleigh_w_full, vct_a=grid_savepoint.vct_a(), @@ -233,7 +241,14 @@ def test_compute_exner_exfac( backend: gtx_typing.Backend, ) -> None: horizontal_start = icon_grid.start_index(cell_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_2)) - exner_expol = 0.333 if experiment == definitions.Experiments.MCH_CH_R04B09 else 0.3333333333333 + match experiment: + case definitions.Experiments.MCH_CH_R04B09: + exner_expol = 0.333 + case definitions.Experiments.WEISMAN_KLEMP_TORUS: + exner_expol = 0.333 + case _: + exner_expol = 1.0 / 3.0 + exner_exfac = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, allocator=backend) max_slp = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, allocator=backend) max_hgtd = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, allocator=backend) @@ -329,7 +344,14 @@ def test_compute_exner_w_implicit_weight_parameter( ) vwind_impl_wgt_ref = metrics_savepoint.vwind_impl_wgt() dual_edge_length = grid_savepoint.dual_edge_length() - vwind_offctr = 0.2 if experiment == definitions.Experiments.MCH_CH_R04B09 else 0.15 + match experiment: + case definitions.Experiments.MCH_CH_R04B09: + vwind_offctr = 0.2 + case definitions.Experiments.WEISMAN_KLEMP_TORUS: + vwind_offctr = 0.2 + case _: + vwind_offctr = 0.15 + xp = data_alloc.import_array_ns(backend) exner_w_implicit_weight_parameter = mf.compute_exner_w_implicit_weight_parameter( c2e=icon_grid.get_connectivity(dims.C2E).ndarray, diff --git a/model/common/tests/common/metrics/unit_tests/test_metrics_factory.py b/model/common/tests/common/metrics/unit_tests/test_metrics_factory.py index bfdcd6bb37..b3ba399319 100644 --- a/model/common/tests/common/metrics/unit_tests/test_metrics_factory.py +++ b/model/common/tests/common/metrics/unit_tests/test_metrics_factory.py @@ -25,7 +25,7 @@ serialbox, test_utils as test_helpers, ) -from icon4py.model.testing.definitions import metrics_config +from icon4py.model.testing.definitions import construct_metrics_config from icon4py.model.testing.fixtures.datatest import ( backend, data_provider, @@ -66,7 +66,7 @@ def _get_metrics_factory( exner_expol, vwind_offctr, rayleigh_type, - ) = metrics_config(experiment) + ) = construct_metrics_config(experiment) vertical_config = v_grid.VerticalGridConfig( geometry.grid.num_levels, @@ -603,6 +603,9 @@ def test_factory_compute_diffusion_mask_and_coef( experiment: definitions.Experiment, backend: gtx_typing.Backend | None, ) -> None: + if experiment == definitions.Experiments.GAUSS3D: + pytest.xfail("TODO") + field_ref_1 = metrics_savepoint.mask_hdiff() field_ref_2 = metrics_savepoint.zd_diffcoef() factory = _get_metrics_factory( @@ -628,6 +631,9 @@ def test_factory_compute_diffusion_intcoeff_and_vertoffset( experiment: definitions.Experiment, backend: gtx_typing.Backend | None, ) -> None: + if experiment == definitions.Experiments.GAUSS3D: + pytest.xfail("TODO") + field_ref_1 = metrics_savepoint.zd_intcoef() field_ref_2 = metrics_savepoint.zd_vertoffset() factory = _get_metrics_factory( diff --git a/model/driver/tests/driver/integration_tests/test_icon4py.py b/model/driver/tests/driver/integration_tests/test_icon4py.py index 97d708818e..f3908eabc1 100644 --- a/model/driver/tests/driver/integration_tests/test_icon4py.py +++ b/model/driver/tests/driver/integration_tests/test_icon4py.py @@ -9,7 +9,6 @@ from typing import TYPE_CHECKING -import click import pytest import icon4py.model.common.grid.states as grid_states @@ -42,7 +41,7 @@ @pytest.mark.embedded_remap_error @pytest.mark.datatest @pytest.mark.parametrize( - "experiment, istep_init, istep_exit, substep_init, substep_exit, timeloop_date_init, timeloop_date_exit, step_date_init, step_date_exit, timeloop_diffusion_linit_init, timeloop_diffusion_linit_exit, vn_only", + "experiment, istep_init, istep_exit, substep_init, substep_exit, timeloop_date_init, timeloop_date_exit, step_date_init, step_date_exit, timeloop_diffusion_linit_init, timeloop_diffusion_linit_exit", [ ( definitions.Experiments.MCH_CH_R04B09, @@ -56,7 +55,6 @@ "2021-06-20T12:00:10.000", True, False, - False, ), ( definitions.Experiments.MCH_CH_R04B09, @@ -70,7 +68,6 @@ "2021-06-20T12:00:20.000", False, False, - True, ), ( definitions.Experiments.GAUSS3D, @@ -84,7 +81,6 @@ "2001-01-01T00:00:04.000", False, False, - False, ), ], ) @@ -93,7 +89,6 @@ def test_run_timeloop_single_step( timeloop_date_init: str, timeloop_date_exit: str, timeloop_diffusion_linit_init: bool, - vn_only: bool, # TODO unused? *, grid_savepoint: sb.IconGridSavepoint, icon_grid: base_grid.Grid, diff --git a/model/testing/src/icon4py/model/testing/definitions.py b/model/testing/src/icon4py/model/testing/definitions.py index 9c78736ac5..3a39edab6e 100644 --- a/model/testing/src/icon4py/model/testing/definitions.py +++ b/model/testing/src/icon4py/model/testing/definitions.py @@ -228,6 +228,10 @@ def construct_diffusion_config( hdiff_temp=True, n_substeps=ndyn_substeps, ) + elif experiment == Experiments.GAUSS3D: + return diffusion.DiffusionConfig( + n_substeps=ndyn_substeps, + ) else: raise NotImplementedError( f"DiffusionConfig for experiment {experiment.name} not implemented." @@ -249,34 +253,58 @@ def construct_nonhydrostatic_config(experiment: Experiment) -> solve_nh.NonHydro rayleigh_coeff=0.1, divdamp_order=dycore_states.DivergenceDampingOrder.COMBINED, # type: ignore[arg-type] # TODO(havogt): typing in `NonHydrostaticConfig` needs to be fixed ) + elif experiment == Experiments.GAUSS3D: + return solve_nh.NonHydrostaticConfig( + fourth_order_divdamp_factor=0.0025, + ) else: raise NotImplementedError( f"NonHydrostaticConfig for experiment {experiment.name} not implemented." ) -def metrics_config(experiment: Experiment) -> tuple: - rayleigh_coeff = 5.0 - lowest_layer_thickness = 50.0 - exner_expol = 0.333 - vwind_offctr = 0.2 - rayleigh_type = 2 - model_top_height = 23500.0 - stretch_factor = 1.0 - damping_height = 45000.0 +def construct_metrics_config(experiment: Experiment) -> tuple: match experiment: case Experiments.MCH_CH_R04B09: lowest_layer_thickness = 20.0 model_top_height = 23000.0 stretch_factor = 0.65 damping_height = 12500.0 + rayleigh_coeff = 5.0 + exner_expol = 0.333 + vwind_offctr = 0.2 + rayleigh_type = 2 case Experiments.EXCLAIM_APE: + lowest_layer_thickness = 50.0 model_top_height = 75000.0 stretch_factor = 0.9 damping_height = 50000.0 rayleigh_coeff = 0.1 exner_expol = 0.3333333333333 vwind_offctr = 0.15 + rayleigh_type = 2 + case Experiments.GAUSS3D: + lowest_layer_thickness = 50.0 + model_top_height = 23500.0 + stretch_factor = 1.0 + damping_height = 45000.0 + rayleigh_coeff = 0.1 + exner_expol = 1.0 / 3.0 + vwind_offctr = 0.15 + rayleigh_type = 2 + case Experiments.WEISMAN_KLEMP_TORUS: + lowest_layer_thickness = 50.0 + model_top_height = 23500.0 + stretch_factor = 1.0 + damping_height = 8000.0 + rayleigh_coeff = 0.75 + exner_expol = 0.333 + vwind_offctr = 0.15 + rayleigh_type = 2 + case _: + raise NotImplementedError( + f"Metrics config for experiment {experiment.name} not implemented." + ) return ( lowest_layer_thickness, 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..9081692f9f 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( 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()