diff --git a/model/common/src/icon4py/model/common/grid/geometry.py b/model/common/src/icon4py/model/common/grid/geometry.py index 79bcd3847b..cd0df333ac 100644 --- a/model/common/src/icon4py/model/common/grid/geometry.py +++ b/model/common/src/icon4py/model/common/grid/geometry.py @@ -29,7 +29,11 @@ horizontal as h_grid, icon, ) -from icon4py.model.common.math import helpers as math_helpers, xp_utils +from icon4py.model.common.math import ( + coordinate_transformations as coord_trans, + math_utils, + vector_operations as vector_ops, +) from icon4py.model.common.states import factory, model, utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc, device_utils @@ -187,7 +191,7 @@ def _inverse_field_provider(self, field_name: str) -> factory.FieldProvider: name = meta["standard_name"] self._attrs.update({name: meta}) provider = factory.ProgramFieldProvider( - func=math_helpers.compute_inverse_on_edges, + func=vector_ops.compute_inverse_on_edges, deps={"f": field_name}, fields={"f_inverse": name}, domain={ @@ -367,7 +371,7 @@ def _register_computed_fields(self) -> None: self.register_provider(mean_dual_cell_area_np) characteristic_length_np = factory.NumpyDataProvider( - func=xp_utils.compute_sqrt, + func=math_utils.compute_sqrt, domain=(), deps={ "input_val": attrs.MEAN_DUAL_AREA, @@ -408,7 +412,7 @@ def _register_normals_and_tangents_icosahedron(self) -> None: # 2. primal_normals: gridfile%zonal_normal_primal_edge - edges%primal_normal%v1, gridfile%meridional_normal_primal_edge - edges%primal_normal%v2, normal_uv = factory.ProgramFieldProvider( - func=math_helpers.compute_zonal_and_meridional_components_on_edges, + func=coord_trans.compute_zonal_and_meridional_components_on_edges, deps={ "lat": attrs.EDGE_LAT, "lon": attrs.EDGE_LON, @@ -431,7 +435,7 @@ def _register_normals_and_tangents_icosahedron(self) -> None: self.register_provider(normal_uv) dual_uv = factory.ProgramFieldProvider( - func=math_helpers.compute_zonal_and_meridional_components_on_edges, + func=coord_trans.compute_zonal_and_meridional_components_on_edges, deps={ "lat": attrs.EDGE_LAT, "lon": attrs.EDGE_LON, @@ -708,7 +712,7 @@ def _register_normals_and_tangents_torus(self) -> None: def _register_cartesian_coordinates_icosahedron(self) -> None: """Register Cartesian coordinate conversions for icosahedron geometry.""" cartesian_vertices = factory.EmbeddedFieldOperatorProvider( - func=math_helpers.geographical_to_cartesian_on_vertices.with_backend(self.backend), + func=coord_trans.geographical_to_cartesian_on_vertices.with_backend(self.backend), domain={ dims.VertexDim: ( h_grid.vertex_domain(h_grid.Zone.LOCAL), @@ -728,7 +732,7 @@ def _register_cartesian_coordinates_icosahedron(self) -> None: ) self.register_provider(cartesian_vertices) cartesian_edge_centers = factory.EmbeddedFieldOperatorProvider( - func=math_helpers.geographical_to_cartesian_on_edges.with_backend(self.backend), + func=coord_trans.geographical_to_cartesian_on_edges.with_backend(self.backend), domain={ dims.EdgeDim: ( h_grid.edge_domain(h_grid.Zone.LOCAL), @@ -748,7 +752,7 @@ def _register_cartesian_coordinates_icosahedron(self) -> None: ) self.register_provider(cartesian_edge_centers) cartesian_cell_centers = factory.EmbeddedFieldOperatorProvider( - func=math_helpers.geographical_to_cartesian_on_cells.with_backend(self.backend), + func=coord_trans.geographical_to_cartesian_on_cells.with_backend(self.backend), domain={ dims.CellDim: ( h_grid.cell_domain(h_grid.Zone.LOCAL), 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 2fc5f78d4e..ab8a4f4a00 100644 --- a/model/common/src/icon4py/model/common/grid/geometry_stencils.py +++ b/model/common/src/icon4py/model/common/grid/geometry_stencils.py @@ -15,15 +15,19 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta from icon4py.model.common.dimension import E2C, E2C2V, E2V, EdgeDim -from icon4py.model.common.math.helpers import ( +from icon4py.model.common.math.coordinate_transformations import ( + geographical_to_cartesian_on_edges, + geographical_to_cartesian_on_vertices, + zonal_and_meridional_components_on_edges, +) +from icon4py.model.common.math.distance 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, +) +from icon4py.model.common.math.vector_operations import ( + cross_product_on_edges, normalize_cartesian_vector_on_edges, - zonal_and_meridional_components_on_edges, ) from icon4py.model.common.utils import data_allocation as data_alloc 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 050c71e38a..24bb83a212 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py @@ -17,14 +17,13 @@ from gt4py.next import where import icon4py.model.common.field_type_aliases as fa -import icon4py.model.common.math.projection as proj import icon4py.model.common.type_alias as ta from icon4py.model.common import dimension as dims from icon4py.model.common.decomposition import definitions as decomposition from icon4py.model.common.dimension import C2E, V2E from icon4py.model.common.grid import base as base_grid, gridfile from icon4py.model.common.grid.geometry_stencils import compute_primal_cart_normal -from icon4py.model.common.math.projection import diff_on_edges_torus_numpy, gnomonic_proj +from icon4py.model.common.math import projection from icon4py.model.common.utils import data_allocation as data_alloc @@ -1057,10 +1056,10 @@ def compute_pos_on_tplane_e_x_y( pos_on_tplane_e_y = array_ns.zeros(e2c.shape) xyloc_plane_n1 = array_ns.zeros([2, e2c.shape[0]]) xyloc_plane_n2 = array_ns.zeros([2, e2c.shape[0]]) - xyloc_plane_n1[0, llb:], xyloc_plane_n1[1, llb:] = proj.gnomonic_proj( + xyloc_plane_n1[0, llb:], xyloc_plane_n1[1, llb:] = projection.gnomonic_proj( edges_lon[llb:], edges_lat[llb:], cells_lon[e2c[llb:, 0]], cells_lat[e2c[llb:, 0]] ) - xyloc_plane_n2[0, llb:], xyloc_plane_n2[1, llb:] = proj.gnomonic_proj( + xyloc_plane_n2[0, llb:], xyloc_plane_n2[1, llb:] = projection.gnomonic_proj( edges_lon[llb:], edges_lat[llb:], cells_lon[e2c[llb:, 1]], cells_lat[e2c[llb:, 1]] ) @@ -1236,7 +1235,7 @@ def compute_lsq_coeffs( case base_grid.GeometryType.ICOSAHEDRON: for js in range(lsq_dim_stencil): z_dist_g[:, js, :] = array_ns.asarray( - gnomonic_proj( + projection.gnomonic_proj( cell_lon, cell_lat, cell_lon[c2e2c[:, js]], @@ -1259,7 +1258,7 @@ def compute_lsq_coeffs( cc_cv = array_ns.asarray((cell_center_x[jc], cell_center_y[jc])) for js in range(lsq_dim_stencil): cc_cell[js, :] = array_ns.asarray( - diff_on_edges_torus_numpy( + projection.diff_on_edges_torus_numpy( cell_center_x[jc], cell_center_y[jc], cell_center_x[ilc_s][js], diff --git a/model/common/src/icon4py/model/common/math/__init__.py b/model/common/src/icon4py/model/common/math/__init__.py index de9850de36..e29e30c57f 100644 --- a/model/common/src/icon4py/model/common/math/__init__.py +++ b/model/common/src/icon4py/model/common/math/__init__.py @@ -5,3 +5,23 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause + +""" +Mathematical operations for ICON grid computations. + +This package provides mathematical operations organized into the following submodules: + +- ``coordinate_transformations``: Conversions between geographical (lat/lon) and cartesian + coordinates, and between zonal/meridional and cartesian vector components. +- ``derivative``: Vertical derivative computations at cell centers. +- ``distance``: Arc length on spheres and distance/difference operations on torus geometries. +- ``gradient``: Finite difference gradient operators (normal and tangential). +- ``math_utils``: General-purpose mathematical utility functions (typed sqrt, etc.). +- ``operators``: Laplacian (nabla²) and difference operators on cell fields. +- ``projection``: Gnomonic projection and NumPy-based torus operations. +- ``smagorinsky``: Smagorinsky diffusion enhancement factor computation. +- ``vector_operations``: Dot product, cross product, norms, normalization and inversion + of vectors on cell, edge and vertex fields. +- ``vertical_operations``: Averaging and differencing of adjacent vertical levels. +- ``stencils``: GT4Py program wrappers for operators with domain control. +""" diff --git a/model/common/src/icon4py/model/common/math/coordinate_transformations.py b/model/common/src/icon4py/model/common/math/coordinate_transformations.py new file mode 100644 index 0000000000..a1bd67fbe2 --- /dev/null +++ b/model/common/src/icon4py/model/common/math/coordinate_transformations.py @@ -0,0 +1,282 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +""" +Coordinate transformation operations on unstructured grid fields. + +Contains conversions between geographical (lat/lon) and cartesian coordinates, +and between zonal/meridional and cartesian components. +""" + +from gt4py import next as gtx +from gt4py.next import cos, sin, sqrt + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.math.vector_operations import norm2_on_cells, norm2_on_edges + + +@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) +def geographical_to_cartesian_on_cells( + lat: fa.CellField[ta.wpfloat], lon: fa.CellField[ta.wpfloat] +) -> tuple[fa.CellField[ta.wpfloat], fa.CellField[ta.wpfloat], fa.CellField[ta.wpfloat]]: + """ + Convert geographical (lat, lon) coordinates to cartesian coordinates on the unit sphere. + + Args: + lat: latitude + lon: longitude + + Returns: + x: x coordinate + y: y coordinate + z: z coordinate + + """ + x = cos(lat) * cos(lon) + y = cos(lat) * sin(lon) + z = sin(lat) + return x, y, z + + +@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) +def geographical_to_cartesian_on_edges( + lat: fa.EdgeField[ta.wpfloat], lon: fa.EdgeField[ta.wpfloat] +) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: + """ + Convert geographical (lat, lon) coordinates to cartesian coordinates on the unit sphere. + + Args: + lat: latitude + lon: longitude + + Returns: + x: x coordinate + y: y coordinate + z: z coordinate + + """ + x = cos(lat) * cos(lon) + y = cos(lat) * sin(lon) + z = sin(lat) + return x, y, z + + +@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) +def geographical_to_cartesian_on_vertices( + lat: fa.VertexField[ta.wpfloat], lon: fa.VertexField[ta.wpfloat] +) -> tuple[fa.VertexField[ta.wpfloat], fa.VertexField[ta.wpfloat], fa.VertexField[ta.wpfloat]]: + """ + Convert geographical (lat, lon) coordinates to cartesian coordinates on the unit sphere. + + Args: + lat: latitude + lon: longitude + + Returns: + x: x coordinate + y: y coordinate + z: z coordinate + + """ + x = cos(lat) * cos(lon) + y = cos(lat) * sin(lon) + z = sin(lat) + return x, y, z + + +@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) +def zonal_and_meridional_components_on_cells( + lat: fa.CellField[ta.wpfloat], + lon: fa.CellField[ta.wpfloat], + x: fa.CellField[ta.wpfloat], + y: fa.CellField[ta.wpfloat], + z: fa.CellField[ta.wpfloat], +) -> tuple[fa.CellField[ta.wpfloat], fa.CellField[ta.wpfloat]]: + """ + Compute normalized zonal and meridional components of a cartesian vector (x, y, z) at point (lat, lon) + + Args: + lat: latitude + lon: longitude + x: x coordinate + y: y coordinate + z: z coordinate + + Returns: + u zonal component + v meridional component + + """ + cos_lat = cos(lat) + sin_lat = sin(lat) + cos_lon = cos(lon) + sin_lon = sin(lon) + u = cos_lon * y - sin_lon * x + + v = cos_lat * z - sin_lat * (cos_lon * x + sin_lon * y) + norm = sqrt(u * u + v * v) + return u / norm, v / norm + + +@gtx.field_operator +def zonal_and_meridional_components_on_edges( + lat: fa.EdgeField[ta.wpfloat], + lon: fa.EdgeField[ta.wpfloat], + x: fa.EdgeField[ta.wpfloat], + y: fa.EdgeField[ta.wpfloat], + z: fa.EdgeField[ta.wpfloat], +) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: + """ + Compute the zonal and meridional component of a vector (x, y, z) at position (lat, lon) + + Args: + lat: latitude + lon: longitude + x: x component of cartesian vector + y: y component of cartesian vector + z: z component of cartesian vector + + Returns: + zonal (eastward) component of (x, y, z) at (lat, lon) + meridional (northward) component of (x, y, z) at (lat, lon) + + """ + cos_lat = cos(lat) + sin_lat = sin(lat) + cos_lon = cos(lon) + sin_lon = sin(lon) + u = cos_lon * y - sin_lon * x + + v = cos_lat * z - sin_lat * (cos_lon * x + sin_lon * y) + norm = sqrt(u * u + v * v) + return u / norm, v / norm + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_zonal_and_meridional_components_on_edges( + lat: fa.EdgeField[ta.wpfloat], + lon: fa.EdgeField[ta.wpfloat], + x: fa.EdgeField[ta.wpfloat], + y: fa.EdgeField[ta.wpfloat], + z: fa.EdgeField[ta.wpfloat], + u: fa.EdgeField[ta.wpfloat], + v: fa.EdgeField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, +): + zonal_and_meridional_components_on_edges( + lat, lon, x, y, z, out=(u, v), domain={dims.EdgeDim: (horizontal_start, horizontal_end)} + ) + + +@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) +def cartesian_coordinates_from_zonal_and_meridional_components_on_edges( + lat: fa.EdgeField[ta.wpfloat], + lon: fa.EdgeField[ta.wpfloat], + u: fa.EdgeField[ta.wpfloat], + v: fa.EdgeField[ta.wpfloat], +) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: + """ + Compute cartesian coordinates from zonal and meridional components at position (lat, lon) + Args: + lat: latitude + lon: longitude + u: zonal component + v: meridional component + + Returns: + x, y, z cartesian components + + """ + cos_lat = cos(lat) + sin_lat = sin(lat) + cos_lon = cos(lon) + sin_lon = sin(lon) + + x = -u * sin_lon - v * sin_lat * cos_lon + y = u * cos_lon - v * sin_lat * sin_lon + z = cos_lat * v + + norm = norm2_on_edges(x, y, z) + return x / norm, y / norm, z / norm + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_cartesian_coordinates_from_zonal_and_meridional_components_on_edges( + edge_lat: fa.EdgeField[ta.wpfloat], + edge_lon: fa.EdgeField[ta.wpfloat], + u: fa.EdgeField[ta.wpfloat], + v: fa.EdgeField[ta.wpfloat], + x: fa.EdgeField[ta.wpfloat], + y: fa.EdgeField[ta.wpfloat], + z: fa.EdgeField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, +): + cartesian_coordinates_from_zonal_and_meridional_components_on_edges( + edge_lat, + edge_lon, + u, + v, + out=(x, y, z), + domain={dims.EdgeDim: (horizontal_start, horizontal_end)}, + ) + + +@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) +def cartesian_coordinates_from_zonal_and_meridional_components_on_cells( + lat: fa.CellField[ta.wpfloat], + lon: fa.CellField[ta.wpfloat], + u: fa.CellField[ta.wpfloat], + v: fa.CellField[ta.wpfloat], +) -> tuple[fa.CellField[ta.wpfloat], fa.CellField[ta.wpfloat], fa.CellField[ta.wpfloat]]: + """ + Compute cartesian coordinates from zonal and meridional components at position (lat, lon) + Args: + lat: latitude + lon: longitude + u: zonal component + v: meridional component + + Returns: + x, y, z cartesian components + + """ + cos_lat = cos(lat) + sin_lat = sin(lat) + cos_lon = cos(lon) + sin_lon = sin(lon) + + x = -u * sin_lon - v * sin_lat * cos_lon + y = u * cos_lon - v * sin_lat * sin_lon + z = cos_lat * v + + norm = norm2_on_cells(x, y, z) + return x / norm, y / norm, z / norm + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_cartesian_coordinates_from_zonal_and_meridional_components_on_cells( + cell_lat: fa.CellField[ta.wpfloat], + cell_lon: fa.CellField[ta.wpfloat], + u: fa.CellField[ta.wpfloat], + v: fa.CellField[ta.wpfloat], + x: fa.CellField[ta.wpfloat], + y: fa.CellField[ta.wpfloat], + z: fa.CellField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, +): + cartesian_coordinates_from_zonal_and_meridional_components_on_cells( + cell_lat, + cell_lon, + u, + v, + out=(x, y, z), + domain={dims.CellDim: (horizontal_start, horizontal_end)}, + ) diff --git a/model/common/src/icon4py/model/common/math/distance.py b/model/common/src/icon4py/model/common/math/distance.py new file mode 100644 index 0000000000..890ef72ca5 --- /dev/null +++ b/model/common/src/icon4py/model/common/math/distance.py @@ -0,0 +1,129 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +""" +Distance and arc length operations on unstructured grid fields. + +Contains arc length computation on spheres and distance/difference operations +on torus geometries. +""" + +from gt4py import next as gtx +from gt4py.next import ( + abs, # noqa: A004 + arccos, + sqrt, + where, +) + +from icon4py.model.common import field_type_aliases as fa, type_alias as ta +from icon4py.model.common.math.vector_operations import dot_product_on_edges + + +@gtx.field_operator +def arc_length_on_edges( + x0: fa.EdgeField[ta.wpfloat], + x1: fa.EdgeField[ta.wpfloat], + y0: fa.EdgeField[ta.wpfloat], + y1: fa.EdgeField[ta.wpfloat], + z0: fa.EdgeField[ta.wpfloat], + z1: fa.EdgeField[ta.wpfloat], + radius: ta.wpfloat, +): + """ + Compute the arc length between two points on the sphere. + + Inputs are cartesian coordinates of the points. + + 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 + z0: z coordinate of point_0 + z1: z coordinate of point_1 + radius: sphere radius + + Returns: + arc length + + """ + return radius * arccos(dot_product_on_edges(x0, x1, y0, y1, z0, z1)) + + +@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) +def diff_on_edges_torus( + x0: fa.EdgeField[ta.wpfloat], + x1: fa.EdgeField[ta.wpfloat], + y0: fa.EdgeField[ta.wpfloat], + y1: fa.EdgeField[ta.wpfloat], + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, +) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: + """ + Compute the difference between two points on the torus. + + Inputs are cartesian coordinates of the points. Z is assumed zero and is + ignored. Distances are computed modulo the domain length and height. + + Args: + x0: x coordinate of point_0 + x1: x coordinate of point_1 + y0: y coordinate of point_0 + y1: y coordinate of point_1 + domain_length: length of the domain + domain_height: height of the domain + + Returns: + dx, dy + + """ + x1 = where( + abs(x1 - x0) <= 0.5 * domain_length, + x1, + where(x0 > x1, x1 + domain_length, x1 - domain_length), + ) + + y1 = where( + abs(y1 - y0) <= 0.5 * domain_height, + y1, + where(y0 > y1, y1 + domain_height, y1 - domain_height), + ) + + return (x1 - x0, y1 - y0) + + +@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) +def distance_on_edges_torus( + x0: fa.EdgeField[ta.wpfloat], + x1: fa.EdgeField[ta.wpfloat], + y0: fa.EdgeField[ta.wpfloat], + y1: fa.EdgeField[ta.wpfloat], + domain_length: ta.wpfloat, + domain_height: ta.wpfloat, +) -> fa.EdgeField[ta.wpfloat]: + """ + Compute the distance between two points on the torus. + + Inputs are cartesian coordinates of the points. Z is assumed zero and is + ignored. Distances are computed modulo the domain length and height. + + Args: + x0: x coordinate of point_0 + x1: x coordinate of point_1 + y0: y coordinate of point_0 + y1: y coordinate of point_1 + domain_length: length of the domain + domain_height: height of the domain + + Returns: + distance + + """ + xdiff, ydiff = diff_on_edges_torus(x0, x1, y0, y1, domain_length, domain_height) + return sqrt(xdiff**2 + ydiff**2) diff --git a/model/common/src/icon4py/model/common/math/gradient.py b/model/common/src/icon4py/model/common/math/gradient.py new file mode 100644 index 0000000000..75a78414ca --- /dev/null +++ b/model/common/src/icon4py/model/common/math/gradient.py @@ -0,0 +1,49 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +""" +Finite difference gradient operators on unstructured grid fields. + +Contains normal and tangential gradient computations on edges using +finite difference approximations. +""" + +from gt4py import next as gtx + +from icon4py.model.common import dimension as dims, field_type_aliases as fa +from icon4py.model.common.dimension import E2C, E2V + + +@gtx.field_operator +def grad_fd_norm( + psi_c: fa.CellKField[float], + inv_dual_edge_length: fa.EdgeField[float], +) -> fa.EdgeKField[float]: + """ + Calculate the gradient value of adjacent interface levels. + + Computes the difference of two offseted values multiplied by a field of the offseted dimension + Args: + psi_c: fa.CellKField[float], + inv_dual_edge_length: Field[Dims[EdgeDim], float], + + Returns: fa.EdgeKField[float] + + """ + grad_norm_psi_e = (psi_c(E2C[1]) - psi_c(E2C[0])) * inv_dual_edge_length + return grad_norm_psi_e + + +@gtx.field_operator +def _grad_fd_tang( + psi_v: gtx.Field[gtx.Dims[dims.VertexDim, dims.KDim], float], + inv_primal_edge_length: fa.EdgeField[float], + tangent_orientation: fa.EdgeField[float], +) -> fa.EdgeKField[float]: + grad_tang_psi_e = tangent_orientation * (psi_v(E2V[1]) - psi_v(E2V[0])) * inv_primal_edge_length + return grad_tang_psi_e diff --git a/model/common/src/icon4py/model/common/math/helpers.py b/model/common/src/icon4py/model/common/math/helpers.py index 02350d8606..de9850de36 100644 --- a/model/common/src/icon4py/model/common/math/helpers.py +++ b/model/common/src/icon4py/model/common/math/helpers.py @@ -5,657 +5,3 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause - -from gt4py import next as gtx -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 -from icon4py.model.common.type_alias import wpfloat - - -@gtx.field_operator -def average_level_plus1_on_cells( - half_level_field: fa.CellKField[wpfloat], -) -> fa.CellKField[wpfloat]: - """ - Calculate the mean value of adjacent interface levels. - - Computes the average of two adjacent interface levels upwards over a cell field for storage - in the corresponding full levels. - Args: - half_level_field: Field[Dims[CellDim, dims.KDim], wpfloat] - - Returns: Field[Dims[CellDim, dims.KDim], wpfloat] full level field - - """ - return 0.5 * (half_level_field + half_level_field(Koff[1])) - - -@gtx.field_operator -def average_level_plus1_on_edges( - half_level_field: fa.EdgeKField[wpfloat], -) -> fa.EdgeKField[wpfloat]: - """ - Calculate the mean value of adjacent interface levels. - - Computes the average of two adjacent interface levels upwards over an edge field for storage - in the corresponding full levels. - Args: - half_level_field: fa.EdgeKField[wpfloat] - - Returns: fa.EdgeKField[wpfloat] full level field - - """ - return 0.5 * (half_level_field + half_level_field(Koff[1])) - - -@gtx.field_operator -def difference_level_plus1_on_cells( - half_level_field: fa.CellKField[wpfloat], -) -> fa.CellKField[wpfloat]: - """ - Calculate the difference value of adjacent interface levels. - - Computes the difference of two adjacent interface levels upwards over a cell field for storage - in the corresponding full levels. - Args: - half_level_field: Field[Dims[CellDim, dims.KDim], wpfloat] - - Returns: Field[Dims[CellDim, dims.KDim], wpfloat] full level field - - """ - return half_level_field - half_level_field(Koff[1]) - - -@gtx.field_operator -def grad_fd_norm( - psi_c: fa.CellKField[float], - inv_dual_edge_length: fa.EdgeField[float], -) -> fa.EdgeKField[float]: - """ - Calculate the gradient value of adjacent interface levels. - - Computes the difference of two offseted values multiplied by a field of the offseted dimension - Args: - psi_c: fa.CellKField[float], - inv_dual_edge_length: Field[Dims[EdgeDim], float], - - Returns: fa.EdgeKField[float] - - """ - grad_norm_psi_e = (psi_c(E2C[1]) - psi_c(E2C[0])) * inv_dual_edge_length - return grad_norm_psi_e - - -@gtx.field_operator -def _grad_fd_tang( - psi_v: gtx.Field[gtx.Dims[dims.VertexDim, dims.KDim], float], - inv_primal_edge_length: fa.EdgeField[float], - tangent_orientation: fa.EdgeField[float], -) -> fa.EdgeKField[float]: - grad_tang_psi_e = tangent_orientation * (psi_v(E2V[1]) - psi_v(E2V[0])) * inv_primal_edge_length - return grad_tang_psi_e - - -@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) -def geographical_to_cartesian_on_cells( - lat: fa.CellField[ta.wpfloat], lon: fa.CellField[ta.wpfloat] -) -> tuple[fa.CellField[ta.wpfloat], fa.CellField[ta.wpfloat], fa.CellField[ta.wpfloat]]: - """ - Convert geographical (lat, lon) coordinates to cartesian coordinates on the unit sphere. - - Args: - lat: latitude - lon: longitude - - Returns: - x: x coordinate - y: y coordinate - z: z coordinate - - """ - x = cos(lat) * cos(lon) - y = cos(lat) * sin(lon) - z = sin(lat) - return x, y, z - - -@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) -def geographical_to_cartesian_on_edges( - lat: fa.EdgeField[ta.wpfloat], lon: fa.EdgeField[ta.wpfloat] -) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: - """ - Convert geographical (lat, lon) coordinates to cartesian coordinates on the unit sphere. - - Args: - lat: latitude - lon: longitude - - Returns: - x: x coordinate - y: y coordinate - z: z coordinate - - """ - x = cos(lat) * cos(lon) - y = cos(lat) * sin(lon) - z = sin(lat) - return x, y, z - - -@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) -def geographical_to_cartesian_on_vertices( - lat: fa.VertexField[ta.wpfloat], lon: fa.VertexField[ta.wpfloat] -) -> tuple[fa.VertexField[ta.wpfloat], fa.VertexField[ta.wpfloat], fa.VertexField[ta.wpfloat]]: - """ - Convert geographical (lat, lon) coordinates to cartesian coordinates on the unit sphere. - - Args: - lat: latitude - lon: longitude - - Returns: - x: x coordinate - y: y coordinate - z: z coordinate - - """ - x = cos(lat) * cos(lon) - y = cos(lat) * sin(lon) - z = sin(lat) - return x, y, z - - -@gtx.field_operator -def dot_product_on_edges( - x1: fa.EdgeField[ta.wpfloat], - x2: fa.EdgeField[ta.wpfloat], - y1: fa.EdgeField[ta.wpfloat], - y2: fa.EdgeField[ta.wpfloat], - z1: fa.EdgeField[ta.wpfloat], - z2: fa.EdgeField[ta.wpfloat], -) -> fa.EdgeField[ta.wpfloat]: - """Compute dot product of cartesian vectors (x1, y1, z1) * (x2, y2, z2)""" - return x1 * x2 + y1 * y2 + z1 * z2 - - -@gtx.field_operator -def dot_product_on_cells( - x1: fa.CellField[ta.wpfloat], - x2: fa.CellField[ta.wpfloat], - y1: fa.CellField[ta.wpfloat], - y2: fa.CellField[ta.wpfloat], - z1: fa.CellField[ta.wpfloat], - z2: fa.CellField[ta.wpfloat], -) -> fa.CellField[ta.wpfloat]: - """Compute dot product of cartesian vectors (x1, y1, z1) * (x2, y2, z2)""" - return x1 * x2 + y1 * y2 + z1 * z2 - - -@gtx.field_operator -def dot_product_on_vertices( - x1: fa.VertexField[ta.wpfloat], - x2: fa.VertexField[ta.wpfloat], - y1: fa.VertexField[ta.wpfloat], - y2: fa.VertexField[ta.wpfloat], - z1: fa.VertexField[ta.wpfloat], - z2: fa.VertexField[ta.wpfloat], -) -> fa.VertexField[ta.wpfloat]: - """Compute dot product of cartesian vectors (x1, y1, z1) * (x2, y2, z2)""" - return x1 * x2 + y1 * y2 + z1 * z2 - - -@gtx.field_operator -def cross_product_on_edges( - x1: fa.EdgeField[ta.wpfloat], - x2: fa.EdgeField[ta.wpfloat], - y1: fa.EdgeField[ta.wpfloat], - y2: fa.EdgeField[ta.wpfloat], - z1: fa.EdgeField[ta.wpfloat], - z2: fa.EdgeField[ta.wpfloat], -) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: - """Compute cross product of cartesian vectors (x1, y1, z1) x (x2, y2, z2)""" - x = y1 * z2 - z1 * y2 - y = z1 * x2 - x1 * z2 - z = x1 * y2 - y1 * x2 - return x, y, z - - -@gtx.field_operator -def norm2_on_edges( - x: fa.EdgeField[ta.wpfloat], y: fa.EdgeField[ta.wpfloat], z: fa.EdgeField[ta.wpfloat] -) -> fa.EdgeField[ta.wpfloat]: - """ - Compute 2 norm of a cartesian vector (x, y, z) - Args: - x: x coordinate - y: y coordinate - z: z coordinate - - Returns: - norma - - """ - return sqrt(dot_product_on_edges(x, x, y, y, z, z)) - - -@gtx.field_operator -def norm2_on_cells( - x: fa.CellField[ta.wpfloat], y: fa.CellField[ta.wpfloat], z: fa.CellField[ta.wpfloat] -) -> fa.CellField[ta.wpfloat]: - """ - Compute 2 norm of a cartesian vector (x, y, z) - Args: - x: x coordinate - y: y coordinate - z: z coordinate - - Returns: - norma - - """ - return sqrt(dot_product_on_cells(x, x, y, y, z, z)) - - -@gtx.field_operator -def norm2_on_vertices( - x: fa.VertexField[ta.wpfloat], y: fa.VertexField[ta.wpfloat], z: fa.VertexField[ta.wpfloat] -) -> fa.VertexField[ta.wpfloat]: - """ - Compute 2 norm of a cartesian vector (x, y, z) - Args: - x: x coordinate - y: y coordinate - z: z coordinate - - Returns: - norma - - """ - return sqrt(dot_product_on_vertices(x, x, y, y, z, z)) - - -@gtx.field_operator -def normalize_cartesian_vector_on_edges( - v_x: fa.EdgeField[ta.wpfloat], v_y: fa.EdgeField[ta.wpfloat], v_z: fa.EdgeField[ta.wpfloat] -) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: - """ - Normalize a cartesian vector. - - Args: - v_x: x coordinate - v_y: y coordinate - v_z: z coordinate - - Returns: - normalized vector - - """ - norm = norm2_on_edges(v_x, v_y, v_z) - return v_x / norm, v_y / norm, v_z / norm - - -@gtx.field_operator -def invert_edge_field(f: fa.EdgeField[ta.wpfloat]) -> fa.EdgeField[ta.wpfloat]: - """ - Invert values. - Args: - f: values - - Returns: - 1/f where f is not zero. - """ - return where(f != 0.0, 1.0 / f, f) - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def compute_inverse_on_edges( - f: fa.EdgeField[ta.wpfloat], - f_inverse: fa.EdgeField[ta.wpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, -): - invert_edge_field(f, out=f_inverse, domain={dims.EdgeDim: (horizontal_start, horizontal_end)}) - - -@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) -def zonal_and_meridional_components_on_cells( - lat: fa.CellField[ta.wpfloat], - lon: fa.CellField[ta.wpfloat], - x: fa.CellField[ta.wpfloat], - y: fa.CellField[ta.wpfloat], - z: fa.CellField[ta.wpfloat], -) -> tuple[fa.CellField[ta.wpfloat], fa.CellField[ta.wpfloat]]: - """ - Compute normalized zonal and meridional components of a cartesian vector (x, y, z) at point (lat, lon) - - Args: - lat: latitude - lon: longitude - x: x coordinate - y: y coordinate - z: z coordinate - - Returns: - u zonal component - v meridional component - - """ - cos_lat = cos(lat) - sin_lat = sin(lat) - cos_lon = cos(lon) - sin_lon = sin(lon) - u = cos_lon * y - sin_lon * x - - v = cos_lat * z - sin_lat * (cos_lon * x + sin_lon * y) - norm = sqrt(u * u + v * v) - return u / norm, v / norm - - -@gtx.field_operator -def zonal_and_meridional_components_on_edges( - lat: fa.EdgeField[ta.wpfloat], - lon: fa.EdgeField[ta.wpfloat], - x: fa.EdgeField[ta.wpfloat], - y: fa.EdgeField[ta.wpfloat], - z: fa.EdgeField[ta.wpfloat], -) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: - """ - Compute the zonal and meridional component of a vector (x, y, z) at position (lat, lon) - - Args: - lat: latitude - lon: longitude - x: x component of cartesian vector - y: y component of cartesian vector - z: z component of cartesian vector - - Returns: - zonal (eastward) component of (x, y, z) at (lat, lon) - meridional (northward) component of (x, y, z) at (lat, lon) - - """ - cos_lat = cos(lat) - sin_lat = sin(lat) - cos_lon = cos(lon) - sin_lon = sin(lon) - u = cos_lon * y - sin_lon * x - - v = cos_lat * z - sin_lat * (cos_lon * x + sin_lon * y) - norm = sqrt(u * u + v * v) - return u / norm, v / norm - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def compute_zonal_and_meridional_components_on_edges( - lat: fa.EdgeField[ta.wpfloat], - lon: fa.EdgeField[ta.wpfloat], - x: fa.EdgeField[ta.wpfloat], - y: fa.EdgeField[ta.wpfloat], - z: fa.EdgeField[ta.wpfloat], - u: fa.EdgeField[ta.wpfloat], - v: fa.EdgeField[ta.wpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, -): - zonal_and_meridional_components_on_edges( - lat, lon, x, y, z, out=(u, v), domain={dims.EdgeDim: (horizontal_start, horizontal_end)} - ) - - -@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) -def cartesian_coordinates_from_zonal_and_meridional_components_on_edges( - lat: fa.EdgeField[ta.wpfloat], - lon: fa.EdgeField[ta.wpfloat], - u: fa.EdgeField[ta.wpfloat], - v: fa.EdgeField[ta.wpfloat], -) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: - """ - Compute cartesian coordinates from zonal an meridional components at position (lat, lon) - Args: - lat: latitude - lon: longitude - u: zonal component - v: meridional component - - Returns: - x, y, z cartesian components - - """ - cos_lat = cos(lat) - sin_lat = sin(lat) - cos_lon = cos(lon) - sin_lon = sin(lon) - - x = -u * sin_lon - v * sin_lat * cos_lon - y = u * cos_lon - v * sin_lat * sin_lon - z = cos_lat * v - - norm = norm2_on_edges(x, y, z) - return x / norm, y / norm, z / norm - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def compute_cartesian_coordinates_from_zonal_and_meridional_components_on_edges( - edge_lat: fa.EdgeField[ta.wpfloat], - edge_lon: fa.EdgeField[ta.wpfloat], - u: fa.EdgeField[ta.wpfloat], - v: fa.EdgeField[ta.wpfloat], - x: fa.EdgeField[ta.wpfloat], - y: fa.EdgeField[ta.wpfloat], - z: fa.EdgeField[ta.wpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, -): - cartesian_coordinates_from_zonal_and_meridional_components_on_edges( - edge_lat, - edge_lon, - u, - v, - out=(x, y, z), - domain={dims.EdgeDim: (horizontal_start, horizontal_end)}, - ) - - -@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) -def cartesian_coordinates_from_zonal_and_meridional_components_on_cells( - lat: fa.CellField[ta.wpfloat], - lon: fa.CellField[ta.wpfloat], - u: fa.CellField[ta.wpfloat], - v: fa.CellField[ta.wpfloat], -) -> tuple[fa.CellField[ta.wpfloat], fa.CellField[ta.wpfloat], fa.CellField[ta.wpfloat]]: - """ - Compute cartesian coordinates form zonal an meridonal components at position (lat, lon) - Args: - lat: latitude - lon: longitude - u: zonal component - v: meridional component - - Returns: - x, y, z cartesian components - - """ - cos_lat = cos(lat) - sin_lat = sin(lat) - cos_lon = cos(lon) - sin_lon = sin(lon) - - x = -u * sin_lon - v * sin_lat * cos_lon - y = u * cos_lon - v * sin_lat * sin_lon - z = cos_lat * v - - norm = norm2_on_cells(x, y, z) - return x / norm, y / norm, z / norm - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def compute_cartesian_coordinates_from_zonal_and_meridional_components_on_cells( - cell_lat: fa.CellField[ta.wpfloat], - cell_lon: fa.CellField[ta.wpfloat], - u: fa.CellField[ta.wpfloat], - v: fa.CellField[ta.wpfloat], - x: fa.CellField[ta.wpfloat], - y: fa.CellField[ta.wpfloat], - z: fa.CellField[ta.wpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, -): - cartesian_coordinates_from_zonal_and_meridional_components_on_cells( - cell_lat, - cell_lon, - u, - v, - out=(x, y, z), - domain={dims.CellDim: (horizontal_start, horizontal_end)}, - ) - - -@gtx.field_operator -def arc_length_on_edges( - x0: fa.EdgeField[ta.wpfloat], - x1: fa.EdgeField[ta.wpfloat], - y0: fa.EdgeField[ta.wpfloat], - y1: fa.EdgeField[ta.wpfloat], - z0: fa.EdgeField[ta.wpfloat], - z1: fa.EdgeField[ta.wpfloat], - radius: ta.wpfloat, -): - """ - Compute the arc length between two points on the sphere. - - Inputs are cartesian coordinates of the points. - - 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 - z0: z coordinate of point_0 - z1: z coordinate of point_1 - radius: sphere radius - - Returns: - arc length - - """ - return radius * arccos(dot_product_on_edges(x0, x1, y0, y1, z0, z1)) - - -@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) -def diff_on_edges_torus( - x0: fa.EdgeField[ta.wpfloat], - x1: fa.EdgeField[ta.wpfloat], - y0: fa.EdgeField[ta.wpfloat], - y1: fa.EdgeField[ta.wpfloat], - domain_length: ta.wpfloat, - domain_height: ta.wpfloat, -) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: - """ - Compute the difference between two points on the torus. - - Inputs are cartesian coordinates of the points. Z is assumed zero and is - ignored. Distances are computed modulo the domain length and height. - - Args: - x0: x coordinate of point_0 - x1: x coordinate of point_1 - y0: y coordinate of point_0 - y1: y coordinate of point_1 - domain_length: length of the domain - domain_height: height of the domain - - Returns: - dx, dy - - """ - x1 = where( - abs(x1 - x0) <= 0.5 * domain_length, - x1, - where(x0 > x1, x1 + domain_length, x1 - domain_length), - ) - - y1 = where( - abs(y1 - y0) <= 0.5 * domain_height, - y1, - where(y0 > y1, y1 + domain_height, y1 - domain_height), - ) - - return (x1 - x0, y1 - y0) - - -@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) -def distance_on_edges_torus( - x0: fa.EdgeField[ta.wpfloat], - x1: fa.EdgeField[ta.wpfloat], - y0: fa.EdgeField[ta.wpfloat], - y1: fa.EdgeField[ta.wpfloat], - domain_length: ta.wpfloat, - domain_height: ta.wpfloat, -) -> fa.EdgeField[ta.wpfloat]: - """ - Compute the distance between two points on the torus. - - Inputs are cartesian coordinates of the points. Z is assumed zero and is - ignored. Distances are computed modulo the domain length and height. - - Args: - x0: x coordinate of point_0 - x1: x coordinate of point_1 - y0: y coordinate of point_0 - y1: y coordinate of point_1 - domain_length: length of the domain - domain_height: height of the domain - - Returns: - distance - - """ - xdiff, ydiff = diff_on_edges_torus(x0, x1, y0, y1, domain_length, domain_height) - return sqrt(xdiff**2 + ydiff**2) - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def average_two_vertical_levels_downwards_on_edges( - input_field: fa.EdgeKField[wpfloat], - average: fa.EdgeKField[wpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -): - average_level_plus1_on_edges( - input_field, - out=average, - domain={ - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end), - }, - ) - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def average_two_vertical_levels_downwards_on_cells( - input_field: fa.CellKField[wpfloat], - average: fa.CellKField[wpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -) -> None: - average_level_plus1_on_cells( - input_field, - out=average, - domain={ - dims.CellDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/common/src/icon4py/model/common/math/math_utils.py b/model/common/src/icon4py/model/common/math/math_utils.py new file mode 100644 index 0000000000..353ebb8784 --- /dev/null +++ b/model/common/src/icon4py/model/common/math/math_utils.py @@ -0,0 +1,28 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +""" +General-purpose mathematical utility functions. + +Contains typed wrappers around standard math operations for use in factories +and validation. +""" + +import math + +import numpy as np + + +def compute_sqrt( + input_val: np.float64, +) -> np.float64: + """ + Compute the square root of input_val. + math.sqrt is not sufficiently typed for the validation happening in the factories. + """ + return np.float64(math.sqrt(input_val)) diff --git a/model/common/src/icon4py/model/common/math/vector_operations.py b/model/common/src/icon4py/model/common/math/vector_operations.py new file mode 100644 index 0000000000..0d16be072d --- /dev/null +++ b/model/common/src/icon4py/model/common/math/vector_operations.py @@ -0,0 +1,171 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +""" +Vector operations on unstructured grid fields. + +Contains dot product, cross product, norm, normalization and inversion operations +for vectors defined on cell, edge and vertex fields. +""" + +from gt4py import next as gtx +from gt4py.next import sqrt, where + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta + + +@gtx.field_operator +def dot_product_on_edges( + x1: fa.EdgeField[ta.wpfloat], + x2: fa.EdgeField[ta.wpfloat], + y1: fa.EdgeField[ta.wpfloat], + y2: fa.EdgeField[ta.wpfloat], + z1: fa.EdgeField[ta.wpfloat], + z2: fa.EdgeField[ta.wpfloat], +) -> fa.EdgeField[ta.wpfloat]: + """Compute dot product of cartesian vectors (x1, y1, z1) * (x2, y2, z2)""" + return x1 * x2 + y1 * y2 + z1 * z2 + + +@gtx.field_operator +def dot_product_on_cells( + x1: fa.CellField[ta.wpfloat], + x2: fa.CellField[ta.wpfloat], + y1: fa.CellField[ta.wpfloat], + y2: fa.CellField[ta.wpfloat], + z1: fa.CellField[ta.wpfloat], + z2: fa.CellField[ta.wpfloat], +) -> fa.CellField[ta.wpfloat]: + """Compute dot product of cartesian vectors (x1, y1, z1) * (x2, y2, z2)""" + return x1 * x2 + y1 * y2 + z1 * z2 + + +@gtx.field_operator +def dot_product_on_vertices( + x1: fa.VertexField[ta.wpfloat], + x2: fa.VertexField[ta.wpfloat], + y1: fa.VertexField[ta.wpfloat], + y2: fa.VertexField[ta.wpfloat], + z1: fa.VertexField[ta.wpfloat], + z2: fa.VertexField[ta.wpfloat], +) -> fa.VertexField[ta.wpfloat]: + """Compute dot product of cartesian vectors (x1, y1, z1) * (x2, y2, z2)""" + return x1 * x2 + y1 * y2 + z1 * z2 + + +@gtx.field_operator +def cross_product_on_edges( + x1: fa.EdgeField[ta.wpfloat], + x2: fa.EdgeField[ta.wpfloat], + y1: fa.EdgeField[ta.wpfloat], + y2: fa.EdgeField[ta.wpfloat], + z1: fa.EdgeField[ta.wpfloat], + z2: fa.EdgeField[ta.wpfloat], +) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: + """Compute cross product of cartesian vectors (x1, y1, z1) x (x2, y2, z2)""" + x = y1 * z2 - z1 * y2 + y = z1 * x2 - x1 * z2 + z = x1 * y2 - y1 * x2 + return x, y, z + + +@gtx.field_operator +def norm2_on_edges( + x: fa.EdgeField[ta.wpfloat], y: fa.EdgeField[ta.wpfloat], z: fa.EdgeField[ta.wpfloat] +) -> fa.EdgeField[ta.wpfloat]: + """ + Compute 2 norm of a cartesian vector (x, y, z) + Args: + x: x coordinate + y: y coordinate + z: z coordinate + + Returns: + norma + + """ + return sqrt(dot_product_on_edges(x, x, y, y, z, z)) + + +@gtx.field_operator +def norm2_on_cells( + x: fa.CellField[ta.wpfloat], y: fa.CellField[ta.wpfloat], z: fa.CellField[ta.wpfloat] +) -> fa.CellField[ta.wpfloat]: + """ + Compute 2 norm of a cartesian vector (x, y, z) + Args: + x: x coordinate + y: y coordinate + z: z coordinate + + Returns: + norma + + """ + return sqrt(dot_product_on_cells(x, x, y, y, z, z)) + + +@gtx.field_operator +def norm2_on_vertices( + x: fa.VertexField[ta.wpfloat], y: fa.VertexField[ta.wpfloat], z: fa.VertexField[ta.wpfloat] +) -> fa.VertexField[ta.wpfloat]: + """ + Compute 2 norm of a cartesian vector (x, y, z) + Args: + x: x coordinate + y: y coordinate + z: z coordinate + + Returns: + norma + + """ + return sqrt(dot_product_on_vertices(x, x, y, y, z, z)) + + +@gtx.field_operator +def normalize_cartesian_vector_on_edges( + v_x: fa.EdgeField[ta.wpfloat], v_y: fa.EdgeField[ta.wpfloat], v_z: fa.EdgeField[ta.wpfloat] +) -> tuple[fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat], fa.EdgeField[ta.wpfloat]]: + """ + Normalize a cartesian vector. + + Args: + v_x: x coordinate + v_y: y coordinate + v_z: z coordinate + + Returns: + normalized vector + + """ + norm = norm2_on_edges(v_x, v_y, v_z) + return v_x / norm, v_y / norm, v_z / norm + + +@gtx.field_operator +def invert_edge_field(f: fa.EdgeField[ta.wpfloat]) -> fa.EdgeField[ta.wpfloat]: + """ + Invert values. + Args: + f: values + + Returns: + 1/f where f is not zero. + """ + return where(f != 0.0, 1.0 / f, f) + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_inverse_on_edges( + f: fa.EdgeField[ta.wpfloat], + f_inverse: fa.EdgeField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, +): + invert_edge_field(f, out=f_inverse, domain={dims.EdgeDim: (horizontal_start, horizontal_end)}) diff --git a/model/common/src/icon4py/model/common/math/vertical_operations.py b/model/common/src/icon4py/model/common/math/vertical_operations.py new file mode 100644 index 0000000000..4a24d5d8fd --- /dev/null +++ b/model/common/src/icon4py/model/common/math/vertical_operations.py @@ -0,0 +1,112 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +""" +Vertical level operations on unstructured grid fields. + +Contains averaging and difference operations between adjacent vertical levels +on cell and edge fields. +""" + +from gt4py import next as gtx + +from icon4py.model.common import dimension as dims, field_type_aliases as fa +from icon4py.model.common.dimension import Koff +from icon4py.model.common.type_alias import wpfloat + + +@gtx.field_operator +def average_level_plus1_on_cells( + half_level_field: fa.CellKField[wpfloat], +) -> fa.CellKField[wpfloat]: + """ + Calculate the mean value of adjacent interface levels. + + Computes the average of two adjacent interface levels upwards over a cell field for storage + in the corresponding full levels. + Args: + half_level_field: Field[Dims[CellDim, dims.KDim], wpfloat] + + Returns: Field[Dims[CellDim, dims.KDim], wpfloat] full level field + + """ + return 0.5 * (half_level_field + half_level_field(Koff[1])) + + +@gtx.field_operator +def average_level_plus1_on_edges( + half_level_field: fa.EdgeKField[wpfloat], +) -> fa.EdgeKField[wpfloat]: + """ + Calculate the mean value of adjacent interface levels. + + Computes the average of two adjacent interface levels upwards over an edge field for storage + in the corresponding full levels. + Args: + half_level_field: fa.EdgeKField[wpfloat] + + Returns: fa.EdgeKField[wpfloat] full level field + + """ + return 0.5 * (half_level_field + half_level_field(Koff[1])) + + +@gtx.field_operator +def difference_level_plus1_on_cells( + half_level_field: fa.CellKField[wpfloat], +) -> fa.CellKField[wpfloat]: + """ + Calculate the difference value of adjacent interface levels. + + Computes the difference of two adjacent interface levels upwards over a cell field for storage + in the corresponding full levels. + Args: + half_level_field: Field[Dims[CellDim, dims.KDim], wpfloat] + + Returns: Field[Dims[CellDim, dims.KDim], wpfloat] full level field + + """ + return half_level_field - half_level_field(Koff[1]) + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def average_two_vertical_levels_downwards_on_edges( + input_field: fa.EdgeKField[wpfloat], + average: fa.EdgeKField[wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + average_level_plus1_on_edges( + input_field, + out=average, + domain={ + dims.EdgeDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def average_two_vertical_levels_downwards_on_cells( + input_field: fa.CellKField[wpfloat], + average: fa.CellKField[wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +) -> None: + average_level_plus1_on_cells( + input_field, + out=average, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/common/src/icon4py/model/common/math/xp_utils.py b/model/common/src/icon4py/model/common/math/xp_utils.py index 8968875984..de9850de36 100644 --- a/model/common/src/icon4py/model/common/math/xp_utils.py +++ b/model/common/src/icon4py/model/common/math/xp_utils.py @@ -5,17 +5,3 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause - -import math - -import numpy as np - - -def compute_sqrt( - input_val: np.float64, -) -> np.float64: - """ - Compute the square root of input_val. - math.sqrt is not sufficiently typed for the validation happening in the factories. - """ - return np.float64(math.sqrt(input_val)) diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index fa7c9a6ff9..9c0e63f5a0 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -40,11 +40,8 @@ from icon4py.model.common.interpolation.stencils.compute_cell_2_vertex_interpolation import ( _compute_cell_2_vertex_interpolation, ) -from icon4py.model.common.math.helpers import ( - _grad_fd_tang, - difference_level_plus1_on_cells, - grad_fd_norm, -) +from icon4py.model.common.math.gradient import _grad_fd_tang, grad_fd_norm +from icon4py.model.common.math.vertical_operations import difference_level_plus1_on_cells from icon4py.model.common.type_alias import vpfloat, wpfloat from icon4py.model.common.utils import data_allocation as data_alloc diff --git a/model/common/src/icon4py/model/common/metrics/metrics_factory.py b/model/common/src/icon4py/model/common/metrics/metrics_factory.py index 8a6ed38dda..48d9711ed3 100644 --- a/model/common/src/icon4py/model/common/metrics/metrics_factory.py +++ b/model/common/src/icon4py/model/common/metrics/metrics_factory.py @@ -12,7 +12,7 @@ import gt4py.next as gtx import gt4py.next.typing as gtx_typing -import icon4py.model.common.math.helpers as math_helpers +from icon4py.model.common.math import vertical_operations as vertical_ops import icon4py.model.common.metrics.compute_weight_factors as weight_factors from icon4py.model.common import ( constants, @@ -187,7 +187,7 @@ def _register_computed_fields(self) -> None: # noqa: PLR0915 [too-many-statemen self.register_provider(vertical_coordinates_on_half_levels) height = factory.ProgramFieldProvider( - func=math_helpers.average_two_vertical_levels_downwards_on_cells.with_backend( + func=vertical_ops.average_two_vertical_levels_downwards_on_cells.with_backend( self._backend ), domain={ @@ -499,7 +499,7 @@ def _register_computed_fields(self) -> None: # noqa: PLR0915 [too-many-statemen # ddxn_z_full is dependent only on attrs.DDXN_Z_HALF_E, which has halo exchange. That's why halo_exchange is set to True compute_ddxn_z_full = factory.ProgramFieldProvider( - func=math_helpers.average_two_vertical_levels_downwards_on_edges.with_backend( + func=vertical_ops.average_two_vertical_levels_downwards_on_edges.with_backend( self._backend ), deps={ @@ -521,7 +521,7 @@ def _register_computed_fields(self) -> None: # noqa: PLR0915 [too-many-statemen self.register_provider(compute_ddxn_z_full) compute_ddxt_z_full = factory.ProgramFieldProvider( - func=math_helpers.average_two_vertical_levels_downwards_on_edges.with_backend( + func=vertical_ops.average_two_vertical_levels_downwards_on_edges.with_backend( self._backend ), deps={ diff --git a/model/common/tests/common/grid/mpi_tests/test_parallel_geometry.py b/model/common/tests/common/grid/mpi_tests/test_parallel_geometry.py index a063fc33be..acd734ff8b 100644 --- a/model/common/tests/common/grid/mpi_tests/test_parallel_geometry.py +++ b/model/common/tests/common/grid/mpi_tests/test_parallel_geometry.py @@ -23,7 +23,7 @@ geometry_attributes as attrs, horizontal as h_grid, ) -from icon4py.model.common.math import helpers as math_helpers +from icon4py.model.common.math import vector_operations as vector_ops from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import definitions as test_defs, parallel_helpers, test_utils @@ -194,7 +194,7 @@ def test_cartesian_geometry_attr_no_halos( norm = data_alloc.zero_field( grid_geometry.grid, dimension, dtype=x_field.dtype, allocator=backend ) - math_helpers.norm2_on_vertices(x_field, z_field, y_field, out=norm, offset_provider={}) + vector_ops.norm2_on_vertices(x_field, z_field, y_field, out=norm, offset_provider={}) assert test_utils.dallclose(norm.asnumpy(), 1.0) case base.GeometryType.TORUS: # on a torus coordinates should be within the domain 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 f918de9210..0334473559 100644 --- a/model/common/tests/common/grid/unit_tests/test_geometry.py +++ b/model/common/tests/common/grid/unit_tests/test_geometry.py @@ -22,7 +22,7 @@ simple, ) from icon4py.model.common.grid.geometry import as_sparse_field -from icon4py.model.common.math import helpers as math_helpers +from icon4py.model.common.math import vector_operations as vector_ops from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import definitions, grid_utils, test_utils from icon4py.model.testing.fixtures import ( @@ -346,7 +346,7 @@ def test_cartesian_centers_edge( 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={}) + vector_ops.norm2_on_edges(x, z, y, out=norm, offset_provider={}) assert test_utils.dallclose(norm.asnumpy(), 1.0) case base.GeometryType.TORUS: # on a torus coordinates should be within the domain @@ -385,7 +385,7 @@ def test_cartesian_centers_cell( 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={}) + vector_ops.norm2_on_cells(x, z, y, out=norm, offset_provider={}) assert test_utils.dallclose(norm.asnumpy(), 1.0) case base.GeometryType.TORUS: # on a torus coordinates should be within the domain @@ -424,7 +424,7 @@ def test_vertex( 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={}) + vector_ops.norm2_on_vertices(x, z, y, out=norm, offset_provider={}) assert test_utils.dallclose(norm.asnumpy(), 1.0) case base.GeometryType.TORUS: # on a torus coordinates should be within the domain diff --git a/model/common/tests/common/math/stencil_tests/test_cell_horizontal_gradients_by_green_gauss_method.py b/model/common/tests/common/math/stencil_tests/test_cell_horizontal_gradients_by_green_gauss_method.py index 6c4ae32a04..21a98e8bab 100644 --- a/model/common/tests/common/math/stencil_tests/test_cell_horizontal_gradients_by_green_gauss_method.py +++ b/model/common/tests/common/math/stencil_tests/test_cell_horizontal_gradients_by_green_gauss_method.py @@ -13,8 +13,8 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.math.stencils.cell_horizontal_gradients_by_green_gauss_method import ( - cell_horizontal_gradients_by_green_gauss_method, +from icon4py.model.common.math.stencils import ( + cell_horizontal_gradients_by_green_gauss_method as green_gauss, ) from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -45,7 +45,7 @@ def cell_horizontal_gradients_by_green_gauss_method_numpy( @pytest.mark.embedded_remap_error class TestMoMathGradientsGradGreenGaussCellDsl(StencilTest): - PROGRAM = cell_horizontal_gradients_by_green_gauss_method + PROGRAM = green_gauss.cell_horizontal_gradients_by_green_gauss_method OUTPUTS = ("out",) @staticmethod diff --git a/model/common/tests/common/math/stencil_tests/test_compute_first_vertical_derivative.py b/model/common/tests/common/math/stencil_tests/test_compute_first_vertical_derivative.py index 5cd46b15ff..651da0b0e5 100644 --- a/model/common/tests/common/math/stencil_tests/test_compute_first_vertical_derivative.py +++ b/model/common/tests/common/math/stencil_tests/test_compute_first_vertical_derivative.py @@ -13,7 +13,7 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.math.derivative import compute_first_vertical_derivative_at_cells +from icon4py.model.common.math import derivative from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat from icon4py.model.common.utils.data_allocation import random_field, zero_field @@ -28,7 +28,7 @@ def compute_first_vertical_derivative_numpy( class TestComputeFirstVerticalDerivative(StencilTest): - PROGRAM = compute_first_vertical_derivative_at_cells + PROGRAM = derivative.compute_first_vertical_derivative_at_cells OUTPUTS = ("first_vertical_derivative",) @staticmethod diff --git a/model/common/tests/common/math/unit_tests/test_helpers.py b/model/common/tests/common/math/unit_tests/test_helpers.py index e1de50194d..c02fd101e0 100644 --- a/model/common/tests/common/math/unit_tests/test_helpers.py +++ b/model/common/tests/common/math/unit_tests/test_helpers.py @@ -15,7 +15,7 @@ import icon4py.model.testing.test_utils from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, simple -from icon4py.model.common.math import helpers +from icon4py.model.common.math import vector_operations as vector_ops, vertical_operations as vertical_ops from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests from icon4py.model.testing.fixtures.datatest import backend, backend_like @@ -34,7 +34,7 @@ def test_cross_product(backend: gtx_typing.Backend) -> None: y = data_alloc.zero_field(mesh, dims.EdgeDim, allocator=backend) z = data_alloc.zero_field(mesh, dims.EdgeDim, allocator=backend) - helpers.cross_product_on_edges.with_backend(backend)( + vector_ops.cross_product_on_edges.with_backend(backend)( x1, x2, y1, y2, z1, z2, out=(x, y, z), offset_provider={} ) a = np.column_stack((x1.asnumpy(), y1.asnumpy(), z1.asnumpy())) @@ -47,7 +47,7 @@ def test_cross_product(backend: gtx_typing.Backend) -> None: class TestAverageTwoVerticalLevelsDownwardsOnEdges(stencil_tests.StencilTest): - PROGRAM = helpers.average_two_vertical_levels_downwards_on_edges + PROGRAM = vertical_ops.average_two_vertical_levels_downwards_on_edges OUTPUTS = ( stencil_tests.Output( "average", @@ -81,7 +81,7 @@ def input_data(self, grid: base.Grid) -> dict: class TestAverageTwoVerticalLevelsDownwardsOnCells(stencil_tests.StencilTest): - PROGRAM = helpers.average_two_vertical_levels_downwards_on_cells + PROGRAM = vertical_ops.average_two_vertical_levels_downwards_on_cells OUTPUTS = ( stencil_tests.Output( "average", diff --git a/model/common/tests/common/math/unit_tests/test_operators.py b/model/common/tests/common/math/unit_tests/test_operators.py index b487d1a50b..505635b783 100644 --- a/model/common/tests/common/math/unit_tests/test_operators.py +++ b/model/common/tests/common/math/unit_tests/test_operators.py @@ -13,8 +13,7 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, base as base_grid -from icon4py.model.common.math.stencils.compute_nabla2_on_cell import compute_nabla2_on_cell -from icon4py.model.common.math.stencils.compute_nabla2_on_cell_k import compute_nabla2_on_cell_k +from icon4py.model.common.math.stencils import compute_nabla2_on_cell, compute_nabla2_on_cell_k from icon4py.model.common.utils.data_allocation import constant_field, zero_field from icon4py.model.testing import reference_funcs from icon4py.model.testing.fixtures.datatest import backend_like @@ -24,7 +23,7 @@ @pytest.mark.embedded_remap_error class TestNabla2OnCell(StencilTest): - PROGRAM = compute_nabla2_on_cell + PROGRAM = compute_nabla2_on_cell.compute_nabla2_on_cell OUTPUTS = ("nabla2_psi_c",) @staticmethod @@ -53,7 +52,7 @@ def input_data(self, grid: base_grid.Grid) -> dict: @pytest.mark.embedded_remap_error class TestNabla2OnCellK(StencilTest): - PROGRAM = compute_nabla2_on_cell_k + PROGRAM = compute_nabla2_on_cell_k.compute_nabla2_on_cell_k OUTPUTS = ("nabla2_psi_c",) @staticmethod diff --git a/model/common/tests/common/math/unit_tests/test_smagorinsky.py b/model/common/tests/common/math/unit_tests/test_smagorinsky.py index 4c0e27e116..b99569a65d 100644 --- a/model/common/tests/common/math/unit_tests/test_smagorinsky.py +++ b/model/common/tests/common/math/unit_tests/test_smagorinsky.py @@ -11,7 +11,7 @@ from icon4py.model.common import dimension as dims, model_backends, model_options from icon4py.model.common.grid import base as base_grid -from icon4py.model.common.math.smagorinsky import en_smag_fac_for_zero_nshift +from icon4py.model.common.math import smagorinsky from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing.fixtures.datatest import backend_like from icon4py.model.testing.fixtures.stencil_tests import grid, grid_manager @@ -28,7 +28,7 @@ def test_init_enh_smag_fac(backend_like: model_backends.BackendLike, grid: base_ z = (0.1, 0.2, 0.3, 0.4) enhanced_smag_fac_np = enhanced_smagorinski_factor_numpy(fac, z, a_vec.asnumpy()) - en_smag_fac_for_zero_nshift.with_backend(backend)( + smagorinsky.en_smag_fac_for_zero_nshift.with_backend(backend)( a_vec, *fac, *z, diff --git a/model/common/tests/common/states/unit_tests/test_factory.py b/model/common/tests/common/states/unit_tests/test_factory.py index 12067e9818..4c96e5f71e 100644 --- a/model/common/tests/common/states/unit_tests/test_factory.py +++ b/model/common/tests/common/states/unit_tests/test_factory.py @@ -18,7 +18,10 @@ from icon4py.model.common import dimension as dims, utils as common_utils from icon4py.model.common.decomposition import definitions as decomposition from icon4py.model.common.grid import horizontal as h_grid, icon, simple, vertical as v_grid -from icon4py.model.common.math import helpers as math_helpers +from icon4py.model.common.math import ( + coordinate_transformations as coord_trans, + vertical_operations as vertical_ops, +) from icon4py.model.common.states import factory, model, utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import definitions, serialbox @@ -154,7 +157,7 @@ def height_coordinate_source( @pytest.mark.datatest def test_field_operator_provider(cell_coordinate_source: SimpleFieldSource) -> None: - field_op = math_helpers.geographical_to_cartesian_on_cells.with_backend(None) + field_op = coord_trans.geographical_to_cartesian_on_cells.with_backend(None) domain = {dims.CellDim: (cell_domain(h_grid.Zone.LOCAL), cell_domain(h_grid.Zone.LOCAL))} deps = {"lat": "lat", "lon": "lon"} @@ -177,7 +180,7 @@ def test_field_operator_provider(cell_coordinate_source: SimpleFieldSource) -> N @pytest.mark.datatest def test_program_provider(height_coordinate_source: SimpleFieldSource) -> None: - program = math_helpers.average_two_vertical_levels_downwards_on_cells + program = vertical_ops.average_two_vertical_levels_downwards_on_cells domain = { dims.CellDim: (cell_domain(h_grid.Zone.LOCAL), cell_domain(h_grid.Zone.LOCAL)), dims.KDim: (k_domain(v_grid.Zone.TOP), k_domain(v_grid.Zone.BOTTOM)), @@ -201,7 +204,7 @@ def test_program_provider(height_coordinate_source: SimpleFieldSource) -> None: @pytest.mark.datatest def test_field_source_raise_error_on_register(cell_coordinate_source: SimpleFieldSource) -> None: - program = math_helpers.average_two_vertical_levels_downwards_on_cells + program = vertical_ops.average_two_vertical_levels_downwards_on_cells domain = { dims.CellDim: (cell_domain(h_grid.Zone.LOCAL), cell_domain(h_grid.Zone.LOCAL)), dims.KDim: (k_domain(v_grid.Zone.TOP), k_domain(v_grid.Zone.BOTTOM)),