diff --git a/dev-requirements.txt b/dev-requirements.txt
index 504a4f29ea69f60d238d8d75c8d91199453a90ab..4735ccd8dba857f3e118b10149302c1c9dbcfd6b 100644
--- a/dev-requirements.txt
+++ b/dev-requirements.txt
@@ -1,6 +1,6 @@
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
diff --git a/hog/ast.py b/hog/ast.py
index 3f1f97a86dd6e55713267f2f57c8f2afd3b0e8ce..aa03b57f0275742a433c2bae387cde18161bd4f5 100644
--- a/hog/ast.py
+++ b/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(
diff --git a/hog/external_functions.py b/hog/external_functions.py
index 0180c9cc6ef912980de4ded5eebcc15989212aab..62da4b055fed35625300489ce6486b76eab021ab 100644
--- a/hog/external_functions.py
+++ b/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):
diff --git a/hog/operator_generation/indexing.py b/hog/operator_generation/indexing.py
index 218a8fef66eb9cf5dcec11e9494d2bcd92e130ba..b8874624eabb81092bf797fc5a40a359c6567140 100644
--- a/hog/operator_generation/indexing.py
+++ b/hog/operator_generation/indexing.py
@@ -15,7 +15,7 @@
# along with this program. If not, see <https://www.gnu.org/licenses/>.
from enum import Enum
-from typing import Any, List, Tuple, Union, Optional
+from typing import Any, List, Tuple, Union
import operator
import sympy as sp
@@ -27,7 +27,7 @@ from hog.symbolizer import Symbolizer
from pystencils.integer_functions import int_div
from pystencils import TypedSymbol
from enum import Enum
-from typing import Tuple, List, Union, cast
+from typing import Tuple, List, TypeAlias, Union, cast
import operator
import sympy as sp
@@ -42,6 +42,12 @@ from pystencils.integer_functions import int_div
USE_SYMPY_INT_DIV = False
+Index1d: TypeAlias = int | sp.Symbol
+Index2d: TypeAlias = Tuple[Index1d, Index1d]
+Index3d: TypeAlias = Tuple[Index1d, Index1d, Index1d]
+Index: TypeAlias = Index2d | Index3d
+
+
class DoFType(Enum):
"""
Used to differentiate the different DoFs defined in HyTeG.
@@ -90,17 +96,6 @@ class CellType(Enum):
GREEN_DOWN = "GREEN_DOWN"
-class Point:
- """
- Resembles a point in 2D or 3D to make index computations more visual.
- """
-
- def _init_(self, x=0, y=0, z=None):
- self.x = x
- self.y = y
- self.z = z
-
-
class VolumeDoFMemoryLayout(Enum):
"""
Array of structs or stuct of array; memory layouts for DG
@@ -115,35 +110,22 @@ def sympy_int_div(x: float, y: float) -> int:
return sp.floor(x / y)
-def generalized_macro_face_index(
- logical_index: Tuple[int, int, int], width: int
-) -> int:
+def generalized_macro_face_index(logical_index: Index2d, width: int) -> int:
"""Indexes macro faces. width depends on the level and quantifies the amount of primitives in one direction of the refined triangle."""
- x, y, _ = logical_index
- if USE_SYMPY_INT_DIV:
- row_offset = y * (width + 1) - sympy_int_div(((y + 1) * y), 2)
- else:
- row_offset = y * (width + 1) - int_div(((y + 1) * y), 2)
+ x, y = logical_index
+ row_offset = y * (width + 1) - num_microvertices_per_face_from_width(y)
return row_offset + x
-def generalized_macro_cell_index(
- logical_index: Tuple[int, int, int], width: int
-) -> int:
+def generalized_macro_cell_index(logical_index: Index3d, width: int) -> int:
"""Indexes macro tetrahedra. width depends on the level and quantifies the amount of primitives in one direction of the refined tetrahedron."""
x, y, z = logical_index
width_minus_slice = width - z
- if USE_SYMPY_INT_DIV:
- slice_offset = int_div(((width + 2) * (width + 1) * width), 6) - sympy_int_div(
- ((width_minus_slice + 2) * (width_minus_slice + 1) * width_minus_slice), 6
- )
- row_offset = y * (width_minus_slice + 1) - sympy_int_div(((y + 1) * y), 2)
- else:
- slice_offset = int_div(((width + 2) * (width + 1) * width), 6) - int_div(
- ((width_minus_slice + 2) * (width_minus_slice + 1) * width_minus_slice), 6
- )
- row_offset = y * (width_minus_slice + 1) - int_div(((y + 1) * y), 2)
+ slice_offset = num_microvertices_per_cell_from_width(
+ width
+ ) - num_microvertices_per_cell_from_width(width_minus_slice)
+ row_offset = y * (width_minus_slice + 1) - num_microvertices_per_face_from_width(y)
return slice_offset + row_offset + x
@@ -151,18 +133,18 @@ def num_microvertices_per_face_from_width(width: int) -> int:
"""Computes the number of microvertices in a refined macro triangle. width depends on the level and quantifies the amount of primitives in one direction of the refined triangle."""
if USE_SYMPY_INT_DIV:
- return sympy_int_div((width * (width + 1)), 2)
+ return sympy_int_div(width * (width + 1), 2)
else:
- return int_div((width * (width + 1)), 2)
+ return int_div(width * (width + 1), 2)
def num_microvertices_per_cell_from_width(width: int) -> int:
"""Computes the number of microvertices in a refined macro tetrahedron. width depends on the level and quantifies the amount of primitives in one direction of the refined triangle."""
if USE_SYMPY_INT_DIV:
- return sympy_int_div(((width + 2) * (width + 1) * width), 6)
+ return sympy_int_div((width + 2) * (width + 1) * width, 6)
else:
- return int_div(((width + 2) * (width + 1) * width), 6)
+ return int_div((width + 2) * (width + 1) * width, 6)
def num_microfaces_per_face(level: int) -> int:
@@ -173,13 +155,6 @@ def num_microcells_per_cell(level: int) -> int:
return 8 ** (level)
-def linear_macro_cell_size(width: int) -> int:
- if USE_SYMPY_INT_DIV:
- return sympy_int_div((width + 2) * (width + 1) * width, 6)
- else:
- return int_div((width + 2) * (width + 1) * width, 6)
-
-
def num_microvertices_per_edge(level: int) -> int:
return 2**level + 1
@@ -188,9 +163,7 @@ def num_microedges_per_edge(level: int) -> int:
return num_microvertices_per_edge(level) - 1
-def num_faces_per_row_by_type(
- level: int, faceType: Union[None, EdgeType, FaceType, CellType]
-) -> int:
+def num_faces_per_row_by_type(level: int, faceType: FaceType) -> int:
if faceType == FaceType.GRAY:
return num_microedges_per_edge(level)
elif faceType == FaceType.BLUE:
@@ -201,7 +174,7 @@ def num_faces_per_row_by_type(
def num_cells_per_row_by_type(
level: int,
- cellType: Union[None, EdgeType, FaceType, CellType],
+ cellType: CellType,
num_microedges_per_edge: sp.Symbol,
) -> int:
if cellType == CellType.WHITE_UP:
@@ -234,74 +207,39 @@ def num_micro_cells_per_macro_cell(
)
-def linear_macro_face_index(width: int, x: int, y: int) -> int:
- if USE_SYMPY_INT_DIV:
- rowOffset = y * ((width) + 1) - sympy_int_div(((y + 1) * (y)), 2)
- else:
- rowOffset = y * ((width) + 1) - int_div(((y + 1) * (y)), 2)
- return rowOffset + x
-
-
-def linear_macro_cell_index(width: int, x: int, y: int, z: int) -> int:
- widthMinusSlice = width - z
- if USE_SYMPY_INT_DIV:
- sliceOffset = linear_macro_cell_size(width) - sympy_int_div(
- (widthMinusSlice + 2) * (widthMinusSlice + 1) * widthMinusSlice, 6
- )
- rowOffset = y * (widthMinusSlice + 1) - sympy_int_div((y + 1) * y, 2)
- else:
- sliceOffset = linear_macro_cell_size(width) - int_div(
- (widthMinusSlice + 2) * (widthMinusSlice + 1) * widthMinusSlice, 6
- )
- rowOffset = y * (widthMinusSlice + 1) - int_div((y + 1) * y, 2)
- return sliceOffset + rowOffset + x
-
-
def facedof_index(
level: int,
- # Hack: see issue #46 for details
- index: Union[Tuple[int, int], Tuple[int, int, int]],
- faceType: Union[None, EdgeType, FaceType, CellType],
+ index: Index2d,
+ faceType: FaceType,
num_microfaces_per_face: sp.Symbol,
num_microedges_per_edge: sp.Symbol,
) -> int:
- """Indexes triangles/faces. Used to compute offsets in volume dof indexing in 2D and AoS layout."""
-
- # second part of hack for #46
- # also see https://github.com/python/mypy/issues/1178
- if len(index) == 3:
- x, y, _ = cast(Tuple[int, int, int], index)
- else:
- x, y = cast(Tuple[int, int], index)
-
- # width = num_faces_per_row_by_type(level, faceType)
if faceType == FaceType.GRAY:
- return linear_macro_face_index(num_microedges_per_edge, x, y)
+ return generalized_macro_face_index(index, num_microedges_per_edge)
elif faceType == FaceType.BLUE:
return num_microvertices_per_face_from_width(
num_microedges_per_edge
- ) + linear_macro_face_index(num_microedges_per_edge - 1, x, y)
+ ) + generalized_macro_face_index(index, num_microedges_per_edge - 1)
else:
raise HOGException(f"Unexpected face type: {faceType}")
def celldof_index(
level: int,
- index: Tuple[int, int, int],
- cellType: Union[None, EdgeType, FaceType, CellType],
+ index: Index3d,
+ cellType: CellType,
num_microedges_per_edge: sp.Symbol,
) -> int:
"""Indexes cells/tetrahedra. Used to compute offsets in volume dof indexing in 3D and AoS layout."""
- x, y, z = index
width = num_cells_per_row_by_type(
level, cellType, num_microedges_per_edge
) # gives expr(level)
if cellType == CellType.WHITE_UP:
- return linear_macro_cell_index(width, x, y, z)
+ return generalized_macro_cell_index(index, width)
elif cellType == CellType.BLUE_UP:
return num_micro_cells_per_macro_cell( # gives expr(level)
level, CellType.WHITE_UP, num_microedges_per_edge
- ) + linear_macro_cell_index(width, x, y, z)
+ ) + generalized_macro_cell_index(index, width)
elif cellType == CellType.GREEN_UP:
return (
num_micro_cells_per_macro_cell(
@@ -310,7 +248,7 @@ def celldof_index(
+ num_micro_cells_per_macro_cell(
level, CellType.BLUE_UP, num_microedges_per_edge
)
- + linear_macro_cell_index(width, x, y, z)
+ + generalized_macro_cell_index(index, width)
)
elif cellType == CellType.WHITE_DOWN:
return (
@@ -323,7 +261,7 @@ def celldof_index(
+ num_micro_cells_per_macro_cell(
level, CellType.GREEN_UP, num_microedges_per_edge
)
- + linear_macro_cell_index(width, x, y, z)
+ + generalized_macro_cell_index(index, width)
)
elif cellType == CellType.BLUE_DOWN:
return (
@@ -339,7 +277,7 @@ def celldof_index(
+ num_micro_cells_per_macro_cell(
level, CellType.WHITE_DOWN, num_microedges_per_edge
)
- + linear_macro_cell_index(width, x, y, z)
+ + generalized_macro_cell_index(index, width)
)
elif cellType == CellType.GREEN_DOWN:
return (
@@ -358,7 +296,7 @@ def celldof_index(
+ num_micro_cells_per_macro_cell(
level, CellType.BLUE_DOWN, num_microedges_per_edge
)
- + linear_macro_cell_index(width, x, y, z)
+ + generalized_macro_cell_index(index, width)
)
else:
raise HOGException()
@@ -408,7 +346,7 @@ class DoFIndex:
def __init__(
self,
- primitive_index: Tuple[int, int, int],
+ primitive_index: Index,
dof_type: DoFType = DoFType.VERTEX,
dof_sub_type: Union[None, EdgeType, FaceType, CellType] = None,
n_dofs_per_primitive: int = 1,
@@ -438,16 +376,28 @@ class DoFIndex:
"""
Computes the array index of the passed DoF.
"""
+ if geometry.dimensions != len(self.primitive_index):
+ raise HOGException(
+ "Geometry dimension does not match length of micro-primitive index:\n"
+ f" geometry dimension = {geometry.dimensions}\n"
+ f" index length = {len(self.primitive_index)}"
+ )
+
if self.dof_type == DoFType.VERTEX:
width = indexing_info.micro_edges_per_macro_edge + 1
if isinstance(geometry, TriangleElement):
- return generalized_macro_face_index(self.primitive_index, width)
+ return generalized_macro_face_index(
+ cast(Index2d, self.primitive_index), width
+ )
elif isinstance(geometry, TetrahedronElement):
- return generalized_macro_cell_index(self.primitive_index, width)
+ return generalized_macro_cell_index(
+ cast(Index3d, self.primitive_index), width
+ )
else:
raise HOGException(
"Indexing function not implemented for this geometry."
)
+
elif self.dof_type == DoFType.EDGE:
width = indexing_info.micro_edges_per_macro_edge
if isinstance(geometry, TriangleElement):
@@ -472,7 +422,7 @@ class DoFIndex:
return order.index(
self.dof_sub_type
) * micro_edges_one_type_per_macro_face + generalized_macro_face_index(
- self.primitive_index, width
+ cast(Index2d, self.primitive_index), width
)
elif isinstance(geometry, TetrahedronElement):
order = [
@@ -489,20 +439,23 @@ class DoFIndex:
) * num_microvertices_per_cell_from_width(
width
) + generalized_macro_cell_index(
- self.primitive_index,
+ cast(Index3d, self.primitive_index),
width - (1 if self.dof_sub_type == EdgeType.XYZ else 0),
)
else:
raise HOGException(
"Indexing function not implemented for this geometry."
)
+
elif self.dof_type == DoFType.VOLUME:
if isinstance(geometry, TriangleElement):
+ assert isinstance(self.dof_sub_type, FaceType)
+
numMicroVolumes = indexing_info.num_microfaces_per_face
microVolume = facedof_index(
indexing_info.level,
- self.primitive_index,
+ cast(Index2d, self.primitive_index),
self.dof_sub_type,
indexing_info.num_microfaces_per_face,
indexing_info.micro_edges_per_macro_edge,
@@ -518,11 +471,13 @@ class DoFIndex:
)
return idx
elif isinstance(geometry, TetrahedronElement):
+ assert isinstance(self.dof_sub_type, CellType)
+
numMicroVolumes = indexing_info.num_microcells_per_cell
microVolume = celldof_index(
indexing_info.level,
- self.primitive_index,
+ cast(Index3d, self.primitive_index),
self.dof_sub_type,
indexing_info.micro_edges_per_macro_edge,
)
@@ -541,6 +496,7 @@ class DoFIndex:
raise HOGException(
"Indexing function not implemented for this geometry."
)
+
else:
raise HOGException(
"Indexing function not implemented for this function space."
@@ -561,30 +517,38 @@ class DoFIndex:
def micro_element_to_vertex_indices(
geometry: ElementGeometry,
element_type: Union[FaceType, CellType],
- element_index: Tuple[int, int, int],
+ element_index: Index,
) -> List[DoFIndex]:
"""
Returns the micro vertex indices of the given micro element, ordered such
that orientations of mesh entities are locally consistent.
"""
+ if geometry.dimensions != len(element_index):
+ raise HOGException(
+ "Geometry dimension does not match length of micro-primitive index:\n"
+ f" geometry dimension = {geometry.dimensions}\n"
+ f" index length = {len(element_index)}"
+ )
+
if isinstance(geometry, TriangleElement):
# fmt: off
if element_type == FaceType.GRAY:
return [
- DoFIndex((element_index[0] , element_index[1] , 0)),
- DoFIndex((element_index[0] + 1, element_index[1] , 0)),
- DoFIndex((element_index[0] , element_index[1] + 1, 0)),
+ DoFIndex((element_index[0] , element_index[1] )),
+ DoFIndex((element_index[0] + 1, element_index[1] )),
+ DoFIndex((element_index[0] , element_index[1] + 1)),
]
elif element_type == FaceType.BLUE:
return [
- DoFIndex((element_index[0] + 1, element_index[1] , 0)),
- DoFIndex((element_index[0] , element_index[1] + 1, 0)),
- DoFIndex((element_index[0] + 1, element_index[1] + 1, 0)),
+ DoFIndex((element_index[0] + 1, element_index[1] )),
+ DoFIndex((element_index[0] , element_index[1] + 1)),
+ DoFIndex((element_index[0] + 1, element_index[1] + 1)),
]
else:
raise HOGException(f"Invalid face type {element_type} for 2D.")
# fmt: on
elif isinstance(geometry, TetrahedronElement):
+ element_index = cast(Index3d, element_index)
# fmt: off
if element_type == CellType.WHITE_UP:
return [
@@ -637,7 +601,7 @@ def micro_element_to_vertex_indices(
def micro_element_to_volume_indices(
primitive_type: Union[FaceType, CellType],
- primitive_index: Tuple[int, int, int],
+ primitive_index: Index,
n_dofs_per_primitive: int,
memory_layout: VolumeDoFMemoryLayout,
) -> List[DoFIndex]:
@@ -663,6 +627,10 @@ def micro_vertex_to_edge_indices(
"""
Returns a collection of correctly oriented edges in form of DoFIndices from a given set of vertex DoFIndices.
"""
+ if any(geometry.dimensions != len(idx.primitive_index) for idx in vertex_indices):
+ raise HOGException(
+ "Geometry dimension does not match length of micro-primitive index."
+ )
edge_indices: List[DoFIndex] = []
@@ -675,22 +643,8 @@ def micro_vertex_to_edge_indices(
raise HOGException("Indexing function not implemented for this geometry.")
for vertex_start, vertex_end in edges:
- vertex_idx_start: Tuple[int, int, int]
- vertex_idx_end: Tuple[int, int, int]
- if isinstance(geometry, TriangleElement):
- vertex_idx_start = (
- vertex_indices[vertex_start][0],
- vertex_indices[vertex_start][1],
- 0,
- )
- vertex_idx_end = (
- vertex_indices[vertex_end][0],
- vertex_indices[vertex_end][1],
- 0,
- )
- elif isinstance(geometry, TetrahedronElement):
- vertex_idx_start = vertex_indices[vertex_start].primitive_index
- vertex_idx_end = vertex_indices[vertex_end].primitive_index
+ vertex_idx_start = vertex_indices[vertex_start].primitive_index
+ vertex_idx_end = vertex_indices[vertex_end].primitive_index
# determine orientation of the edge between the two vertices
edge_orientation = calc_edge_orientation(vertex_idx_start, vertex_idx_end)
@@ -712,7 +666,7 @@ def array_access_from_dof_index():
def element_vertex_coordinates(
geometry: ElementGeometry,
- element_index: Tuple[int, int, int],
+ element_index: Index,
element_type: Union[FaceType, CellType],
micro_edges_per_macro_edge: sp.Symbol,
macro_vertex_coordinates: List[sp.Matrix],
@@ -755,15 +709,17 @@ def element_vertex_coordinates(
return vertex_coordinates
-def calc_edge_orientation(
- vertex_idx_start: Tuple[int, int, int], vertex_idx_end: Tuple[int, int, int]
-) -> EdgeType:
+def calc_edge_orientation(vertex_idx_start: Index, vertex_idx_end: Index) -> EdgeType:
"""
This function computes the edge type/edge orientation between two given vertex dof indices.
"""
x = vertex_idx_end[0] - vertex_idx_start[0]
y = vertex_idx_end[1] - vertex_idx_start[1]
- z = vertex_idx_end[2] - vertex_idx_start[2]
+
+ if len(vertex_idx_start) > 2 and len(vertex_idx_end) > 2:
+ z = vertex_idx_end[2] - vertex_idx_start[2]
+ else:
+ z = 0
if x == 1 and y == 0 and z == 0:
return EdgeType.X
@@ -804,10 +760,10 @@ def calc_edge_orientation(
def get_edge_dof_index(
- vertex_idx_start: Tuple[int, int, int],
- vertex_idx_end: Tuple[int, int, int],
+ vertex_idx_start: Index,
+ vertex_idx_end: Index,
orientation: EdgeType,
-) -> Tuple[int, int, int]:
+) -> Index:
"""
This function maps the edge between the vertex dof indices vertex_idx_start and
vertex_idx_end to the corresponding edge dof index.
@@ -824,12 +780,6 @@ def get_edge_dof_index(
if vertex_idx_start[1] < vertex_idx_end[1]
else vertex_idx_end
)
- elif orientation == EdgeType.Z:
- return (
- vertex_idx_start
- if vertex_idx_start[2] < vertex_idx_end[2]
- else vertex_idx_end
- )
elif orientation == EdgeType.XY:
return tuple(
map(
@@ -841,7 +791,17 @@ def get_edge_dof_index(
),
(0, -1, 0),
)
- ) # type: ignore
+ )
+
+ if len(vertex_idx_start) != 3 or len(vertex_idx_end) != 3:
+ raise HOGException(f"Got 3d edge orientation {orientation} but not 3d indices.")
+
+ elif orientation == EdgeType.Z:
+ return (
+ vertex_idx_start
+ if vertex_idx_start[2] < vertex_idx_end[2]
+ else vertex_idx_end
+ )
elif orientation == EdgeType.XZ:
return tuple(
map(
@@ -853,7 +813,7 @@ def get_edge_dof_index(
),
(0, 0, -1),
)
- ) # type: ignore
+ )
elif orientation == EdgeType.YZ:
return tuple(
map(
@@ -865,7 +825,7 @@ def get_edge_dof_index(
),
(0, 0, -1),
)
- ) # type: ignore
+ )
elif orientation == EdgeType.XYZ:
return tuple(
map(
@@ -877,6 +837,6 @@ def get_edge_dof_index(
),
(0, -1, 0),
)
- ) # type: ignore
+ )
else:
raise HOGException("Unexpected orientation.")
diff --git a/hog/operator_generation/operators.py b/hog/operator_generation/operators.py
index ff104c33a346c2c2f48e5a831f0d43c0d03106a2..1e950927c1d9da152090195f1fd1fb6d31a161d6 100644
--- a/hog/operator_generation/operators.py
+++ b/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,
diff --git a/hog_tests/operator_generation/test_indexing.py b/hog_tests/operator_generation/test_indexing.py
index cdebad9453fd873c4f039af47fabe6524cb36f9a..f4a20aabdfab680744b39bd175c8f09b1b15300b 100644
--- a/hog_tests/operator_generation/test_indexing.py
+++ b/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,