From 30c163f051fc3b373bb75a003d055e429be3d0c8 Mon Sep 17 00:00:00 2001
From: Marcus Mohr <marcus.mohr@lmu.de>
Date: Mon, 2 Dec 2024 18:01:25 +0100
Subject: [PATCH] First unsuccessful approach to split up
function_space_impls.py
---
.../n1e1_space_impl.py | 141 +++
.../p0_space_impl.py | 89 ++
.../p1_space_impl.py | 304 +++++++
.../p2_space_impl.py | 344 ++++++++
.../function_space_impls.py | 812 +-----------------
5 files changed, 883 insertions(+), 807 deletions(-)
create mode 100644 hog/operator_generation/function_space_implementations/n1e1_space_impl.py
create mode 100644 hog/operator_generation/function_space_implementations/p0_space_impl.py
create mode 100644 hog/operator_generation/function_space_implementations/p1_space_impl.py
create mode 100644 hog/operator_generation/function_space_implementations/p2_space_impl.py
diff --git a/hog/operator_generation/function_space_implementations/n1e1_space_impl.py b/hog/operator_generation/function_space_implementations/n1e1_space_impl.py
new file mode 100644
index 0000000..c0cdd5a
--- /dev/null
+++ b/hog/operator_generation/function_space_implementations/n1e1_space_impl.py
@@ -0,0 +1,141 @@
+# HyTeG Operator Generator
+# Copyright (C) 2024 HyTeG Team
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+
+import hog.operator_generation.function_space_impls
+
+
+class N1E1FunctionSpaceImpl(FunctionSpaceImpl):
+ def __init__(
+ self,
+ fe_space: FunctionSpace,
+ name: str,
+ type_descriptor: HOGType,
+ is_pointer: bool = False,
+ ):
+ super().__init__(fe_space, name, type_descriptor, is_pointer)
+ self.field = self._create_field(name)
+
+ def pre_communication(self, dim: int) -> str:
+ if dim == 2:
+ return ""
+ else:
+ return (
+ f"{self._deref()}.communicate< Face, Cell >( level );\n"
+ f"{self._deref()}.communicate< Edge, Cell >( level );"
+ )
+
+ def zero_halos(self, dim: int) -> str:
+ if dim == 2:
+ return ""
+ else:
+ return f"edgedof::macrocell::setBoundaryToZero( level, cell, {self._deref()}.getDoFs()->getCellDataID() );"
+
+ def post_communication(
+ self, dim: int, params: str, transform_basis: bool = True
+ ) -> str:
+ if dim == 2:
+ return ""
+ else:
+ dofs = ""
+ if not transform_basis:
+ dofs = "getDoFs()->"
+ return (
+ f"{self._deref()}.{dofs}communicateAdditively< Cell, Face >( {params} );\n"
+ f"{self._deref()}.{dofs}communicateAdditively< Cell, Edge >( {params} );"
+ )
+
+ def pointer_retrieval(self, dim: int) -> str:
+ """C++ code for retrieving pointers to the numerical data stored in the macro primitives `face` (2d) or `cell` (3d)."""
+ Macro = {2: "Face", 3: "Cell"}[dim]
+ macro = {2: "face", 3: "cell"}[dim]
+
+ return f"{self.type_descriptor.pystencils_type}* _data_{self.name} = {macro}.getData( {self._deref()}.getDoFs()->get{Macro}DataID() )->getPointer( level );"
+
+ def invert_elementwise(self, dim: int) -> str:
+ return f"{self._deref()}.getDoFs()->invertElementwise( level );"
+
+ def local_dofs(
+ self,
+ geometry: ElementGeometry,
+ element_index: Tuple[int, int, int],
+ element_type: Union[FaceType, CellType],
+ indexing_info: IndexingInfo,
+ element_vertex_ordering: List[int],
+ ) -> List[Field.Access]:
+ vertex_dof_indices = micro_element_to_vertex_indices(
+ geometry, element_type, element_index
+ )
+
+ vertex_dof_indices = [vertex_dof_indices[i] for i in element_vertex_ordering]
+
+ edge_dof_indices = micro_vertex_to_edge_indices(geometry, vertex_dof_indices)
+ edge_array_indices = [
+ dof_idx.array_index(geometry, indexing_info) for dof_idx in edge_dof_indices
+ ]
+
+ return [self.field.absolute_access((idx,), (0,)) for idx in edge_array_indices]
+
+ def func_type_string(self) -> str:
+ return f"n1e1::N1E1VectorFunction< {self.type_descriptor.pystencils_type} >"
+
+ def includes(self) -> Set[str]:
+ return {
+ "hyteg/n1e1functionspace/N1E1VectorFunction.hpp",
+ "hyteg/n1e1functionspace/N1E1MacroCell.hpp",
+ }
+
+ def dof_transformation(
+ self,
+ geometry: ElementGeometry,
+ element_index: Tuple[int, int, int],
+ element_type: Union[FaceType, CellType],
+ element_vertex_ordering: List[int],
+ ) -> Tuple[CustomCodeNode, sp.MatrixBase]:
+ if element_vertex_ordering != [0, 1, 2, 3]:
+ raise HOGException(
+ "Element vertex re-ordering not supported for Nédélec elements (yet)."
+ )
+
+ Macro = {2: "Face", 3: "Cell"}[geometry.dimensions]
+ macro = {2: "face", 3: "cell"}[geometry.dimensions]
+ name = "basisTransformation"
+ n_dofs = self.fe_space.num_dofs(geometry)
+
+ # NOTE: Types are added manually because pystencils won't add types to accessed/
+ # defined symbols of custom code nodes. As a result the symbols
+ # in the matrix will be typed but not their definition, leading to
+ # undefined symbols.
+ # NOTE: The type is `real_t` (not `self._type_descriptor.pystencils_type`)
+ # because this function is implemented manually in HyTeG with
+ # this signature. Passing `np.float64` is not ideal (if `real_t !=
+ # double`) but it makes sure that casts are inserted if necessary
+ # (though some might be superfluous).
+ symbols = [
+ ps.TypedSymbol(f"{name}.diagonal()({i})", np.float64) for i in range(n_dofs)
+ ]
+ return (
+ CustomCodeNode(
+ f"const Eigen::DiagonalMatrix< real_t, {n_dofs} > {name} = "
+ f"n1e1::macro{macro}::basisTransformation( level, {macro}, {{{element_index[0]}, {element_index[1]}, {element_index[2]}}}, {macro}dof::{Macro}Type::{element_type.name} );",
+ [
+ ps.TypedSymbol("level", "const uint_t"),
+ ps.TypedSymbol(f"{macro}", f"const {Macro}&"),
+ ],
+ symbols,
+ ),
+ sp.diag(*symbols),
+ )
diff --git a/hog/operator_generation/function_space_implementations/p0_space_impl.py b/hog/operator_generation/function_space_implementations/p0_space_impl.py
new file mode 100644
index 0000000..15180b1
--- /dev/null
+++ b/hog/operator_generation/function_space_implementations/p0_space_impl.py
@@ -0,0 +1,89 @@
+# HyTeG Operator Generator
+# Copyright (C) 2024 HyTeG Team
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+import hog.operator_generation.function_space_impls
+
+
+class P0FunctionSpaceImpl(FunctionSpaceImpl):
+ def __init__(
+ self,
+ fe_space: FunctionSpace,
+ name: str,
+ type_descriptor: HOGType,
+ is_pointer: bool = False,
+ ):
+ super().__init__(fe_space, name, type_descriptor, is_pointer)
+ self.field = self._create_field(name)
+
+ def pre_communication(self, dim: int) -> str:
+ return ""
+
+ def zero_halos(self, dim: int) -> str:
+ return ""
+
+ def post_communication(
+ self, dim: int, params: str, transform_basis: bool = True
+ ) -> str:
+ return ""
+
+ def pointer_retrieval(self, dim: int) -> str:
+ """C++ code for retrieving pointers to the numerical data stored in the macro primitives `face` (2d) or `cell` (3d)."""
+ Macro = {2: "Face", 3: "Cell"}[dim]
+ macro = {2: "face", 3: "cell"}[dim]
+
+ return f"{self.type_descriptor.pystencils_type}* _data_{self.name} = {self._deref()}.getDGFunction()->volumeDoFFunction()->dofMemory( it.first, level );"
+
+ def invert_elementwise(self, dim: int) -> str:
+ return f"{self._deref()}.invertElementwise( level );"
+
+ def local_dofs(
+ self,
+ geometry: ElementGeometry,
+ element_index: Tuple[int, int, int],
+ element_type: Union[FaceType, CellType],
+ indexing_info: IndexingInfo,
+ element_vertex_ordering: List[int],
+ ) -> List[Field.Access]:
+ volume_dof_indices = micro_element_to_volume_indices(
+ element_type, element_index, 1, VolumeDoFMemoryLayout.SoA
+ )
+
+ vertex_array_indices = [
+ dof_idx.array_index(geometry, indexing_info)
+ for dof_idx in volume_dof_indices
+ ]
+
+ return [
+ self.field.absolute_access((idx,), (0,)) for idx in vertex_array_indices
+ ]
+
+ def func_type_string(self) -> str:
+ return f"P0Function< {self.type_descriptor.pystencils_type} >"
+
+ def includes(self) -> Set[str]:
+ return {f"hyteg/p0functionspace/P0Function.hpp"}
+
+ def dof_transformation(
+ self,
+ geometry: ElementGeometry,
+ element_index: Tuple[int, int, int],
+ element_type: Union[FaceType, CellType],
+ element_vertex_ordering: List[int],
+ ) -> Tuple[CustomCodeNode, sp.MatrixBase]:
+ return (
+ CustomCodeNode("", [], []),
+ sp.Identity(self.fe_space.num_dofs(geometry)),
+ )
diff --git a/hog/operator_generation/function_space_implementations/p1_space_impl.py b/hog/operator_generation/function_space_implementations/p1_space_impl.py
new file mode 100644
index 0000000..cb851e6
--- /dev/null
+++ b/hog/operator_generation/function_space_implementations/p1_space_impl.py
@@ -0,0 +1,304 @@
+# HyTeG Operator Generator
+# Copyright (C) 2024 HyTeG Team
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+import hog.operator_generation.function_space_impls
+
+
+class P1FunctionSpaceImpl(FunctionSpaceImpl):
+ def __init__(
+ self,
+ fe_space: FunctionSpace,
+ name: str,
+ type_descriptor: HOGType,
+ is_pointer: bool = False,
+ ):
+ super().__init__(fe_space, name, type_descriptor, is_pointer)
+ self.field = self._create_field(name)
+
+ def pre_communication(self, dim: int) -> str:
+ if dim == 2:
+ return f"communication::syncFunctionBetweenPrimitives( {self.name}, level, communication::syncDirection_t::LOW2HIGH );"
+ else:
+ return (
+ f"{self._deref()}.communicate< Face, Cell >( level );\n"
+ f"{self._deref()}.communicate< Edge, Cell >( level );\n"
+ f"{self._deref()}.communicate< Vertex, Cell >( level );"
+ )
+
+ def zero_halos(self, dim: int) -> str:
+ if dim == 2:
+ return (
+ f"for ( const auto& idx : vertexdof::macroface::Iterator( level ) ) {{\n"
+ f" if ( vertexdof::macroface::isVertexOnBoundary( level, idx ) ) {{\n"
+ f" auto arrayIdx = vertexdof::macroface::index( level, idx.x(), idx.y() );\n"
+ f" _data_{self.name}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
+ f" }}\n"
+ f"}}"
+ )
+ else:
+ return (
+ f"for ( const auto& idx : vertexdof::macrocell::Iterator( level ) ) {{\n"
+ f" if ( !vertexdof::macrocell::isOnCellFace( idx, level ).empty() ) {{\n"
+ f" auto arrayIdx = vertexdof::macrocell::index( level, idx.x(), idx.y(), idx.z() );\n"
+ f" _data_{self.name}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
+ f" }}\n"
+ f"}}"
+ )
+
+ def post_communication(
+ self, dim: int, params: str, transform_basis: bool = True
+ ) -> str:
+ if dim == 2:
+ return (
+ f"{self._deref()}.communicateAdditively < Face, Edge > ( {params} );\n"
+ f"{self._deref()}.communicateAdditively < Face, Vertex > ( {params} );"
+ )
+ else:
+ return (
+ f"{self._deref()}.communicateAdditively< Cell, Face >( {params} );\n"
+ f"{self._deref()}.communicateAdditively< Cell, Edge >( {params} );\n"
+ f"{self._deref()}.communicateAdditively< Cell, Vertex >( {params} );"
+ )
+
+ def pointer_retrieval(self, dim: int) -> str:
+ """C++ code for retrieving pointers to the numerical data stored in the macro primitives `face` (2d) or `cell` (3d)."""
+ Macro = {2: "Face", 3: "Cell"}[dim]
+ macro = {2: "face", 3: "cell"}[dim]
+
+ return f"{self.type_descriptor.pystencils_type}* _data_{self.name} = {macro}.getData( {self._deref()}.get{Macro}DataID() )->getPointer( level );"
+
+ def invert_elementwise(self, dim: int) -> str:
+ return f"{self._deref()}.invertElementwise( level );"
+
+ def local_dofs(
+ self,
+ geometry: ElementGeometry,
+ element_index: Tuple[int, int, int],
+ element_type: Union[FaceType, CellType],
+ indexing_info: IndexingInfo,
+ element_vertex_ordering: List[int],
+ ) -> List[Field.Access]:
+ vertex_dof_indices = micro_element_to_vertex_indices(
+ geometry, element_type, element_index
+ )
+
+ vertex_dof_indices = [vertex_dof_indices[i] for i in element_vertex_ordering]
+
+ vertex_array_indices = [
+ dof_idx.array_index(geometry, indexing_info)
+ for dof_idx in vertex_dof_indices
+ ]
+
+ return [
+ self.field.absolute_access((idx,), (0,)) for idx in vertex_array_indices
+ ]
+
+ def func_type_string(self) -> str:
+ return f"P1Function< {self.type_descriptor.pystencils_type} >"
+
+ def includes(self) -> Set[str]:
+ return {f"hyteg/p1functionspace/P1Function.hpp"}
+
+ def dof_transformation(
+ self,
+ geometry: ElementGeometry,
+ element_index: Tuple[int, int, int],
+ element_type: Union[FaceType, CellType],
+ element_vertex_ordering: List[int],
+ ) -> Tuple[CustomCodeNode, sp.MatrixBase]:
+ return (
+ CustomCodeNode("", [], []),
+ sp.Identity(self.fe_space.num_dofs(geometry)),
+ )
+
+
+class P1VectorFunctionSpaceImpl(FunctionSpaceImpl):
+ def __init__(
+ self,
+ fe_space: FunctionSpace,
+ name: str,
+ type_descriptor: HOGType,
+ is_pointer: bool = False,
+ ):
+ super().__init__(fe_space, name, type_descriptor, is_pointer)
+ self.fields: Dict[int, Field] = {}
+
+ def _field_name(self, component: int) -> str:
+ return self.name + f"_{component}"
+
+ def _raw_pointer_name(self, component: int) -> str:
+ return f"_data_" + self._field_name(component)
+
+ def pre_communication(self, dim: int) -> str:
+ if dim == 2:
+ return f"communication::syncVectorFunctionBetweenPrimitives( {self.name}, level, communication::syncDirection_t::LOW2HIGH );"
+ else:
+ ret_str = ""
+ for i in range(dim):
+ ret_str += (
+ f"{self._deref()}[{i}].communicate< Face, Cell >( level );\n"
+ f"{self._deref()}[{i}].communicate< Edge, Cell >( level );\n"
+ f"{self._deref()}[{i}].communicate< Vertex, Cell >( level );"
+ )
+ return ret_str
+
+ def zero_halos(self, dim: int) -> str:
+ if dim == 2:
+ return (
+ f"for ( const auto& idx : vertexdof::macroface::Iterator( level ) ) {{\n"
+ f" if ( vertexdof::macroface::isVertexOnBoundary( level, idx ) ) {{\n"
+ f" auto arrayIdx = vertexdof::macroface::index( level, idx.x(), idx.y() );\n"
+ f" {self._raw_pointer_name(0)}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
+ f" {self._raw_pointer_name(1)}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
+ f" }}\n"
+ f"}}"
+ )
+ else:
+ return (
+ f"for ( const auto& idx : vertexdof::macrocell::Iterator( level ) ) {{\n"
+ f" if ( !vertexdof::macrocell::isOnCellFace( idx, level ).empty() ) {{\n"
+ f" auto arrayIdx = vertexdof::macrocell::index( level, idx.x(), idx.y(), idx.z() );\n"
+ f" {self._raw_pointer_name(0)}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
+ f" {self._raw_pointer_name(1)}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
+ f" {self._raw_pointer_name(2)}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
+ f" }}\n"
+ f"}}"
+ )
+
+ def post_communication(
+ self, dim: int, params: str, transform_basis: bool = True
+ ) -> str:
+ if dim == 2:
+ ret_str = ""
+ for i in range(dim):
+ ret_str += (
+ f"{self._deref()}[{i}].communicateAdditively < Face, Edge > ( {params} );\n"
+ f"{self._deref()}[{i}].communicateAdditively < Face, Vertex > ( {params} );\n"
+ )
+ return ret_str
+ else:
+ ret_str = ""
+ for i in range(dim):
+ ret_str += (
+ f"{self._deref()}[{i}].communicateAdditively< Cell, Face >( {params} );\n"
+ f"{self._deref()}[{i}].communicateAdditively< Cell, Edge >( {params} );\n"
+ f"{self._deref()}[{i}].communicateAdditively< Cell, Vertex >( {params} );"
+ )
+ return ret_str
+
+ def pointer_retrieval(self, dim: int) -> str:
+ """C++ code for retrieving pointers to the numerical data stored in the macro primitives `face` (2d) or `cell` (3d)."""
+ Macro = {2: "Face", 3: "Cell"}[dim]
+ macro = {2: "face", 3: "cell"}[dim]
+
+ ret_str = ""
+ for i in range(dim):
+ ret_str += f"{self.type_descriptor.pystencils_type}* {self._raw_pointer_name(i)} = {macro}.getData( {self._deref()}[{i}].get{Macro}DataID() )->getPointer( level );\n"
+ return ret_str
+
+ def invert_elementwise(self, dim: int) -> str:
+ ret_str = ""
+ for i in range(dim):
+ ret_str += f"{self._deref()}[{i}].invertElementwise( level );\n"
+ return ret_str
+
+ def local_dofs(
+ self,
+ geometry: ElementGeometry,
+ element_index: Tuple[int, int, int],
+ element_type: Union[FaceType, CellType],
+ indexing_info: IndexingInfo,
+ element_vertex_ordering: List[int],
+ ) -> List[Field.Access]:
+ """
+ Returns the element-local DoFs (field accesses) in a list (i.e., linearized).
+
+ The order here has to match that of the function space implementation.
+ TODO: This is a little concerning since that order has to match in two very different parts of this
+ implementation. Here, and in the TensorialVectorFunctionSpace. Maybe we can define this order in a single
+ location.
+ We return them in "SoA" order: first all DoFs corresponding to the function with the first component != 0,
+ then all DoFs corresponding to the function with the second component != 0 etc.
+
+ For instance, in 2D we get the DoFs corresponding to the 6 shape functions in the following order:
+
+ First component != 0
+
+ list[0]: ⎡-x_ref_0 - x_ref_1 + 1⎤
+ ⎢ ⎥
+ ⎣ 0 ⎦
+
+ list[1]: ⎡x_ref_0⎤
+ ⎢ ⎥
+ ⎣ 0 ⎦
+
+ list[2]: ⎡x_ref_1⎤
+ ⎢ ⎥
+ ⎣ 0 ⎦
+
+ Second component != 0
+
+ list[3]: ⎡ 0 ⎤
+ ⎢ ⎥
+ ⎣-x_ref_0 - x_ref_1 + 1⎦
+
+ list[4]: ⎡ 0 ⎤
+ ⎢ ⎥
+ ⎣x_ref_0⎦
+
+ list[5]: ⎡ 0 ⎤
+ ⎢ ⎥
+ ⎣x_ref_1⎦
+ """
+
+ for c in range(geometry.dimensions):
+ if c not in self.fields:
+ self.fields[c] = self._create_field(self._field_name(c))
+
+ vertex_dof_indices = micro_element_to_vertex_indices(
+ geometry, element_type, element_index
+ )
+
+ vertex_dof_indices = [vertex_dof_indices[i] for i in element_vertex_ordering]
+
+ vertex_array_indices = [
+ dof_idx.array_index(geometry, indexing_info)
+ for dof_idx in vertex_dof_indices
+ ]
+
+ return [
+ self.fields[c].absolute_access((idx,), (0,))
+ for c in range(geometry.dimensions)
+ for idx in vertex_array_indices
+ ]
+
+ def func_type_string(self) -> str:
+ return f"P1VectorFunction< {self.type_descriptor.pystencils_type} >"
+
+ def includes(self) -> Set[str]:
+ return {f"hyteg/p1functionspace/P1VectorFunction.hpp"}
+
+ def dof_transformation(
+ self,
+ geometry: ElementGeometry,
+ element_index: Tuple[int, int, int],
+ element_type: Union[FaceType, CellType],
+ element_vertex_ordering: List[int],
+ ) -> Tuple[CustomCodeNode, sp.MatrixBase]:
+ return (
+ CustomCodeNode("", [], []),
+ sp.Identity(self.fe_space.num_dofs(geometry)),
+ )
diff --git a/hog/operator_generation/function_space_implementations/p2_space_impl.py b/hog/operator_generation/function_space_implementations/p2_space_impl.py
new file mode 100644
index 0000000..f56bca5
--- /dev/null
+++ b/hog/operator_generation/function_space_implementations/p2_space_impl.py
@@ -0,0 +1,344 @@
+# HyTeG Operator Generator
+# Copyright (C) 2024 HyTeG Team
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+
+import hog.operator_generation.function_space_impls
+
+
+class P2FunctionSpaceImpl(FunctionSpaceImpl):
+ def __init__(
+ self,
+ fe_space: FunctionSpace,
+ name: str,
+ type_descriptor: HOGType,
+ is_pointer: bool = False,
+ ):
+ super().__init__(fe_space, name, type_descriptor, is_pointer)
+ self.f_vertex = self._create_field(name + "Vertex")
+ self.f_edge = self._create_field(name + "Edge")
+
+ def pre_communication(self, dim: int) -> str:
+ if dim == 2:
+ return f"communication::syncFunctionBetweenPrimitives( {self.name}, level, communication::syncDirection_t::LOW2HIGH );"
+ else:
+ return (
+ f"{self._deref()}.communicate< Face, Cell >( level );\n"
+ f"{self._deref()}.communicate< Edge, Cell >( level );\n"
+ f"{self._deref()}.communicate< Vertex, Cell >( level );"
+ )
+
+ def zero_halos(self, dim: int) -> str:
+ if dim == 2:
+ return (
+ f"for ( const auto& idx : vertexdof::macroface::Iterator( level ) )\n"
+ f"{{\n"
+ f" if ( vertexdof::macroface::isVertexOnBoundary( level, idx ) )\n"
+ f" {{\n"
+ f" auto arrayIdx = vertexdof::macroface::index( level, idx.x(), idx.y() );\n"
+ f" _data_{self.name }Vertex[arrayIdx] = walberla::numeric_cast< {self.type_descriptor.pystencils_type} >( 0 );\n"
+ f" }}\n"
+ f"}}\n"
+ f"for ( const auto& idx : edgedof::macroface::Iterator( level ) )\n"
+ f"{{\n"
+ f" for ( const auto& orientation : edgedof::faceLocalEdgeDoFOrientations )\n"
+ f" {{\n"
+ f" if ( !edgedof::macroface::isInnerEdgeDoF( level, idx, orientation ) )\n"
+ f" {{\n"
+ f" auto arrayIdx = edgedof::macroface::index( level, idx.x(), idx.y(), orientation );\n"
+ f" _data_{self.name }Edge[arrayIdx] = walberla::numeric_cast< {self.type_descriptor.pystencils_type} >( 0 );\n"
+ f" }}\n"
+ f" }}\n"
+ f"}}"
+ )
+ else:
+ return (
+ f"for ( const auto& idx : vertexdof::macrocell::Iterator( level ) )\n"
+ f"{{\n"
+ f"if ( !vertexdof::macrocell::isOnCellFace( idx, level ).empty() )\n"
+ f" {{\n"
+ f" auto arrayIdx = vertexdof::macrocell::index( level, idx.x(), idx.y(), idx.z() );\n"
+ f" _data_{self.name}Vertex[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
+ f" }}\n"
+ f"}}\n"
+ f"edgedof::macrocell::setBoundaryToZero( level, cell, {self._deref()}.getEdgeDoFFunction().getCellDataID() );"
+ )
+
+ def post_communication(
+ self, dim: int, params: str, transform_basis: bool = True
+ ) -> str:
+ if dim == 2:
+ return (
+ f"{self._deref()}.getVertexDoFFunction().communicateAdditively< Face, Edge >( {params} );\n"
+ f"{self._deref()}.getVertexDoFFunction().communicateAdditively< Face, Vertex >( {params} );\n"
+ f"{self._deref()}.getEdgeDoFFunction().communicateAdditively< Face, Edge >( {params} );"
+ )
+ else:
+ return (
+ f"{self._deref()}.getVertexDoFFunction().communicateAdditively< Cell, Face >( {params} );\n"
+ f"{self._deref()}.getVertexDoFFunction().communicateAdditively< Cell, Edge >( {params} );\n"
+ f"{self._deref()}.getVertexDoFFunction().communicateAdditively< Cell, Vertex >( {params} );\n"
+ f"{self._deref()}.getEdgeDoFFunction().communicateAdditively< Cell, Face >( {params} );\n"
+ f"{self._deref()}.getEdgeDoFFunction().communicateAdditively< Cell, Edge >( {params} );"
+ )
+
+ def pointer_retrieval(self, dim: int) -> str:
+ """C++ code for retrieving pointers to the numerical data stored in the macro primitives `face` (2d) or `cell` (3d)."""
+ Macro = {2: "Face", 3: "Cell"}[dim]
+ macro = {2: "face", 3: "cell"}[dim]
+
+ return (
+ f"{self.type_descriptor.pystencils_type}* _data_{self.name}Vertex = {macro}.getData( {self._deref()}.getVertexDoFFunction().get{Macro}DataID() )->getPointer( level );\n"
+ f"{self.type_descriptor.pystencils_type}* _data_{self.name}Edge = {macro}.getData( {self._deref()}.getEdgeDoFFunction().get{Macro}DataID() )->getPointer( level );"
+ )
+
+ def invert_elementwise(self, dim: int) -> str:
+ return f"{self._deref()}.invertElementwise( level );"
+
+ def local_dofs(
+ self,
+ geometry: ElementGeometry,
+ element_index: Tuple[int, int, int],
+ element_type: Union[FaceType, CellType],
+ indexing_info: IndexingInfo,
+ element_vertex_ordering: List[int],
+ ) -> List[Field.Access]:
+ vertex_dof_indices = micro_element_to_vertex_indices(
+ geometry, element_type, element_index
+ )
+
+ vertex_dof_indices = [vertex_dof_indices[i] for i in element_vertex_ordering]
+
+ edge_dof_indices = micro_vertex_to_edge_indices(geometry, vertex_dof_indices)
+
+ vrtx_array_idcs = [
+ dof_idx.array_index(geometry, indexing_info)
+ for dof_idx in vertex_dof_indices
+ ]
+ edge_array_idcs = [
+ dof_idx.array_index(geometry, indexing_info) for dof_idx in edge_dof_indices
+ ]
+
+ # NOTE: The order of DoFs must match the order of shape functions
+ # defined in the FunctionSpace. Holds within the individual lists
+ # and overall.
+ return [
+ self.f_vertex.absolute_access((idx,), (0,)) for idx in vrtx_array_idcs
+ ] + [self.f_edge.absolute_access((idx,), (0,)) for idx in edge_array_idcs]
+
+ def func_type_string(self) -> str:
+ return f"P2Function< {self.type_descriptor.pystencils_type} >"
+
+ def includes(self) -> Set[str]:
+ return {f"hyteg/p2functionspace/P2Function.hpp"}
+
+ def dof_transformation(
+ self,
+ geometry: ElementGeometry,
+ element_index: Tuple[int, int, int],
+ element_type: Union[FaceType, CellType],
+ element_vertex_ordering: List[int],
+ ) -> Tuple[CustomCodeNode, sp.MatrixBase]:
+ return (
+ CustomCodeNode("", [], []),
+ sp.Identity(self.fe_space.num_dofs(geometry)),
+ )
+
+
+class P2VectorFunctionSpaceImpl(FunctionSpaceImpl):
+ def __init__(
+ self,
+ fe_space: FunctionSpace,
+ name: str,
+ type_descriptor: HOGType,
+ is_pointer: bool = False,
+ ):
+ super().__init__(fe_space, name, type_descriptor, is_pointer)
+ self.fields: Dict[Tuple[str, int], Field] = {}
+
+ def _field_name(self, component: int, dof_type: str) -> str:
+ """dof_type should either be 'vertex' or 'edge'"""
+ return self.name + f"_{dof_type}_{component}"
+
+ def _raw_pointer_name(self, component: int, dof_type: str) -> str:
+ """dof_type should either be 'vertex' or 'edge'"""
+ return f"_data_" + self._field_name(component, dof_type)
+
+ def pre_communication(self, dim: int) -> str:
+ if dim == 2:
+ return f"communication::syncVectorFunctionBetweenPrimitives( {self.name}, level, communication::syncDirection_t::LOW2HIGH );"
+ else:
+ ret_str = ""
+ for i in range(dim):
+ ret_str += (
+ f"{self._deref()}[{i}].communicate< Face, Cell >( level );\n"
+ f"{self._deref()}[{i}].communicate< Edge, Cell >( level );\n"
+ f"{self._deref()}[{i}].communicate< Vertex, Cell >( level );"
+ )
+ return ret_str
+
+ def zero_halos(self, dim: int) -> str:
+ if dim == 2:
+ return (
+ f"for ( const auto& idx : vertexdof::macroface::Iterator( level ) )\n"
+ f"{{\n"
+ f" if ( vertexdof::macroface::isVertexOnBoundary( level, idx ) )\n"
+ f" {{\n"
+ f" auto arrayIdx = vertexdof::macroface::index( level, idx.x(), idx.y() );\n"
+ f" {self._raw_pointer_name(0, 'vertex')}[arrayIdx] = walberla::numeric_cast< {self.type_descriptor.pystencils_type} >( 0 );\n"
+ f" {self._raw_pointer_name(1, 'vertex')}[arrayIdx] = walberla::numeric_cast< {self.type_descriptor.pystencils_type} >( 0 );\n"
+ f" }}\n"
+ f"}}\n"
+ f"for ( const auto& idx : edgedof::macroface::Iterator( level ) )\n"
+ f"{{\n"
+ f" for ( const auto& orientation : edgedof::faceLocalEdgeDoFOrientations )\n"
+ f" {{\n"
+ f" if ( !edgedof::macroface::isInnerEdgeDoF( level, idx, orientation ) )\n"
+ f" {{\n"
+ f" auto arrayIdx = edgedof::macroface::index( level, idx.x(), idx.y(), orientation );\n"
+ f" {self._raw_pointer_name(0, 'edge')}[arrayIdx] = walberla::numeric_cast< {self.type_descriptor.pystencils_type} >( 0 );\n"
+ f" {self._raw_pointer_name(1, 'edge')}[arrayIdx] = walberla::numeric_cast< {self.type_descriptor.pystencils_type} >( 0 );\n"
+ f" }}\n"
+ f" }}\n"
+ f"}}"
+ )
+ else:
+ return (
+ f"for ( const auto& idx : vertexdof::macrocell::Iterator( level ) )\n"
+ f"{{\n"
+ f"if ( !vertexdof::macrocell::isOnCellFace( idx, level ).empty() )\n"
+ f" {{\n"
+ f" auto arrayIdx = vertexdof::macrocell::index( level, idx.x(), idx.y(), idx.z() );\n"
+ f" {self._raw_pointer_name(0, 'vertex')}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
+ f" {self._raw_pointer_name(1, 'vertex')}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
+ f" {self._raw_pointer_name(2, 'vertex')}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
+ f" }}\n"
+ f"}}\n"
+ f"edgedof::macrocell::setBoundaryToZero( level, cell, {self._deref()}[0].getEdgeDoFFunction().getCellDataID() );\n"
+ f"edgedof::macrocell::setBoundaryToZero( level, cell, {self._deref()}[1].getEdgeDoFFunction().getCellDataID() );\n"
+ f"edgedof::macrocell::setBoundaryToZero( level, cell, {self._deref()}[2].getEdgeDoFFunction().getCellDataID() );\n"
+ )
+
+ def post_communication(
+ self, dim: int, params: str, transform_basis: bool = True
+ ) -> str:
+ if dim == 2:
+ ret_str = ""
+ for i in range(dim):
+ ret_str += (
+ f"{self._deref()}[{i}].getVertexDoFFunction().communicateAdditively< Face, Edge >( {params} );\n"
+ f"{self._deref()}[{i}].getVertexDoFFunction().communicateAdditively< Face, Vertex >( {params} );\n"
+ f"{self._deref()}[{i}].getEdgeDoFFunction().communicateAdditively< Face, Edge >( {params} );"
+ )
+ return ret_str
+ else:
+ ret_str = ""
+ for i in range(dim):
+ ret_str += (
+ f"{self._deref()}[{i}].getVertexDoFFunction().communicateAdditively< Cell, Face >( {params} );\n"
+ f"{self._deref()}[{i}].getVertexDoFFunction().communicateAdditively< Cell, Edge >( {params} );\n"
+ f"{self._deref()}[{i}].getVertexDoFFunction().communicateAdditively< Cell, Vertex >( {params} );\n"
+ f"{self._deref()}[{i}].getEdgeDoFFunction().communicateAdditively< Cell, Face >( {params} );\n"
+ f"{self._deref()}[{i}].getEdgeDoFFunction().communicateAdditively< Cell, Edge >( {params} );"
+ )
+ return ret_str
+
+ def pointer_retrieval(self, dim: int) -> str:
+ """C++ code for retrieving pointers to the numerical data stored in the macro primitives `face` (2d) or `cell` (3d)."""
+ Macro = {2: "Face", 3: "Cell"}[dim]
+ macro = {2: "face", 3: "cell"}[dim]
+
+ ret_str = ""
+ for i in range(dim):
+ ret_str += f"{self.type_descriptor.pystencils_type}* {self._raw_pointer_name(i, 'vertex')} = {macro}.getData( {self._deref()}[{i}].getVertexDoFFunction().get{Macro}DataID() )->getPointer( level );\n"
+ ret_str += f"{self.type_descriptor.pystencils_type}* {self._raw_pointer_name(i, 'edge')} = {macro}.getData( {self._deref()}[{i}].getEdgeDoFFunction().get{Macro}DataID() )->getPointer( level );\n"
+ return ret_str
+
+ def invert_elementwise(self, dim: int) -> str:
+ ret_str = ""
+ for i in range(dim):
+ ret_str += f"{self._deref()}[{i}].invertElementwise( level );\n"
+ return ret_str
+
+ def local_dofs(
+ self,
+ geometry: ElementGeometry,
+ element_index: Tuple[int, int, int],
+ element_type: Union[FaceType, CellType],
+ indexing_info: IndexingInfo,
+ element_vertex_ordering: List[int],
+ ) -> List[Field.Access]:
+ """
+ Returns the element-local DoFs (field accesses) in a list (i.e., linearized).
+
+ See P1VectorFunctionSpaceImpl::local_dofs() for details.
+ """
+
+ for dt in ["vertex", "edge"]:
+ for c in range(geometry.dimensions):
+ if (dt, c) not in self.fields:
+ self.fields[(dt, c)] = self._create_field(self._field_name(c, dt))
+
+ vertex_dof_indices = micro_element_to_vertex_indices(
+ geometry, element_type, element_index
+ )
+
+ vertex_dof_indices = [vertex_dof_indices[i] for i in element_vertex_ordering]
+
+ edge_dof_indices = micro_vertex_to_edge_indices(geometry, vertex_dof_indices)
+
+ vertex_array_indices = [
+ dof_idx.array_index(geometry, indexing_info)
+ for dof_idx in vertex_dof_indices
+ ]
+
+ edge_array_indices = [
+ dof_idx.array_index(geometry, indexing_info) for dof_idx in edge_dof_indices
+ ]
+
+ loc_dofs = []
+
+ for c in range(geometry.dimensions):
+ loc_dofs += [
+ self.fields[("vertex", c)].absolute_access((idx,), (0,))
+ for idx in vertex_array_indices
+ ]
+ loc_dofs += [
+ self.fields[("edge", c)].absolute_access((idx,), (0,))
+ for idx in edge_array_indices
+ ]
+
+ return loc_dofs
+
+ def func_type_string(self) -> str:
+ return f"P2VectorFunction< {self.type_descriptor.pystencils_type} >"
+
+ def includes(self) -> Set[str]:
+ return {
+ f"hyteg/p2functionspace/P2VectorFunction.hpp",
+ f"hyteg/p2functionspace/P2Function.hpp",
+ }
+
+ def dof_transformation(
+ self,
+ geometry: ElementGeometry,
+ element_index: Tuple[int, int, int],
+ element_type: Union[FaceType, CellType],
+ element_vertex_ordering: List[int],
+ ) -> Tuple[CustomCodeNode, sp.MatrixBase]:
+ return (
+ CustomCodeNode("", [], []),
+ sp.Identity(self.fe_space.num_dofs(geometry)),
+ )
diff --git a/hog/operator_generation/function_space_impls.py b/hog/operator_generation/function_space_impls.py
index 956001b..67ebd43 100644
--- a/hog/operator_generation/function_space_impls.py
+++ b/hog/operator_generation/function_space_impls.py
@@ -95,6 +95,11 @@ class FunctionSpaceImpl(ABC):
"""
impl_class: type
+ import hog.operator_generation.function_space_implementations.p0_space_impl
+ import hog.operator_generation.function_space_implementations.p1_space_impl
+ import hog.operator_generation.function_space_implementations.p2_space_impl
+ import hog.operator_generation.function_space_implementations.n1e1_space_impl
+
if isinstance(func_space, LagrangianFunctionSpace):
if func_space.degree == 2:
impl_class = P2FunctionSpaceImpl
@@ -244,810 +249,3 @@ class FunctionSpaceImpl(ABC):
return f"(*{self.name})"
else:
return self.name
-
-class P0FunctionSpaceImpl(FunctionSpaceImpl):
- def __init__(
- self,
- fe_space: FunctionSpace,
- name: str,
- type_descriptor: HOGType,
- is_pointer: bool = False,
- ):
- super().__init__(fe_space, name, type_descriptor, is_pointer)
- self.field = self._create_field(name)
-
- def pre_communication(self, dim: int) -> str:
- return ""
-
- def zero_halos(self, dim: int) -> str:
- return ""
-
- def post_communication(
- self, dim: int, params: str, transform_basis: bool = True
- ) -> str:
- return ""
-
- def pointer_retrieval(self, dim: int) -> str:
- """C++ code for retrieving pointers to the numerical data stored in the macro primitives `face` (2d) or `cell` (3d)."""
- Macro = {2: "Face", 3: "Cell"}[dim]
- macro = {2: "face", 3: "cell"}[dim]
-
- return f"{self.type_descriptor.pystencils_type}* _data_{self.name} = {self._deref()}.getDGFunction()->volumeDoFFunction()->dofMemory( it.first, level );"
-
- def invert_elementwise(self, dim: int) -> str:
- return f"{self._deref()}.invertElementwise( level );"
-
- def local_dofs(
- self,
- geometry: ElementGeometry,
- element_index: Tuple[int, int, int],
- element_type: Union[FaceType, CellType],
- indexing_info: IndexingInfo,
- element_vertex_ordering: List[int],
- ) -> List[Field.Access]:
- volume_dof_indices = micro_element_to_volume_indices(
- element_type, element_index, 1, VolumeDoFMemoryLayout.SoA
- )
-
- vertex_array_indices = [
- dof_idx.array_index(geometry, indexing_info)
- for dof_idx in volume_dof_indices
- ]
-
- return [
- self.field.absolute_access((idx,), (0,)) for idx in vertex_array_indices
- ]
-
- def func_type_string(self) -> str:
- return f"P0Function< {self.type_descriptor.pystencils_type} >"
-
- def includes(self) -> Set[str]:
- return {f"hyteg/p0functionspace/P0Function.hpp"}
-
- def dof_transformation(
- self,
- geometry: ElementGeometry,
- element_index: Tuple[int, int, int],
- element_type: Union[FaceType, CellType],
- element_vertex_ordering: List[int],
- ) -> Tuple[CustomCodeNode, sp.MatrixBase]:
- return (
- CustomCodeNode("", [], []),
- sp.Identity(self.fe_space.num_dofs(geometry)),
- )
-
-class P1FunctionSpaceImpl(FunctionSpaceImpl):
- def __init__(
- self,
- fe_space: FunctionSpace,
- name: str,
- type_descriptor: HOGType,
- is_pointer: bool = False,
- ):
- super().__init__(fe_space, name, type_descriptor, is_pointer)
- self.field = self._create_field(name)
-
- def pre_communication(self, dim: int) -> str:
- if dim == 2:
- return f"communication::syncFunctionBetweenPrimitives( {self.name}, level, communication::syncDirection_t::LOW2HIGH );"
- else:
- return (
- f"{self._deref()}.communicate< Face, Cell >( level );\n"
- f"{self._deref()}.communicate< Edge, Cell >( level );\n"
- f"{self._deref()}.communicate< Vertex, Cell >( level );"
- )
-
- def zero_halos(self, dim: int) -> str:
- if dim == 2:
- return (
- f"for ( const auto& idx : vertexdof::macroface::Iterator( level ) ) {{\n"
- f" if ( vertexdof::macroface::isVertexOnBoundary( level, idx ) ) {{\n"
- f" auto arrayIdx = vertexdof::macroface::index( level, idx.x(), idx.y() );\n"
- f" _data_{self.name}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
- f" }}\n"
- f"}}"
- )
- else:
- return (
- f"for ( const auto& idx : vertexdof::macrocell::Iterator( level ) ) {{\n"
- f" if ( !vertexdof::macrocell::isOnCellFace( idx, level ).empty() ) {{\n"
- f" auto arrayIdx = vertexdof::macrocell::index( level, idx.x(), idx.y(), idx.z() );\n"
- f" _data_{self.name}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
- f" }}\n"
- f"}}"
- )
-
- def post_communication(
- self, dim: int, params: str, transform_basis: bool = True
- ) -> str:
- if dim == 2:
- return (
- f"{self._deref()}.communicateAdditively < Face, Edge > ( {params} );\n"
- f"{self._deref()}.communicateAdditively < Face, Vertex > ( {params} );"
- )
- else:
- return (
- f"{self._deref()}.communicateAdditively< Cell, Face >( {params} );\n"
- f"{self._deref()}.communicateAdditively< Cell, Edge >( {params} );\n"
- f"{self._deref()}.communicateAdditively< Cell, Vertex >( {params} );"
- )
-
- def pointer_retrieval(self, dim: int) -> str:
- """C++ code for retrieving pointers to the numerical data stored in the macro primitives `face` (2d) or `cell` (3d)."""
- Macro = {2: "Face", 3: "Cell"}[dim]
- macro = {2: "face", 3: "cell"}[dim]
-
- return f"{self.type_descriptor.pystencils_type}* _data_{self.name} = {macro}.getData( {self._deref()}.get{Macro}DataID() )->getPointer( level );"
-
- def invert_elementwise(self, dim: int) -> str:
- return f"{self._deref()}.invertElementwise( level );"
-
- def local_dofs(
- self,
- geometry: ElementGeometry,
- element_index: Tuple[int, int, int],
- element_type: Union[FaceType, CellType],
- indexing_info: IndexingInfo,
- element_vertex_ordering: List[int],
- ) -> List[Field.Access]:
- vertex_dof_indices = micro_element_to_vertex_indices(
- geometry, element_type, element_index
- )
-
- vertex_dof_indices = [vertex_dof_indices[i] for i in element_vertex_ordering]
-
- vertex_array_indices = [
- dof_idx.array_index(geometry, indexing_info)
- for dof_idx in vertex_dof_indices
- ]
-
- return [
- self.field.absolute_access((idx,), (0,)) for idx in vertex_array_indices
- ]
-
- def func_type_string(self) -> str:
- return f"P1Function< {self.type_descriptor.pystencils_type} >"
-
- def includes(self) -> Set[str]:
- return {f"hyteg/p1functionspace/P1Function.hpp"}
-
- def dof_transformation(
- self,
- geometry: ElementGeometry,
- element_index: Tuple[int, int, int],
- element_type: Union[FaceType, CellType],
- element_vertex_ordering: List[int],
- ) -> Tuple[CustomCodeNode, sp.MatrixBase]:
- return (
- CustomCodeNode("", [], []),
- sp.Identity(self.fe_space.num_dofs(geometry)),
- )
-
-
-class P2FunctionSpaceImpl(FunctionSpaceImpl):
- def __init__(
- self,
- fe_space: FunctionSpace,
- name: str,
- type_descriptor: HOGType,
- is_pointer: bool = False,
- ):
- super().__init__(fe_space, name, type_descriptor, is_pointer)
- self.f_vertex = self._create_field(name + "Vertex")
- self.f_edge = self._create_field(name + "Edge")
-
- def pre_communication(self, dim: int) -> str:
- if dim == 2:
- return f"communication::syncFunctionBetweenPrimitives( {self.name}, level, communication::syncDirection_t::LOW2HIGH );"
- else:
- return (
- f"{self._deref()}.communicate< Face, Cell >( level );\n"
- f"{self._deref()}.communicate< Edge, Cell >( level );\n"
- f"{self._deref()}.communicate< Vertex, Cell >( level );"
- )
-
- def zero_halos(self, dim: int) -> str:
- if dim == 2:
- return (
- f"for ( const auto& idx : vertexdof::macroface::Iterator( level ) )\n"
- f"{{\n"
- f" if ( vertexdof::macroface::isVertexOnBoundary( level, idx ) )\n"
- f" {{\n"
- f" auto arrayIdx = vertexdof::macroface::index( level, idx.x(), idx.y() );\n"
- f" _data_{self.name }Vertex[arrayIdx] = walberla::numeric_cast< {self.type_descriptor.pystencils_type} >( 0 );\n"
- f" }}\n"
- f"}}\n"
- f"for ( const auto& idx : edgedof::macroface::Iterator( level ) )\n"
- f"{{\n"
- f" for ( const auto& orientation : edgedof::faceLocalEdgeDoFOrientations )\n"
- f" {{\n"
- f" if ( !edgedof::macroface::isInnerEdgeDoF( level, idx, orientation ) )\n"
- f" {{\n"
- f" auto arrayIdx = edgedof::macroface::index( level, idx.x(), idx.y(), orientation );\n"
- f" _data_{self.name }Edge[arrayIdx] = walberla::numeric_cast< {self.type_descriptor.pystencils_type} >( 0 );\n"
- f" }}\n"
- f" }}\n"
- f"}}"
- )
- else:
- return (
- f"for ( const auto& idx : vertexdof::macrocell::Iterator( level ) )\n"
- f"{{\n"
- f"if ( !vertexdof::macrocell::isOnCellFace( idx, level ).empty() )\n"
- f" {{\n"
- f" auto arrayIdx = vertexdof::macrocell::index( level, idx.x(), idx.y(), idx.z() );\n"
- f" _data_{self.name}Vertex[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
- f" }}\n"
- f"}}\n"
- f"edgedof::macrocell::setBoundaryToZero( level, cell, {self._deref()}.getEdgeDoFFunction().getCellDataID() );"
- )
-
- def post_communication(
- self, dim: int, params: str, transform_basis: bool = True
- ) -> str:
- if dim == 2:
- return (
- f"{self._deref()}.getVertexDoFFunction().communicateAdditively< Face, Edge >( {params} );\n"
- f"{self._deref()}.getVertexDoFFunction().communicateAdditively< Face, Vertex >( {params} );\n"
- f"{self._deref()}.getEdgeDoFFunction().communicateAdditively< Face, Edge >( {params} );"
- )
- else:
- return (
- f"{self._deref()}.getVertexDoFFunction().communicateAdditively< Cell, Face >( {params} );\n"
- f"{self._deref()}.getVertexDoFFunction().communicateAdditively< Cell, Edge >( {params} );\n"
- f"{self._deref()}.getVertexDoFFunction().communicateAdditively< Cell, Vertex >( {params} );\n"
- f"{self._deref()}.getEdgeDoFFunction().communicateAdditively< Cell, Face >( {params} );\n"
- f"{self._deref()}.getEdgeDoFFunction().communicateAdditively< Cell, Edge >( {params} );"
- )
-
- def pointer_retrieval(self, dim: int) -> str:
- """C++ code for retrieving pointers to the numerical data stored in the macro primitives `face` (2d) or `cell` (3d)."""
- Macro = {2: "Face", 3: "Cell"}[dim]
- macro = {2: "face", 3: "cell"}[dim]
-
- return (
- f"{self.type_descriptor.pystencils_type}* _data_{self.name}Vertex = {macro}.getData( {self._deref()}.getVertexDoFFunction().get{Macro}DataID() )->getPointer( level );\n"
- f"{self.type_descriptor.pystencils_type}* _data_{self.name}Edge = {macro}.getData( {self._deref()}.getEdgeDoFFunction().get{Macro}DataID() )->getPointer( level );"
- )
-
- def invert_elementwise(self, dim: int) -> str:
- return f"{self._deref()}.invertElementwise( level );"
-
- def local_dofs(
- self,
- geometry: ElementGeometry,
- element_index: Tuple[int, int, int],
- element_type: Union[FaceType, CellType],
- indexing_info: IndexingInfo,
- element_vertex_ordering: List[int],
- ) -> List[Field.Access]:
- vertex_dof_indices = micro_element_to_vertex_indices(
- geometry, element_type, element_index
- )
-
- vertex_dof_indices = [vertex_dof_indices[i] for i in element_vertex_ordering]
-
- edge_dof_indices = micro_vertex_to_edge_indices(geometry, vertex_dof_indices)
-
- vrtx_array_idcs = [
- dof_idx.array_index(geometry, indexing_info)
- for dof_idx in vertex_dof_indices
- ]
- edge_array_idcs = [
- dof_idx.array_index(geometry, indexing_info) for dof_idx in edge_dof_indices
- ]
-
- # NOTE: The order of DoFs must match the order of shape functions
- # defined in the FunctionSpace. Holds within the individual lists
- # and overall.
- return [
- self.f_vertex.absolute_access((idx,), (0,)) for idx in vrtx_array_idcs
- ] + [self.f_edge.absolute_access((idx,), (0,)) for idx in edge_array_idcs]
-
- def func_type_string(self) -> str:
- return f"P2Function< {self.type_descriptor.pystencils_type} >"
-
- def includes(self) -> Set[str]:
- return {f"hyteg/p2functionspace/P2Function.hpp"}
-
- def dof_transformation(
- self,
- geometry: ElementGeometry,
- element_index: Tuple[int, int, int],
- element_type: Union[FaceType, CellType],
- element_vertex_ordering: List[int],
- ) -> Tuple[CustomCodeNode, sp.MatrixBase]:
- return (
- CustomCodeNode("", [], []),
- sp.Identity(self.fe_space.num_dofs(geometry)),
- )
-
-
-class P1VectorFunctionSpaceImpl(FunctionSpaceImpl):
- def __init__(
- self,
- fe_space: FunctionSpace,
- name: str,
- type_descriptor: HOGType,
- is_pointer: bool = False,
- ):
- super().__init__(fe_space, name, type_descriptor, is_pointer)
- self.fields: Dict[int, Field] = {}
-
- def _field_name(self, component: int) -> str:
- return self.name + f"_{component}"
-
- def _raw_pointer_name(self, component: int) -> str:
- return f"_data_" + self._field_name(component)
-
- def pre_communication(self, dim: int) -> str:
- if dim == 2:
- return f"communication::syncVectorFunctionBetweenPrimitives( {self.name}, level, communication::syncDirection_t::LOW2HIGH );"
- else:
- ret_str = ""
- for i in range(dim):
- ret_str += (
- f"{self._deref()}[{i}].communicate< Face, Cell >( level );\n"
- f"{self._deref()}[{i}].communicate< Edge, Cell >( level );\n"
- f"{self._deref()}[{i}].communicate< Vertex, Cell >( level );"
- )
- return ret_str
-
- def zero_halos(self, dim: int) -> str:
- if dim == 2:
- return (
- f"for ( const auto& idx : vertexdof::macroface::Iterator( level ) ) {{\n"
- f" if ( vertexdof::macroface::isVertexOnBoundary( level, idx ) ) {{\n"
- f" auto arrayIdx = vertexdof::macroface::index( level, idx.x(), idx.y() );\n"
- f" {self._raw_pointer_name(0)}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
- f" {self._raw_pointer_name(1)}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
- f" }}\n"
- f"}}"
- )
- else:
- return (
- f"for ( const auto& idx : vertexdof::macrocell::Iterator( level ) ) {{\n"
- f" if ( !vertexdof::macrocell::isOnCellFace( idx, level ).empty() ) {{\n"
- f" auto arrayIdx = vertexdof::macrocell::index( level, idx.x(), idx.y(), idx.z() );\n"
- f" {self._raw_pointer_name(0)}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
- f" {self._raw_pointer_name(1)}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
- f" {self._raw_pointer_name(2)}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
- f" }}\n"
- f"}}"
- )
-
- def post_communication(
- self, dim: int, params: str, transform_basis: bool = True
- ) -> str:
- if dim == 2:
- ret_str = ""
- for i in range(dim):
- ret_str += (
- f"{self._deref()}[{i}].communicateAdditively < Face, Edge > ( {params} );\n"
- f"{self._deref()}[{i}].communicateAdditively < Face, Vertex > ( {params} );\n"
- )
- return ret_str
- else:
- ret_str = ""
- for i in range(dim):
- ret_str += (
- f"{self._deref()}[{i}].communicateAdditively< Cell, Face >( {params} );\n"
- f"{self._deref()}[{i}].communicateAdditively< Cell, Edge >( {params} );\n"
- f"{self._deref()}[{i}].communicateAdditively< Cell, Vertex >( {params} );"
- )
- return ret_str
-
- def pointer_retrieval(self, dim: int) -> str:
- """C++ code for retrieving pointers to the numerical data stored in the macro primitives `face` (2d) or `cell` (3d)."""
- Macro = {2: "Face", 3: "Cell"}[dim]
- macro = {2: "face", 3: "cell"}[dim]
-
- ret_str = ""
- for i in range(dim):
- ret_str += f"{self.type_descriptor.pystencils_type}* {self._raw_pointer_name(i)} = {macro}.getData( {self._deref()}[{i}].get{Macro}DataID() )->getPointer( level );\n"
- return ret_str
-
- def invert_elementwise(self, dim: int) -> str:
- ret_str = ""
- for i in range(dim):
- ret_str += f"{self._deref()}[{i}].invertElementwise( level );\n"
- return ret_str
-
- def local_dofs(
- self,
- geometry: ElementGeometry,
- element_index: Tuple[int, int, int],
- element_type: Union[FaceType, CellType],
- indexing_info: IndexingInfo,
- element_vertex_ordering: List[int],
- ) -> List[Field.Access]:
- """
- Returns the element-local DoFs (field accesses) in a list (i.e., linearized).
-
- The order here has to match that of the function space implementation.
- TODO: This is a little concerning since that order has to match in two very different parts of this
- implementation. Here, and in the TensorialVectorFunctionSpace. Maybe we can define this order in a single
- location.
- We return them in "SoA" order: first all DoFs corresponding to the function with the first component != 0,
- then all DoFs corresponding to the function with the second component != 0 etc.
-
- For instance, in 2D we get the DoFs corresponding to the 6 shape functions in the following order:
-
- First component != 0
-
- list[0]: ⎡-x_ref_0 - x_ref_1 + 1⎤
- ⎢ ⎥
- ⎣ 0 ⎦
-
- list[1]: ⎡x_ref_0⎤
- ⎢ ⎥
- ⎣ 0 ⎦
-
- list[2]: ⎡x_ref_1⎤
- ⎢ ⎥
- ⎣ 0 ⎦
-
- Second component != 0
-
- list[3]: ⎡ 0 ⎤
- ⎢ ⎥
- ⎣-x_ref_0 - x_ref_1 + 1⎦
-
- list[4]: ⎡ 0 ⎤
- ⎢ ⎥
- ⎣x_ref_0⎦
-
- list[5]: ⎡ 0 ⎤
- ⎢ ⎥
- ⎣x_ref_1⎦
- """
-
- for c in range(geometry.dimensions):
- if c not in self.fields:
- self.fields[c] = self._create_field(self._field_name(c))
-
- vertex_dof_indices = micro_element_to_vertex_indices(
- geometry, element_type, element_index
- )
-
- vertex_dof_indices = [vertex_dof_indices[i] for i in element_vertex_ordering]
-
- vertex_array_indices = [
- dof_idx.array_index(geometry, indexing_info)
- for dof_idx in vertex_dof_indices
- ]
-
- return [
- self.fields[c].absolute_access((idx,), (0,))
- for c in range(geometry.dimensions)
- for idx in vertex_array_indices
- ]
-
- def func_type_string(self) -> str:
- return f"P1VectorFunction< {self.type_descriptor.pystencils_type} >"
-
- def includes(self) -> Set[str]:
- return {f"hyteg/p1functionspace/P1VectorFunction.hpp"}
-
- def dof_transformation(
- self,
- geometry: ElementGeometry,
- element_index: Tuple[int, int, int],
- element_type: Union[FaceType, CellType],
- element_vertex_ordering: List[int],
- ) -> Tuple[CustomCodeNode, sp.MatrixBase]:
- return (
- CustomCodeNode("", [], []),
- sp.Identity(self.fe_space.num_dofs(geometry)),
- )
-
-
-class P2VectorFunctionSpaceImpl(FunctionSpaceImpl):
- def __init__(
- self,
- fe_space: FunctionSpace,
- name: str,
- type_descriptor: HOGType,
- is_pointer: bool = False,
- ):
- super().__init__(fe_space, name, type_descriptor, is_pointer)
- self.fields: Dict[Tuple[str, int], Field] = {}
-
- def _field_name(self, component: int, dof_type: str) -> str:
- """dof_type should either be 'vertex' or 'edge'"""
- return self.name + f"_{dof_type}_{component}"
-
- def _raw_pointer_name(self, component: int, dof_type: str) -> str:
- """dof_type should either be 'vertex' or 'edge'"""
- return f"_data_" + self._field_name(component, dof_type)
-
- def pre_communication(self, dim: int) -> str:
- if dim == 2:
- return f"communication::syncVectorFunctionBetweenPrimitives( {self.name}, level, communication::syncDirection_t::LOW2HIGH );"
- else:
- ret_str = ""
- for i in range(dim):
- ret_str += (
- f"{self._deref()}[{i}].communicate< Face, Cell >( level );\n"
- f"{self._deref()}[{i}].communicate< Edge, Cell >( level );\n"
- f"{self._deref()}[{i}].communicate< Vertex, Cell >( level );"
- )
- return ret_str
-
- def zero_halos(self, dim: int) -> str:
- if dim == 2:
- return (
- f"for ( const auto& idx : vertexdof::macroface::Iterator( level ) )\n"
- f"{{\n"
- f" if ( vertexdof::macroface::isVertexOnBoundary( level, idx ) )\n"
- f" {{\n"
- f" auto arrayIdx = vertexdof::macroface::index( level, idx.x(), idx.y() );\n"
- f" {self._raw_pointer_name(0, 'vertex')}[arrayIdx] = walberla::numeric_cast< {self.type_descriptor.pystencils_type} >( 0 );\n"
- f" {self._raw_pointer_name(1, 'vertex')}[arrayIdx] = walberla::numeric_cast< {self.type_descriptor.pystencils_type} >( 0 );\n"
- f" }}\n"
- f"}}\n"
- f"for ( const auto& idx : edgedof::macroface::Iterator( level ) )\n"
- f"{{\n"
- f" for ( const auto& orientation : edgedof::faceLocalEdgeDoFOrientations )\n"
- f" {{\n"
- f" if ( !edgedof::macroface::isInnerEdgeDoF( level, idx, orientation ) )\n"
- f" {{\n"
- f" auto arrayIdx = edgedof::macroface::index( level, idx.x(), idx.y(), orientation );\n"
- f" {self._raw_pointer_name(0, 'edge')}[arrayIdx] = walberla::numeric_cast< {self.type_descriptor.pystencils_type} >( 0 );\n"
- f" {self._raw_pointer_name(1, 'edge')}[arrayIdx] = walberla::numeric_cast< {self.type_descriptor.pystencils_type} >( 0 );\n"
- f" }}\n"
- f" }}\n"
- f"}}"
- )
- else:
- return (
- f"for ( const auto& idx : vertexdof::macrocell::Iterator( level ) )\n"
- f"{{\n"
- f"if ( !vertexdof::macrocell::isOnCellFace( idx, level ).empty() )\n"
- f" {{\n"
- f" auto arrayIdx = vertexdof::macrocell::index( level, idx.x(), idx.y(), idx.z() );\n"
- f" {self._raw_pointer_name(0, 'vertex')}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
- f" {self._raw_pointer_name(1, 'vertex')}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
- f" {self._raw_pointer_name(2, 'vertex')}[arrayIdx] = {self.type_descriptor.pystencils_type}( 0 );\n"
- f" }}\n"
- f"}}\n"
- f"edgedof::macrocell::setBoundaryToZero( level, cell, {self._deref()}[0].getEdgeDoFFunction().getCellDataID() );\n"
- f"edgedof::macrocell::setBoundaryToZero( level, cell, {self._deref()}[1].getEdgeDoFFunction().getCellDataID() );\n"
- f"edgedof::macrocell::setBoundaryToZero( level, cell, {self._deref()}[2].getEdgeDoFFunction().getCellDataID() );\n"
- )
-
- def post_communication(
- self, dim: int, params: str, transform_basis: bool = True
- ) -> str:
- if dim == 2:
- ret_str = ""
- for i in range(dim):
- ret_str += (
- f"{self._deref()}[{i}].getVertexDoFFunction().communicateAdditively< Face, Edge >( {params} );\n"
- f"{self._deref()}[{i}].getVertexDoFFunction().communicateAdditively< Face, Vertex >( {params} );\n"
- f"{self._deref()}[{i}].getEdgeDoFFunction().communicateAdditively< Face, Edge >( {params} );"
- )
- return ret_str
- else:
- ret_str = ""
- for i in range(dim):
- ret_str += (
- f"{self._deref()}[{i}].getVertexDoFFunction().communicateAdditively< Cell, Face >( {params} );\n"
- f"{self._deref()}[{i}].getVertexDoFFunction().communicateAdditively< Cell, Edge >( {params} );\n"
- f"{self._deref()}[{i}].getVertexDoFFunction().communicateAdditively< Cell, Vertex >( {params} );\n"
- f"{self._deref()}[{i}].getEdgeDoFFunction().communicateAdditively< Cell, Face >( {params} );\n"
- f"{self._deref()}[{i}].getEdgeDoFFunction().communicateAdditively< Cell, Edge >( {params} );"
- )
- return ret_str
-
- def pointer_retrieval(self, dim: int) -> str:
- """C++ code for retrieving pointers to the numerical data stored in the macro primitives `face` (2d) or `cell` (3d)."""
- Macro = {2: "Face", 3: "Cell"}[dim]
- macro = {2: "face", 3: "cell"}[dim]
-
- ret_str = ""
- for i in range(dim):
- ret_str += f"{self.type_descriptor.pystencils_type}* {self._raw_pointer_name(i, 'vertex')} = {macro}.getData( {self._deref()}[{i}].getVertexDoFFunction().get{Macro}DataID() )->getPointer( level );\n"
- ret_str += f"{self.type_descriptor.pystencils_type}* {self._raw_pointer_name(i, 'edge')} = {macro}.getData( {self._deref()}[{i}].getEdgeDoFFunction().get{Macro}DataID() )->getPointer( level );\n"
- return ret_str
-
- def invert_elementwise(self, dim: int) -> str:
- ret_str = ""
- for i in range(dim):
- ret_str += f"{self._deref()}[{i}].invertElementwise( level );\n"
- return ret_str
-
- def local_dofs(
- self,
- geometry: ElementGeometry,
- element_index: Tuple[int, int, int],
- element_type: Union[FaceType, CellType],
- indexing_info: IndexingInfo,
- element_vertex_ordering: List[int],
- ) -> List[Field.Access]:
- """
- Returns the element-local DoFs (field accesses) in a list (i.e., linearized).
-
- See P1VectorFunctionSpaceImpl::local_dofs() for details.
- """
-
- for dt in ["vertex", "edge"]:
- for c in range(geometry.dimensions):
- if (dt, c) not in self.fields:
- self.fields[(dt, c)] = self._create_field(self._field_name(c, dt))
-
- vertex_dof_indices = micro_element_to_vertex_indices(
- geometry, element_type, element_index
- )
-
- vertex_dof_indices = [vertex_dof_indices[i] for i in element_vertex_ordering]
-
- edge_dof_indices = micro_vertex_to_edge_indices(geometry, vertex_dof_indices)
-
- vertex_array_indices = [
- dof_idx.array_index(geometry, indexing_info)
- for dof_idx in vertex_dof_indices
- ]
-
- edge_array_indices = [
- dof_idx.array_index(geometry, indexing_info) for dof_idx in edge_dof_indices
- ]
-
- loc_dofs = []
-
- for c in range(geometry.dimensions):
- loc_dofs += [
- self.fields[("vertex", c)].absolute_access((idx,), (0,))
- for idx in vertex_array_indices
- ]
- loc_dofs += [
- self.fields[("edge", c)].absolute_access((idx,), (0,))
- for idx in edge_array_indices
- ]
-
- return loc_dofs
-
- def func_type_string(self) -> str:
- return f"P2VectorFunction< {self.type_descriptor.pystencils_type} >"
-
- def includes(self) -> Set[str]:
- return {
- f"hyteg/p2functionspace/P2VectorFunction.hpp",
- f"hyteg/p2functionspace/P2Function.hpp",
- }
-
- def dof_transformation(
- self,
- geometry: ElementGeometry,
- element_index: Tuple[int, int, int],
- element_type: Union[FaceType, CellType],
- element_vertex_ordering: List[int],
- ) -> Tuple[CustomCodeNode, sp.MatrixBase]:
- return (
- CustomCodeNode("", [], []),
- sp.Identity(self.fe_space.num_dofs(geometry)),
- )
-
-
-class N1E1FunctionSpaceImpl(FunctionSpaceImpl):
- def __init__(
- self,
- fe_space: FunctionSpace,
- name: str,
- type_descriptor: HOGType,
- is_pointer: bool = False,
- ):
- super().__init__(fe_space, name, type_descriptor, is_pointer)
- self.field = self._create_field(name)
-
- def pre_communication(self, dim: int) -> str:
- if dim == 2:
- return ""
- else:
- return (
- f"{self._deref()}.communicate< Face, Cell >( level );\n"
- f"{self._deref()}.communicate< Edge, Cell >( level );"
- )
-
- def zero_halos(self, dim: int) -> str:
- if dim == 2:
- return ""
- else:
- return f"edgedof::macrocell::setBoundaryToZero( level, cell, {self._deref()}.getDoFs()->getCellDataID() );"
-
- def post_communication(
- self, dim: int, params: str, transform_basis: bool = True
- ) -> str:
- if dim == 2:
- return ""
- else:
- dofs = ""
- if not transform_basis:
- dofs = "getDoFs()->"
- return (
- f"{self._deref()}.{dofs}communicateAdditively< Cell, Face >( {params} );\n"
- f"{self._deref()}.{dofs}communicateAdditively< Cell, Edge >( {params} );"
- )
-
- def pointer_retrieval(self, dim: int) -> str:
- """C++ code for retrieving pointers to the numerical data stored in the macro primitives `face` (2d) or `cell` (3d)."""
- Macro = {2: "Face", 3: "Cell"}[dim]
- macro = {2: "face", 3: "cell"}[dim]
-
- return f"{self.type_descriptor.pystencils_type}* _data_{self.name} = {macro}.getData( {self._deref()}.getDoFs()->get{Macro}DataID() )->getPointer( level );"
-
- def invert_elementwise(self, dim: int) -> str:
- return f"{self._deref()}.getDoFs()->invertElementwise( level );"
-
- def local_dofs(
- self,
- geometry: ElementGeometry,
- element_index: Tuple[int, int, int],
- element_type: Union[FaceType, CellType],
- indexing_info: IndexingInfo,
- element_vertex_ordering: List[int],
- ) -> List[Field.Access]:
- vertex_dof_indices = micro_element_to_vertex_indices(
- geometry, element_type, element_index
- )
-
- vertex_dof_indices = [vertex_dof_indices[i] for i in element_vertex_ordering]
-
- edge_dof_indices = micro_vertex_to_edge_indices(geometry, vertex_dof_indices)
- edge_array_indices = [
- dof_idx.array_index(geometry, indexing_info) for dof_idx in edge_dof_indices
- ]
-
- return [self.field.absolute_access((idx,), (0,)) for idx in edge_array_indices]
-
- def func_type_string(self) -> str:
- return f"n1e1::N1E1VectorFunction< {self.type_descriptor.pystencils_type} >"
-
- def includes(self) -> Set[str]:
- return {
- "hyteg/n1e1functionspace/N1E1VectorFunction.hpp",
- "hyteg/n1e1functionspace/N1E1MacroCell.hpp",
- }
-
- def dof_transformation(
- self,
- geometry: ElementGeometry,
- element_index: Tuple[int, int, int],
- element_type: Union[FaceType, CellType],
- element_vertex_ordering: List[int],
- ) -> Tuple[CustomCodeNode, sp.MatrixBase]:
-
- if element_vertex_ordering != [0, 1, 2, 3]:
- raise HOGException(
- "Element vertex re-ordering not supported for Nédélec elements (yet)."
- )
-
- Macro = {2: "Face", 3: "Cell"}[geometry.dimensions]
- macro = {2: "face", 3: "cell"}[geometry.dimensions]
- name = "basisTransformation"
- n_dofs = self.fe_space.num_dofs(geometry)
-
- # NOTE: Types are added manually because pystencils won't add types to accessed/
- # defined symbols of custom code nodes. As a result the symbols
- # in the matrix will be typed but not their definition, leading to
- # undefined symbols.
- # NOTE: The type is `real_t` (not `self._type_descriptor.pystencils_type`)
- # because this function is implemented manually in HyTeG with
- # this signature. Passing `np.float64` is not ideal (if `real_t !=
- # double`) but it makes sure that casts are inserted if necessary
- # (though some might be superfluous).
- symbols = [
- ps.TypedSymbol(f"{name}.diagonal()({i})", np.float64) for i in range(n_dofs)
- ]
- return (
- CustomCodeNode(
- f"const Eigen::DiagonalMatrix< real_t, {n_dofs} > {name} = "
- f"n1e1::macro{macro}::basisTransformation( level, {macro}, {{{element_index[0]}, {element_index[1]}, {element_index[2]}}}, {macro}dof::{Macro}Type::{element_type.name} );",
- [
- ps.TypedSymbol("level", "const uint_t"),
- ps.TypedSymbol(f"{macro}", f"const {Macro}&"),
- ],
- symbols,
- ),
- sp.diag(*symbols),
- )
--
GitLab