diff --git a/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_diffusion.py b/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_diffusion.py index 630c560ff5..2900d88eae 100644 --- a/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_diffusion.py +++ b/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_diffusion.py @@ -165,6 +165,7 @@ def test_smagorinski_factor_diffusion_type_5(experiment): [ (definitions.Experiments.MCH_CH_R04B09, "2021-06-20T12:00:10.000"), (definitions.Experiments.MCH_CH_R04B09, "2021-06-20T12:00:20.000"), + (definitions.Experiments.GAUSS3D, "2001-01-01T00:00:04.000"), ], ) def test_diffusion_init( @@ -298,6 +299,7 @@ def _verify_init_values_against_savepoint( (definitions.Experiments.MCH_CH_R04B09, "2021-06-20T12:00:20.000"), (definitions.Experiments.EXCLAIM_APE, "2000-01-01T00:00:02.000"), (definitions.Experiments.EXCLAIM_APE, "2000-01-01T00:00:04.000"), + (definitions.Experiments.GAUSS3D, "2001-01-01T00:00:04.000"), ], ) @pytest.mark.parametrize("ndyn_substeps", (2,)) @@ -315,6 +317,9 @@ def test_verify_diffusion_init_against_savepoint( ndyn_substeps, backend, ): + if experiment == definitions.Experiments.GAUSS3D: + pytest.xfail("wrong results") + grid = get_grid_for_experiment(experiment, backend) cell_params = get_cell_geometry_for_experiment(experiment, backend) edge_params = get_edge_geometry_for_experiment(experiment, backend) @@ -365,6 +370,11 @@ def test_verify_diffusion_init_against_savepoint( "2000-01-01T00:00:02.000", "2000-01-01T00:00:02.000", ), + ( + definitions.Experiments.GAUSS3D, + "2001-01-01T00:00:04.000", + "2001-01-01T00:00:04.000", + ), ], ) @pytest.mark.parametrize("ndyn_substeps", [2]) @@ -390,6 +400,9 @@ def test_run_diffusion_single_step( if orchestration and not test_utils.is_dace(backend): pytest.skip("Orchestration test requires a dace backend.") + if experiment == definitions.Experiments.GAUSS3D: + pytest.xfail("wrong results") + grid = get_grid_for_experiment(experiment, backend) cell_geometry = get_cell_geometry_for_experiment(experiment, backend) edge_geometry = get_edge_geometry_for_experiment(experiment, backend) @@ -451,6 +464,11 @@ def test_run_diffusion_single_step( "2021-06-20T12:00:10.000", "2021-06-20T12:00:10.000", ), + ( + definitions.Experiments.GAUSS3D, + "2001-01-01T00:00:04.000", + "2001-01-01T00:00:04.000", + ), ], ) @pytest.mark.parametrize("ndyn_substeps", (2,)) @@ -475,6 +493,9 @@ def test_run_diffusion_multiple_steps( pytest.skip("dace orchestration broken by precompiled programs") if not test_utils.is_dace(backend): raise pytest.skip("This test is only executed for dace backends") + + if experiment == definitions.Experiments.GAUSS3D: + pytest.xfail("probably wrong results? not tested") ###################################################################### # Diffusion initialization ###################################################################### @@ -574,8 +595,23 @@ def test_run_diffusion_multiple_steps( @pytest.mark.datatest @pytest.mark.embedded_remap_error -@pytest.mark.parametrize("experiment", [definitions.Experiments.MCH_CH_R04B09]) -@pytest.mark.parametrize("linit", [True]) +@pytest.mark.parametrize( + "experiment,step_date_init,step_date_exit,linit", + [ + ( + definitions.Experiments.MCH_CH_R04B09, + "2021-06-20T12:00:10.000", + "2021-06-20T12:00:10.000", + True, + ), + ( + definitions.Experiments.GAUSS3D, + "2001-01-01T00:00:04.000", + "2001-01-01T00:00:04.000", + False, + ), + ], +) # TODO(): Enable dace orchestration, currently broken by precompiled programs @pytest.mark.parametrize("orchestration", [False]) def test_run_diffusion_initial_step( @@ -594,6 +630,10 @@ def test_run_diffusion_initial_step( ): if orchestration and not test_utils.is_dace(backend): pytest.skip("Orchestration test requires a dace backend.") + + if experiment == definitions.Experiments.GAUSS3D: + pytest.xfail("wrong results") + grid = get_grid_for_experiment(experiment, backend) cell_geometry = get_cell_geometry_for_experiment(experiment, backend) edge_geometry = get_edge_geometry_for_experiment(experiment, backend) @@ -652,17 +692,20 @@ def test_run_diffusion_initial_step( @pytest.mark.datatest -@pytest.mark.parametrize("linit", [True]) # TODO(havogt): Remove custom `experiment` parametrization @pytest.mark.parametrize( - "experiment,step_date_init", + "experiment,step_date_init,linit", [ - (definitions.Experiments.MCH_CH_R04B09, "2021-06-20T12:00:10.000"), + (definitions.Experiments.MCH_CH_R04B09, "2021-06-20T12:00:10.000", True), + (definitions.Experiments.GAUSS3D, "2001-01-01T00:00:04.000", False), ], ) def test_verify_special_diffusion_inital_step_values_against_initial_savepoint( savepoint_diffusion_init, experiment, icon_grid, linit, ndyn_substeps, backend ): + if experiment == definitions.Experiments.GAUSS3D: + pytest.xfail("wrong results") + savepoint = savepoint_diffusion_init config = definitions.construct_diffusion_config(experiment, ndyn_substeps=ndyn_substeps) diff --git a/model/common/src/icon4py/model/common/grid/geometry.py b/model/common/src/icon4py/model/common/grid/geometry.py index df18139e89..dd966aeee5 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 @@ -80,33 +80,116 @@ def __init__( grid: icon.IconGrid, decomposition_info: definitions.DecompositionInfo, backend: gtx_typing.Backend | None, - coordinates: gm.CoordinateDict, - extra_fields: gm.GeometryDict, metadata: dict[str, model.FieldMetaData], - ): + ) -> None: """ Args: grid: IconGrid the grid topology decomposition_info: data structure containing owner masks for field dimensions backend: backend used for memory allocation and computation - coordinates: dictionary containing geographical coordinates for grid cells, edges and vertices, - extra_fields: fields that are not computed but directly read off the grid file, - currently only the edge_system_orientation cell_area. Should eventually disappear. + # TODO(msimberg): Should eventually disappear? Using for now to pass cartesian coordinates for torus. metadata: a dictionary of FieldMetaData for all fields computed in GridGeometry. """ self._providers = {} self._backend = backend + self._xp = data_alloc.import_array_ns(backend) self._allocator = gtx.constructors.zeros.partial(allocator=backend) self._grid = grid self._decomposition_info = decomposition_info self._attrs = metadata - self._geometry_type: base.GeometryType | None = grid.global_properties.geometry_type + self._geometry_type: base.GeometryType = grid.global_properties.geometry_type self._edge_domain = h_grid.domain(dims.EdgeDim) log.info( f"initialized geometry for backend = '{self._backend_name()}' and grid = '{self._grid}'" ) + @staticmethod + def with_geometry_type( + grid: icon.IconGrid, + decomposition_info: definitions.DecompositionInfo, + backend: gtx_typing.Backend | None, + coordinates: gm.CoordinateDict, + extra_fields: gm.GeometryDict, + metadata: dict[str, model.FieldMetaData], + ) -> "GridGeometry": + match grid.global_properties.geometry_type: + case base.GeometryType.ICOSAHEDRON: + return IcosahedronGridGeometry( + grid, + decomposition_info, + backend, + coordinates, + extra_fields, + metadata, + ) + case base.GeometryType.TORUS: + return TorusGridGeometry( + grid, + decomposition_info, + backend, + coordinates, + extra_fields, + metadata, + ) + case _: + # TODO(msimberg): Warn? Error? Fall back to icosahedron? + raise ValueError( + f"Geometry type {gm.grid.global_properties.geometry_type} not supported." + ) + + def _inverse_field_provider(self, field_name: str) -> factory.FieldProvider: + meta = attrs.metadata_for_inverse(attrs.attrs[field_name]) + name = meta["standard_name"] + self._attrs.update({name: meta}) + provider = factory.ProgramFieldProvider( + func=math_helpers.compute_inverse_on_edges, + deps={"f": field_name}, + fields={"f_inverse": name}, + domain={ + dims.EdgeDim: ( + self._edge_domain(h_grid.Zone.LOCAL), + self._edge_domain(h_grid.Zone.LOCAL), + ) + }, + ) + return provider + + def __repr__(self) -> str: + geometry_name = self._geometry_type._name_ if self._geometry_type else "" + return ( + f"{self.__class__.__name__} for geometry_type={geometry_name} (grid={self._grid.id!r})" + ) + + @property + def metadata(self) -> dict[str, model.FieldMetaData]: + return self._attrs + + @property + def backend(self) -> gtx_typing.Backend: + return self._backend + + @property + def grid(self) -> icon.IconGrid: + return self._grid + + @property + def vertical_grid(self) -> None: + return None + + +class IcosahedronGridGeometry(GridGeometry): + def __init__( + self, + grid: icon.IconGrid, + decomposition_info: definitions.DecompositionInfo, + backend: gtx_typing.Backend | None, + coordinates: gm.CoordinateDict, + extra_fields: gm.GeometryDict, + metadata: dict[str, model.FieldMetaData], + ) -> None: + super().__init__(grid, decomposition_info, backend, metadata) + ( edge_orientation0_lat, edge_orientation0_lon, @@ -174,9 +257,11 @@ def __init__( } ) self.register_provider(input_fields_provider) - self._register_computed_fields() + self._register_computed_fields(input_fields_provider) - def _register_computed_fields(self) -> None: + def _register_computed_fields( + self, input_fields_provider: factory.PrecomputedFieldProvider + ) -> None: meta = attrs.metadata_for_inverse(attrs.attrs[attrs.EDGE_LENGTH]) name = meta["standard_name"] self._attrs.update({name: meta}) @@ -501,50 +586,258 @@ 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]) + +# TODO(msimberg): Don't do separate class but just initialize differently? Then +# caller doesn't have to know which one it is. IconGrid knows anyway which type +# it is. Or add factory method. +class TorusGridGeometry(GridGeometry): + def __init__( + self, + grid: icon.IconGrid, + decomposition_info: definitions.DecompositionInfo, + backend: gtx_typing.Backend | None, + coordinates: gm.CoordinateDict, + extra_fields: gm.GeometryDict, + metadata: dict[str, model.FieldMetaData], + ) -> None: + super().__init__(grid, decomposition_info, backend, metadata) + + log.info( + f"initialized geometry for backend = '{self._backend_name()}' and grid = '{self._grid}'" + ) + + coordinates_ = { + attrs.CELL_LAT: coordinates[dims.CellDim]["lat"], + attrs.CELL_LON: coordinates[dims.CellDim]["lon"], + attrs.CELL_CENTER_X: coordinates[dims.CellDim]["x"], + attrs.CELL_CENTER_Y: coordinates[dims.CellDim]["y"], + attrs.CELL_CENTER_Z: coordinates[dims.CellDim]["z"], + attrs.VERTEX_LAT: coordinates[dims.VertexDim]["lat"], + attrs.VERTEX_LON: coordinates[dims.VertexDim]["lon"], + attrs.VERTEX_X: coordinates[dims.VertexDim]["x"], + attrs.VERTEX_Y: coordinates[dims.VertexDim]["y"], + attrs.VERTEX_Z: coordinates[dims.VertexDim]["z"], + attrs.EDGE_LON: coordinates[dims.EdgeDim]["lon"], + attrs.EDGE_LAT: coordinates[dims.EdgeDim]["lat"], + attrs.EDGE_CENTER_X: coordinates[dims.EdgeDim]["x"], + attrs.EDGE_CENTER_Y: coordinates[dims.EdgeDim]["y"], + attrs.EDGE_CENTER_Z: coordinates[dims.EdgeDim]["z"], + } + coodinate_provider = factory.PrecomputedFieldProvider(coordinates_) + self.register_provider(coodinate_provider) + + input_fields_provider = factory.PrecomputedFieldProvider( + { + attrs.EDGE_LENGTH: extra_fields[gridfile.GeometryName.EDGE_LENGTH], + attrs.DUAL_EDGE_LENGTH: extra_fields[gridfile.GeometryName.DUAL_EDGE_LENGTH], + attrs.EDGE_CELL_DISTANCE: extra_fields[gridfile.GeometryName.EDGE_CELL_DISTANCE], + attrs.EDGE_VERTEX_DISTANCE: extra_fields[ + gridfile.GeometryName.EDGE_VERTEX_DISTANCE + ], + attrs.CELL_AREA: extra_fields[gridfile.GeometryName.CELL_AREA], + attrs.DUAL_AREA: extra_fields[gridfile.GeometryName.DUAL_AREA], + attrs.TANGENT_ORIENTATION: extra_fields[gridfile.GeometryName.TANGENT_ORIENTATION], + "edge_owner_mask": gtx.as_field( + (dims.EdgeDim,), + decomposition_info.owner_mask(dims.EdgeDim), + dtype=bool, + allocator=self._backend, + ), + attrs.CELL_NORMAL_ORIENTATION: extra_fields[ + gridfile.GeometryName.CELL_NORMAL_ORIENTATION + ], + attrs.VERTEX_EDGE_ORIENTATION: extra_fields[ + gridfile.GeometryName.EDGE_ORIENTATION_ON_VERTEX + ], + "vertex_owner_mask": gtx.as_field( + (dims.VertexDim,), + decomposition_info.owner_mask(dims.VertexDim), + allocator=self._backend, + dtype=bool, + ), + "cell_owner_mask": gtx.as_field( + (dims.CellDim,), + decomposition_info.owner_mask(dims.CellDim), + allocator=self._backend, + dtype=bool, + ), + } + ) + self.register_provider(input_fields_provider) + self._register_computed_fields(input_fields_provider) + + def _register_computed_fields( + self, input_fields_provider: factory.PrecomputedFieldProvider + ) -> None: + meta = attrs.metadata_for_inverse(attrs.attrs[attrs.EDGE_LENGTH]) name = meta["standard_name"] self._attrs.update({name: meta}) - provider = factory.ProgramFieldProvider( - func=math_helpers.compute_inverse_on_edges, - deps={"f": field_name}, - fields={"f_inverse": name}, + inverse_edge_length = self._inverse_field_provider(attrs.EDGE_LENGTH) + self.register_provider(inverse_edge_length) + + inverse_dual_edge_length = self._inverse_field_provider(attrs.DUAL_EDGE_LENGTH) + self.register_provider(inverse_dual_edge_length) + + # TODO(msimberg): Is this needed for torus? Even the serialized data + # doesn't seem to make sense. Earth radius is used even for the torus. + vertex_vertex_distance = factory.ProgramFieldProvider( + func=stencils.compute_distance_of_far_edges_in_diamond_torus, domain={ dims.EdgeDim: ( - self._edge_domain(h_grid.Zone.LOCAL), + self._edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2), self._edge_domain(h_grid.Zone.LOCAL), ) }, + fields={"far_vertex_distance": attrs.VERTEX_VERTEX_LENGTH}, + deps={ + "vertex_x": attrs.VERTEX_X, + "vertex_y": attrs.VERTEX_Y, + }, + params={ + "domain_length": self._grid.global_properties.domain_length, + "domain_height": self._grid.global_properties.domain_height, + }, ) - return provider - - def __repr__(self) -> str: - geometry_name = self._geometry_type._name_ if self._geometry_type else "" - return ( - f"{self.__class__.__name__} for geometry_type={geometry_name} (grid={self._grid.id!r})" + self.register_provider(vertex_vertex_distance) + inverse_far_edge_distance_provider = self._inverse_field_provider( + attrs.VERTEX_VERTEX_LENGTH ) + self.register_provider(inverse_far_edge_distance_provider) - @property - def metadata(self) -> dict[str, model.FieldMetaData]: - return self._attrs - - @property - def backend(self) -> gtx_typing.Backend | None: - return self._backend + edge_areas = factory.ProgramFieldProvider( + func=stencils.compute_edge_area, + deps={ + "owner_mask": "edge_owner_mask", + "primal_edge_length": attrs.EDGE_LENGTH, + "dual_edge_length": attrs.DUAL_EDGE_LENGTH, + }, + fields={"area": attrs.EDGE_AREA}, + domain={ + dims.EdgeDim: ( + self._edge_domain(h_grid.Zone.LOCAL), + self._edge_domain(h_grid.Zone.END), + ) + }, + ) + self.register_provider(edge_areas) + coriolis_params = factory.PrecomputedFieldProvider( + { + "coriolis_parameter": gtx.as_field( + (dims.EdgeDim,), + self._xp.zeros( + self._grid.start_index(self._edge_domain(h_grid.Zone.END)) + - self._grid.start_index(self._edge_domain(h_grid.Zone.LOCAL)) + ), + dtype=ta.wpfloat, + allocator=self._backend, + ) + } + ) + self.register_provider(coriolis_params) - @property - def grid(self) -> icon.IconGrid: - return self._grid + # normals and tangents + tangent_normal_coordinates = factory.ProgramFieldProvider( + func=stencils.compute_cartesian_coordinates_of_edge_tangent_and_normal_torus, + deps={ + "vertex_x": attrs.VERTEX_X, + "vertex_y": attrs.VERTEX_Y, + "edge_x": attrs.EDGE_CENTER_X, + "edge_y": attrs.EDGE_CENTER_Y, + "edge_orientation": attrs.TANGENT_ORIENTATION, + }, + fields={ + "tangent_x": attrs.EDGE_TANGENT_X, + "tangent_y": attrs.EDGE_TANGENT_Y, + "tangent_z": attrs.EDGE_TANGENT_Z, + "tangent_u": attrs.EDGE_DUAL_U, + "tangent_v": attrs.EDGE_DUAL_V, + "normal_x": attrs.EDGE_NORMAL_X, + "normal_y": attrs.EDGE_NORMAL_Y, + "normal_z": attrs.EDGE_NORMAL_Z, + "normal_u": attrs.EDGE_NORMAL_U, + "normal_v": attrs.EDGE_NORMAL_V, + }, + domain={ + dims.EdgeDim: ( + self._edge_domain(h_grid.Zone.LOCAL), + self._edge_domain(h_grid.Zone.END), + ) + }, + params={ + "domain_length": self._grid.global_properties.domain_length, + "domain_height": self._grid.global_properties.domain_height, + }, + ) + self.register_provider(tangent_normal_coordinates) - @property - def vertical_grid(self) -> None: - return None + # 3. primal_normal_vert, primal_normal_cell + normal_vert_wrapper = SparseFieldProviderWrapper( + tangent_normal_coordinates, + target_dims=attrs.attrs[attrs.EDGE_NORMAL_VERTEX_U]["dims"], + fields=(attrs.EDGE_NORMAL_VERTEX_U, attrs.EDGE_NORMAL_VERTEX_V), + pairs=( + ( + attrs.EDGE_NORMAL_X, + attrs.EDGE_NORMAL_X, + attrs.EDGE_NORMAL_X, + attrs.EDGE_NORMAL_X, + ), + ( + attrs.EDGE_NORMAL_Y, + attrs.EDGE_NORMAL_Y, + attrs.EDGE_NORMAL_Y, + attrs.EDGE_NORMAL_Y, + ), + ), + ) + self.register_provider(normal_vert_wrapper) + normal_cell_wrapper = SparseFieldProviderWrapper( + tangent_normal_coordinates, + target_dims=attrs.attrs[attrs.EDGE_NORMAL_CELL_U]["dims"], + fields=(attrs.EDGE_NORMAL_CELL_U, attrs.EDGE_NORMAL_CELL_V), + pairs=( + (attrs.EDGE_NORMAL_X, attrs.EDGE_NORMAL_X), + (attrs.EDGE_NORMAL_Y, attrs.EDGE_NORMAL_Y), + ), + ) + self.register_provider(normal_cell_wrapper) + # 3. dual normals: the dual normals are the edge tangents + tangent_vert_wrapper = SparseFieldProviderWrapper( + tangent_normal_coordinates, + target_dims=attrs.attrs[attrs.EDGE_TANGENT_VERTEX_U]["dims"], + fields=(attrs.EDGE_TANGENT_VERTEX_U, attrs.EDGE_TANGENT_VERTEX_V), + pairs=( + ( + attrs.EDGE_TANGENT_X, + attrs.EDGE_TANGENT_X, + attrs.EDGE_TANGENT_X, + attrs.EDGE_TANGENT_X, + ), + ( + attrs.EDGE_TANGENT_Y, + attrs.EDGE_TANGENT_Y, + attrs.EDGE_TANGENT_Y, + attrs.EDGE_TANGENT_Y, + ), + ), + ) + self.register_provider(tangent_vert_wrapper) + tangent_cell_wrapper = SparseFieldProviderWrapper( + tangent_normal_coordinates, + target_dims=attrs.attrs[attrs.EDGE_TANGENT_CELL_U]["dims"], + fields=(attrs.EDGE_TANGENT_CELL_U, attrs.EDGE_TANGENT_CELL_V), + pairs=( + (attrs.EDGE_TANGENT_X, attrs.EDGE_TANGENT_X), + (attrs.EDGE_TANGENT_Y, attrs.EDGE_TANGENT_Y), + ), + ) + self.register_provider(tangent_cell_wrapper) class SparseFieldProviderWrapper(factory.FieldProvider): def __init__( self, - field_provider: factory.ProgramFieldProvider, + field_provider: factory.FieldProvider, target_dims: Sequence[gtx.Dimension], fields: Sequence[str], pairs: Sequence[tuple[str, ...]], diff --git a/model/common/src/icon4py/model/common/grid/geometry_attributes.py b/model/common/src/icon4py/model/common/grid/geometry_attributes.py index ca5d1587b0..047e25267e 100644 --- a/model/common/src/icon4py/model/common/grid/geometry_attributes.py +++ b/model/common/src/icon4py/model/common/grid/geometry_attributes.py @@ -50,9 +50,6 @@ EDGE_LENGTH: Final[str] = "edge_length" DUAL_EDGE_LENGTH: Final[str] = "length_of_dual_edge" -EDGE_TANGENT_X: Final[str] = "x_component_of_edge_tangential_unit_vector" -EDGE_TANGENT_Y: Final[str] = "y_component_of_edge_tangential_unit_vector" -EDGE_TANGENT_Z: Final[str] = "z_component_of_edge_tangential_unit_vector" EDGE_TANGENT_VERTEX_U: Final[str] = "eastward_component_of_edge_tangent_on_vertex" EDGE_TANGENT_VERTEX_V: Final[str] = "northward_component_of_edge_tangent_on_vertex" EDGE_TANGENT_CELL_U: Final[str] = "eastward_component_of_edge_tangent_on_cell" @@ -63,6 +60,9 @@ EDGE_NORMAL_Z: Final[str] = "z_component_of_edge_normal_unit_vector" EDGE_NORMAL_U: Final[str] = "eastward_component_of_edge_normal" EDGE_NORMAL_V: Final[str] = "northward_component_of_edge_normal" +EDGE_TANGENT_X: Final[str] = "x_component_of_edge_tangential_unit_vector" +EDGE_TANGENT_Y: Final[str] = "y_component_of_edge_tangential_unit_vector" +EDGE_TANGENT_Z: Final[str] = "z_component_of_edge_tangential_unit_vector" EDGE_DUAL_U: Final[str] = "eastward_component_of_edge_tangent" EDGE_DUAL_V: Final[str] = "northward_component_of_edge_tangent" EDGE_NORMAL_VERTEX_U: Final[str] = "eastward_component_of_edge_normal_on_vertex" diff --git a/model/common/src/icon4py/model/common/grid/geometry_stencils.py b/model/common/src/icon4py/model/common/grid/geometry_stencils.py index 90addf46df..9ef7cbf681 100644 --- a/model/common/src/icon4py/model/common/grid/geometry_stencils.py +++ b/model/common/src/icon4py/model/common/grid/geometry_stencils.py @@ -17,6 +17,8 @@ from icon4py.model.common.math.helpers import ( arc_length_on_edges, cross_product_on_edges, + diff_on_edges_torus, + distance_on_edges_torus, geographical_to_cartesian_on_edges, geographical_to_cartesian_on_vertices, normalize_cartesian_vector_on_edges, @@ -56,6 +58,49 @@ def cartesian_coordinates_of_edge_tangent( return normalize_cartesian_vector_on_edges(x, y, z) +@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) +def cartesian_coordinates_of_edge_tangent_torus( + vertex_x: fa.VertexField[ta.wpfloat], + vertex_y: fa.VertexField[ta.wpfloat], + edge_orientation: fa.EdgeField[ta.wpfloat], + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, +) -> tuple[ + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], +]: + """ + Compute normalized cartesian vector tangential to an edge on a torus grid. + + That is: computes the distance between the two vertices adjacent to the edge: + t = d(v1, v2) + + Args: + vertex_x: x coordinates of vertices + vertex_y: y coordinates of vertices + edge_orientation: encoding of the edge orientation: (-1, +1) depending on whether the + edge is directed from first to second neighbor of vice versa. + Returns: + x: x coordinate of normalized tangent vector + y: y coordinate of normalized tangent vector + z: z coordinate of normalized tangent vector + """ + xdiff, ydiff = diff_on_edges_torus( + vertex_x(E2V[0]), + vertex_x(E2V[1]), + vertex_y(E2V[0]), + vertex_y(E2V[1]), + domain_length, + domain_height, + ) + x = edge_orientation * xdiff + y = edge_orientation * ydiff + z = 0.0 * x # TODO(msimberg): zeros + + return normalize_cartesian_vector_on_edges(x, y, z) + + @gtx.field_operator def cartesian_coordinates_of_edge_normal( edge_lat: fa.EdgeField[ta.wpfloat], @@ -93,6 +138,30 @@ def cartesian_coordinates_of_edge_normal( return normalize_cartesian_vector_on_edges(x, y, z) +@gtx.field_operator +def cartesian_coordinates_of_edge_normal_torus( + edge_tangent_x: fa.EdgeField[ta.wpfloat], + edge_tangent_y: fa.EdgeField[ta.wpfloat], +) -> tuple[ + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], +]: + """ + Compute the normal to the edge tangent vector on a torus grid. + + Args: + edge_tangent_x: x coordinate of the tangent + edge_tangent_y: y coordinate of the tangent + Returns: + edge_normal_x: x coordinate of the normal + edge_normal_y: y coordinate of the normal + edge_normal_z: y coordinate of the normal + """ + z = 0.0 * edge_tangent_x # TODO(msimberg): zeros + return normalize_cartesian_vector_on_edges(-edge_tangent_y, edge_tangent_x, z) + + @gtx.field_operator def cartesian_coordinates_edge_tangent_and_normal( vertex_lat: fa.VertexField[ta.wpfloat], @@ -123,6 +192,59 @@ def cartesian_coordinates_edge_tangent_and_normal( return tangent_x, tangent_y, tangent_z, normal_x, normal_y, normal_z +@gtx.field_operator +def cartesian_coordinates_edge_tangent_and_normal_torus( + vertex_x: fa.VertexField[ta.wpfloat], + vertex_y: fa.VertexField[ta.wpfloat], + edge_x: fa.EdgeField[ta.wpfloat], + edge_y: fa.EdgeField[ta.wpfloat], + edge_orientation: fa.EdgeField[ta.wpfloat], + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, +) -> tuple[ + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], + fa.EdgeField[ta.wpfloat], +]: + """Compute normalized cartesian vectors of edge tangent and edge normal.""" + tangent_x, tangent_y, tangent_z = cartesian_coordinates_of_edge_tangent_torus( + vertex_x, + vertex_y, + edge_orientation, + domain_length, + domain_height, + ) + tangent_u = tangent_x + tangent_v = tangent_y + + normal_x, normal_y, normal_z = cartesian_coordinates_of_edge_normal_torus( + tangent_x, + tangent_y, + ) + normal_u = normal_x + normal_v = normal_y + + return ( + tangent_x, + tangent_y, + tangent_z, + tangent_u, + tangent_v, + normal_x, + normal_y, + normal_z, + normal_u, + normal_v, + ) + + @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_cartesian_coordinates_of_edge_tangent_and_normal( vertex_lat: fa.VertexField[ta.wpfloat], @@ -150,6 +272,52 @@ def compute_cartesian_coordinates_of_edge_tangent_and_normal( ) +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_cartesian_coordinates_of_edge_tangent_and_normal_torus( + vertex_x: fa.VertexField[ta.wpfloat], + vertex_y: fa.VertexField[ta.wpfloat], + edge_x: fa.EdgeField[ta.wpfloat], + edge_y: fa.EdgeField[ta.wpfloat], + edge_orientation: fa.EdgeField[ta.wpfloat], + tangent_x: fa.EdgeField[ta.wpfloat], + tangent_y: fa.EdgeField[ta.wpfloat], + tangent_z: fa.EdgeField[ta.wpfloat], + tangent_u: fa.EdgeField[ta.wpfloat], + tangent_v: fa.EdgeField[ta.wpfloat], + normal_x: fa.EdgeField[ta.wpfloat], + normal_y: fa.EdgeField[ta.wpfloat], + normal_z: fa.EdgeField[ta.wpfloat], + normal_u: fa.EdgeField[ta.wpfloat], + normal_v: fa.EdgeField[ta.wpfloat], + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, +): + cartesian_coordinates_edge_tangent_and_normal_torus( + vertex_x, + vertex_y, + edge_x, + edge_y, + edge_orientation, + domain_length, + domain_height, + out=( + tangent_x, + tangent_y, + tangent_z, + tangent_u, + tangent_v, + normal_x, + normal_y, + normal_z, + normal_u, + normal_v, + ), + domain={dims.EdgeDim: (horizontal_start, horizontal_end)}, + ) + + @gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) def zonal_and_meridional_component_of_edge_field_at_vertex( vertex_lat: fa.VertexField[ta.wpfloat], @@ -410,6 +578,38 @@ def arc_distance_of_far_edges_in_diamond( return far_vertex_vertex_length +@gtx.field_operator +def distance_of_far_edges_in_diamond_torus( + vertex_x: fa.VertexField[ta.wpfloat], + vertex_y: fa.VertexField[ta.wpfloat], + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, +) -> fa.EdgeField[ta.wpfloat]: + """ + Compute the distance between the "far" vertices of an edge. + + See arc_distance_of_far_edges_in_diamond for details. + + Args: + vertex_x: x coordinate of vertices + vertex_y: y coordinate of vertices + domain_length: length of the domain + domain_height: height of the domain + + Returns: + distance between the "far" vertices in the diamond. + + """ + return distance_on_edges_torus( + vertex_x(E2C2V[2]), + vertex_x(E2C2V[3]), + vertex_y(E2C2V[2]), + vertex_y(E2C2V[3]), + domain_length, + domain_height, + ) + + @gtx.field_operator def edge_length( vertex_lat: fa.VertexField[ta.wpfloat], @@ -502,6 +702,26 @@ def compute_arc_distance_of_far_edges_in_diamond( ) +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_distance_of_far_edges_in_diamond_torus( + vertex_x: fa.VertexField[ta.wpfloat], + vertex_y: fa.VertexField[ta.wpfloat], + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, + far_vertex_distance: fa.EdgeField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, +): + distance_of_far_edges_in_diamond_torus( + vertex_x, + vertex_y, + domain_length, + domain_height, + out=far_vertex_distance, + domain={dims.EdgeDim: (horizontal_start, horizontal_end)}, + ) + + @gtx.field_operator def edge_area( owner_mask: fa.EdgeField[bool], diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py index cdb39bf3b9..2e01df220f 100644 --- a/model/common/src/icon4py/model/common/grid/grid_manager.py +++ b/model/common/src/icon4py/model/common/grid/grid_manager.py @@ -63,7 +63,10 @@ def __call__(self, array: data_alloc.NDArray): ) -CoordinateDict: TypeAlias = dict[gtx.Dimension, dict[Literal["lat", "lon"], gtx.Field]] +# TODO(msimberg): x, y, z added temporarily for torus grids, but should they be in a separate dict? +CoordinateDict: TypeAlias = dict[ + gtx.Dimension, dict[Literal["lat", "lon", "x", "y", "z"], gtx.Field] +] # TODO (halungge): use a TypeDict for that GeometryDict: TypeAlias = dict[gridfile.GeometryName, gtx.Field] @@ -120,25 +123,37 @@ def __exit__(self, exc_type, exc_val, exc_tb): def __call__(self, allocator: gtx_typing.FieldBufferAllocationUtil, keep_skip_values: bool): if not self._reader: self.open() - self._geometry = self._read_geometry_fields(allocator) - self._grid = self._construct_grid(allocator=allocator, with_skip_values=keep_skip_values) - self._coordinates = self._read_coordinates(allocator) + + if geometry_type := self._reader.try_attribute(gridfile.MPIMPropertyName.GEOMETRY): + geometry_type = base.GeometryType(geometry_type) + + self._geometry = self._read_geometry_fields(allocator, geometry_type) + self._grid = self._construct_grid( + allocator=allocator, + with_skip_values=keep_skip_values, + geometry_type=geometry_type, + ) + self._coordinates = self._read_coordinates(allocator, geometry_type) self.close() - def _read_coordinates(self, backend: gtx_typing.Backend | None) -> CoordinateDict: - return { + def _read_coordinates( + self, + allocator: gtx_typing.FieldBufferAllocationUtil | None, + geometry_type: base.GeometryType | None, + ) -> CoordinateDict: + coordinates = { dims.CellDim: { "lat": gtx.as_field( (dims.CellDim,), self._reader.variable(gridfile.CoordinateName.CELL_LATITUDE), dtype=ta.wpfloat, - allocator=backend, + allocator=allocator, ), "lon": gtx.as_field( (dims.CellDim,), self._reader.variable(gridfile.CoordinateName.CELL_LONGITUDE), dtype=ta.wpfloat, - allocator=backend, + allocator=allocator, ), }, dims.EdgeDim: { @@ -146,32 +161,94 @@ def _read_coordinates(self, backend: gtx_typing.Backend | None) -> CoordinateDic (dims.EdgeDim,), self._reader.variable(gridfile.CoordinateName.EDGE_LATITUDE), dtype=ta.wpfloat, - allocator=backend, + allocator=allocator, ), "lon": gtx.as_field( (dims.EdgeDim,), self._reader.variable(gridfile.CoordinateName.EDGE_LONGITUDE), dtype=ta.wpfloat, - allocator=backend, + allocator=allocator, ), }, dims.VertexDim: { "lat": gtx.as_field( (dims.VertexDim,), self._reader.variable(gridfile.CoordinateName.VERTEX_LATITUDE), - allocator=backend, + allocator=allocator, dtype=ta.wpfloat, ), "lon": gtx.as_field( (dims.VertexDim,), self._reader.variable(gridfile.CoordinateName.VERTEX_LONGITUDE), - allocator=backend, + allocator=allocator, dtype=ta.wpfloat, ), }, } - def _read_geometry_fields(self, allocator: gtx_typing.FieldBufferAllocationUtil): + if geometry_type == base.GeometryType.TORUS: + coordinates[dims.CellDim]["x"] = gtx.as_field( + (dims.CellDim,), + self._reader.variable(gridfile.CoordinateName.CELL_X), + dtype=ta.wpfloat, + allocator=allocator, + ) + coordinates[dims.CellDim]["y"] = gtx.as_field( + (dims.CellDim,), + self._reader.variable(gridfile.CoordinateName.CELL_Y), + dtype=ta.wpfloat, + allocator=allocator, + ) + coordinates[dims.CellDim]["z"] = gtx.as_field( + (dims.CellDim,), + self._reader.variable(gridfile.CoordinateName.CELL_Z), + dtype=ta.wpfloat, + allocator=allocator, + ) + coordinates[dims.EdgeDim]["x"] = gtx.as_field( + (dims.EdgeDim,), + self._reader.variable(gridfile.CoordinateName.EDGE_X), + dtype=ta.wpfloat, + allocator=allocator, + ) + coordinates[dims.EdgeDim]["y"] = gtx.as_field( + (dims.EdgeDim,), + self._reader.variable(gridfile.CoordinateName.EDGE_Y), + dtype=ta.wpfloat, + allocator=allocator, + ) + coordinates[dims.EdgeDim]["z"] = gtx.as_field( + (dims.EdgeDim,), + self._reader.variable(gridfile.CoordinateName.EDGE_Z), + dtype=ta.wpfloat, + allocator=allocator, + ) + coordinates[dims.VertexDim]["x"] = gtx.as_field( + (dims.VertexDim,), + self._reader.variable(gridfile.CoordinateName.VERTEX_X), + dtype=ta.wpfloat, + allocator=allocator, + ) + coordinates[dims.VertexDim]["y"] = gtx.as_field( + (dims.VertexDim,), + self._reader.variable(gridfile.CoordinateName.VERTEX_Y), + dtype=ta.wpfloat, + allocator=allocator, + ) + coordinates[dims.VertexDim]["z"] = gtx.as_field( + (dims.VertexDim,), + self._reader.variable(gridfile.CoordinateName.VERTEX_Z), + dtype=ta.wpfloat, + allocator=allocator, + ) + + return coordinates + + def _read_geometry_fields( + self, + allocator: gtx_typing.FieldBufferAllocationUtil, + geometry_type: base.GeometryType | None, + ) -> GeometryDict: return { # TODO(halungge): still needs to ported, values from "our" grid files contains (wrong) values: # based on bug in generator fixed with this [PR40](https://gitlab.dkrz.de/dwd-sw/dwd_icon_tools/-/merge_requests/40) . @@ -242,7 +319,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 +351,10 @@ def coordinates(self) -> CoordinateDict: return self._coordinates def _construct_grid( - self, allocator: gtx_typing.FieldBufferAllocationUtil, with_skip_values: bool + self, + allocator: gtx_typing.FieldBufferAllocationUtil | None, + with_skip_values: bool, + geometry_type: base.GeometryType | None, ) -> icon.IconGrid: """Construct the grid topology from the icon grid file. @@ -359,7 +439,9 @@ def _construct_grid( } neighbor_tables.update(_get_derived_connectivities(neighbor_tables, array_ns=xp)) domain_bounds_constructor = functools.partial( - refinement.compute_domain_bounds, refinement_fields=refinement_fields, array_ns=xp + refinement.compute_domain_bounds, + refinement_fields=refinement_fields, + array_ns=xp, ) start_index, end_index = icon.get_start_and_end_index(domain_bounds_constructor) @@ -382,7 +464,8 @@ def _get_index_field(self, field: gridfile.GridFileName, transpose=True, apply_o def _get_derived_connectivities( - neighbor_tables: dict[gtx.FieldOffset, data_alloc.NDArray], array_ns: ModuleType = np + neighbor_tables: dict[gtx.FieldOffset, data_alloc.NDArray], + array_ns: ModuleType = np, ) -> dict[gtx.FieldOffset, data_alloc.NDArray]: e2v_table = neighbor_tables[dims.E2V] c2v_table = neighbor_tables[dims.C2V] diff --git a/model/common/src/icon4py/model/common/grid/gridfile.py b/model/common/src/icon4py/model/common/grid/gridfile.py index 0eaeee8c2b..5873089d67 100644 --- a/model/common/src/icon4py/model/common/grid/gridfile.py +++ b/model/common/src/icon4py/model/common/grid/gridfile.py @@ -187,6 +187,12 @@ class GeometryName(FieldName): # TODO(halungge): compute from coordinates EDGE_CELL_DISTANCE = "edge_cell_distance" EDGE_VERTEX_DISTANCE = "edge_vert_distance" + EDGE_NORMAL_X = "edge_primal_normal_cartesian_x" + EDGE_NORMAL_Y = "edge_primal_normal_cartesian_y" + EDGE_NORMAL_Z = "edge_primal_normal_cartesian_z" + EDGE_TANGENT_X = "edge_dual_normal_cartesian_x" + EDGE_TANGENT_Y = "edge_dual_normal_cartesian_y" + EDGE_TANGENT_Z = "edge_dual_normal_cartesian_z" class CoordinateName(FieldName): @@ -202,6 +208,16 @@ class CoordinateName(FieldName): VERTEX_LONGITUDE = "vlon" VERTEX_LATITUDE = "vlat" + CELL_X = "cell_circumcenter_cartesian_x" + CELL_Y = "cell_circumcenter_cartesian_y" + CELL_Z = "cell_circumcenter_cartesian_z" + EDGE_X = "edge_middle_cartesian_x" + EDGE_Y = "edge_middle_cartesian_y" + EDGE_Z = "edge_middle_cartesian_z" + VERTEX_X = "cartesian_x_vertices" + VERTEX_Y = "cartesian_y_vertices" + VERTEX_Z = "cartesian_z_vertices" + class GridRefinementName(FieldName): """Names of arrays in grid file defining the grid control, definition of boundaries layers, start and end indices of horizontal zones.""" diff --git a/model/common/src/icon4py/model/common/grid/icon.py b/model/common/src/icon4py/model/common/grid/icon.py index 8382b8ea72..40ac367f4b 100644 --- a/model/common/src/icon4py/model/common/grid/icon.py +++ b/model/common/src/icon4py/model/common/grid/icon.py @@ -166,8 +166,8 @@ def __post_init__(self) -> None: object.__setattr__(self, "characteristic_length", math.sqrt(self.mean_cell_area)) @property - def geometry_type(self) -> base.GeometryType | None: - return self.grid_shape.geometry_type if self.grid_shape else None + def geometry_type(self) -> base.GeometryType: + return self.grid_shape.geometry_type @property def subdivision(self) -> GridSubdivision | None: diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py index cb154ea4ed..480cb9d048 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py @@ -15,6 +15,7 @@ from icon4py.model.common import constants, dimension as dims from icon4py.model.common.decomposition import definitions from icon4py.model.common.grid import ( + base, geometry, geometry_attributes as geometry_attrs, grid_refinement as refinement, @@ -54,7 +55,9 @@ def __init__( self._attrs = metadata self._providers: dict[str, factory.FieldProvider] = {} self._geometry = geometry_source + geometry_type = self._grid.global_properties.geometry_type characteristic_length = self._grid.global_properties.characteristic_length + mean_dual_edge_length = self._grid.global_properties.mean_dual_edge_length # TODO @halungge: Dummy config dict - to be replaced by real configuration self._config = { "divavg_cntrwgt": 0.5, @@ -66,13 +69,13 @@ def __init__( "rbf_kernel_edge": rbf.DEFAULT_RBF_KERNEL[rbf.RBFDimension.EDGE], "rbf_kernel_vertex": rbf.DEFAULT_RBF_KERNEL[rbf.RBFDimension.VERTEX], "rbf_scale_cell": rbf.compute_default_rbf_scale( - characteristic_length, rbf.RBFDimension.CELL + geometry_type, characteristic_length, mean_dual_edge_length, rbf.RBFDimension.CELL ), "rbf_scale_edge": rbf.compute_default_rbf_scale( - characteristic_length, rbf.RBFDimension.EDGE + geometry_type, characteristic_length, mean_dual_edge_length, rbf.RBFDimension.EDGE ), "rbf_scale_vertex": rbf.compute_default_rbf_scale( - characteristic_length, rbf.RBFDimension.VERTEX + geometry_type, characteristic_length, mean_dual_edge_length, rbf.RBFDimension.VERTEX ), } log.info( @@ -182,31 +185,59 @@ def _register_computed_fields(self) -> None: self.register_provider(geofac_grdiv) - cell_average_weight = factory.NumpyDataProvider( - func=functools.partial( - interpolation_fields.compute_mass_conserving_bilinear_cell_average_weight, - array_ns=self._xp, - ), - fields=(attrs.C_BLN_AVG,), - domain=(dims.CellDim, dims.C2E2CODim), - deps={ - "lat": geometry_attrs.CELL_LAT, - "lon": geometry_attrs.CELL_LON, - "cell_areas": geometry_attrs.CELL_AREA, - "cell_owner_mask": "cell_owner_mask", - }, - connectivities={"c2e2c0": dims.C2E2CODim}, - params={ - "horizontal_start": self.grid.start_index( - cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) - ), - "horizontal_start_level_3": self.grid.start_index( - cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) - ), - "divavg_cntrwgt": self._config["divavg_cntrwgt"], - }, - ) - self.register_provider(cell_average_weight) + match self.grid.global_properties.geometry_type: + case base.GeometryType.ICOSAHEDRON: + cell_average_weight = factory.NumpyDataProvider( + func=functools.partial( + interpolation_fields.compute_mass_conserving_bilinear_cell_average_weight, + array_ns=self._xp, + ), + fields=(attrs.C_BLN_AVG,), + domain=(dims.CellDim, dims.C2E2CODim), + deps={ + "lat": geometry_attrs.CELL_LAT, + "lon": geometry_attrs.CELL_LON, + "cell_areas": geometry_attrs.CELL_AREA, + "cell_owner_mask": "cell_owner_mask", + }, + connectivities={"c2e2c0": dims.C2E2CODim}, + params={ + "horizontal_start": self.grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ), + "horizontal_start_level_3": self.grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) + ), + "divavg_cntrwgt": self._config["divavg_cntrwgt"], + }, + ) + self.register_provider(cell_average_weight) + case base.GeometryType.TORUS: + cell_average_weight = factory.NumpyDataProvider( + func=functools.partial( + interpolation_fields.compute_mass_conserving_bilinear_cell_average_weight_torus, + array_ns=self._xp, + ), + fields=(attrs.C_BLN_AVG,), + domain=(dims.CellDim, dims.C2E2CODim), + deps={ + "cell_areas": geometry_attrs.CELL_AREA, + "cell_owner_mask": "cell_owner_mask", + }, + connectivities={"c2e2c0": dims.C2E2CODim}, + params={ + "horizontal_start": self.grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ), + "horizontal_start_level_3": self.grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) + ), + "divavg_cntrwgt": self._config["divavg_cntrwgt"], + }, + ) + self.register_provider(cell_average_weight) + case _: + raise ValueError("TODO") c_lin_e = factory.NumpyDataProvider( func=functools.partial(interpolation_fields.compute_c_lin_e, array_ns=self._xp), @@ -274,47 +305,83 @@ def _register_computed_fields(self) -> None: ) self.register_provider(e_flx_avg) - e_bln_c_s = factory.NumpyDataProvider( - func=functools.partial(interpolation_fields.compute_e_bln_c_s, array_ns=self._xp), - fields=(attrs.E_BLN_C_S,), - domain=(dims.CellDim, dims.C2EDim), - deps={ - "cells_lat": geometry_attrs.CELL_LAT, - "cells_lon": geometry_attrs.CELL_LON, - "edges_lat": geometry_attrs.EDGE_LAT, - "edges_lon": geometry_attrs.EDGE_LON, - }, - connectivities={"c2e": dims.C2EDim}, - params={"weighting_factor": self._config["weighting_factor"]}, - ) - self.register_provider(e_bln_c_s) + match self.grid.global_properties.geometry_type: + case base.GeometryType.ICOSAHEDRON: + e_bln_c_s = factory.NumpyDataProvider( + func=functools.partial( + interpolation_fields.compute_e_bln_c_s, array_ns=self._xp + ), + fields=(attrs.E_BLN_C_S,), + domain=(dims.CellDim, dims.C2EDim), + deps={ + "cells_lat": geometry_attrs.CELL_LAT, + "cells_lon": geometry_attrs.CELL_LON, + "edges_lat": geometry_attrs.EDGE_LAT, + "edges_lon": geometry_attrs.EDGE_LON, + }, + connectivities={"c2e": dims.C2EDim}, + params={"weighting_factor": self._config["weighting_factor"]}, + ) + self.register_provider(e_bln_c_s) + case base.GeometryType.TORUS: + e_bln_c_s = factory.NumpyDataProvider( + func=functools.partial( + interpolation_fields.compute_e_bln_c_s_torus, array_ns=self._xp + ), + fields=(attrs.E_BLN_C_S,), + domain=(dims.CellDim, dims.C2EDim), + deps={}, + connectivities={"c2e": dims.C2EDim}, + params={}, + ) + self.register_provider(e_bln_c_s) + case _: + raise ValueError("TODO") - pos_on_tplane_e_x_y = factory.NumpyDataProvider( - func=functools.partial( - interpolation_fields.compute_pos_on_tplane_e_x_y, array_ns=self._xp - ), - fields=(attrs.POS_ON_TPLANE_E_X, attrs.POS_ON_TPLANE_E_Y), - domain=(dims.EdgeDim, dims.E2CDim), - deps={ - "primal_normal_v1": geometry_attrs.EDGE_NORMAL_U, - "primal_normal_v2": geometry_attrs.EDGE_NORMAL_V, - "dual_normal_v1": geometry_attrs.EDGE_DUAL_U, - "dual_normal_v2": geometry_attrs.EDGE_DUAL_V, - "cells_lon": geometry_attrs.CELL_LON, - "cells_lat": geometry_attrs.CELL_LAT, - "edges_lon": geometry_attrs.EDGE_LON, - "edges_lat": geometry_attrs.EDGE_LAT, - "owner_mask": "edge_owner_mask", - }, - connectivities={"e2c": dims.E2CDim}, - params={ - "grid_sphere_radius": constants.EARTH_RADIUS, - "horizontal_start": self.grid.start_index( - edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) - ), - }, - ) - self.register_provider(pos_on_tplane_e_x_y) + match self.grid.global_properties.geometry_type: + case base.GeometryType.ICOSAHEDRON: + pos_on_tplane_e_x_y = factory.NumpyDataProvider( + func=functools.partial( + interpolation_fields.compute_pos_on_tplane_e_x_y, array_ns=self._xp + ), + fields=(attrs.POS_ON_TPLANE_E_X, attrs.POS_ON_TPLANE_E_Y), + domain=(dims.EdgeDim, dims.E2CDim), + deps={ + "primal_normal_v1": geometry_attrs.EDGE_NORMAL_U, + "primal_normal_v2": geometry_attrs.EDGE_NORMAL_V, + "dual_normal_v1": geometry_attrs.EDGE_DUAL_U, + "dual_normal_v2": geometry_attrs.EDGE_DUAL_V, + "cells_lon": geometry_attrs.CELL_LON, + "cells_lat": geometry_attrs.CELL_LAT, + "edges_lon": geometry_attrs.EDGE_LON, + "edges_lat": geometry_attrs.EDGE_LAT, + "owner_mask": "edge_owner_mask", + }, + connectivities={"e2c": dims.E2CDim}, + params={ + "grid_sphere_radius": constants.EARTH_RADIUS, + "horizontal_start": self.grid.start_index( + edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ), + }, + ) + self.register_provider(pos_on_tplane_e_x_y) + case base.GeometryType.TORUS: + pos_on_tplane_e_x_y = factory.NumpyDataProvider( + func=functools.partial( + interpolation_fields.compute_pos_on_tplane_e_x_y_torus, array_ns=self._xp + ), + fields=(attrs.POS_ON_TPLANE_E_X, attrs.POS_ON_TPLANE_E_Y), + domain=(dims.EdgeDim, dims.E2CDim), + deps={ + "dual_edge_length": geometry_attrs.DUAL_EDGE_LENGTH, + }, + connectivities={"e2c": dims.E2CDim}, + params={}, + ) + self.register_provider(pos_on_tplane_e_x_y) + case _: + raise ValueError("TODO") cells_aw_verts = factory.NumpyDataProvider( func=functools.partial(interpolation_fields.compute_cells_aw_verts, array_ns=self._xp), @@ -359,10 +426,17 @@ def _register_computed_fields(self) -> None: connectivities={"rbf_offset": dims.C2E2C2EDim}, params={ "rbf_kernel": self._config["rbf_kernel_cell"].value, + "geometry_type": self._grid.global_properties.geometry_type.value, "scale_factor": self._config["rbf_scale_cell"], "horizontal_start": self.grid.start_index( cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) ), + "domain_length": self._grid.global_properties.domain_length + if self._grid.global_properties.domain_length + else -1.0, + "domain_height": self._grid.global_properties.domain_height + if self._grid.global_properties.domain_height + else -1.0, }, ) self.register_provider(rbf_vec_coeff_c) @@ -386,10 +460,17 @@ def _register_computed_fields(self) -> None: connectivities={"rbf_offset": dims.E2C2EDim}, params={ "rbf_kernel": self._config["rbf_kernel_edge"].value, + "geometry_type": self._grid.global_properties.geometry_type.value, "scale_factor": self._config["rbf_scale_edge"], "horizontal_start": self.grid.start_index( edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) ), + "domain_length": self._grid.global_properties.domain_length + if self._grid.global_properties.domain_length + else -1.0, + "domain_height": self._grid.global_properties.domain_height + if self._grid.global_properties.domain_height + else -1.0, }, ) self.register_provider(rbf_vec_coeff_e) @@ -414,10 +495,17 @@ def _register_computed_fields(self) -> None: connectivities={"rbf_offset": dims.V2EDim}, params={ "rbf_kernel": self._config["rbf_kernel_vertex"].value, + "geometry_type": self._grid.global_properties.geometry_type.value, "scale_factor": self._config["rbf_scale_vertex"], "horizontal_start": self.grid.start_index( vertex_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) ), + "domain_length": self._grid.global_properties.domain_length + if self._grid.global_properties.domain_length + else -1.0, + "domain_height": self._grid.global_properties.domain_height + if self._grid.global_properties.domain_height + else -1.0, }, ) self.register_provider(rbf_vec_coeff_v) diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py index c39ed10956..982e95fe09 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py @@ -622,6 +622,33 @@ def _enforce_mass_conservation( return c_bln_avg +def _compute_uniform_c_bln_avg( + c2e2c: data_alloc.NDArray, + divavg_cntrwgt: ta.wpfloat, + horizontal_start: gtx.int32, + array_ns: ModuleType = np, +) -> data_alloc.NDArray: + """ + Compute bilinear cell average weight for a torus grid. + + Args: + divavg_cntrwgt: + c2e2c: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, C2E2CDim], gtx.int32] + horizontal_start: + + Returns: + c_bln_avg: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], ta.wpfloat] + """ + local_weight = divavg_cntrwgt + neighbor_weight = (1.0 - divavg_cntrwgt) / 3.0 + c_bln_avg = array_ns.full( + (c2e2c.shape[0], c2e2c.shape[1] + 1), + [local_weight, neighbor_weight, neighbor_weight, neighbor_weight], + ) + + return c_bln_avg + + def compute_mass_conserving_bilinear_cell_average_weight( c2e2c0: data_alloc.NDArray, lat: data_alloc.NDArray, @@ -647,6 +674,31 @@ def compute_mass_conserving_bilinear_cell_average_weight( ) +def compute_mass_conserving_bilinear_cell_average_weight_torus( + c2e2c0: data_alloc.NDArray, + cell_areas: data_alloc.NDArray, + cell_owner_mask: data_alloc.NDArray, + divavg_cntrwgt: ta.wpfloat, + horizontal_start: gtx.int32, + horizontal_start_level_3: gtx.int32, + array_ns: ModuleType = np, +) -> data_alloc.NDArray: + c_bln_avg = _compute_uniform_c_bln_avg( + c2e2c0[:, 1:], divavg_cntrwgt, horizontal_start, array_ns + ) + # TODO(msimberg): Exact result for torus without the following. 1e-16 error + # with the the following. Is it needed? + return _force_mass_conservation_to_c_bln_avg( + c2e2c0, + c_bln_avg, + cell_areas, + cell_owner_mask, + divavg_cntrwgt, + horizontal_start_level_3, + array_ns, + ) + + def _create_inverse_neighbor_index( source_offset, inverse_offset, array_ns: ModuleType ) -> data_alloc.NDArray: @@ -956,6 +1008,22 @@ def compute_e_bln_c_s( return e_bln_c_s +def compute_e_bln_c_s_torus( + c2e: data_alloc.NDArray, + array_ns: ModuleType = np, +) -> data_alloc.NDArray: + """ + Compute e_bln_c_s. + + Args: + c2e: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], gtx.int32] + + Returns: + e_bln_c_s: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], ta.wpfloat] + """ + return array_ns.full_like(c2e, 1.0 / 3.0, dtype=ta.wpfloat) + + def compute_pos_on_tplane_e_x_y( grid_sphere_radius: ta.wpfloat, primal_normal_v1: data_alloc.NDArray, @@ -1047,3 +1115,44 @@ def compute_pos_on_tplane_e_x_y( ) return pos_on_tplane_e[:, :, 0], pos_on_tplane_e[:, :, 1] + + +def compute_pos_on_tplane_e_x_y_torus( + dual_edge_length: data_alloc.NDArray, + e2c: data_alloc.NDArray, + array_ns: ModuleType = np, +) -> data_alloc.NDArray: + """ + Compute pos_on_tplane_e_x_y. + get geographical coordinates of edge midpoint + get line and block indices of neighbour cells + get geographical coordinates of first cell center + projection first cell center into local \\lambda-\\Phi-system + get geographical coordinates of second cell center + projection second cell center into local \\lambda-\\Phi-system + + Args: + dual_edge_length: numpy_array, representing a gtx.Field[gtx.Dims[EdgeDim], ta.wpfloat] + e2c: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], gtx.int32] + + Returns: + pos_on_tplane_e_x: \\ numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], ta.wpfloat] + pos_on_tplane_e_y: // + """ + # The implementation makes the simplifying assumptions that: + # - The torus grid consists of equilateral triangles, which means that the + # neighboring cell centers must always be at 0.5 * dual_edge_length from + # the edge center (the edge lies symmetrically perpendicular to the primal + # edge). + # - The neighboring cell centers are exactly along the primal normal/dual + # tangent direction, which means the x component in the local coordinate + # system is always zero, and the y component is always 0.5 * + # dual_edge_length. + # - The first neighbor cell is in the opposite direction of the primal + # normal and the second neighbor is in the direction of the primal normal. + half_dual_edge_length = 0.5 * dual_edge_length[0] + num_edges = e2c.shape[0] + return ( + array_ns.full([num_edges, 2], [-half_dual_edge_length, half_dual_edge_length]), + array_ns.zeros([num_edges, 2]), + ) diff --git a/model/common/src/icon4py/model/common/interpolation/rbf_interpolation.py b/model/common/src/icon4py/model/common/interpolation/rbf_interpolation.py index 38b9cf0912..9a00aaadcc 100644 --- a/model/common/src/icon4py/model/common/interpolation/rbf_interpolation.py +++ b/model/common/src/icon4py/model/common/interpolation/rbf_interpolation.py @@ -47,26 +47,38 @@ class InterpolationKernel(enum.Enum): } -def compute_default_rbf_scale(mean_characteristic_length: ta.wpfloat, dim: RBFDimension): - """Compute the default RBF scale factor. This assumes that the Gaussian - kernel is used for vertices and cells, and that the inverse multiquadratic - kernel is used for edges.""" - - threshold = 2.5 if dim == RBFDimension.CELL else 2.0 - c1 = 0.4 if dim == RBFDimension.EDGE else 1.8 - if dim == RBFDimension.CELL: - c2 = 3.75 - c3 = 0.9 - elif dim == RBFDimension.VERTEX: - c2 = 3.0 - c3 = 0.96 - else: - c2 = 2.0 - c3 = 0.325 - - resol = mean_characteristic_length / 1000.0 - scale = 0.5 / (1.0 + c1 * math.log(threshold / resol) ** c2) if resol < threshold else 0.5 - return scale * (resol / 0.125) ** c3 if resol <= 0.125 else scale +def compute_default_rbf_scale( + geometry_type: base_grid.GeometryType, + mean_characteristic_length: ta.wpfloat, + mean_dual_edge_length: ta.wpfloat, + dim: RBFDimension, +): + """Compute the default RBF scale factor. This assumes that the default + interpolation kernels are used for each dimension""" + + match geometry_type: + case base_grid.GeometryType.ICOSAHEDRON: + threshold = 2.5 if dim == RBFDimension.CELL else 2.0 + c1 = 0.4 if dim == RBFDimension.EDGE else 1.8 + if dim == RBFDimension.CELL: + c2 = 3.75 + c3 = 0.9 + elif dim == RBFDimension.VERTEX: + c2 = 3.0 + c3 = 0.96 + else: + c2 = 2.0 + c3 = 0.325 + + resol = mean_characteristic_length / 1000.0 + scale = ( + 0.5 / (1.0 + c1 * math.log(threshold / resol) ** c2) if resol < threshold else 0.5 + ) + return scale * (resol / 0.125) ** c3 if resol <= 0.125 else scale + case base_grid.GeometryType.TORUS: + return mean_dual_edge_length + case _: + raise ValueError(f"Unsupported geometry type: {geometry_type}") def construct_rbf_matrix_offsets_tables_for_cells( @@ -105,7 +117,15 @@ def _dot_product( return array_ns.matmul(v1, v2_tilde) -def _arc_length_pairwise(v: data_alloc.NDArray, array_ns: ModuleType = np) -> data_alloc.NDArray: +# TODO(msimberg): Rename again. This is arc length on a unit sphere for +# icosahedral grids and real distance for torus grids. +def _distance_pairwise( + geometry_type: base_grid.GeometryType, + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, + v: data_alloc.NDArray, + array_ns: ModuleType = np, +) -> data_alloc.NDArray: """ Compute the pairwise arc lengths between points in each row of v. @@ -115,23 +135,41 @@ def _arc_length_pairwise(v: data_alloc.NDArray, array_ns: ModuleType = np) -> da dimension of the points. array_ns: numpy or cupy module to use for computations. """ - # For pairs of points p1 and p2 compute: - # arccos(dot(p1, p2) / (norm(p1) * norm(p2))) noqa: ERA001 - # Compute all pairs of dot products - arc_lengths = _dot_product(v, v, array_ns=array_ns) - # Use the dot product of the diagonals to get the norm of each point - norms = array_ns.sqrt(array_ns.diagonal(arc_lengths, axis1=1, axis2=2)) - # Divide the dot products by the broadcasted norms - array_ns.divide(arc_lengths, norms[:, :, array_ns.newaxis], out=arc_lengths) - array_ns.divide(arc_lengths, norms[:, array_ns.newaxis, :], out=arc_lengths) - # Ensure all points are within [-1.0, 1.0] (may be outside due to numerical - # inaccuracies) - array_ns.clip(arc_lengths, -1.0, 1.0, out=arc_lengths) - return array_ns.arccos(arc_lengths) - - -def _arc_length_vector_matrix( - v1: data_alloc.NDArray, v2: data_alloc.NDArray, array_ns: ModuleType = np + match geometry_type: + case base_grid.GeometryType.ICOSAHEDRON: + # For pairs of points p1 and p2 compute: + # arccos(dot(p1, p2) / (norm(p1) * norm(p2))) noqa: ERA001 + # Compute all pairs of dot products + arc_lengths = _dot_product(v, v, array_ns=array_ns) + # Use the dot product of the diagonals to get the norm of each point + norms = array_ns.sqrt(array_ns.diagonal(arc_lengths, axis1=1, axis2=2)) + # Divide the dot products by the broadcasted norms + array_ns.divide(arc_lengths, norms[:, :, array_ns.newaxis], out=arc_lengths) + array_ns.divide(arc_lengths, norms[:, array_ns.newaxis, :], out=arc_lengths) + # Ensure all points are within [-1.0, 1.0] (may be outside due to numerical + # inaccuracies) + array_ns.clip(arc_lengths, -1.0, 1.0, out=arc_lengths) + return array_ns.arccos(arc_lengths) + case base_grid.GeometryType.TORUS: + # For pairs of points p1 and p2 compute: + # norm(p1 - p2), taking into account the periodic boundaries noqa: ERA001 + diff = array_ns.abs(v[:, :, array_ns.newaxis, :] - v[:, array_ns.newaxis, :, :]) + domain_size = array_ns.asarray([domain_length, domain_height, ta.wpfloat(0.0)]) + domain_size_expanded = domain_size[array_ns.newaxis, array_ns.newaxis, :] + inverted_diff = array_ns.subtract(domain_size_expanded, diff) + array_ns.minimum(diff, inverted_diff, out=diff) + return array_ns.linalg.norm(diff, axis=-1) + case _: + raise ValueError(f"Unsupported geometry type: {geometry_type}") + + +def _distance_vector_matrix( + geometry_type: base_grid.GeometryType, + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, + v1: data_alloc.NDArray, + v2: data_alloc.NDArray, + array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ Compute the arc lengths between each point in v1 and the points in v2 at the same row. @@ -144,19 +182,32 @@ def _arc_length_vector_matrix( of the points. array_ns: numpy or cupy module to use for computations. """ - # For pairs of points p1 and p2 compute: - # arccos(dot(p1, p2) / (norm(p1) * norm(p2))) noqa: ERA001 - # Compute all pairs of dot products - arc_lengths = _dot_product(v1, v2, array_ns=array_ns) - v1_norm = array_ns.linalg.norm(v1, axis=-1) - v2_norm = array_ns.linalg.norm(v2, axis=-1) - # Divide the dot products by the broadcasted norms - array_ns.divide(arc_lengths, v1_norm[:, :, array_ns.newaxis], out=arc_lengths) - array_ns.divide(arc_lengths, v2_norm[:, array_ns.newaxis, :], out=arc_lengths) - # Ensure all points are within [-1.0, 1.0] (may be outside due to numerical - # inaccuracies) - array_ns.clip(arc_lengths, -1.0, 1.0, out=arc_lengths) - return array_ns.squeeze(array_ns.arccos(arc_lengths), axis=1) + match geometry_type: + case base_grid.GeometryType.ICOSAHEDRON: + # For pairs of points p1 and p2 compute: + # arccos(dot(p1, p2) / (norm(p1) * norm(p2))) noqa: ERA001 + # Compute all pairs of dot products + arc_lengths = _dot_product(v1, v2, array_ns=array_ns) + v1_norm = array_ns.linalg.norm(v1, axis=-1) + v2_norm = array_ns.linalg.norm(v2, axis=-1) + # Divide the dot products by the broadcasted norms + array_ns.divide(arc_lengths, v1_norm[:, :, array_ns.newaxis], out=arc_lengths) + array_ns.divide(arc_lengths, v2_norm[:, array_ns.newaxis, :], out=arc_lengths) + # Ensure all points are within [-1.0, 1.0] (may be outside due to numerical + # inaccuracies) + array_ns.clip(arc_lengths, -1.0, 1.0, out=arc_lengths) + return array_ns.squeeze(array_ns.arccos(arc_lengths), axis=1) + case base_grid.GeometryType.TORUS: + # For pairs of points p1 and p2 compute: + # norm(p1 - p2) noqa: ERA001 + diff = np.abs(v1 - v2) + domain_size = array_ns.asarray([domain_length, domain_height, ta.wpfloat(0.0)]) + domain_size_expanded = domain_size[array_ns.newaxis, array_ns.newaxis, :] + inverted_diff = array_ns.subtract(domain_size_expanded, diff) + diff = array_ns.minimum(diff, inverted_diff, out=diff) + return array_ns.linalg.norm(diff, axis=-1) + case _: + raise ValueError(f"Unsupported geometry type: {geometry_type}") def _gaussian( @@ -191,22 +242,29 @@ def _kernel( def _cartesian_coordinates_from_zonal_and_meridional_components( + geometry_type: base_grid.GeometryType, lat: data_alloc.NDArray, lon: data_alloc.NDArray, u: data_alloc.NDArray, v: data_alloc.NDArray, array_ns: ModuleType = np, ) -> tuple[data_alloc.NDArray, data_alloc.NDArray, data_alloc.NDArray]: - cos_lat = array_ns.cos(lat) - sin_lat = array_ns.sin(lat) - cos_lon = array_ns.cos(lon) - sin_lon = array_ns.sin(lon) - - x = -u * sin_lon - v * sin_lat * cos_lon - y = u * cos_lon - v * sin_lat * sin_lon - z = cos_lat * v - - return x, y, z + match geometry_type: + case base_grid.GeometryType.ICOSAHEDRON: + cos_lat = array_ns.cos(lat) + sin_lat = array_ns.sin(lat) + cos_lon = array_ns.cos(lon) + sin_lon = array_ns.sin(lon) + + x = -u * sin_lon - v * sin_lat * cos_lon + y = u * cos_lon - v * sin_lat * sin_lon + z = cos_lat * v + + return x, y, z + case base_grid.GeometryType.TORUS: + return u, v, array_ns.zeros_like(u) + case _: + raise ValueError(f"Unsupported geometry type: {geometry_type}") def _compute_rbf_interpolation_coeffs( @@ -224,8 +282,11 @@ def _compute_rbf_interpolation_coeffs( uv: tuple[tuple[data_alloc.NDArray, data_alloc.NDArray], ...], rbf_offset: data_alloc.NDArray, rbf_kernel: InterpolationKernel, + geometry_type: base_grid.GeometryType, scale_factor: ta.wpfloat, horizontal_start: gtx.int32, + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, array_ns: ModuleType = np, ) -> tuple[data_alloc.NDArray, ...]: rbf_offset_shape_full = rbf_offset.shape @@ -271,8 +332,13 @@ def index_offset(f): axis=-1, ) assert element_center.shape == (rbf_offset.shape[0], 3) - vector_dist = _arc_length_vector_matrix( - element_center[:, array_ns.newaxis, :], edge_center, array_ns=array_ns + vector_dist = _distance_vector_matrix( + geometry_type, + domain_length, + domain_height, + element_center[:, array_ns.newaxis, :], + edge_center, + array_ns=array_ns, ) assert vector_dist.shape == rbf_offset.shape rbf_val = _kernel(rbf_kernel, vector_dist, scale_factor, array_ns=array_ns) @@ -287,6 +353,7 @@ def index_offset(f): assert 1 <= num_zonal_meridional_components <= 2 for i in range(num_zonal_meridional_components): z_nx_x, z_nx_y, z_nx_z = _cartesian_coordinates_from_zonal_and_meridional_components( + geometry_type, element_center_lat[horizontal_start:], element_center_lon[horizontal_start:], uv[i][0][horizontal_start:], @@ -311,7 +378,9 @@ def index_offset(f): ) # Distance between edge midpoints for RBF interpolation matrix - z_dist = _arc_length_pairwise(edge_center, array_ns=array_ns) + z_dist = _distance_pairwise( + geometry_type, domain_length, domain_height, edge_center, array_ns=array_ns + ) assert z_dist.shape == ( rbf_offset.shape[0], rbf_offset.shape[1], @@ -333,7 +402,7 @@ def index_offset(f): # https://github.com/cupy/cupy/pull/9116. rbf_vec_coeff_np = [ np.zeros(rbf_offset_shape_full, dtype=ta.wpfloat) - for j in range(num_zonal_meridional_components) + for _ in range(num_zonal_meridional_components) ] rbf_offset_np = data_alloc.as_numpy(rbf_offset) z_rbfmat_np = data_alloc.as_numpy(z_rbfmat) @@ -372,8 +441,11 @@ def compute_rbf_interpolation_coeffs_cell( rbf_offset: data_alloc.NDArray, # TODO(): Can't pass enum as "params" in NumpyFieldsProvider? rbf_kernel: int, + geometry_type: int, scale_factor: ta.wpfloat, horizontal_start: gtx.int32, + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, array_ns: ModuleType = np, ) -> tuple[data_alloc.NDArray]: zeros = array_ns.zeros(rbf_offset.shape[0], dtype=ta.wpfloat) @@ -394,8 +466,11 @@ def compute_rbf_interpolation_coeffs_cell( ((ones, zeros), (zeros, ones)), rbf_offset, InterpolationKernel(rbf_kernel), + base_grid.GeometryType(geometry_type), scale_factor, horizontal_start, + domain_length, + domain_height, array_ns=array_ns, ) @@ -413,13 +488,16 @@ def compute_rbf_interpolation_coeffs_edge( edge_dual_normal_v: data_alloc.NDArray, rbf_offset: data_alloc.NDArray, rbf_kernel: int, + geometry_type: int, scale_factor: ta.wpfloat, horizontal_start: gtx.int32, + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, array_ns: ModuleType = np, ) -> data_alloc.NDArray: return _compute_rbf_interpolation_coeffs( - edge_lat, - edge_lon, + edge_lat, # these should be ignored for torus + edge_lon, # these should be ignored for torus edge_center_x, edge_center_y, edge_center_z, @@ -432,8 +510,11 @@ def compute_rbf_interpolation_coeffs_edge( ((edge_dual_normal_u, edge_dual_normal_v),), rbf_offset, InterpolationKernel(rbf_kernel), + base_grid.GeometryType(geometry_type), scale_factor, horizontal_start, + domain_length, + domain_height, array_ns=array_ns, )[0] @@ -452,8 +533,11 @@ def compute_rbf_interpolation_coeffs_vertex( edge_normal_z: data_alloc.NDArray, rbf_offset: data_alloc.NDArray, rbf_kernel: int, + geometry_type: int, scale_factor: ta.wpfloat, horizontal_start: gtx.int32, + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, array_ns: ModuleType = np, ) -> tuple[data_alloc.NDArray, data_alloc.NDArray]: zeros = array_ns.zeros(rbf_offset.shape[0], dtype=ta.wpfloat) @@ -474,7 +558,10 @@ def compute_rbf_interpolation_coeffs_vertex( ((ones, zeros), (zeros, ones)), rbf_offset, InterpolationKernel(rbf_kernel), + base_grid.GeometryType(geometry_type), scale_factor, horizontal_start, + domain_length, + domain_height, array_ns=array_ns, ) diff --git a/model/common/src/icon4py/model/common/math/helpers.py b/model/common/src/icon4py/model/common/math/helpers.py index 8d47cf0a16..a46f9209df 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, +): + """ + 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, +): + """ + Compute the distance between two points on the torus. + + Inputs are cartesian coordinates of the points. Z is assumed zero and is + ignored. Distances are computed modulo the domain length and height. + + Args: + x0: x coordinate of point_0 + x1: x coordinate of point_1 + y0: y coordinate of point_0 + y1: y coordinate of point_1 + domain_length: length of the domain + domain_height: height of the domain + + Returns: + distance + + """ + xdiff, ydiff = diff_on_edges_torus(x0, x1, y0, y1, domain_length, domain_height) + return sqrt(xdiff**2 + ydiff**2) + + @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def average_two_vertical_levels_downwards_on_edges( input_field: fa.EdgeKField[wpfloat], diff --git a/model/common/tests/common/grid/unit_tests/test_geometry.py b/model/common/tests/common/grid/unit_tests/test_geometry.py index 003c3e0cc1..c6893701d6 100644 --- a/model/common/tests/common/grid/unit_tests/test_geometry.py +++ b/model/common/tests/common/grid/unit_tests/test_geometry.py @@ -13,8 +13,9 @@ import numpy as np import pytest -from icon4py.model.common import dimension as dims +from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import ( + base, geometry, geometry_attributes as attrs, horizontal as h_grid, @@ -41,8 +42,11 @@ from icon4py.model.testing import serialbox as sb -def test_geometry_raises_for_unknown_field(backend: gtx_typing.Backend) -> None: - geometry = grid_utils.get_grid_geometry(backend, definitions.Experiments.EXCLAIM_APE) +@pytest.mark.datatest +def test_geometry_raises_for_unknown_field( + backend: gtx_typing.Backend, experiment: definitions.Experiment +) -> None: + geometry = grid_utils.get_grid_geometry(backend, experiment) with pytest.raises(ValueError) as e: geometry.get("foo") assert "'foo'" in e.value # type: ignore[operator] @@ -54,6 +58,7 @@ def test_geometry_raises_for_unknown_field(backend: gtx_typing.Backend) -> None: [ (definitions.Experiments.MCH_CH_R04B09, 1e-7), (definitions.Experiments.EXCLAIM_APE, 3e-12), + (definitions.Experiments.GAUSS3D, 1e-13), ], ) @pytest.mark.datatest @@ -87,6 +92,7 @@ def test_coriolis_parameter( [ (definitions.Experiments.MCH_CH_R04B09, 1e-9), (definitions.Experiments.EXCLAIM_APE, 1e-12), + (definitions.Experiments.GAUSS3D, 1e-13), ], ) @pytest.mark.datatest @@ -107,6 +113,7 @@ def test_compute_edge_length( [ (definitions.Experiments.MCH_CH_R04B09, 1e-9), (definitions.Experiments.EXCLAIM_APE, 1e-12), + (definitions.Experiments.GAUSS3D, 1e-13), ], ) @pytest.mark.datatest @@ -128,6 +135,7 @@ def test_compute_inverse_edge_length( [ (definitions.Experiments.MCH_CH_R04B09, 1e-7), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 1e-13), ], ) @pytest.mark.datatest @@ -149,6 +157,7 @@ def test_compute_dual_edge_length( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 1e-13), ], ) @pytest.mark.datatest @@ -175,6 +184,7 @@ def test_compute_inverse_dual_edge_length( [ (definitions.Experiments.MCH_CH_R04B09, 5e-10), (definitions.Experiments.EXCLAIM_APE, 1e-12), + (definitions.Experiments.GAUSS3D, 1e-14), ], ) @pytest.mark.datatest @@ -186,9 +196,13 @@ 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): icon fortran multiplies sphere radius in even for + # torus grids. Fix upstream. + result = result / constants.EARTH_RADIUS + assert test_utils.dallclose(result, expected, rtol=rtol) @pytest.mark.datatest @@ -338,6 +352,7 @@ def test_dual_normal_vert( assert test_utils.dallclose(dual_normal_vert_v.asnumpy(), dual_normal_vert_v_ref, atol=1e-12) +@pytest.mark.datatest def test_cartesian_centers_edge( backend: gtx_typing.Backend, experiment: definitions.Experiment ) -> None: @@ -349,12 +364,22 @@ 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: + assert all(x.asnumpy() >= 0.0) + assert all(x.asnumpy() <= grid.global_properties.domain_length) + assert all(y.asnumpy() >= 0.0) + assert all(y.asnumpy() <= grid.global_properties.domain_height) + assert all(z.asnumpy() == 0.0) + +@pytest.mark.datatest def test_cartesian_centers_cell( backend: gtx_typing.Backend, experiment: definitions.Experiment ) -> None: @@ -366,12 +391,22 @@ 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: + assert all(x.asnumpy() >= 0.0) + assert all(x.asnumpy() <= grid.global_properties.domain_length) + assert all(y.asnumpy() >= 0.0) + assert all(y.asnumpy() <= grid.global_properties.domain_height) + assert all(z.asnumpy() == 0.0) +@pytest.mark.datatest def test_vertex(backend: gtx_typing.Backend, experiment: definitions.Experiment) -> None: grid_geometry = grid_utils.get_grid_geometry(backend, experiment) grid = grid_geometry.grid @@ -381,10 +416,19 @@ 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: + 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_icon.py b/model/common/tests/common/grid/unit_tests/test_icon.py index 84656a9faf..024179aead 100644 --- a/model/common/tests/common/grid/unit_tests/test_icon.py +++ b/model/common/tests/common/grid/unit_tests/test_icon.py @@ -231,6 +231,7 @@ def _sphere_area(radius: float) -> float: return 4.0 * math.pi * radius**2.0 +# TODO(msimberg): Test construction from fields. @pytest.mark.parametrize( "geometry_type,grid_root,grid_level,global_num_cells,num_cells,mean_cell_area,expected_global_num_cells,expected_num_cells,expected_mean_cell_area", [ diff --git a/model/common/tests/common/interpolation/unit_tests/test_interpolation_factory.py b/model/common/tests/common/interpolation/unit_tests/test_interpolation_factory.py index 30bb81024a..179a982334 100644 --- a/model/common/tests/common/interpolation/unit_tests/test_interpolation_factory.py +++ b/model/common/tests/common/interpolation/unit_tests/test_interpolation_factory.py @@ -101,6 +101,7 @@ def test_factory_raises_error_on_unknown_field( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -124,6 +125,7 @@ def test_get_c_lin_e( [ (definitions.Experiments.MCH_CH_R04B09, 1e-9), (definitions.Experiments.EXCLAIM_APE, 1e-12), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -149,6 +151,7 @@ def test_get_geofac_div( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -183,6 +186,7 @@ def assert_reordered(val: np.ndarray, ref: np.ndarray, **kwargs: Any) -> None: [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -209,6 +213,7 @@ def test_get_geofac_rot( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -262,6 +267,7 @@ def test_get_geofac_grg( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 1e-15), ], ) @pytest.mark.datatest @@ -301,6 +307,7 @@ def test_e_flx_avg( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -324,6 +331,7 @@ def test_e_bln_c_s( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 1e-13), ], ) @pytest.mark.datatest @@ -348,6 +356,7 @@ def test_pos_on_tplane_e_x_y( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -372,6 +381,7 @@ def test_cells_aw_verts( [ (definitions.Experiments.MCH_CH_R04B09, 5e-9), (definitions.Experiments.EXCLAIM_APE, 1e-11), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -394,6 +404,7 @@ def test_nudgecoeffs( [ (definitions.Experiments.EXCLAIM_APE, 3e-9), (definitions.Experiments.MCH_CH_R04B09, 4e-2), + (definitions.Experiments.GAUSS3D, 1e-14), ], ) @pytest.mark.datatest @@ -427,6 +438,7 @@ def test_rbf_interpolation_coeffs_cell( [ (definitions.Experiments.EXCLAIM_APE, 8e-14), (definitions.Experiments.MCH_CH_R04B09, 2e-9), + (definitions.Experiments.GAUSS3D, 0), ], ) @pytest.mark.datatest @@ -455,6 +467,7 @@ def test_rbf_interpolation_coeffs_edge( [ (definitions.Experiments.EXCLAIM_APE, 3e-10), (definitions.Experiments.MCH_CH_R04B09, 3e-3), + (definitions.Experiments.GAUSS3D, 1e-15), ], ) @pytest.mark.datatest diff --git a/model/common/tests/common/interpolation/unit_tests/test_interpolation_fields.py b/model/common/tests/common/interpolation/unit_tests/test_interpolation_fields.py index b623edc8f2..55ef66a918 100644 --- a/model/common/tests/common/interpolation/unit_tests/test_interpolation_fields.py +++ b/model/common/tests/common/interpolation/unit_tests/test_interpolation_fields.py @@ -20,6 +20,7 @@ compute_c_lin_e, compute_cells_aw_verts, compute_e_bln_c_s, + compute_e_bln_c_s_torus, compute_e_flx_avg, compute_geofac_div, compute_geofac_grdiv, @@ -27,7 +28,9 @@ compute_geofac_n2s, compute_geofac_rot, compute_mass_conserving_bilinear_cell_average_weight, + compute_mass_conserving_bilinear_cell_average_weight_torus, compute_pos_on_tplane_e_x_y, + compute_pos_on_tplane_e_x_y_torus, ) from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import definitions, serialbox as sb @@ -245,7 +248,11 @@ def test_compute_geofac_grdiv( @pytest.mark.datatest @pytest.mark.parametrize( "experiment, atol", - [(definitions.Experiments.MCH_CH_R04B09, 1e-10), (definitions.Experiments.EXCLAIM_APE, 1e-10)], + [ + (definitions.Experiments.MCH_CH_R04B09, 1e-10), + (definitions.Experiments.EXCLAIM_APE, 1e-10), + (definitions.Experiments.GAUSS3D, 1e-15), + ], ) def test_compute_c_bln_avg( grid_savepoint: sb.IconGridSavepoint, @@ -268,18 +275,30 @@ def test_compute_c_bln_avg( c2e2c0 = icon_grid.get_connectivity(dims.C2E2CO).ndarray - c_bln_avg = functools.partial( - compute_mass_conserving_bilinear_cell_average_weight, array_ns=xp - )( - c2e2c0, - lat, - lon, - cell_areas, - cell_owner_mask, - divavg_cntrwgt, - horizontal_start, - horizontal_start_p2, - ) + match icon_grid.global_properties.geometry_type: + case base_grid.GeometryType.ICOSAHEDRON: + c_bln_avg = compute_mass_conserving_bilinear_cell_average_weight( + c2e2c0, + lat, + lon, + cell_areas, + cell_owner_mask, + divavg_cntrwgt, + horizontal_start, + horizontal_start_p2, + array_ns=xp, + ) + case base_grid.GeometryType.TORUS: + c_bln_avg = compute_mass_conserving_bilinear_cell_average_weight_torus( + c2e2c0, + cell_areas, + cell_owner_mask, + divavg_cntrwgt, + horizontal_start, + horizontal_start_p2, + array_ns=xp, + ) + assert test_helpers.dallclose(data_alloc.as_numpy(c_bln_avg), c_bln_avg_ref, atol=atol) @@ -373,9 +392,14 @@ def test_compute_e_bln_c_s( edges_lon = grid_savepoint.edges_center_lon().ndarray xp = data_alloc.import_array_ns(backend) - e_bln_c_s = functools.partial(compute_e_bln_c_s, array_ns=xp)( - c2e, cells_lat, cells_lon, edges_lat, edges_lon, 0.0 - ) + # TODO(msimberg): pass geometry type to compute_pos_on_tplane_e_x_y? + match icon_grid.global_properties.geometry_type: + case base_grid.GeometryType.ICOSAHEDRON: + e_bln_c_s = compute_e_bln_c_s( + c2e, cells_lat, cells_lon, edges_lat, edges_lon, 0.0, array_ns=xp + ) + case base_grid.GeometryType.TORUS: + e_bln_c_s = compute_e_bln_c_s_torus(c2e, array_ns=xp) assert test_helpers.dallclose( data_alloc.as_numpy(e_bln_c_s), e_bln_c_s_ref.asnumpy(), atol=1e-6, rtol=1e-7 ) @@ -393,6 +417,7 @@ def test_compute_pos_on_tplane_e( pos_on_tplane_e_x_ref = interpolation_savepoint.pos_on_tplane_e_x().asnumpy() pos_on_tplane_e_y_ref = interpolation_savepoint.pos_on_tplane_e_y().asnumpy() sphere_radius = constants.EARTH_RADIUS + dual_edge_length = grid_savepoint.dual_edge_length().ndarray primal_normal_v1 = grid_savepoint.primal_normal_v1().ndarray primal_normal_v2 = grid_savepoint.primal_normal_v2().ndarray dual_normal_v1 = grid_savepoint.dual_normal_v1().ndarray @@ -404,20 +429,29 @@ def test_compute_pos_on_tplane_e( edges_lat = grid_savepoint.edges_center_lat().ndarray e2c = icon_grid.get_connectivity(dims.E2C).ndarray horizontal_start = icon_grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)) - pos_on_tplane_e_x, pos_on_tplane_e_y = compute_pos_on_tplane_e_x_y( - sphere_radius, - primal_normal_v1, - primal_normal_v2, - dual_normal_v1, - dual_normal_v2, - cells_lon, - cells_lat, - edges_lon, - edges_lat, - owner_mask, - e2c, - horizontal_start, - array_ns=xp, - ) + # TODO(msimberg): pass geometry type to compute_pos_on_tplane_e_x_y? + match icon_grid.global_properties.geometry_type: + case base_grid.GeometryType.ICOSAHEDRON: + pos_on_tplane_e_x, pos_on_tplane_e_y = compute_pos_on_tplane_e_x_y( + sphere_radius, + primal_normal_v1, + primal_normal_v2, + dual_normal_v1, + dual_normal_v2, + cells_lon, + cells_lat, + edges_lon, + edges_lat, + owner_mask, + e2c, + horizontal_start, + array_ns=xp, + ) + case base_grid.GeometryType.TORUS: + pos_on_tplane_e_x, pos_on_tplane_e_y = compute_pos_on_tplane_e_x_y_torus( + dual_edge_length, + e2c, + array_ns=xp, + ) assert test_helpers.dallclose(pos_on_tplane_e_x, pos_on_tplane_e_x_ref, atol=1e-6, rtol=1e-7) assert test_helpers.dallclose(pos_on_tplane_e_y, pos_on_tplane_e_y_ref, atol=1e-6, rtol=1e-7) diff --git a/model/common/tests/common/interpolation/unit_tests/test_rbf_interpolation.py b/model/common/tests/common/interpolation/unit_tests/test_rbf_interpolation.py index 9912cc3b94..4e172e82b2 100644 --- a/model/common/tests/common/interpolation/unit_tests/test_rbf_interpolation.py +++ b/model/common/tests/common/interpolation/unit_tests/test_rbf_interpolation.py @@ -145,6 +145,7 @@ def test_construct_rbf_matrix_offsets_tables_for_vertices( [ (definitions.Experiments.EXCLAIM_APE, 3e-9), (definitions.Experiments.MCH_CH_R04B09, 3e-2), + (definitions.Experiments.GAUSS3D, 1e-14), ], ) def test_rbf_interpolation_coeffs_cell( @@ -177,8 +178,16 @@ def test_rbf_interpolation_coeffs_cell( geometry.get(geometry_attrs.EDGE_NORMAL_Z).ndarray, rbf.construct_rbf_matrix_offsets_tables_for_cells(grid), rbf.DEFAULT_RBF_KERNEL[rbf_dim], - rbf.compute_default_rbf_scale(math.sqrt(grid_savepoint.mean_cell_area()), rbf_dim), + grid.global_properties.geometry_type.value, + rbf.compute_default_rbf_scale( + grid.global_properties.geometry_type, + grid.global_properties.characteristic_length, + grid.global_properties.mean_dual_edge_length, + rbf_dim, + ), horizontal_start, + grid.global_properties.domain_length, + grid.global_properties.domain_height, array_ns=data_alloc.import_array_ns(backend), ) @@ -214,6 +223,7 @@ def test_rbf_interpolation_coeffs_cell( [ (definitions.Experiments.EXCLAIM_APE, 3e-10), (definitions.Experiments.MCH_CH_R04B09, 3e-3), + (definitions.Experiments.GAUSS3D, 1e-15), ], ) def test_rbf_interpolation_coeffs_vertex( @@ -246,8 +256,16 @@ def test_rbf_interpolation_coeffs_vertex( geometry.get(geometry_attrs.EDGE_NORMAL_Z).ndarray, rbf.construct_rbf_matrix_offsets_tables_for_vertices(grid), rbf.DEFAULT_RBF_KERNEL[rbf_dim], - rbf.compute_default_rbf_scale(math.sqrt(grid_savepoint.mean_cell_area()), rbf_dim), + grid.global_properties.geometry_type.value, + rbf.compute_default_rbf_scale( + grid.global_properties.geometry_type, + grid.global_properties.characteristic_length, + grid.global_properties.mean_dual_edge_length, + rbf_dim, + ), horizontal_start, + grid.global_properties.domain_length, + grid.global_properties.domain_height, array_ns=data_alloc.import_array_ns(backend), ) @@ -283,6 +301,7 @@ def test_rbf_interpolation_coeffs_vertex( [ (definitions.Experiments.EXCLAIM_APE, 8e-14), (definitions.Experiments.MCH_CH_R04B09, 2e-9), + (definitions.Experiments.GAUSS3D, 0), ], ) def test_rbf_interpolation_coeffs_edge( @@ -317,8 +336,16 @@ def test_rbf_interpolation_coeffs_edge( # coefficients in savepoint. grid_savepoint.e2c2e(), rbf.DEFAULT_RBF_KERNEL[rbf_dim], - rbf.compute_default_rbf_scale(math.sqrt(grid_savepoint.mean_cell_area()), rbf_dim), + grid.global_properties.geometry_type.value, + rbf.compute_default_rbf_scale( + grid.global_properties.geometry_type, + grid.global_properties.characteristic_length, + grid.global_properties.mean_dual_edge_length, + rbf_dim, + ), horizontal_start, + grid.global_properties.domain_length, + grid.global_properties.domain_height, array_ns=data_alloc.import_array_ns(backend), ) diff --git a/model/common/tests/common/metrics/unit_tests/test_metric_fields.py b/model/common/tests/common/metrics/unit_tests/test_metric_fields.py index 48c397334c..6978ceab64 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 @@ -146,8 +146,19 @@ def test_compute_rayleigh_w( 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 + match experiment: + case definitions.Experiments.EXCLAIM_APE: + rayleigh_coeff = 0.1 + damping_height = 50000.0 + case definitions.Experiments.GAUSS3D: + rayleigh_coeff = 0.1 + damping_height = 45000.0 + case definitions.Experiments.WEISMAN_KLEMP_TORUS: + rayleigh_coeff = 0.75 + damping_height = 8000.0 + case _: + rayleigh_coeff = 5.0 + damping_height = 12500.0 mf.compute_rayleigh_w.with_backend(backend=backend)( rayleigh_w=rayleigh_w_full, vct_a=grid_savepoint.vct_a(), @@ -231,7 +242,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) exner_exfac_ref = metrics_savepoint.exner_exfac() mf.compute_exner_exfac.with_backend(backend)( @@ -314,7 +332,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 7a1befd7c8..8f3a1b204c 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 @@ -15,7 +15,7 @@ import pytest from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import vertical as v_grid +from icon4py.model.common.grid import base, vertical as v_grid from icon4py.model.common.interpolation import interpolation_attributes, interpolation_factory from icon4py.model.common.metrics import metrics_attributes as attrs, metrics_factory from icon4py.model.common.utils import data_allocation as data_alloc @@ -46,7 +46,13 @@ def metrics_config(experiment: definitions.Experiment) -> tuple: rayleigh_coeff = 5.0 lowest_layer_thickness = 50.0 exner_expol = 0.333 - vwind_offctr = 0.2 + 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 rayleigh_type = 2 model_top_height = 23500.0 stretch_factor = 1.0 @@ -64,6 +70,13 @@ def metrics_config(experiment: definitions.Experiment) -> tuple: rayleigh_coeff = 0.1 exner_expol = 0.3333333333333 vwind_offctr = 0.15 + case definitions.Experiments.GAUSS3D: + rayleigh_coeff = 0.1 + damping_height = 45000.0 + exner_expol = 1.0 / 3.0 + case definitions.Experiments.WEISMAN_KLEMP_TORUS: + rayleigh_coeff = 0.75 + damping_height = 8000.0 return ( lowest_layer_thickness, diff --git a/model/testing/src/icon4py/model/testing/definitions.py b/model/testing/src/icon4py/model/testing/definitions.py index 8818c487c5..323fa7cd04 100644 --- a/model/testing/src/icon4py/model/testing/definitions.py +++ b/model/testing/src/icon4py/model/testing/definitions.py @@ -228,6 +228,29 @@ def construct_diffusion_config( hdiff_temp=True, n_substeps=ndyn_substeps, ) + elif experiment == Experiments.GAUSS3D: + return diffusion.DiffusionConfig( + diffusion_type=diffusion.DiffusionType.SMAGORINSKY_4TH_ORDER, # TODO(msimberg): check + hdiff_w=True, + hdiff_vn=True, + hdiff_temp=True, + type_vn_diffu=1, + smag_3d=False, + type_t_diffu=2, + hdiff_efdt_ratio=36.0, + hdiff_w_efdt_ratio=15.0, + smagorinski_scaling_factor=0.015, + n_substeps=ndyn_substeps, + zdiffu_t=True, + thslp_zdiffu=0.025, + thhgtd_zdiffu=200.0, + velocity_boundary_diffusion_denom=200.0, # TODO(msimberg): check + temperature_boundary_diffusion_denom=135.0, # TODO(msimberg): check + max_nudging_coefficient=0.375, # TODO(msimberg): check + nudging_decay_rate=2.0, # TODO(msimberg): check + shear_type=diffusion.TurbulenceShearForcingType.VERTICAL_HORIZONTAL_OF_HORIZONTAL_VERTICAL_WIND, # TODO(msimberg): check + ltkeshs=True, # TODO(msimberg): check + ) else: raise NotImplementedError( f"DiffusionConfig for experiment {experiment.name} not implemented." diff --git a/model/testing/src/icon4py/model/testing/fixtures/benchmark.py b/model/testing/src/icon4py/model/testing/fixtures/benchmark.py index f230f43e82..658dc605ac 100644 --- a/model/testing/src/icon4py/model/testing/fixtures/benchmark.py +++ b/model/testing/src/icon4py/model/testing/fixtures/benchmark.py @@ -43,7 +43,7 @@ def geometry_field_source( generic_concrete_backend = model_options.customize_backend(None, backend_like) decomposition_info = grid_utils.construct_decomposition_info(mesh, allocator) - geometry_field_source = grid_geometry.GridGeometry( + geometry_field_source = grid_geometry.GridGeometry.with_geometry_type( grid=mesh, decomposition_info=decomposition_info, backend=generic_concrete_backend, diff --git a/model/testing/src/icon4py/model/testing/fixtures/datatest.py b/model/testing/src/icon4py/model/testing/fixtures/datatest.py index 123884f464..3058fc7210 100644 --- a/model/testing/src/icon4py/model/testing/fixtures/datatest.py +++ b/model/testing/src/icon4py/model/testing/fixtures/datatest.py @@ -88,7 +88,11 @@ def cpu_allocator() -> gtx_typing.FieldBufferAllocationUtil: @pytest.fixture( - params=[definitions.Experiments.MCH_CH_R04B09, definitions.Experiments.EXCLAIM_APE], + params=[ + definitions.Experiments.MCH_CH_R04B09, + definitions.Experiments.EXCLAIM_APE, + definitions.Experiments.GAUSS3D, + ], ids=lambda r: r.name, ) def experiment(request: pytest.FixtureRequest) -> definitions.Experiment: diff --git a/model/testing/src/icon4py/model/testing/grid_utils.py b/model/testing/src/icon4py/model/testing/grid_utils.py index 5d2dc9696d..4e5f920538 100644 --- a/model/testing/src/icon4py/model/testing/grid_utils.py +++ b/model/testing/src/icon4py/model/testing/grid_utils.py @@ -135,7 +135,7 @@ def _construct_grid_geometry() -> geometry.GridGeometry: ) grid = gm.grid decomposition_info = construct_decomposition_info(grid, backend) - geometry_source = geometry.GridGeometry( + return geometry.GridGeometry.with_geometry_type( grid, decomposition_info, backend, @@ -143,7 +143,6 @@ def _construct_grid_geometry() -> geometry.GridGeometry: gm.geometry_fields, geometry_attrs.attrs, ) - return geometry_source if not grid_geometries.get(register_name): grid_geometries[register_name] = _construct_grid_geometry()