Handle tensor/vector repr in 2D and axisymmetric cases

Author: bleyerj

Diff for: bindings/python/mgis/fenics/gradient_flux.py

@@ -17,21 +17,22 @@
     get_quadrature_element,
     vector_to_tensor,
 )
+import mgis.behaviour as mgis_bv
 
 
 class QuadratureFunction:
     """An abstract class for functions defined at quadrature points"""
-    def __init__(self, name, shape):
+
+    def __init__(self, name, shape, hypothesis):
         self.shape = shape
         self.name = name
+        self.hypothesis = hypothesis
 
     def initialize_function(self, mesh, quadrature_degree):
         self.quadrature_degree = quadrature_degree
         self.mesh = mesh
-        self.dx = Measure("dx",
-                          metadata={"quadrature_degree": quadrature_degree})
-        We = get_quadrature_element(mesh.ufl_cell(), quadrature_degree,
-                                    self.shape)
+        self.dx = Measure("dx", metadata={"quadrature_degree": quadrature_degree})
+        We = get_quadrature_element(mesh.ufl_cell(), quadrature_degree, self.shape)
         self.function_space = FunctionSpace(mesh, We)
         self.function = Function(self.function_space, name=self.name)
 
@@ -64,12 +65,13 @@ def project_on(self, space, degree, as_tensor=False, **kwargs):
             V = VectorFunctionSpace(self.mesh, space, degree, dim=self.shape)
         v = Function(V, name=self.name)
         v.assign(
-            project(fun,
-                    V,
-                    form_compiler_parameters={
-                        "quadrature_degree": self.quadrature_degree
-                    },
-                    **kwargs))
+            project(
+                fun,
+                V,
+                form_compiler_parameters={"quadrature_degree": self.quadrature_degree},
+                **kwargs
+            )
+        )
         return v
 
 
@@ -87,26 +89,45 @@ class Gradient(QuadratureFunction):
     This class is intended for internal use only. Gradient objects must be
     declared by the user using the registration concept.
     """
-    def __init__(self, variable, expression, name, symmetric=None):
+
+    def __init__(self, variable, expression, name, hypothesis, symmetric=None):
         self.variable = variable
         if symmetric is None:
             self.expression = expression
         # TODO: treat axisymmetric case
-        elif symmetric:
-            if ufl.shape(expression) == (2, 2):
-                self.expression = as_vector([
-                    symmetric_tensor_to_vector(expression)[i] for i in range(4)
-                ])
-            else:
-                self.expression = symmetric_tensor_to_vector(expression)
         else:
-            if ufl.shape(expression) == (2, 2):
-                self.expression = as_vector([
-                    nonsymmetric_tensor_to_vector(expression)[i]
-                    for i in range(5)
-                ])
+            if symmetric:
+                converter = symmetric_tensor_to_vector
             else:
-                self.expression = nonsymmetric_tensor_to_vector(expression)
+                converter = nonsymmetric_tensor_to_vector
+            if hypothesis in [
+                mgis_bv.Hypothesis.PlaneStrain,
+                mgis_bv.Hypothesis.Axisymmetrical,
+            ]:
+                if ufl.shape(expression) == (2, 2):
+                    T22 = 0
+                    expression_2d = expression
+                elif ufl.shape(expression) == (3, 3):
+                    T22 = expression[2, 2]
+                    expression_2d = ufl.as_tensor(
+                        [[expression[i, j] for j in range(2)] for i in range(2)]
+                    )
+                self.expression = converter(expression_2d, T22)
+            else:
+                self.expression = converter(expression)
+
+        #     self.expression = as_vector(
+        #         [converter(expression)[i] for i in range(4)]
+        #     )
+        #     else:
+        #         self.expression = symmetric_tensor_to_vector(expression)
+        # else:
+        #     if ufl.shape(expression) == (2, 2):
+        #         self.expression = as_vector(
+        #             [nonsymmetric_tensor_to_vector(expression)[i] for i in range(5)]
+        #         )
+        #     else:
+        #         self.expression = nonsymmetric_tensor_to_vector(expression)
         shape = ufl.shape(self.expression)
         if len(shape) == 1:
             self.shape = shape[0]
@@ -115,17 +136,20 @@ def __init__(self, variable, expression, name, symmetric=None):
         else:
             self.shape = shape
         self.name = name
+        self.hypothesis = hypothesis
 
     def __call__(self, v):
         return ufl.replace(self.expression, {self.variable: v})
 
     def variation(self, v):
         """ Directional derivative in direction v """
         # return ufl.algorithms.expand_derivatives(ufl.derivative(self.expression, self.variable, v))
-        deriv = sum([
-            ufl.derivative(self.expression, var, v_)
-            for (var, v_) in zip(split(self.variable), split(v))
-        ])
+        deriv = sum(
+            [
+                ufl.derivative(self.expression, var, v_)
+                for (var, v_) in zip(split(self.variable), split(v))
+            ]
+        )
         return ufl.algorithms.expand_derivatives(deriv)
 
     def initialize_function(self, mesh, quadrature_degree):
@@ -148,15 +172,17 @@ def _evaluate_at_quadrature_points(self, x):
 
 class Var(Gradient):
     """ A simple variable """
-    def __init__(self, variable, expression, name):
-        Gradient.__init__(self, variable, expression, name)
+
+    def __init__(self, variable, expression, name, hypothesis):
+        Gradient.__init__(self, variable, expression, name, hypothesis)
 
     def _evaluate_at_quadrature_points(self, x):
         local_project(x, self.function_space, self.dx, self.function)
 
 
 class QuadratureFunctionTangentBlocks(QuadratureFunction):
     """An abstract class for Flux and InternalStateVariables"""
+
     def initialize_tangent_blocks(self, variables):
         self.variables = variables
         values = [
@@ -170,7 +196,8 @@ def initialize_tangent_blocks(self, variables):
                     ),
                 ),
                 name="d{}_d{}".format(self.name, v.name),
-            ) for v in self.variables
+            )
+            for v in self.variables
         ]
         keys = [v.name for v in self.variables]
         self.tangent_blocks = dict(zip(keys, values))