Merge branch 'bauerd/indexing' into 'main'

Clean up indexing, Pt. 1

Closes #46

See merge request !43

dev-requirements.txt

 black ~= 23.10
-mypy ~= 1.2.0
+mypy ~= 1.15.0
 types-tabulate == 0.8.7
 Sphinx == 7.2.6
 furo == 2024.1.29
-myst-parser == 2.0.0
\ No newline at end of file
+myst-parser == 2.0.0

hog/ast.py

@@ -36,7 +36,7 @@ class Operations:
 
     def to_table(self) -> str:
         d = vars(self)
-        return tabulate.tabulate([d.values()], headers=d.keys())  # type: ignore
+        return tabulate.tabulate([d.values()], headers=list(d.keys()))
 
 
 def count_operations(

hog/external_functions.py

@@ -21,10 +21,12 @@ class BlendingFTriangle(MultiAssignment):
     def function_name(self):
         return "Blending_F_Triangle"
 
-    def num_input_args(self):
+    @classmethod
+    def num_input_args(cls):
         return 2
 
-    def num_output_args(self):
+    @classmethod
+    def num_output_args(cls):
         return 2
 
     def implementation(self):
@@ -39,10 +41,12 @@ class BlendingFEmbeddedTriangle(MultiAssignment):
     def function_name(self):
         return "Blending_F_EmbeddedTriangle"
 
-    def num_input_args(self):
+    @classmethod
+    def num_input_args(cls):
         return 3
 
-    def num_output_args(self):
+    @classmethod
+    def num_output_args(cls):
         return 3
 
     def implementation(self):
@@ -58,10 +62,12 @@ class BlendingFTetrahedron(MultiAssignment):
     def function_name(self):
         return "Blending_F_Tetrahedron"
 
-    def num_input_args(self):
+    @classmethod
+    def num_input_args(cls):
         return 3
 
-    def num_output_args(self):
+    @classmethod
+    def num_output_args(cls):
         return 3
 
     def implementation(self):

hog/operator_generation/operators.py

@@ -90,6 +90,7 @@ from hog.exception import HOGException
 from hog.operator_generation.indexing import (
     all_element_types,
     element_vertex_coordinates,
+    Index,
     IndexingInfo,
     FaceType,
     CellType,
@@ -907,9 +908,9 @@ class HyTeGElementwiseOperator:
     def _compute_micro_element_coordinates(
         self,
         integration_info: IntegrationInfo,
-        element_index: List[sp.Symbol],
+        element_index: Tuple[sp.Symbol, sp.Symbol, sp.Symbol],
         geometry: ElementGeometry,
-    ) -> Tuple[List[int | sp.Symbol | Field.Access], List[CustomCodeNode]]:
+    ) -> Tuple[Index, List[CustomCodeNode]]:
         """
         Computes coordinates of the micro-element. This is _not_ to be confused with the coordinates of the element's
         vertices!
@@ -947,7 +948,7 @@ class HyTeGElementwiseOperator:
         ):
             # The Jacobians are loop-counter dependent, and we do not care about vectorization.
             # So we just use the indices. pystencils will handle casting them to float.
-            el_matrix_element_index = element_index.copy()
+            el_matrix_element_index = element_index
 
         else:
             # Vectorization and loop-counter dependencies
@@ -1000,12 +1001,12 @@ class HyTeGElementwiseOperator:
             # are _not_ array accesses. We assign them to the first entry of the float loop counter array. The
             # vectorizer can automatically assign multiple variables here.
             # Note that the first, i.e., innermost loop counter is chosen for the array access for all dimensions!
-            el_matrix_element_index = [
+            el_matrix_element_index = tuple(
                 float_loop_ctr_arrays[d].absolute_access(
                     (element_index[0] - phantom_ctr,), (0,)
                 )
                 for d in range(geometry.dimensions)
-            ]
+            )
 
             # Let's fill the array.
             float_ctr_array_size = (
@@ -1175,9 +1176,9 @@ class HyTeGElementwiseOperator:
         ]
 
         # create loop according to loop strategy
-        element_index = [
+        element_index = tuple(
             LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(dim)
-        ]
+        )
         loop = integration_info.loop_strategy.create_loop(
             dim, element_index, indexing_info.micro_edges_per_macro_edge
         )
@@ -1235,7 +1236,7 @@ class HyTeGElementwiseOperator:
             src_vecs_accesses = [
                 src_field.local_dofs(
                     geometry,
-                    element_index,  # type: ignore[arg-type] # list of sympy expressions also works
+                    element_index,
                     element_type,
                     indexing_info,
                     element_vertex_order,
@@ -1245,7 +1246,7 @@ class HyTeGElementwiseOperator:
             dst_vecs_accesses = [
                 dst_field.local_dofs(
                     geometry,
-                    element_index,  # type: ignore[arg-type] # list of sympy expressions also works
+                    element_index,
                     element_type,
                     indexing_info,
                     element_vertex_order,
@@ -1262,7 +1263,7 @@ class HyTeGElementwiseOperator:
 
             kernel_op_post_assignments = kernel_type.kernel_post_operation(
                 geometry,
-                element_index,  # type: ignore[arg-type] # list of sympy expressions also works
+                element_index,
                 element_type,
                 src_vecs_accesses,
                 dst_vecs_accesses,
@@ -1298,7 +1299,7 @@ class HyTeGElementwiseOperator:
                         sp.Symbol(dof_symbol.name),
                         coeffs[dof_symbol.function_id].local_dofs(
                             geometry,
-                            element_index,  # type: ignore[arg-type] # list of sympy expressions also works
+                            element_index,
                             element_type,
                             indexing_info,
                             element_vertex_order,
@@ -1322,7 +1323,7 @@ class HyTeGElementwiseOperator:
 
             el_vertex_coordinates = element_vertex_coordinates(
                 geometry,
-                el_matrix_element_index,  # type: ignore[arg-type] # list of sympy expressions also works
+                el_matrix_element_index,
                 element_type,
                 indexing_info.micro_edges_per_macro_edge_float,
                 macro_vertex_coordinates,

hog_tests/operator_generation/test_indexing.py

@@ -45,19 +45,19 @@ def test_micro_element_to_vertex_indices():
     clear_cache()
 
     assert indexing.micro_element_to_vertex_indices(
-        TriangleElement(), FaceType.GRAY, (0, 0, 0)
+        TriangleElement(), FaceType.GRAY, (0, 0)
     ) == [
-        DoFIndex((0, 0, 0), DoFType.VERTEX),
-        DoFIndex((1, 0, 0), DoFType.VERTEX),
-        DoFIndex((0, 1, 0), DoFType.VERTEX),
+        DoFIndex((0, 0), DoFType.VERTEX),
+        DoFIndex((1, 0), DoFType.VERTEX),
+        DoFIndex((0, 1), DoFType.VERTEX),
     ]
 
     assert indexing.micro_element_to_vertex_indices(
-        TriangleElement(), FaceType.BLUE, (2, 4, 0)
+        TriangleElement(), FaceType.BLUE, (2, 4)
     ) == [
-        DoFIndex((3, 4, 0), DoFType.VERTEX),
-        DoFIndex((2, 5, 0), DoFType.VERTEX),
-        DoFIndex((3, 5, 0), DoFType.VERTEX),
+        DoFIndex((3, 4), DoFType.VERTEX),
+        DoFIndex((2, 5), DoFType.VERTEX),
+        DoFIndex((3, 5), DoFType.VERTEX),
     ]
 
     assert indexing.micro_element_to_vertex_indices(
@@ -122,28 +122,28 @@ def test_micro_vertex_to_edge_indices():
     assert indexing.micro_vertex_to_edge_indices(
         TriangleElement(),
         [
-            DoFIndex((5, 4, 0), DoFType.VERTEX),
-            DoFIndex((6, 4, 0), DoFType.VERTEX),
-            DoFIndex((5, 5, 0), DoFType.VERTEX),
+            DoFIndex((5, 4), DoFType.VERTEX),
+            DoFIndex((6, 4), DoFType.VERTEX),
+            DoFIndex((5, 5), DoFType.VERTEX),
         ],
     ) == [
-        DoFIndex((5, 4, 0), DoFType.EDGE, EdgeType.XY),
-        DoFIndex((5, 4, 0), DoFType.EDGE, EdgeType.Y),
-        DoFIndex((5, 4, 0), DoFType.EDGE, EdgeType.X),
+        DoFIndex((5, 4), DoFType.EDGE, EdgeType.XY),
+        DoFIndex((5, 4), DoFType.EDGE, EdgeType.Y),
+        DoFIndex((5, 4), DoFType.EDGE, EdgeType.X),
     ]
 
     # BLUE
     assert indexing.micro_vertex_to_edge_indices(
         TriangleElement(),
         [
-            DoFIndex((6, 4, 0), DoFType.VERTEX),
-            DoFIndex((5, 5, 0), DoFType.VERTEX),
-            DoFIndex((6, 5, 0), DoFType.VERTEX),
+            DoFIndex((6, 4), DoFType.VERTEX),
+            DoFIndex((5, 5), DoFType.VERTEX),
+            DoFIndex((6, 5), DoFType.VERTEX),
         ],
     ) == [
-        DoFIndex((5, 5, 0), DoFType.EDGE, EdgeType.X),
-        DoFIndex((6, 4, 0), DoFType.EDGE, EdgeType.Y),
-        DoFIndex((5, 4, 0), DoFType.EDGE, EdgeType.XY),
+        DoFIndex((5, 5), DoFType.EDGE, EdgeType.X),
+        DoFIndex((6, 4), DoFType.EDGE, EdgeType.Y),
+        DoFIndex((5, 4), DoFType.EDGE, EdgeType.XY),
     ]
 
     # WHITE UP
@@ -266,7 +266,7 @@ def test_micro_volume_to_volume_indices():
         indexing_info: IndexingInfo,
         n_dofs_per_primitive: int,
         primitive_type: Union[FaceType, CellType],
-        primitive_index: Tuple[int, int, int],
+        primitive_index: Tuple[int, int] | Tuple[int, int, int],
         target_array_index: int,
         intra_primitive_index: int = 0,
         memory_layout: VolumeDoFMemoryLayout = VolumeDoFMemoryLayout.AoS,
@@ -286,16 +286,16 @@ def test_micro_volume_to_volume_indices():
 
     # 2D, P0:
     test_element_type_on_level(
-        TriangleElement(), 2, indexingInfo, 1, FaceType.GRAY, (2, 0, 0), 2
+        TriangleElement(), 2, indexingInfo, 1, FaceType.GRAY, (2, 0), 2
     )
     test_element_type_on_level(
-        TriangleElement(), 2, indexingInfo, 1, FaceType.GRAY, (1, 2, 0), 8
+        TriangleElement(), 2, indexingInfo, 1, FaceType.GRAY, (1, 2), 8
     )
     test_element_type_on_level(
-        TriangleElement(), 2, indexingInfo, 1, FaceType.BLUE, (1, 1, 0), 14
+        TriangleElement(), 2, indexingInfo, 1, FaceType.BLUE, (1, 1), 14
     )
     test_element_type_on_level(
-        TriangleElement(), 2, indexingInfo, 1, FaceType.BLUE, (0, 2, 0), 15
+        TriangleElement(), 2, indexingInfo, 1, FaceType.BLUE, (0, 2), 15
     )
 
     # 2D, P1:
@@ -310,7 +310,7 @@ def test_micro_volume_to_volume_indices():
         indexingInfo,
         3,
         FaceType.GRAY,
-        (2, 0, 0),
+        (2, 0),
         7,
         intra_primitive_index=1,
     )
@@ -320,7 +320,7 @@ def test_micro_volume_to_volume_indices():
         indexingInfo,
         3,
         FaceType.GRAY,
-        (1, 2, 0),
+        (1, 2),
         26,
         intra_primitive_index=2,
     )
@@ -330,7 +330,7 @@ def test_micro_volume_to_volume_indices():
         indexingInfo,
         3,
         FaceType.BLUE,
-        (1, 1, 0),
+        (1, 1),
         30 + 13,
         intra_primitive_index=1,
     )
@@ -340,7 +340,7 @@ def test_micro_volume_to_volume_indices():
         indexingInfo,
         3,
         FaceType.BLUE,
-        (0, 2, 0),
+        (0, 2),
         30 + 17,
         intra_primitive_index=2,
     )
@@ -355,7 +355,7 @@ def test_micro_volume_to_volume_indices():
         indexingInfo,
         3,
         FaceType.GRAY,
-        (2, 0, 0),
+        (2, 0),
         18,
         intra_primitive_index=1,
         memory_layout=VolumeDoFMemoryLayout.SoA,
@@ -366,7 +366,7 @@ def test_micro_volume_to_volume_indices():
         indexingInfo,
         3,
         FaceType.GRAY,
-        (1, 1, 0),
+        (1, 1),
         32 + 5,
         intra_primitive_index=2,
         memory_layout=VolumeDoFMemoryLayout.SoA,
@@ -377,7 +377,7 @@ def test_micro_volume_to_volume_indices():
         indexingInfo,
         3,
         FaceType.BLUE,
-        (1, 1, 0),
+        (1, 1),
         26 + 4,
         intra_primitive_index=1,
         memory_layout=VolumeDoFMemoryLayout.SoA,
@@ -388,7 +388,7 @@ def test_micro_volume_to_volume_indices():
         indexingInfo,
         3,
         FaceType.BLUE,
-        (0, 2, 0),
+        (0, 2),
         42 + 5,
         intra_primitive_index=2,
         memory_layout=VolumeDoFMemoryLayout.SoA,