TensorProductElement: rename sub_elements factor_elements

Author: ksagiyam

firedrake/embedding.py

@@ -59,7 +59,7 @@ def get_embedding_method_for_checkpointing(element):
                               'S', 'DPC', 'Real']:
         return "interpolate"
     elif isinstance(element, finat.ufl.TensorProductElement):
-        methods = [get_embedding_method_for_checkpointing(elem) for elem in element.sub_elements]
+        methods = [get_embedding_method_for_checkpointing(elem) for elem in element.factor_elements]
         if any(method == "project" for method in methods):
             return "project"
         else:

firedrake/extrusion_utils.py

@@ -50,7 +50,7 @@ def make_extruded_coords(extruded_topology, base_coords, ext_coords,
     coordinates on the extruded cell (to write to), the fixed layer
     height, and the current cell layer.
     """
-    _, vert_space = ext_coords.function_space().ufl_element().sub_elements[0].sub_elements
+    _, vert_space = ext_coords.function_space().ufl_element().sub_elements[0].factor_elements
     if kernel is None and not (vert_space.degree() == 1
                                and vert_space.family() in ['Lagrange',
                                                            'Discontinuous Lagrange']):

firedrake/functionspaceimpl.py

@@ -60,7 +60,7 @@ def check_element(element, top=True):
     elif type(element) is finat.ufl.EnrichedElement:
         inner = element._elements
     elif type(element) is finat.ufl.TensorProductElement:
-        inner = element.sub_elements
+        inner = element.factor_elements
     elif isinstance(element, finat.ufl.MixedElement):
         if not top:
             raise ValueError(f"{type(element).__name__} modifier must be outermost")

firedrake/interpolation.py

@@ -543,9 +543,7 @@ def __init__(
         # input ordering VOM will only contain the points on rank 0!
         # QUESTION: Should any of the below have annotation turned off?
         ufl_scalar_element = V_dest.ufl_element()
-        if ufl_scalar_element.num_sub_elements and not isinstance(
-            ufl_scalar_element, finat.ufl.TensorProductElement
-        ):
+        if isinstance(ufl_scalar_element, finat.ufl.MixedElement):
             if all(
                 ufl_scalar_element.sub_elements[0] == e
                 for e in ufl_scalar_element.sub_elements
@@ -558,7 +556,7 @@ def __init__(
                     raise NotImplementedError(
                         "Can't yet cross-mesh interpolate onto function spaces made from VectorElements or TensorElements made from sub elements with value shape other than ()."
                     )
-            elif type(ufl_scalar_element) is finat.ufl.MixedElement:
+            else:
                 # Build and save an interpolator for each sub-element
                 # separately for MixedFunctionSpaces. NOTE: since we can't have
                 # expressions for MixedFunctionSpaces we know that the input
@@ -597,10 +595,6 @@ def __init__(
                     )
                     self.sub_interpolators.append(sub_interpolator)
                 return
-            else:
-                raise NotImplementedError(
-                    f"Unhandled cross-mesh interpolation ufl element type: {repr(ufl_scalar_element)}"
-                )
 
         from firedrake.assemble import assemble
         V_dest_vec = firedrake.VectorFunctionSpace(dest_mesh, ufl_scalar_element)

firedrake/mg/embedded.py

@@ -70,8 +70,11 @@ def __init__(self, *, native_transfers=None, use_averaging=True):
     def is_native(self, element, op):
         if element in self.native_transfers.keys():
             return self.native_transfers[element][op] is not None
-        if isinstance(element.cell, ufl.TensorProductCell) and len(element.sub_elements) > 0:
-            return all(self.is_native(e, op) for e in element.sub_elements)
+        if isinstance(element.cell, ufl.TensorProductCell):
+            if isinstance(element, finat.ufl.TensorProductElement):
+                return all(self.is_native(e, op) for e in element.factor_elements)
+            elif isinstance(element, finat.ufl.MixedElement):
+                return all(self.is_native(e, op) for e in element.sub_elements)
         return (element.family() in native_families) and not (element.variant() in non_native_variants)
 
     def _native_transfer(self, element, op):

firedrake/preconditioners/hiptmair.py

@@ -250,7 +250,7 @@ def div_to_curl(ele):
     elif isinstance(ele, finat.ufl.EnrichedElement):
         return type(ele)(*(div_to_curl(e) for e in reversed(ele._elements)))
     elif isinstance(ele, finat.ufl.TensorProductElement):
-        return type(ele)(*(div_to_curl(e) for e in ele.sub_elements), cell=ele.cell)
+        return type(ele)(*(div_to_curl(e) for e in ele.factor_elements), cell=ele.cell)
     elif isinstance(ele, finat.ufl.WithMapping):
         return type(ele)(div_to_curl(ele.wrapee), ele.mapping())
     elif isinstance(ele, finat.ufl.BrokenElement):

firedrake/preconditioners/pmg.py

@@ -423,7 +423,7 @@ def reconstruct_degree(ele, degree):
             return type(ele)(*(PMGBase.reconstruct_degree(e, PMGBase.max_degree(e) + shift) for e in ele._elements))
         elif isinstance(ele, finat.ufl.TensorProductElement):
             shift = degree - PMGBase.max_degree(ele)
-            return type(ele)(*(PMGBase.reconstruct_degree(e, PMGBase.max_degree(e) + shift) for e in ele.sub_elements), cell=ele.cell)
+            return type(ele)(*(PMGBase.reconstruct_degree(e, PMGBase.max_degree(e) + shift) for e in ele.factor_elements), cell=ele.cell)
         elif isinstance(ele, finat.ufl.MixedElement):
             shift = degree - PMGBase.max_degree(ele)
             return type(ele)(*(PMGBase.reconstruct_degree(e, PMGBase.max_degree(e) + shift) for e in ele.sub_elements))