FInAT: Dual evaluation for compositions of TensorProductElement + EnrichedElement

finat/discontinuous.py

@@ -75,6 +75,9 @@ def point_evaluation(self, order, refcoords, entity=None, coordinate_mapping=Non
     def dual_basis(self):
         return self.element.dual_basis
 
+    def dual_evaluation(self, fn, coordinate_mapping=None):
+        return self.element.dual_evaluation(fn, coordinate_mapping)
+
     @property
     def mapping(self):
         return self.element.mapping

finat/enriched.py

@@ -1,4 +1,4 @@
-from functools import partial
+from functools import partial, singledispatch
 from itertools import chain
 from operator import add, methodcaller
 
@@ -8,8 +8,13 @@
 from gem.interpreter import evaluate
 from gem.utils import cached_property
 
-from finat.finiteelementbase import FiniteElementBase
-from finat.hdivcurl import HCurlElement, HDivElement
+from finat.cube import FlattenedDimensions
+from finat.discontinuous import DiscontinuousElement
+from finat.finiteelementbase import FiniteElementBase, broadcast_tensor
+from finat.hdivcurl import HCurlElement, HDivElement, WrapperElementBase
+from finat.point_set import UnionPointSet
+from finat.tensor_product import TensorProductElement
+from finat.tensorfiniteelement import TensorFiniteElement
 
 
 class EnrichedElement(FiniteElementBase):
@@ -160,34 +165,164 @@ def mapping(self):
             result, = mappings
             return result
 
-    def dual_evaluation(self, argument, coordinate_mapping=None):
+    @cached_property
+    def _summands(self):
+        """The summands that are not themselves direct sums, in basis order.
+
+        An element is brought out as a direct sum one level at a time, so a
+        summand may be a direct sum in turn.  These are the elements that
+        evaluate their dual basis on their own points, and whose points make
+        up the union that :attr:`dual_basis` works against.
+        """
+        summands = []
+        for element in self.elements:
+            expanded = as_enriched(element)
+            summands.extend(expanded._summands if expanded is not None
+                            else [element])
+        return tuple(summands)
+
+    @property
+    def dual_basis(self):
+        """The weights of the dual basis, on the union of the summands' points.
+
+        A summand's functionals evaluate only on that summand's points, so the
+        weights are block diagonal: each summand's weights sit at its own
+        offset in the union, and are zero against every other summand's
+        points.  :meth:`_dual_evaluation` contracts each summand on its own
+        points instead of carrying those zeros.
+        """
+        duals = [element.dual_basis for element in self._summands]
+        x = UnionPointSet([xk for _, xk in duals])
+        p, = x.indices
+        zeta = self.get_value_indices()
+        # The natural shape of each summand's own points, in the order they
+        # occupy the union.
+        shapes = [tuple(i.extent for i in xk.indices) for _, xk in duals]
+
+        blocks = []
+        for k, (element, (Q, xk)) in enumerate(zip(self._summands, duals)):
+            alpha = element.get_indices()
+            # Turn this summand's point indices into a shape, so that its
+            # weights can be embedded at its own offset in the union.
+            own = gem.ComponentTensor(gem.Indexed(Q, alpha + zeta), xk.indices)
+            pieces = [own if j == k else gem.Zero(shape)
+                      for j, shape in enumerate(shapes)]
+            weights = gem.Indexed(gem.Concatenate(*pieces), (p,))
+            blocks.append(gem.ComponentTensor(weights, alpha))
+
+        beta, = self.get_indices()
+        Q = gem.Indexed(gem.Concatenate(*blocks), (beta,))
+        return gem.ComponentTensor(Q, (beta,) + zeta), x
+
+    def _dual_evaluation(self, fn, coordinate_mapping=None):
+        """Dual evaluate each summand on its own points.
+
+        :arg fn: Callable representing the function to dual evaluate.
+                 Callable should take in an :class:`AbstractPointSet` and
+                 return a GEM expression for evaluation of the function at
+                 those points.
+        :arg coordinate_mapping: a
+           :class:`~.physically_mapped.PhysicalGeometry` object that
+           provides physical geometry callbacks (may be None).
+        :returns: an ``(evaluation, point_indices, basis_indices)`` triple, as
+           :meth:`~finat.finiteelementbase.FiniteElementBase.dual_evaluation`
+           returns.  The points are contracted here, so ``point_indices`` is
+           empty.
+
+        The summands do not share their points, so each one contracts on its
+        own, and the results stack along the basis index.  Concatenating over
+        a free index is what :func:`~gem.unconcatenate.unconcatenate` splits
+        downstream; a concatenation over the contracted points could not be.
+        """
         if not self.is_nodal_enriched:
             raise NotImplementedError(
-                f"Dual evaluation not defined for element {type(self).__name__}"
+                f"Dual evaluation not defined for non-nodal {type(self).__name__}"
             )
-        # Gather results from all sub-elements
-        # Each sub_result is (eval_expr, point_indices, local_indices)
-        sub_results = [sub.dual_evaluation(argument, coordinate_mapping=coordinate_mapping)
-                       for sub in self.elements]
-
-        # Extract the evaluation sub-expressions, contracting each over its own points.
-        # We must ensure that all subindices are in the free indices of subexpr
-        # before wrapping in ComponentTensor. If some are missing (e.g. if the
-        # expression simplified to a constant), we multiply by a dummy ones tensor.
         evals = []
-        for subexpr, point_indices, subindices in sub_results:
-            subexpr = gem.IndexSum(subexpr, point_indices)
-            missing_indices = tuple(idx for idx in subindices if idx not in subexpr.free_indices)
-            if missing_indices:
-                shape = tuple(idx.extent for idx in missing_indices)
-                ones = gem.Literal(numpy.ones(shape))
-                dummy = gem.Indexed(ones, missing_indices)
-                subexpr = gem.Product(subexpr, dummy)
-            evals.append(gem.ComponentTensor(subexpr, subindices))
+        for element in self.elements:
+            expr, point_indices, indices = element.dual_evaluation(
+                fn, coordinate_mapping=coordinate_mapping)
+            evals.append(broadcast_tensor(gem.IndexSum(expr, point_indices), indices))
 
         beta = self.get_indices()
-        expr = gem.Indexed(gem.Concatenate(*evals), beta)
-        return expr, (), beta
+        return gem.Indexed(gem.Concatenate(*evals), beta), (), beta
+
+
+@singledispatch
+def as_enriched(element):
+    """Rewrite an element as a direct sum, bringing the sum outermost.
+
+    :arg element: the :class:`~finat.finiteelementbase.FiniteElementBase` to
+                  rewrite.
+    :returns: an :class:`EnrichedElement` with the same basis functions in the
+              same order as ``element``, or ``None`` if ``element`` is not a
+              direct sum.
+
+    Only a direct sum can be dual evaluated summand by summand.  The sum
+    commutes with everything that can hold it: a pullback is linear, a tensor
+    product distributes over a sum in its first factor, and the wrappers that
+    leave the dual basis alone are transparent.
+    """
+    return None
+
+
+@as_enriched.register(EnrichedElement)
+def as_enriched_enriched(element):
+    return element
+
+
+@as_enriched.register(FlattenedDimensions)
+def as_enriched_flattened(element):
+    return as_enriched(element.product)
+
+
+@as_enriched.register(DiscontinuousElement)
+def as_enriched_discontinuous(element):
+    return as_enriched(element.element)
+
+
+@as_enriched.register(WrapperElementBase)
+def as_enriched_wrapper(element):
+    """Distribute the pullback over the sum the wrapped element is."""
+    summands = as_enriched(element.wrappee)
+    if summands is None:
+        return None
+    return EnrichedElement([type(element)(e) for e in summands.elements],
+                           is_nodal_enriched=summands.is_nodal_enriched)
+
+
+@as_enriched.register(TensorFiniteElement)
+def as_enriched_tensor_finite_element(element):
+    """Distribute the vector/tensor wrapper over the sum its base element is."""
+    summands = as_enriched(element.base_element)
+    if summands is None:
+        return None
+    return EnrichedElement(
+        [TensorFiniteElement(e, element._shape, element._transpose) for e in summands.elements],
+        is_nodal_enriched=summands.is_nodal_enriched)
+
+
+@as_enriched.register(TensorProductElement)
+def as_enriched_tensor_product(element):
+    """Distribute the product over the sum its first factor is.
+
+    The summands of a sum in the first factor own a contiguous range of the
+    flat basis index, so they stack in the order the product already numbers
+    them.  A sum in any later factor would interleave with the factors before
+    it, and stacking would renumber the degrees of freedom.
+    """
+    first, *rest = element.factors
+    if any(as_enriched(factor) is not None for factor in rest):
+        raise NotImplementedError(
+            "Only the first factor of a TensorProductElement may be a direct"
+            " sum, as the degrees of freedom of a later one do not stack"
+        )
+    summands = as_enriched(first)
+    if summands is None:
+        return None
+    return EnrichedElement(
+        [TensorProductElement((e, *rest)) for e in summands.elements],
+        is_nodal_enriched=summands.is_nodal_enriched)
 
 
 def tree_map(f, *args):

finat/finiteelementbase.py

@@ -9,6 +9,24 @@
 from finat.quadrature import make_quadrature
 
 
+def broadcast_tensor(expression, multiindex):
+    """Reshape a tensor expression, broadcasting over the indices it lacks.
+
+    :arg expression: an indexed tensor.
+    :arg multiindex: the indices to turn into a shape.
+    :returns: ``expression`` as a tensor of the extents of ``multiindex``.
+
+    :class:`gem.ComponentTensor` requires every index to be free in the
+    expression, which a cellwise constant, say, does not satisfy.  Multiply
+    by ones to broadcast over the missing indices.
+    """
+    missing = tuple(i for i in multiindex if i not in expression.free_indices)
+    if missing:
+        ones = gem.Literal(numpy.ones(tuple(i.extent for i in missing)))
+        expression = gem.Product(expression, gem.Indexed(ones, missing))
+    return gem.ComponentTensor(expression, multiindex)
+
+
 class FiniteElementBase(metaclass=ABCMeta):
 
     @abstractproperty
@@ -163,6 +181,39 @@ def basis_evaluation(self, order, ps, entity=None, coordinate_mapping=None):
            provides physical geometry callbacks (may be None).
         '''
 
+    def _stack_tabulations(self, order, ps, entity=None, coordinate_mapping=None):
+        """Tabulate on each point set of a union, stacking on the point index.
+
+        :arg order: return derivatives up to this order.
+        :arg ps: the :class:`~finat.point_set.UnionPointSet` to tabulate on.
+        :arg entity: the cell entity on which to tabulate.
+        :arg coordinate_mapping: a
+           :class:`~.physically_mapped.PhysicalGeometry` object that
+           provides physical geometry callbacks (may be None).
+        :returns: the tabulation on the whole of ``ps``, as
+           :meth:`basis_evaluation` returns.
+
+        A union of points has no structure of its own, so structured elements
+        tabulate on each point set in turn and stack the tabulations here.
+        """
+        tables = [self.basis_evaluation(order, sub, entity,
+                                        coordinate_mapping=coordinate_mapping)
+                  for sub in ps.point_sets]
+        keys, = set(map(frozenset, tables))
+        p, = ps.indices
+        multiindex = tuple(chain(self.get_indices(), self.get_value_indices()))
+
+        def concatenate(alpha):
+            # The point indices are free in each table, so promote them to a
+            # shape before concatenating the tables along it.
+            pieces = [broadcast_tensor(gem.Indexed(table[alpha], multiindex),
+                                       sub.indices)
+                      for table, sub in zip(tables, ps.point_sets)]
+            return gem.ComponentTensor(
+                gem.Indexed(gem.Concatenate(*pieces), (p,)), multiindex)
+
+        return {alpha: concatenate(alpha) for alpha in keys}
+
     @abstractmethod
     def point_evaluation(self, order, refcoords, entity=None, coordinate_mapping=None):
         '''Return code for evaluating the element at an arbitrary points on
@@ -262,6 +313,30 @@ def dual_evaluation(self, fn, coordinate_mapping=None):
                   is compiled from ``evaluation`` (alongside any argument
                   multiindices already encoded within ``fn``)
         '''
+        # Only a direct sum can dual evaluate each summand on its own points,
+        # so bring the sum outermost first.
+        from finat.enriched import as_enriched  # Avoid circular import
+        summands = as_enriched(self)
+        if summands is not None and summands is not self:
+            return summands.dual_evaluation(fn, coordinate_mapping=coordinate_mapping)
+        return self._dual_evaluation(fn, coordinate_mapping=coordinate_mapping)
+
+    def _dual_evaluation(self, fn, coordinate_mapping=None):
+        """Dual evaluate an element that is not a direct sum.
+
+        :arg fn: Callable representing the function to dual evaluate.
+                 Callable should take in an :class:`AbstractPointSet` and
+                 return a GEM expression for evaluation of the function at
+                 those points.
+        :arg coordinate_mapping: a
+           :class:`~.physically_mapped.PhysicalGeometry` object that
+           provides physical geometry callbacks (may be None).
+        :returns: an ``(evaluation, point_indices, basis_indices)`` triple, as
+           :meth:`dual_evaluation` returns.
+
+        :meth:`dual_evaluation` rewrites an element as a direct sum before
+        calling this, so the element here has a single set of points.
+        """
         Q, x = self.dual_basis
         Q = self.dual_transformation(Q, coordinate_mapping=coordinate_mapping)
 

finat/point_set.py

@@ -233,6 +233,29 @@ def almost_equal(self, other, tolerance=1e-12):
                 for s, o in zip(self.factors, other.factors))
 
 
+class UnionPointSet(PointSet):
+    """All of the points of several point sets, along a single point index.
+
+    :arg point_sets: the point sets to take the union of, in the order they
+        occupy the point index.  A union of unions is flattened.
+
+    These are the points of a dual basis that evaluates summand by summand,
+    named so that a function space can be built on them.  The point sets are
+    kept because an element that needs their structure to tabulate -- a tensor
+    product, which cannot factor the union -- tabulates on each in turn.
+    """
+
+    def __init__(self, point_sets):
+        self.point_sets = tuple(chain(*(
+            ps.point_sets if isinstance(ps, UnionPointSet) else (ps,)
+            for ps in point_sets)))
+        super().__init__(numpy.concatenate([ps.points
+                                            for ps in self.point_sets]))
+
+    def __repr__(self):
+        return f"{type(self).__name__}({self.point_sets!r})"
+
+
 class FacetPointSet(AbstractPointSet):
     """A point set on facets.