Merged in mapdes/fiat/tpc-trace (pull request #34)

Tensor product support for the trace element

Approved-by: Miklós Homolya <[email protected]>

FIAT/hdiv_trace.py

@@ -16,14 +16,20 @@
 # along with FIAT. If not, see <http://www.gnu.org/licenses/>.
 
 from __future__ import absolute_import, print_function, division
+from six import iteritems, itervalues
+from six.moves import range
 
 import numpy as np
 from FIAT.discontinuous_lagrange import DiscontinuousLagrange
-from FIAT.reference_element import ufc_simplex
+from FIAT.dual_set import DualSet
+from FIAT.finite_element import FiniteElement
 from FIAT.functional import PointEvaluation
 from FIAT.polynomial_set import mis
-from FIAT import FiniteElement
-from FIAT import dual_set
+from FIAT.reference_element import (ufc_simplex, POINT,
+                                    LINE, QUADRILATERAL,
+                                    TRIANGLE, TETRAHEDRON,
+                                    TENSORPRODUCT)
+from FIAT.tensor_product import TensorProductElement
 
 # Numerical tolerance for facet-entity identifications
 epsilon = 1e-10
@@ -49,47 +55,89 @@ class HDivTrace(FiniteElement):
     """
 
     def __init__(self, ref_el, degree):
+        """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.
+        """
         sd = ref_el.get_spatial_dimension()
         if sd in (0, 1):
-            raise ValueError("Cannot use this trace class on a %d-dimensional cell." % sd)
-
-        # Constructing facet element as a discontinuous Lagrange element
-        dglagrange = DiscontinuousLagrange(ufc_simplex(sd - 1), degree)
-
-        # Construct entity ids (assigning top. dim. and initializing as empty)
+            raise ValueError("Cannot take the trace of a %d-dim cell." % sd)
+
+        # Store the degrees if on a tensor product cell
+        if ref_el.get_shape() == TENSORPRODUCT:
+            try:
+                degree = tuple(degree)
+            except TypeError:
+                degree = (degree,) * len(ref_el.cells)
+
+            assert len(ref_el.cells) == len(degree), (
+                "Number of specified degrees must be equal to the number of cells."
+            )
+        else:
+            if ref_el.get_shape() not in [TRIANGLE, TETRAHEDRON, QUADRILATERAL]:
+                raise NotImplementedError(
+                    "Trace element on a %s not implemented" % type(ref_el)
+                )
+            # Cannot have varying degrees for these reference cells
+            if isinstance(degree, tuple):
+                raise ValueError("Must have a tensor product cell if providing multiple degrees")
+
+        # Initialize entity dofs and construct the DG elements
+        # for the facets
+        facet_sd = sd - 1
+        dg_elements = {}
         entity_dofs = {}
-
-        # Looping over dictionary of cell topology to construct the empty
-        # dictionary for entity ids of the trace element
         topology = ref_el.get_topology()
-        for top_dim, entities in topology.items():
+        for top_dim, entities in iteritems(topology):
+            cell = ref_el.construct_subelement(top_dim)
             entity_dofs[top_dim] = {}
+
+            # We have a facet entity!
+            if cell.get_spatial_dimension() == facet_sd:
+                dg_elements[top_dim] = construct_dg_element(cell, degree)
+            # Initialize
             for entity in entities:
                 entity_dofs[top_dim][entity] = []
 
-        # Filling in entity ids and generating points for dual basis
-        nf = dglagrange.space_dimension()
-        points = []
-        num_facets = sd + 1
-        for f in range(num_facets):
-            entity_dofs[sd - 1][f] = range(f * nf, (f + 1) * nf)
-
-            for dof in dglagrange.dual_basis():
-                facet_point = list(dof.get_point_dict().keys())[0]
-                transform = ref_el.get_entity_transform(sd - 1, f)
-                points.append(tuple(transform(facet_point)))
+        # Compute the dof numbering for all facet entities
+        # and extract nodes
+        offset = 0
+        pts = []
+        for facet_dim in sorted(dg_elements):
+            element = dg_elements[facet_dim]
+            nf = element.space_dimension()
+            num_facets = len(topology[facet_dim])
+
+            for i in range(num_facets):
+                entity_dofs[facet_dim][i] = list(range(offset, offset + nf))
+                offset += nf
+
+                # Run over nodes and collect the points for point evaluations
+                for dof in element.dual_basis():
+                    facet_pt, = dof.get_point_dict()
+                    transform = ref_el.get_entity_transform(facet_dim, i)
+                    pts.append(tuple(transform(facet_pt)))
 
         # Setting up dual basis - only point evaluations
-        nodes = [PointEvaluation(ref_el, pt) for pt in points]
-        dual = dual_set.DualSet(nodes, ref_el, entity_dofs)
+        nodes = [PointEvaluation(ref_el, pt) for pt in pts]
+        dual = DualSet(nodes, ref_el, entity_dofs)
+
+        # Degree of the element
+        deg = max([e.degree() for e in dg_elements.values()])
+
+        super(HDivTrace, self).__init__(ref_el, dual, order=deg,
+                                        formdegree=facet_sd,
+                                        mapping="affine")
 
-        super(HDivTrace, self).__init__(ref_el, dual, dglagrange.get_order(),
-                                        dglagrange.get_formdegree(), dglagrange.mapping()[0])
-        # Set up facet element
-        self.facet_element = dglagrange
+        # Set up facet elements
+        self.dg_elements = dg_elements
 
-        # degree for quadrature rule
-        self.polydegree = degree
+        # Degree for quadrature rule
+        self.polydegree = deg
 
     def degree(self):
         """Return the degree of the (embedding) polynomial space."""
@@ -119,71 +167,104 @@ def tabulate(self, order, points, entity=None):
 
         .. note ::
 
-        Performing illegal tabulations on this element will result in either
-        a tabulation table of `numpy.nan` arrays (`entity=None` case), or
-        insertions of the `TraceError` exception class. This is due to the
-        fact that performing cell-wise tabulations, or asking for any order
-        of derivative evaluations, are not mathematically well-defined.
+           Performing illegal tabulations on this element will result in either
+           a tabulation table of `numpy.nan` arrays (`entity=None` case), or
+           insertions of the `TraceError` exception class. This is due to the
+           fact that performing cell-wise tabulations, or asking for any order
+           of derivative evaluations, are not mathematically well-defined.
         """
-        facet_dim = self.ref_el.get_spatial_dimension() - 1
-        sdim = self.space_dimension()
-        nf = self.facet_element.space_dimension()
+        sd = self.ref_el.get_spatial_dimension()
+        facet_sd = sd - 1
 
         # Initializing dictionary with zeros
         phivals = {}
         for i in range(order + 1):
-            alphas = mis(self.ref_el.get_spatial_dimension(), i)
+            alphas = mis(sd, i)
             for alpha in alphas:
-                phivals[alpha] = np.zeros(shape=(sdim, len(points)))
-        evalkey = (0,) * (facet_dim + 1)
+                phivals[alpha] = np.zeros(shape=(self.space_dimension(), len(points)))
+
+        evalkey = (0,) * sd
 
         # If entity is None, identify facet using numerical tolerance and
         # return the tabulated values
         if entity is None:
+            # NOTE: Numerical approximation of the facet id is currently only
+            # implemented for simplex reference cells.
+            if self.ref_el.get_shape() not in [TRIANGLE, TETRAHEDRON]:
+                raise NotImplementedError(
+                    "Tabulating this element on a %s cell without providing "
+                    "an entity is not currently supported." % type(self.ref_el)
+                )
+
             # Attempt to identify which facet (if any) the given points are on
             vertices = self.ref_el.vertices
             coordinates = barycentric_coordinates(points, vertices)
-            (unique_facet, success) = extract_unique_facet(coordinates)
+            unique_facet, success = extract_unique_facet(coordinates)
 
-            # If successful, insert evaluations
-            if success:
-                # Map points to the reference facet
-                new_points = map_to_reference_facet(points, vertices, unique_facet)
+            # If not successful, return NaNs
+            if not success:
+                for key in phivals:
+                    phivals[key] = np.full(shape=(sd, len(points)), fill_value=np.nan)
+
+                return phivals
 
-                # Retrieve values by tabulating the DiscontinuousLagrange element
-                nonzerovals = list(self.facet_element.tabulate(order, new_points).values())[0]
-                phivals[evalkey][nf*unique_facet:nf*(unique_facet + 1), :] = nonzerovals
-            # Otherwise, return NaNs
+            # Otherwise, extract non-zero values and insertion indices
             else:
-                for key in phivals.keys():
-                    phivals[key] = np.full(shape=(sdim, len(points)), fill_value=np.nan)
+                # Map points to the reference facet
+                new_points = map_to_reference_facet(points, vertices, unique_facet)
 
-            return phivals
+                # Retrieve values by tabulating the DG element
+                element = self.dg_elements[facet_sd]
+                nf = element.space_dimension()
+                nonzerovals, = itervalues(element.tabulate(order, new_points))
+                indices = slice(nf * unique_facet, nf * (unique_facet + 1))
 
-        entity_dim, entity_id = entity
+        else:
+            entity_dim, _ = entity
 
-        # If the user is directly specifying cell-wise tabulation, return TraceErrors in dict for
-        # appropriate handling in the form compiler
-        if entity_dim != facet_dim:
-            for key in phivals.keys():
-                phivals[key] = TraceError("Attempting to tabulate a %d-entity. Expecting a %d-entitiy" % (entity_dim, facet_dim))
-            return phivals
+            # If the user is directly specifying cell-wise tabulation, return
+            # TraceErrors in dict for appropriate handling in the form compiler
+            if entity_dim not in self.dg_elements:
+                for key in phivals:
+                    msg = "The HDivTrace element can only be tabulated on facets."
+                    phivals[key] = TraceError(msg)
 
-        else:
-            # Retrieve function evaluations (order = 0 case)
-            nonzerovals = list(self.facet_element.tabulate(0, points).values())[0]
-            phivals[evalkey][nf*entity_id:nf*(entity_id + 1), :] = nonzerovals
+                return phivals
 
-            # If asking for gradient evaluations, insert TraceError in gradient evaluations
-            if order > 0:
-                for key in phivals.keys():
-                    if key != evalkey:
-                        phivals[key] = TraceError("Gradient evaluations are illegal on trace elements.")
-            return phivals
+            else:
+                # Retrieve function evaluations (order = 0 case)
+                offset = 0
+                for facet_dim in sorted(self.dg_elements):
+                    element = self.dg_elements[facet_dim]
+                    nf = element.space_dimension()
+                    num_facets = len(self.ref_el.get_topology()[facet_dim])
+
+                    # Loop over the number of facets until we find a facet
+                    # with matching dimension and id
+                    for i in range(num_facets):
+                        # Found it! Grab insertion indices
+                        if (facet_dim, i) == entity:
+                            nonzerovals, = itervalues(element.tabulate(0, points))
+                            indices = slice(offset, offset + nf)
+
+                        offset += nf
+
+        # If asking for gradient evaluations, insert TraceError in
+        # gradient slots
+        if order > 0:
+            msg = "Gradients on trace elements are not well-defined."
+            for key in phivals:
+                if key != evalkey:
+                    phivals[key] = TraceError(msg)
+
+        # Insert non-zero values in appropriate place
+        phivals[evalkey][indices, :] = nonzerovals
+
+        return phivals
 
     def value_shape(self):
         """Return the value shape of the finite element functions."""
-        return self.facet_element.value_shape()
+        return ()
 
     def dmats(self):
         """Return dmats: expansion coefficients for basis function
@@ -195,6 +276,43 @@ def get_num_members(self, arg):
         raise NotImplementedError("get_num_members not implemented for the trace element.")
 
 
+def construct_dg_element(ref_el, degree):
+    """Constructs a discontinuous galerkin element of a given degree
+    on a particular reference cell.
+    """
+    if ref_el.get_shape() in [LINE, TRIANGLE]:
+        dg_element = DiscontinuousLagrange(ref_el, degree)
+
+    # Quadrilateral facets could be on a FiredrakeQuadrilateral.
+    # In this case, we treat this as an interval x interval cell:
+    elif ref_el.get_shape() == QUADRILATERAL:
+        dg_a = DiscontinuousLagrange(ufc_simplex(1), degree)
+        dg_b = DiscontinuousLagrange(ufc_simplex(1), degree)
+        dg_element = TensorProductElement(dg_a, dg_b)
+
+    # This handles the more general case for facets:
+    elif ref_el.get_shape() == TENSORPRODUCT:
+        assert len(degree) == len(ref_el.cells), (
+            "Must provide the same number of degrees as the number "
+            "of cells that make up the tensor product cell."
+        )
+        sub_elements = [construct_dg_element(c, d)
+                        for c, d in zip(ref_el.cells, degree)
+                        if c.get_shape() != POINT]
+
+        if len(sub_elements) > 1:
+            dg_element = TensorProductElement(*sub_elements)
+        else:
+            dg_element, = sub_elements
+
+    else:
+        raise NotImplementedError(
+            "Reference cells of type %s not currently supported" % type(ref_el)
+        )
+
+    return dg_element
+
+
 # The following functions are credited to Marie E. Rognes:
 def extract_unique_facet(coordinates, tolerance=epsilon):
     """Determines whether a set of points (described in its barycentric coordinates)

doc/sphinx/source/releases/next.rst

@@ -11,6 +11,13 @@ Summary of changes
           be published (and renamed) to list the most important
           changes in the new release.
 
+- Extended the discontinuous trace element ``HDivTrace`` to support tensor
+  product reference cells. Tabulating the trace defined on a tensor product
+  cell relies on the argument ``entity`` to specify a facet of the cell. The
+  backwards compatibility case ``entity=None`` does not support tensor product
+  tabulation as a result. Tabulating the trace of triangles or tetrahedron
+  remains unaffected and works as usual with or without an entity argument.
+
 Detailed changes
 ================