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
@@ -723,7 +702,10 @@ def compute_cell_point_map(ref_el, pts, unique=True, tol=1E-12):
         return {cell: Ellipsis for cell in sorted(top[sd])}
 
     # The distance to the nearest cell is equal to the distance to the parent cell
-    best = ref_el.get_parent().distance_to_point_l1(pts, rescale=True)
+    parent = ref_el
+    while parent.get_parent() is not None:
+        parent = parent.get_parent()
+    best = parent.distance_to_point_l1(pts, rescale=True)
     tol = best + tol
 
     cell_point_map = {}
@@ -779,3 +761,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

FIAT/hdiv_trace.py

@@ -100,11 +100,8 @@ class HDivTrace(DiscontinuousLagrange):
     def __new__(cls, ref_el, degree, variant="equispaced_interior"):
         """Constructor for the HDivTrace element.
 
-        :arg ref_el: A reference element, which may be a tensor product
-                     cell.
-        :arg degree: The degree of approximation. If on a tensor product
-                     cell, then provide a tuple of degrees if you want
-                     varying degrees.
+        :arg ref_el: A reference element.
+        :arg degree: The degree of approximation.
         :arg variant: The point distribution variant passed on to recursivenodes.
         """
         facets = TraceSimplicialComplex(ref_el)

FIAT/reference_element.py

@@ -618,7 +618,7 @@ def compute_barycentric_coordinates(self, points, entity=None, rescale=False):
         if len(points) == 0:
             return points
         if entity is None:
-            entity = (self.get_spatial_dimension(), 0)
+            entity = (self.get_dimension(), 0)
         entity_dim, entity_id = entity
         sd = len(self.vertices[0])
 
@@ -628,10 +628,9 @@ def compute_barycentric_coordinates(self, points, entity=None, rescale=False):
         top = self.get_topology()
         subcomplex = top[entity_dim][entity_id]
         if entity_dim != sd:
-            parent = self.get_parent() if self.is_trace() else self
-            cell_id = parent.connectivity[(entity_dim, sd)][0][0]
-            top = parent.get_topology()
-            subcell = top[sd][cell_id]
+            parent = self.get_parent_complex() if self.is_trace() else self
+            cell_id = min(parent.connectivity[(entity_dim, sd)][entity_id])
+            subcell = parent.topology[sd][cell_id]
             while len(subcell) > sd + 1:
                 # construct a simplex if we have a hypercube
                 k = max(set(subcell) - set(subcomplex))
@@ -648,7 +647,7 @@ def compute_barycentric_coordinates(self, points, entity=None, rescale=False):
         A, b = make_affine_mapping(cell_verts, ref_verts)
         A, b = A[indices], b[indices]
         if rescale:
-            # rescale barycentric coordinates by the height wrt. to the facet
+            # rescale barycentric coordinates by the height w.r.t. the facet
             h = 1 / numpy.linalg.norm(A, axis=1)
             b *= h
             A *= h[:, None]

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

@@ -2,6 +2,7 @@
 from FIAT.reference_element import LINE
 
 import gem
+import numpy
 from gem.utils import cached_property
 from finat.finiteelementbase import FiniteElementBase
 from finat.tensor_product import TensorProductElement
@@ -58,7 +59,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 +73,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 +93,59 @@ 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 or sum(entity_dim) != sum(self.support_dim):
+            return core_eval
+
+        def zero_failure(expr):
+            if isinstance(expr, gem.Failure):
+                return gem.Literal(numpy.zeros(expr.shape, expr.dtype))
+            return expr.reconstruct(*map(zero_failure, expr.children))
+
+        # Create a zero tabulation with the same tensor-product structure
+        return {alpha: zero_failure(table) 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 = \
@@ -96,18 +97,21 @@ def __new__(cls,
 
             elif family == "HDiv Trace":
                 cell_h, cell_v = cell.sub_cells()
-                cell_h, cell_v = cell.sub_cells()
+                if not isinstance(degree, tuple):
+                    degree = (degree, degree)
+                hdegree, vdegree = degree
+
+                dg_family = lambda cell: "DG" if cell.is_simplex() else "DQ"
+                is_interval = lambda cell: cell.cellname() == "interval"
 
-                hdegree = 0 if cell_h.cellname() == "interval" else degree
-                vdegree = 0 if cell_v.cellname() == "interval" else degree
-                tr_h = FiniteElement("HDiv Trace", cell_h, hdegree, variant=variant)
-                tr_v = FiniteElement("HDiv Trace", cell_v, vdegree, variant=variant)
+                dg_h = FiniteElement(dg_family(cell_h), cell_h, hdegree, variant=variant)
+                tr_h = FiniteElement("HDiv Trace", cell_h, 0 if is_interval(cell_h) else hdegree, variant=variant)
 
-                dg_h = FiniteElement("DG", cell_h, degree, variant=variant)
-                dg_v = FiniteElement("DG", cell_v, degree, variant=variant)
+                dg_v = FiniteElement(dg_family(cell_v), cell_v, vdegree, variant=variant)
+                tr_v = FiniteElement("HDiv Trace", cell_v, 0 if is_interval(cell_v) else vdegree, 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)