Better torus support

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"

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

@@ -17,12 +17,14 @@
 from icon4py.model.common.math.helpers import (
     arc_length_on_edges,
     cross_product_on_edges,
+    diff_on_edges_torus,
+    distance_on_edges_torus,
     geographical_to_cartesian_on_edges,
     geographical_to_cartesian_on_vertices,
     normalize_cartesian_vector_on_edges,
     zonal_and_meridional_components_on_edges,
 )
-from icon4py.model.common.utils import data_allocation as alloc
+from icon4py.model.common.utils import data_allocation as data_alloc
 
 
 @gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED)
@@ -56,6 +58,51 @@ def cartesian_coordinates_of_edge_tangent(
     return normalize_cartesian_vector_on_edges(x, y, z)
 
 
+@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED)
+def cartesian_coordinates_of_edge_tangent_torus(
+    vertex_x: fa.VertexField[ta.wpfloat],
+    vertex_y: fa.VertexField[ta.wpfloat],
+    edge_orientation: fa.EdgeField[ta.wpfloat],
+    domain_length: ta.wpfloat,
+    domain_height: ta.wpfloat,
+) -> tuple[
+    fa.EdgeField[ta.wpfloat],
+    fa.EdgeField[ta.wpfloat],
+    fa.EdgeField[ta.wpfloat],
+]:
+    """
+    Compute normalized cartesian vector tangential to an edge on a torus grid.
+
+    That is: computes the delta between the two vertices adjacent to the edge:
+    t = d(v1, v2)
+
+    Args:
+        vertex_x: x coordinates of vertices
+        vertex_y: y coordinates of vertices
+        edge_orientation: encoding of the edge orientation: (-1, +1) depending on whether the
+            edge is directed from first to second neighbor of vice versa.
+    Returns:
+          x: x coordinate of normalized tangent vector
+          y: y coordinate of normalized tangent vector
+          z: z coordinate of normalized tangent vector
+    """
+    xdiff, ydiff = diff_on_edges_torus(
+        vertex_x(E2V[0]),
+        vertex_x(E2V[1]),
+        vertex_y(E2V[0]),
+        vertex_y(E2V[1]),
+        domain_length,
+        domain_height,
+    )
+    x = edge_orientation * xdiff
+    y = edge_orientation * ydiff
+    z = 0.0 * x
+    # TODO(msimberg): This should use something like numpy.zeros_like if and
+    # when that becomes available in gt4py.
+
+    return normalize_cartesian_vector_on_edges(x, y, z)
+
+
 @gtx.field_operator
 def cartesian_coordinates_of_edge_normal(
     edge_lat: fa.EdgeField[ta.wpfloat],
@@ -93,6 +140,32 @@ def cartesian_coordinates_of_edge_normal(
     return normalize_cartesian_vector_on_edges(x, y, z)
 
 
+@gtx.field_operator
+def cartesian_coordinates_of_edge_normal_torus(
+    edge_tangent_x: fa.EdgeField[ta.wpfloat],
+    edge_tangent_y: fa.EdgeField[ta.wpfloat],
+) -> tuple[
+    fa.EdgeField[ta.wpfloat],
+    fa.EdgeField[ta.wpfloat],
+    fa.EdgeField[ta.wpfloat],
+]:
+    """
+    Compute the normal to the edge tangent vector on a torus grid.
+
+    Args:
+        edge_tangent_x: x coordinate of the tangent
+        edge_tangent_y: y coordinate of the tangent
+    Returns:
+        edge_normal_x: x coordinate of the normal
+        edge_normal_y: y coordinate of the normal
+        edge_normal_z: y coordinate of the normal
+    """
+    z = 0.0 * edge_tangent_x
+    # TODO(msimberg): This should use something like numpy.zeros_like if and
+    # when that becomes available in gt4py.
+    return normalize_cartesian_vector_on_edges(-edge_tangent_y, edge_tangent_x, z)
+
+
 @gtx.field_operator
 def cartesian_coordinates_edge_tangent_and_normal(
     vertex_lat: fa.VertexField[ta.wpfloat],
@@ -123,6 +196,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 +276,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 +582,38 @@ def arc_distance_of_far_edges_in_diamond(
     return far_vertex_vertex_length
 
 
+@gtx.field_operator
+def distance_of_far_edges_in_diamond_torus(
+    vertex_x: fa.VertexField[ta.wpfloat],
+    vertex_y: fa.VertexField[ta.wpfloat],
+    domain_length: ta.wpfloat,
+    domain_height: ta.wpfloat,
+) -> fa.EdgeField[ta.wpfloat]:
+    """
+    Compute the distance between the "far" vertices of an edge on a torus grid.
+
+    See arc_distance_of_far_edges_in_diamond for details.
+
+    Args:
+        vertex_x: x coordinate of vertices
+        vertex_y: y coordinate of vertices
+        domain_length: length of the domain
+        domain_height: height of the domain
+
+    Returns:
+        distance between the "far" vertices in the diamond.
+
+    """
+    return distance_on_edges_torus(
+        vertex_x(E2C2V[2]),
+        vertex_x(E2C2V[3]),
+        vertex_y(E2C2V[2]),
+        vertex_y(E2C2V[3]),
+        domain_length,
+        domain_height,
+    )
+
+
 @gtx.field_operator
 def edge_length(
     vertex_lat: fa.VertexField[ta.wpfloat],
@@ -502,6 +706,26 @@ def compute_arc_distance_of_far_edges_in_diamond(
     )
 
 
+@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
+def compute_distance_of_far_edges_in_diamond_torus(
+    vertex_x: fa.VertexField[ta.wpfloat],
+    vertex_y: fa.VertexField[ta.wpfloat],
+    domain_length: ta.wpfloat,
+    domain_height: ta.wpfloat,
+    far_vertex_distance: fa.EdgeField[ta.wpfloat],
+    horizontal_start: gtx.int32,
+    horizontal_end: gtx.int32,
+):
+    distance_of_far_edges_in_diamond_torus(
+        vertex_x,
+        vertex_y,
+        domain_length,
+        domain_height,
+        out=far_vertex_distance,
+        domain={dims.EdgeDim: (horizontal_start, horizontal_end)},
+    )
+
+
 @gtx.field_operator
 def edge_area(
     owner_mask: fa.EdgeField[bool],
@@ -574,11 +798,11 @@ def compute_coriolis_parameter_on_edges(
 
 
 def compute_primal_cart_normal(
-    primal_cart_normal_x: alloc.NDArray,
-    primal_cart_normal_y: alloc.NDArray,
-    primal_cart_normal_z: alloc.NDArray,
+    primal_cart_normal_x: data_alloc.NDArray,
+    primal_cart_normal_y: data_alloc.NDArray,
+    primal_cart_normal_z: data_alloc.NDArray,
     array_ns: ModuleType = np,
-) -> alloc.NDArray:
+) -> data_alloc.NDArray:
     primal_cart_normal = array_ns.transpose(
         array_ns.stack(
             (