Add rotation wrapper for epsilon operator

Fix return branch without wrapper

hog/forms.py

@@ -33,7 +33,7 @@ from hog.fem_helpers import (
     fem_function_on_element,
     fem_function_gradient_on_element,
 )
-from hog.function_space import FunctionSpace, N1E1Space, TrialSpace, TestSpace
+from hog.function_space import FunctionSpace, N1E1Space, TrialSpace, TestSpace, LagrangianFunctionSpace
 from hog.math_helpers import dot, inv, abs, det, double_contraction
 from hog.quadrature import Quadrature, Tabulation
 from hog.symbolizer import Symbolizer
@@ -321,6 +321,8 @@ def epsilon(
     component_test: int = 0,
     variable_viscosity: bool = True,
     coefficient_function_space: Optional[FunctionSpace] = None,
+    rotation_wrapper: bool = False,
+    normal_fspace: Optional[FunctionSpace] = None,
 ) -> Form:
     docstring = f"""
 "Epsilon" operator.
@@ -353,6 +355,8 @@ where
     if blending == IdentityMap():
         integr = integrand_affine
 
+
+
     return process_integrand(
         integr,
         trial,
@@ -362,7 +366,13 @@ where
         blending=blending,
         is_symmetric=trial == test,
         docstring=docstring,
-        fe_coefficients={"mu": coefficient_function_space},
+        fe_coefficients={
+            "mu": coefficient_function_space, 
+            "nx": normal_fspace, 
+            "ny": normal_fspace, 
+            "nz": normal_fspace
+        } if rotation_wrapper else {"mu": coefficient_function_space},
+        rotation_wrapper = rotation_wrapper,
     )
 
 

hog/integrand.py

@@ -73,6 +73,13 @@ from hog.fem_helpers import (
     trafo_ref_to_physical,
 )
 from hog.math_helpers import inv, det
+from enum import Enum
+
+
+class RotationType(Enum):
+    PRE_MULTIPLY = 1
+    POST_MULTIPLY = 2
+    PRE_AND_POST_MULTIPLY = 3
 
 
 @dataclass
@@ -87,6 +94,8 @@ class Form:
     symmetric: bool
     free_symbols: List[sp.Symbol] = field(default_factory=lambda: list())
     docstring: str = ""
+    rotmat: sp.MatrixBase = sp.Matrix([[0]])
+    rot_type: RotationType = RotationType.PRE_AND_POST_MULTIPLY
 
     def integrate(self, quad: Quadrature, symbolizer: Symbolizer) -> sp.Matrix:
         """Integrates the form using the passed quadrature directly, i.e. without tabulations or loops."""
@@ -230,6 +239,7 @@ def process_integrand(
     boundary_geometry: ElementGeometry | None = None,
     fe_coefficients: Dict[str, Union[FunctionSpace, None]] | None = None,
     is_symmetric: bool = False,
+    rotation_wrapper: bool = False,
     docstring: str = "",
 ) -> Form:
     """
@@ -375,10 +385,10 @@ def process_integrand(
                 f"You cannot use the name {special_name_of_micromesh_coeff} for your FE coefficient."
                 f"It is reserved."
             )
-        fe_coefficients_modified[special_name_of_micromesh_coeff] = (
-            TensorialVectorFunctionSpace(
-                LagrangianFunctionSpace(blending.degree, symbolizer)
-            )
+        fe_coefficients_modified[
+            special_name_of_micromesh_coeff
+        ] = TensorialVectorFunctionSpace(
+            LagrangianFunctionSpace(blending.degree, symbolizer)
         )
 
     s.k = dict()
@@ -491,10 +501,58 @@ def process_integrand(
 
     free_symbols_sorted = sorted(list(free_symbols), key=lambda x: str(x))
 
+    if rotation_wrapper:
+        if not trial.is_vectorial:
+            raise HOGException("Nope")
+
+        from hog.recipes.integrands.volume.rotation import rotation_matrix
+
+        normal_fspace = fe_coefficients_modified["nx"]
+
+        # phi_eval_symbols = tabulation.register_phi_evals(
+        #     normal_fspace.shape(volume_geometry)
+        # )
+
+        normals = ["nx", "ny"] if volume_geometry.dimensions == 2 else ["nx", "ny", "nz"]
+
+        n_dof_symbols = []
+
+
+        for normal in normals:
+            if normal_fspace is None:
+                raise HOGException("Invalid normal function space")
+            else:
+                _, nc_dof_symbols = fem_function_on_element(
+                    normal_fspace,
+                    volume_geometry,
+                    symbolizer,
+                    domain="reference",
+                    function_id=normal,
+                    # basis_eval=phi_eval_symbols,
+                )
+
+                n_dof_symbols.append(nc_dof_symbols)
+
+        rotmat = rotation_matrix(
+            trial.num_dofs(volume_geometry),
+            int(trial.num_dofs(volume_geometry) / volume_geometry.dimensions),
+            n_dof_symbols,
+            volume_geometry,
+        )
+
+        return Form(
+            mat,
+            tabulation,
+            symmetric=is_symmetric,
+            free_symbols=free_symbols_sorted,
+            docstring=docstring,
+            rotmat=rotmat,
+        )
+    
     return Form(
         mat,
         tabulation,
         symmetric=is_symmetric,
         free_symbols=free_symbols_sorted,
         docstring=docstring,
-    )
+    )
\ No newline at end of file

hog/operator_generation/operators.py

@@ -101,6 +101,7 @@ from hog.operator_generation.optimizer import Optimizer, Opts
 from hog.quadrature import QuadLoop, Quadrature
 from hog.symbolizer import Symbolizer
 from hog.operator_generation.types import HOGType
+from hog.integrand import RotationType
 
 
 class MacroIntegrationDomain(Enum):
@@ -397,6 +398,17 @@ class HyTeGElementwiseOperator:
                                 mat[row, col], self.symbolizer, blending
                             )
 
+        if not form.rotmat.is_zero_matrix:
+            if form.rot_type == RotationType.PRE_AND_POST_MULTIPLY:
+                mat = form.rotmat * mat * form.rotmat.T
+            elif form.rot_type == RotationType.PRE_MULTIPLY:
+                mat = form.rotmat * mat
+            elif form.rot_type == RotationType.POST_MULTIPLY:
+                mat = mat * form.rotmat.T
+            else:
+                raise HOGException("Not implemented")
+
+
         if volume_geometry.space_dimension not in self.integration_infos:
             self.integration_infos[volume_geometry.space_dimension] = []
 

hog/recipes/integrands/volume/rotation.py

+# 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/>.
+
+from hog.recipes.common import *
+from hog.exception import HOGException
+
+
+def rotation_matrix(
+    mat_size,
+    mat_comp_size,
+    n_dof_symbols,
+    geometry
+):
+
+    rotmat = sp.zeros(mat_size, mat_size)
+
+    for row in range(mat_comp_size):
+        for col in range(mat_comp_size):
+            if row != col:
+                continue
+
+            nx_ = n_dof_symbols[0][row]
+            ny_ = n_dof_symbols[1][row]
+
+            if geometry.dimensions == 2:
+                n_ = sp.Matrix([[nx_], [ny_]])
+                nnorm = sp.sqrt(nx_ * nx_ + ny_ * ny_)
+                rotmat_ = sp.Matrix([[-ny_, nx_], [nx_, ny_]])
+            elif geometry.dimensions == 3:
+                nz_ = n_dof_symbols[2][row]
+                n_ = sp.Matrix([[nx_], [ny_], [nz_]])
+                nnorm = sp.sqrt(nx_ * nx_ + ny_ * ny_ + nz_ * nz_)
+
+                ex = sp.Matrix([[1.0], [0.0], [0.0]])
+                ey = sp.Matrix([[0.0], [1.0], [0.0]])
+                ez = sp.Matrix([[0.0], [0.0], [1.0]])
+
+                ncross = [n_.cross(ex), n_.cross(ey), n_.cross(ez)]
+                ncrossnorm = [nc_.norm() for nc_ in ncross]
+
+                t1all = [
+                    n_.cross(ex).normalized(),
+                    n_.cross(ey).normalized(),
+                    n_.cross(ez).normalized(),
+                ]
+
+                getrotmat = lambda t1: (t1.row_join(n_.cross(t1)).row_join(n_)).T
+
+                machine_eps = 1e-10
+                rotmat_ = sp.eye(geometry.dimensions)
+                for iRot in range(geometry.dimensions):
+                    for jRot in range(geometry.dimensions):
+                        rotmat_[iRot, jRot] = sp.Piecewise(
+                            (
+                                getrotmat(t1all[0])[iRot, jRot],
+                                sp.And(
+                                    ncrossnorm[0] + machine_eps > ncrossnorm[1],
+                                    ncrossnorm[0] + machine_eps > ncrossnorm[2],
+                                ),
+                            ),
+                            (
+                                getrotmat(t1all[1])[iRot, jRot],
+                                sp.And(
+                                    ncrossnorm[1] + machine_eps > ncrossnorm[0],
+                                    ncrossnorm[1] + machine_eps > ncrossnorm[2],
+                                ),
+                            ),
+                            (getrotmat(t1all[2])[iRot, jRot], True),
+                        )
+
+                # print("rotmat_ = ", rotmat_.subs([(nx_, 0.0), (ny_, 0.0), (nz_, 1.0)]))
+                # print("rotmat_ = ", rotmat_[0, 0])
+                # exit()
+            else:
+                raise HOGException(
+                    "We have not reached your dimensions, supreme being or little one"
+                )
+
+
+            rotmatid_ = sp.eye(geometry.dimensions)
+
+            ndofs = mat_comp_size
+
+            for iRot in range(geometry.dimensions):
+                for jRot in range(geometry.dimensions):
+                    rotmat[row + iRot * ndofs, col + jRot * ndofs] = sp.Piecewise(
+                        (rotmat_[iRot, jRot], nnorm > 0.5),
+                        (rotmatid_[iRot, jRot], True),
+                    )
+
+    return rotmat