finat.ufl.HDivTraceElement, full support for tensor product elements

Author: pbrubeck

FIAT/barycentric_interpolation.py

@@ -75,7 +75,7 @@ def __init__(self, ref_el, pts):
             x = pts[ibfs]
             if ref_el.is_trace():
                 verts = ref_el.get_vertices_of_subcomplex(ref_el.topology[1][cell])
-                A = numpy.diff(verts, axis=0)[0]/2
+                A, = numpy.diff(verts, axis=0)
                 A /= numpy.linalg.norm(A)
                 b = -numpy.dot(numpy.sum(verts, axis=0)/2, A.T)
                 self.affine_mappings[cell] = (A, b)

FIAT/expansions.py

@@ -345,8 +345,6 @@ def distance(alpha, beta):
 
     def _tabulate(self, n, pts, order=0):
         """A version of tabulate() that also works for a single point."""
-        from FIAT.hdiv_trace import TraceError
-        from FIAT.polynomial_set import mis
         pts = numpy.asarray(pts)
         unique = self.continuity is not None and order == 0
         cell_point_map = compute_cell_point_map(self.ref_el, pts, unique=unique)
@@ -357,26 +355,7 @@ def _tabulate(self, n, pts, order=0):
             return phis[0]
 
         if self.ref_el.is_trace():
-            parent = self.ref_el.get_parent()
-            tdim = self.ref_el.get_spatial_dimension()
-            gdim = parent.get_spatial_dimension()
-            for cell in phis:
-                # Promote facet keys to cell keys
-                phi = phis[cell][(0,) * tdim]
-                # Raise TraceError on gradient tabulations
-                msg = "Gradients on trace elements are not well-defined."
-                phis[cell] = {alpha: phi if sum(alpha) == 0 else TraceError(msg)
-                              for i in range(order+1)
-                              for alpha in mis(gdim, i)}
-            if tdim == 0 and len(phis) == 0:
-                # Hack for TensorProduct HDivTrace: do not raise TraceError on the interval
-                for cell in parent.topology[gdim]:
-                    phis[cell] = {(0,)*gdim: numpy.zeros(())}
-            elif sum(len(cell_point_map[cell]) for cell in cell_point_map) < len(pts):
-                # Raise TraceError when interior points fail to be binned on facets
-                for cell in parent.topology[gdim]:
-                    msg = "The HDivTrace element can only be tabulated on facets."
-                    phis[cell] = {(0,)*gdim: TraceError(msg)}
+            phis = trace_tabulation(self.ref_el, cell_point_map, order, pts, phis)
 
         if pts.dtype == object:
             # If binning is undefined, scale by the characteristic function of each subcell
@@ -402,7 +381,7 @@ def _tabulate(self, n, pts, order=0):
         result = {}
         base_phi = tuple(phis.values())[0]
         for alpha in base_phi:
-            if isinstance(base_phi[alpha], TraceError):
+            if isinstance(base_phi[alpha], Exception):
                 result[alpha] = base_phi[alpha]
                 continue
             dtype = base_phi[alpha].dtype
@@ -779,3 +758,30 @@ def compute_partition_of_unity(ref_el, pt, unique=True, tol=1E-12):
         mult = sum(masks)
         masks = [m / mult for m in masks]
     return masks
+
+
+def trace_tabulation(ref_el, cell_point_map, order, pts, phis):
+    """Lift trace tabulations into the cells and raise TraceError on invalid tabulations."""
+    from FIAT.polynomial_set import mis
+    from FIAT.hdiv_trace import TraceError
+    parent = ref_el.get_parent()
+    tdim = ref_el.get_spatial_dimension()
+    gdim = parent.get_spatial_dimension()
+    facet_key = (0, ) * tdim
+    cell_key = (0, ) * gdim
+
+    for cell in phis:
+        # Lift facet keys to cell keys
+        phi = phis[cell][facet_key]
+        # Raise TraceError on gradient tabulations
+        msg = "Gradients on trace elements are not well-defined."
+        phis[cell] = {alpha: phi if sum(alpha) == 0 else TraceError(msg)
+                      for i in range(order+1)
+                      for alpha in mis(gdim, i)}
+
+    if sum(len(cell_point_map[cell]) for cell in cell_point_map) < len(pts):
+        # Raise TraceError when interior points fail to be binned on facets
+        for cell in parent.topology[gdim]:
+            msg = "The HDivTrace element can only be tabulated on facets."
+            phis[cell] = {cell_key: TraceError(msg)}
+    return phis

finat/__init__.py

@@ -32,7 +32,7 @@
 from .cube import FlattenedDimensions                              # noqa: F401
 from .discontinuous import DiscontinuousElement                    # noqa: F401
 from .enriched import EnrichedElement                              # noqa: F401
-from .hdivcurl import HCurlElement, HDivElement                    # noqa: F401
+from .hdivcurl import HCurlElement, HDivElement, HDivTraceElement  # noqa: F401
 from .mixed import MixedElement                                    # noqa: F401
 from .nodal_enriched import NodalEnrichedElement                   # noqa: F401
 from .quadrature_element import QuadratureElement, make_quadrature_element  # noqa: F401

finat/element_factory.py

@@ -287,6 +287,12 @@ def convert_tensorproductelement(element, **kwargs):
     return finat.TensorProductElement(elements), deps
 
 
+@convert.register(finat.ufl.HDivTraceElement)
+def convert_hdivtraceelement(element, **kwargs):
+    finat_elem, deps = _create_element(element._element, **kwargs)
+    return finat.HDivTraceElement(finat_elem), deps
+
+
 @convert.register(finat.ufl.HDivElement)
 def convert_hdivelement(element, **kwargs):
     finat_elem, deps = _create_element(element._element, **kwargs)

finat/fiat_elements.py

@@ -112,7 +112,8 @@ def basis_evaluation(self, order, ps, entity=None, coordinate_mapping=None):
         index_shape = (self._element.space_dimension(),)
         for alpha, fiat_table in fiat_result.items():
             if isinstance(fiat_table, Exception):
-                result[alpha] = gem.Failure(self.index_shape + self.value_shape, fiat_table)
+                shape = ps.points.shape[:-1] + self.index_shape + self.value_shape
+                result[alpha] = gem.partial_indexed(gem.Failure(shape, fiat_table), ps.indices)
                 continue
 
             derivative = sum(alpha)

finat/hdivcurl.py

@@ -58,7 +58,7 @@ def index_shape(self):
     def value_shape(self):
         return (self.cell.get_spatial_dimension(),)
 
-    def _transform_evaluation(self, core_eval):
+    def _transform_evaluation(self, core_eval, entity=None):
         beta = self.get_indices()
         zeta = self.get_value_indices()
 
@@ -72,11 +72,11 @@ def promote(table):
 
     def basis_evaluation(self, order, ps, entity=None, coordinate_mapping=None):
         core_eval = self.wrappee.basis_evaluation(order, ps, entity)
-        return self._transform_evaluation(core_eval)
+        return self._transform_evaluation(core_eval, entity)
 
     def point_evaluation(self, order, refcoords, entity=None, coordinate_mapping=None):
         core_eval = self.wrappee.point_evaluation(order, refcoords, entity)
-        return self._transform_evaluation(core_eval)
+        return self._transform_evaluation(core_eval, entity)
 
     @property
     def dual_basis(self):
@@ -92,6 +92,53 @@ def dual_basis(self):
         return gem.ComponentTensor(Q[zeta], beta + zeta), x
 
 
+class HDivTraceElement(WrapperElementBase):
+    """HDiv Trace wrapper element for tensor product elements."""
+
+    def __init__(self, wrappee):
+        assert isinstance(wrappee, TensorProductElement)
+        if any(fe.formdegree is None for fe in wrappee.factors):
+            raise ValueError("Form degree of subelement is None, cannot HDiv Trace!")
+
+        formdegree = sum(fe.formdegree for fe in wrappee.factors)
+        if formdegree != wrappee.cell.get_spatial_dimension() - 1:
+            raise ValueError("HDiv Trace requires (n-1)-form element!")
+
+        self.support_dim = tuple(fe.formdegree for fe in wrappee.factors)
+
+        super().__init__(wrappee, None)
+
+    def _transform_evaluation(self, core_eval, entity):
+        if entity is None:
+            entity = (self.cell.get_dimension(), 0)
+        entity_dim, entity_id = entity
+        if entity_dim == self.support_dim:
+            return core_eval
+
+        # Create a zero tabulation with the same tensor-product structure
+        return {alpha: table.zero() for alpha, table in core_eval.items()}
+
+    @property
+    def formdegree(self):
+        return self.wrappee.formdegree
+
+    @property
+    def value_shape(self):
+        return self.wrappee.value_shape
+
+    @cached_property
+    def fiat_equivalent(self):
+        return self.wrappee.fiat_equivalent
+
+    @property
+    def mapping(self):
+        return self.wrappee.mapping
+
+    @property
+    def dual_basis(self):
+        return self.wrappee.dual_basis
+
+
 class HDivElement(WrapperElementBase):
     """H(div) wrapper element for tensor product elements."""
 

finat/ufl/__init__.py

@@ -15,7 +15,7 @@
 from finat.ufl.enrichedelement import EnrichedElement, NodalEnrichedElement  # noqa: F401
 from finat.ufl.finiteelement import FiniteElement  # noqa: F401
 from finat.ufl.finiteelementbase import FiniteElementBase  # noqa: F401
-from finat.ufl.hdivcurl import HCurlElement, HDivElement, WithMapping, HDiv, HCurl  # noqa: F401
+from finat.ufl.hdivcurl import HCurlElement, HDivElement, HDivTraceElement, WithMapping, HDiv, HCurl  # noqa: F401
 from finat.ufl.mixedelement import MixedElement, TensorElement, VectorElement  # noqa: F401
 from finat.ufl.restrictedelement import RestrictedElement  # noqa: F401
 from finat.ufl.tensorproductelement import TensorProductElement  # noqa: F401

finat/ufl/finiteelement.py

@@ -39,6 +39,7 @@ def __new__(cls,
             from finat.ufl.enrichedelement import EnrichedElement
             from finat.ufl.hdivcurl import HCurlElement as HCurl
             from finat.ufl.hdivcurl import HDivElement as HDiv
+            from finat.ufl.hdivcurl import HDivTraceElement as HDivTrace
             from finat.ufl.tensorproductelement import TensorProductElement
 
             family, short_name, degree, reference_value_shape, sobolev_space, mapping, embedded_degree = \
@@ -106,8 +107,8 @@ def __new__(cls,
                 dg_h = FiniteElement("DG", cell_h, degree, variant=variant)
                 dg_v = FiniteElement("DG", cell_v, degree, variant=variant)
 
-                return EnrichedElement(TensorProductElement(tr_h, dg_v, cell=cell),
-                                       TensorProductElement(dg_h, tr_v, cell=cell))
+                return EnrichedElement(HDivTrace(TensorProductElement(tr_h, dg_v, cell=cell)),
+                                       HDivTrace(TensorProductElement(dg_h, tr_v, cell=cell)))
 
             elif family == "Q":
                 return TensorProductElement(*[FiniteElement("CG", c, degree, variant=variant)

finat/ufl/hdivcurl.py

@@ -35,6 +35,62 @@ def __call__(self, element):
 HDiv = CallableSobolevSpace(HDivSobolevSpace.name, HDivSobolevSpace.parents)
 
 
+class HDivTraceElement(FiniteElementBase):
+    """A HDiv Trace version of an outer product element, assuming this makes mathematical sense."""
+    __slots__ = ("_element", )
+
+    def __init__(self, element):
+        """Doc."""
+        self._element = element
+
+        family = "TensorProductElement"
+        cell = element.cell
+        degree = element.degree()
+        quad_scheme = element.quadrature_scheme()
+        reference_value_shape = ()
+
+        super().__init__(family, cell, degree, quad_scheme, reference_value_shape)
+
+    def __repr__(self):
+        """Doc."""
+        return f"HDivTraceElement({repr(self._element)})"
+
+    def mapping(self):
+        """Doc."""
+        return "affine"
+
+    @property
+    def sobolev_space(self):
+        """Return the underlying Sobolev space."""
+        return L2
+
+    def reconstruct(self, **kwargs):
+        """Doc."""
+        return HDivTraceElement(self._element.reconstruct(**kwargs))
+
+    def variant(self):
+        """Doc."""
+        return self._element.variant()
+
+    def __str__(self):
+        """Doc."""
+        return f"HDivTraceElement({repr(self._element)})"
+
+    def shortstr(self):
+        """Format as string for pretty printing."""
+        return f"HDivTraceElement({self._element.shortstr()})"
+
+    @property
+    def embedded_subdegree(self):
+        """Return embedded subdegree."""
+        return self._element.embedded_subdegree
+
+    @property
+    def embedded_superdegree(self):
+        """Return embedded superdegree."""
+        return self._element.embedded_superdegree
+
+
 class HDivElement(FiniteElementBase):
     """A div-conforming version of an outer product element, assuming this makes mathematical sense."""
     __slots__ = ("_element", )

gem/gem.py

@@ -166,6 +166,9 @@ def inherit_dtype_from_children(children):
         else:
             return numpy.result_type(*(c.dtype for c in children))
 
+    def zero(self):
+        return type(self)(*(c.zero() for c in self.children))
+
 
 class Terminal(Node):
     """Abstract class for terminal GEM nodes."""
@@ -197,6 +200,9 @@ def dtype(self):
         else:
             raise AttributeError(f"Must set _dtype on terminal node, {type(self)}")
 
+    def zero(self):
+        return Literal(numpy.zeros(self.shape, dtype=self.dtype))
+
 
 class Scalar(Node):
     """Abstract class for scalar-valued GEM nodes."""
@@ -717,6 +723,9 @@ def index_ordering(self):
                 free_indices.extend(i.expression.free_indices)
         return tuple(free_indices)
 
+    def zero(self):
+        return type(self)(*(c.zero() for c in self.children), self.multiindex)
+
 
 class FlexiblyIndexed(Scalar):
     """Flexible indexing of :py:class:`Variable`s to implement views and
@@ -847,6 +856,9 @@ def __new__(cls, expression, multiindex):
 
         return self
 
+    def zero(self):
+        return type(self)(*(c.zero() for c in self.children), self.multiindex)
+
 
 class IndexSum(Scalar):
     __slots__ = ('children', 'multiindex')