add some cell version of math stencils (#603)

* add some cell version of math stencils
fix wrong return value in gvec2cvec

model/common/src/icon4py/model/common/grid/geometry.py

@@ -429,7 +429,7 @@ def _inverse_field_provider(self, field_name: str):
         name = meta["standard_name"]
         self._attrs.update({name: meta})
         provider = factory.ProgramFieldProvider(
-            func=math_helpers.compute_inverse,
+            func=math_helpers.compute_inverse_on_edges,
             deps={"f": field_name},
             fields={"f_inverse": name},
             domain={

model/common/src/icon4py/model/common/grid/geometry_stencils.py

@@ -13,10 +13,10 @@
 from icon4py.model.common.dimension import E2C, E2C2V, E2V, EdgeDim
 from icon4py.model.common.math.helpers import (
     arc_length,
-    cross_product,
+    cross_product_on_edges,
     geographical_to_cartesian_on_edges,
     geographical_to_cartesian_on_vertex,
-    normalize_cartesian_vector,
+    normalize_cartesian_vector_on_edges,
     zonal_and_meridional_components_on_edges,
 )
 
@@ -49,7 +49,7 @@ def cartesian_coordinates_of_edge_tangent(
     y = edge_orientation * (vertex_y(E2V[1]) - vertex_y(E2V[0]))
     z = edge_orientation * (vertex_z(E2V[1]) - vertex_z(E2V[0]))
 
-    return normalize_cartesian_vector(x, y, z)
+    return normalize_cartesian_vector_on_edges(x, y, z)
 
 
 @gtx.field_operator
@@ -83,10 +83,10 @@ def cartesian_coordinates_of_edge_normal(
     edge_center_x, edge_center_y, edge_center_z = geographical_to_cartesian_on_edges(
         edge_lat, edge_lon
     )
-    x, y, z = cross_product(
+    x, y, z = cross_product_on_edges(
         edge_center_x, edge_tangent_x, edge_center_y, edge_tangent_y, edge_center_z, edge_tangent_z
     )
-    return normalize_cartesian_vector(x, y, z)
+    return normalize_cartesian_vector_on_edges(x, y, z)
 
 
 @gtx.field_operator

model/common/src/icon4py/model/common/math/helpers.py

@@ -187,7 +187,7 @@ def geographical_to_cartesian_on_vertex(
 
 
 @gtx.field_operator
-def dot_product(
+def dot_product_on_edges(
     x1: fa.EdgeField[ta.wpfloat],
     x2: fa.EdgeField[ta.wpfloat],
     y1: fa.EdgeField[ta.wpfloat],
@@ -200,7 +200,20 @@ def dot_product(
 
 
 @gtx.field_operator
-def cross_product(
+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 cross_product_on_edges(
     x1: fa.EdgeField[ta.wpfloat],
     x2: fa.EdgeField[ta.wpfloat],
     y1: fa.EdgeField[ta.wpfloat],
@@ -216,7 +229,7 @@ def cross_product(
 
 
 @gtx.field_operator
-def norm2(
+def norm2_on_edges(
     x: fa.EdgeField[ta.wpfloat], y: fa.EdgeField[ta.wpfloat], z: fa.EdgeField[ta.wpfloat]
 ) -> fa.EdgeField[ta.wpfloat]:
     """
@@ -230,11 +243,29 @@ def norm2(
         norma
 
     """
-    return sqrt(dot_product(x, x, y, y, z, z))
+    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 normalize_cartesian_vector(
+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]]:
     """
@@ -249,12 +280,12 @@ def normalize_cartesian_vector(
         normalized vector
 
     """
-    norm = norm2(v_x, v_y, v_z)
+    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(f: fa.EdgeField[ta.wpfloat]) -> fa.EdgeField[ta.wpfloat]:
+def invert_edge_field(f: fa.EdgeField[ta.wpfloat]) -> fa.EdgeField[ta.wpfloat]:
     """
     Invert values.
     Args:
@@ -267,13 +298,13 @@ def invert(f: fa.EdgeField[ta.wpfloat]) -> fa.EdgeField[ta.wpfloat]:
 
 
 @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
-def compute_inverse(
+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(f, out=f_inverse, domain={dims.EdgeDim: (horizontal_start, horizontal_end)})
+    invert_edge_field(f, out=f_inverse, domain={dims.EdgeDim: (horizontal_start, horizontal_end)})
 
 
 @gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED)
@@ -389,8 +420,8 @@ def cartesian_coordinates_from_zonal_and_meridional_components_on_edges(
     y = u * cos_lon - v * sin_lat * sin_lon
     z = cos_lat * v
 
-    norm = norm2(x, y, z)
-    return x / norm, y / norm, y / norm
+    norm = norm2_on_edges(x, y, z)
+    return x / norm, y / norm, z / norm
 
 
 @gtx.program
@@ -415,6 +446,60 @@ def compute_cartesian_coordinates_from_zonal_and_meridional_components_on_edges(
     )
 
 
+@gtx.field_operator
+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
+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(
     x0: fa.EdgeField[ta.wpfloat],
@@ -443,4 +528,4 @@ def arc_length(
         arc length
 
     """
-    return radius * arccos(dot_product(x0, x1, y0, y1, z0, z1))
+    return radius * arccos(dot_product_on_edges(x0, x1, y0, y1, z0, z1))

model/common/tests/math_tests/test_helpers.py

@@ -25,7 +25,7 @@ def test_cross_product(backend):
     y = test_helpers.zero_field(mesh, dims.EdgeDim)
     z = test_helpers.zero_field(mesh, dims.EdgeDim)
 
-    helpers.cross_product.with_backend(backend)(
+    helpers.cross_product_on_edges.with_backend(backend)(
         x1, x2, y1, y2, z1, z2, out=(x, y, z), offset_provider={}
     )
     a = xp.column_stack((x1.ndarray, y1.ndarray, z1.ndarray))