diff --git a/finat/discontinuous.py b/finat/discontinuous.py index abfb3782..f00a36d9 100644 --- a/finat/discontinuous.py +++ b/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 diff --git a/finat/enriched.py b/finat/enriched.py index 8c94d4a6..3012159b 100644 --- a/finat/enriched.py +++ b/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.quadrature_element import QuadratureElement +from finat.tensor_product import TensorProductElement class EnrichedElement(FiniteElementBase): @@ -160,34 +165,159 @@ 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. + """ + 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(QuadratureElement) +def as_enriched_quadrature_element(element): + """Rewrite a rule on a union of point sets as a sum of one rule each.""" + rules = element._summand_rules + if not rules: + return None + return EnrichedElement( + [QuadratureElement(element.cell, rule) for rule in rules], + is_nodal_enriched=True) + + +@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): diff --git a/finat/finiteelementbase.py b/finat/finiteelementbase.py index 8f21c552..b1f0e28c 100644 --- a/finat/finiteelementbase.py +++ b/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) diff --git a/finat/point_set.py b/finat/point_set.py index 068f7659..f1b7af66 100644 --- a/finat/point_set.py +++ b/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. diff --git a/finat/quadrature_element.py b/finat/quadrature_element.py index 7ba5633d..edd8e8e0 100644 --- a/finat/quadrature_element.py +++ b/finat/quadrature_element.py @@ -1,4 +1,4 @@ -from finat.point_set import UnknownPointSet, FacetPointSet +from finat.point_set import UnknownPointSet, FacetPointSet, UnionPointSet import numpy @@ -9,7 +9,7 @@ from gem.utils import cached_property from finat.finiteelementbase import FiniteElementBase -from finat.quadrature import make_quadrature, AbstractQuadratureRule +from finat.quadrature import make_quadrature, AbstractQuadratureRule, QuadratureRule def make_quadrature_element(fiat_ref_cell, degree, scheme="default", codim=0): @@ -113,18 +113,21 @@ def value_shape(self): return () @cached_property - def fiat_equivalent(self): - ps = self._point_set - if isinstance(ps, UnknownPointSet): - raise ValueError("A quadrature element with rule with runtime points has no fiat equivalent!") + def _weights(self): weights = getattr(self._rule, 'weights', None) if weights is None: # we need the weights. weights, = evaluate([self._rule.weight_expression]) weights = weights.arr.flatten() self._rule.weights = weights + return weights - return FIAT.QuadratureElement(self.cell, ps.points, weights) + @cached_property + def fiat_equivalent(self): + ps = self._point_set + if isinstance(ps, UnknownPointSet): + raise ValueError("A quadrature element with rule with runtime points has no fiat equivalent!") + return FIAT.QuadratureElement(self.cell, ps.points, self._weights) def basis_evaluation(self, order, ps, entity=None, coordinate_mapping=None): '''Return code for evaluating the element at known points on the @@ -148,15 +151,35 @@ def basis_evaluation(self, order, ps, entity=None, coordinate_mapping=None): if order: raise ValueError("Derivatives are not defined on a QuadratureElement.") - if not self._rule.point_set.almost_equal(ps): - raise ValueError("Mismatch of quadrature points!") + # A union of points has no structure of its own to tabulate on. + if isinstance(ps, UnionPointSet): + return self._stack_tabulations(order, ps, entity, coordinate_mapping=coordinate_mapping) - # Return an outer product of identity matrices basis_indices = self.get_indices() - point_indices = ps.indices - if len(basis_indices) > len(point_indices): - point_indices = (entity_id, *point_indices) - delta = gem.Delta(point_indices, basis_indices) + ps_indices = ps.indices + if isinstance(self._point_set, FacetPointSet): + # A FacetPointSet carries a facet index, absent from the rule's. + ps_indices = (entity_id, *ps_indices) + + rule_ps = self._rule.point_set + blocks = rule_ps.point_sets if isinstance(rule_ps, UnionPointSet) else (rule_ps,) + matches = [k for k, block in enumerate(blocks) if block.almost_equal(ps)] + if not matches: + raise ValueError("Mismatch of quadrature points!") + k, = matches + + if isinstance(rule_ps, UnionPointSet): + # `ps` is one point set of the union: tabulate onto the rows of the + # identity it owns, and zero onto the others. Concatenating along + # the basis index lets the contraction with a coefficient split. + beta = tuple(gem.Index(extent=index.extent) for index in ps_indices) + own = gem.ComponentTensor(gem.Delta(ps_indices, beta), beta) + branches = [own if j == k else gem.Zero(tuple(i.extent for i in block.indices)) + for j, block in enumerate(blocks)] + delta = gem.Indexed(gem.Concatenate(*branches), basis_indices) + else: + # Return an outer product of identity matrices + delta = gem.Delta(ps_indices, basis_indices) sd = self.cell.get_spatial_dimension() return {(0,) * sd: gem.ComponentTensor(delta, basis_indices)} @@ -174,6 +197,37 @@ def dual_basis(self): Q = gem.ComponentTensor(Q, multiindex) return Q, ps + @cached_property + def _summand_rules(self): + """The rules of the summands this element is a direct sum of. + + Returns + ------- + tuple + One :class:`~finat.quadrature.QuadratureRule` for each point set of + a :class:`~finat.point_set.UnionPointSet` rule, in the order they + are stacked, or an empty tuple if the rule has a single point set. + + Notes + ----- + A union of point sets is how the points of a direct sum are stacked, so + splitting it back up recovers a rule for each summand, each on the + points its own functionals evaluate on. + + """ + rule_ps = self._rule.point_set + if not isinstance(rule_ps, UnionPointSet): + return () + + rules = [] + offset = 0 + for ps in rule_ps.point_sets: + n = len(ps.points) + rules.append(QuadratureRule(ps, self._weights[offset:offset + n], + ref_el=self._rule.ref_el)) + offset += n + return tuple(rules) + @property def mapping(self): return "affine" diff --git a/finat/restricted.py b/finat/restricted.py index 6b4e291e..b0ca0d21 100644 --- a/finat/restricted.py +++ b/finat/restricted.py @@ -1,4 +1,4 @@ -from functools import singledispatch +from functools import partial, singledispatch from itertools import chain import FIAT @@ -120,7 +120,10 @@ def restrict_enriched(element, domain, take_closure): elif not any(isinstance(e, finat.mixed.MixedSubElement) for e in element.elements): elements = tuple(restrict(e, domain, take_closure) for e in element.elements) - reconstruct = finat.EnrichedElement + # Restriction selects disjoint subsets of the DoFs, so the restricted + # subelements are nodal whenever the original ones are. + reconstruct = partial(finat.EnrichedElement, + is_nodal_enriched=element.is_nodal_enriched) else: raise NotImplementedError("Not expecting enriched with mixture of MixedSubElement and others") @@ -131,30 +134,23 @@ def restrict_enriched(element, domain, take_closure): return null_element -@restrict.register(finat.HCurlElement) -def restrict_hcurl(element, domain, take_closure): - restricted = restrict(element.wrappee, domain, take_closure) - if restricted is null_element: - return null_element - else: - if isinstance(restricted, finat.EnrichedElement): - return finat.EnrichedElement(finat.HCurlElement(e) - for e in restricted.elements) - else: - return finat.HCurlElement(restricted) - +@restrict.register(finat.hdivcurl.WrapperElementBase) +def restrict_hdivcurl(element, domain, take_closure): + """Restrict an element that wraps another with a pullback. -@restrict.register(finat.HDivElement) -def restrict_hdiv(element, domain, take_closure): + The pullback acts on each subelement separately, so it preserves + whichever subelements the restriction selects, and with them the + nodality of the restricted element. + """ restricted = restrict(element.wrappee, domain, take_closure) if restricted is null_element: return null_element + elif isinstance(restricted, finat.EnrichedElement): + return finat.EnrichedElement( + [type(element)(e) for e in restricted.elements], + is_nodal_enriched=restricted.is_nodal_enriched) else: - if isinstance(restricted, finat.EnrichedElement): - return finat.EnrichedElement(finat.HDivElement(e) - for e in restricted.elements) - else: - return finat.HDivElement(restricted) + return type(element)(restricted) @restrict.register(finat.mixed.MixedSubElement) diff --git a/finat/tensor_product.py b/finat/tensor_product.py index 2672fa01..274d05e4 100644 --- a/finat/tensor_product.py +++ b/finat/tensor_product.py @@ -12,7 +12,7 @@ from gem.utils import cached_property from finat.finiteelementbase import FiniteElementBase -from finat.point_set import PointSingleton, PointSet, TensorPointSet +from finat.point_set import PointSingleton, PointSet, TensorPointSet, UnionPointSet class TensorProductElement(FiniteElementBase): @@ -133,6 +133,13 @@ def _merge_evaluations(self, factor_results): return result def basis_evaluation(self, order, ps, entity=None, coordinate_mapping=None): + if isinstance(ps, UnionPointSet): + # A union of points does not factor, so tabulate on each of its + # point sets, where the product structure survives to be sum + # factorised. + return self._stack_tabulations(order, ps, entity, + coordinate_mapping=coordinate_mapping) + entities = self._factor_entity(entity) entity_dim, _ = zip(*entities) diff --git a/finat/tensorfiniteelement.py b/finat/tensorfiniteelement.py index 8293f8fa..85f7bd25 100644 --- a/finat/tensorfiniteelement.py +++ b/finat/tensorfiniteelement.py @@ -170,29 +170,48 @@ def dual_basis(self): tQ = gem.ComponentTensor(Qi*deltas, index_ordering) return tQ, points - def dual_evaluation(self, fn, coordinate_mapping=None): - tQ, x = self.dual_basis - tQ = self._base_element.dual_transformation(tQ, coordinate_mapping) - - expr = fn(x) - # NOTE: any shape indices in the expression are because the - # expression is tensor valued. - assert expr.shape == self.value_shape - - scalar_i = self.base_element.get_indices() - scalar_vi = self.base_element.get_value_indices() + def _dual_evaluation(self, fn, coordinate_mapping=None): + """Dual evaluate the base element on one tensor component at a time. + + :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 tensor wrapper couples a basis function to a value component + through an identity, so the components share the base element's dual + basis and the base element keeps whatever structure it dual evaluates + with -- a direct sum among it. + """ + base = self.base_element + base_vi = base.get_value_indices() tensor_i = tuple(gem.Index(extent=d) for d in self._shape) - tensor_vi = tuple(gem.Index(extent=d) for d in self._shape) + def component(ps): + expr = fn(ps) + assert expr.shape == self.value_shape + return gem.ComponentTensor(gem.Indexed(expr, tensor_i + base_vi), base_vi) + + evaluation, point_indices, scalar_i = base.dual_evaluation( + component, coordinate_mapping=coordinate_mapping) + # Couple the component the base element saw to the basis function that + # carries it. Contracting the pair keeps the component index out of + # the basis indices, where a free index into a ListTensor would defeat + # argument factorisation. + tensor_vi = tuple(gem.Index(extent=d) for d in self._shape) + evaluation = gem.IndexSum( + gem.Product(gem.Delta(tensor_i, tensor_vi), evaluation), tensor_i) if self._transpose: - index_ordering = tensor_i + scalar_i + tensor_vi + scalar_vi + basis_indices = tensor_vi + scalar_i else: - index_ordering = scalar_i + tensor_i + tensor_vi + scalar_vi - - tQi = tQ[index_ordering] - expri = expr[tensor_i + scalar_vi] - evaluation = gem.IndexSum(tQi * expri, scalar_vi + tensor_i) - return evaluation, x.indices, scalar_i + tensor_vi + basis_indices = scalar_i + tensor_vi + return evaluation, point_indices, basis_indices @property def mapping(self): diff --git a/gem/unconcatenate.py b/gem/unconcatenate.py index ce6e30b3..74b2171d 100644 --- a/gem/unconcatenate.py +++ b/gem/unconcatenate.py @@ -66,7 +66,7 @@ __all__ = ['flatten', 'unconcatenate'] -def find_group(expressions): +def find_group(expressions, free_indices): """Finds a full set of indexed Concatenate nodes with the same free index, if any such node exists. @@ -74,10 +74,12 @@ def find_group(expressions): must be removed. :arg expressions: a multi-root GEM expression DAG + :arg free_indices: the indices that may be split along, that is, those + carried by the assignment variables. A Concatenate + indexed by anything else has nothing to be split + against, and is left alone. :returns: a list of GEM nodes, or None """ - free_indices = set().union(chain(*[e.free_indices for e in expressions])) - # Result variables index = None nodes = [] @@ -100,8 +102,7 @@ def find_group(expressions): child, = node.children if isinstance(child, Concatenate): i, = node.multiindex - assert i in free_indices - if (index or i) == i: + if i in free_indices and (index or i) == i: index = i nodes.append(node) # Skip adding children @@ -178,7 +179,9 @@ def replace_node(expression, mapping, cut=None): def _unconcatenate(cache, pairs): # Tail-call recursive core of unconcatenate. # Assumes that input has already been sanitised. - concat_group = find_group([e for v, e in pairs]) + # Only an index carried by an assignment variable can be split against it. + splittable = set().union(chain(*[v.free_indices for v, e in pairs])) + concat_group = find_group([e for v, e in pairs], splittable) if concat_group is None: return pairs diff --git a/test/finat/test_dual_basis.py b/test/finat/test_dual_basis.py index c995cfa9..1737730a 100644 --- a/test/finat/test_dual_basis.py +++ b/test/finat/test_dual_basis.py @@ -1,9 +1,18 @@ +from itertools import chain + import pytest import numpy import finat import gem -from FIAT import ufc_simplex +import ufl +import finat.ufl +from finat.element_factory import create_element +from finat.enriched import as_enriched +from finat.point_set import UnionPointSet +from finat.quadrature import QuadratureRule +from finat.quadrature_element import QuadratureElement from gem.interpreter import evaluate +from FIAT import ufc_simplex @pytest.mark.parametrize("dim", (2, 3)) @@ -31,6 +40,103 @@ def test_collapse_repeated_points(dim): assert len(points) == expected +def check_nodal(element): + """Assert that applying the dual basis to the primal basis is the identity.""" + j = element.get_indices() + zeta = element.get_value_indices() + dim = element.cell.get_spatial_dimension() + + def tabulate(ps): + table = element.basis_evaluation(0, ps)[(0,) * dim] + return gem.ComponentTensor(gem.Indexed(table, j + zeta), zeta) + + expr, point_indices, indices = element.dual_evaluation(tabulate) + if point_indices: + expr = gem.IndexSum(expr, point_indices) + result, = evaluate([gem.ComponentTensor(expr, indices + j)]) + n = element.space_dimension() + assert numpy.allclose(result.arr.reshape(n, n), numpy.eye(n)) + + +def check_dual_basis(element): + """Assert that contracting the dual weights with the primal basis is the identity.""" + Q, x = element.dual_basis + assert Q.shape == element.index_shape + element.value_shape + assert set(Q.free_indices) == set(x.indices) + summands = as_enriched(element) + if summands is not None: + assert len(x.points) == sum(len(e.dual_basis[1].points) + for e in summands._summands) + + i = element.get_indices() + j = element.get_indices() + zeta = element.get_value_indices() + dim = element.cell.get_spatial_dimension() + table = element.basis_evaluation(0, x)[(0,) * dim] + expr = gem.IndexSum(gem.Product(gem.Indexed(Q, i + zeta), + gem.Indexed(table, j + zeta)), + x.indices + zeta) + result, = evaluate([gem.ComponentTensor(expr, i + j)]) + n = element.space_dimension() + assert numpy.allclose(result.arr.reshape(n, n), numpy.eye(n)) + + +def test_enriched_element_dual_basis(): + # The weights of a direct sum are block diagonal: each summand's weights + # sit at its own offset in the union of the points, and are zero against + # every other summand's points. + cell = ufc_simplex(2) + fe = finat.Lagrange(cell, 3) + enriched = finat.EnrichedElement( + [finat.RestrictedElement(fe, restriction_domain=domain) + for domain in ("interior", "facet")], is_nodal_enriched=True) + + assert isinstance(enriched.dual_basis[1], UnionPointSet) + check_dual_basis(enriched) + + +def test_quadrature_element_on_union_of_points(): + # Firedrake interpolates through a quadrature space on the points of the + # target's dual basis, which for a direct sum is a union. That element + # has to evaluate on each point set of the union, just as the sum does. + cell = ufc_simplex(2) + fe = finat.Lagrange(cell, 3) + enriched = finat.EnrichedElement( + [finat.RestrictedElement(fe, restriction_domain=domain) + for domain in ("interior", "facet")], is_nodal_enriched=True) + + _, ps = enriched.dual_basis + # The weights are not used, this quadrature scheme is not for integration. + rule = QuadratureRule(ps, numpy.full(len(ps.points), numpy.nan), ref_el=cell) + element = QuadratureElement(cell, rule) + # A vector-valued target interpolates through a wrapper of that element, + # which is distributed over the sum just the same. + vector = finat.TensorFiniteElement(element, (cell.get_spatial_dimension(),)) + for e in (element, vector): + check_nodal(e) + # Each summand evaluates on its own points. Handed the union instead, + # the callable tabulates into a Concatenate over the points, which is + # contracted away before anything downstream can split along it. + seen = [] + + def fn(point_set, e=e): + seen.append(point_set) + return gem.Literal(numpy.zeros(e.value_shape)) + + e.dual_evaluation(fn) + assert seen == list(ps.point_sets) + + +@pytest.mark.parametrize("family", ("RTCE", "RTCF", "NCE", "NCF")) +@pytest.mark.parametrize("degree", (1, 2)) +def test_hdivcurl_dual_basis(family, degree): + # A union of points is a point set like any other, so a tensor product + # tabulates on it by splitting the coordinates and sharing the point + # index, and the weights contract against that tabulation. + element = create_element(finat.ufl.FiniteElement(family, hdivcurl_cell(family), degree)) + check_dual_basis(element) + + def test_enriched_element_dual_evaluation(): cell = ufc_simplex(2) fe = finat.Lagrange(cell, 3) @@ -39,7 +145,6 @@ def test_enriched_element_dual_evaluation(): fe2 = finat.RestrictedElement(fe, restriction_domain="facet") enriched = finat.EnrichedElement([fe1, fe2], is_nodal_enriched=True) - # Check that calling dual_evaluation returns a valid Indexed expression fn = lambda x: gem.Literal(1.0) expr, point_indices, basis_indices = enriched.dual_evaluation(fn) assert isinstance(expr, gem.Indexed) @@ -47,6 +152,26 @@ def test_enriched_element_dual_evaluation(): assert len(basis_indices) == 1 assert basis_indices[0].extent == enriched.space_dimension() + check_nodal(enriched) + + +def test_enriched_element_as_tensor_product_factor(): + # Restricting an element on a tensor product cell to its facets makes + # the restriction of each factor a factor of the result. Those factors + # are themselves EnrichedElements, so the tensor product is the direct + # sum of the products of their summands. + interval = ufc_simplex(1) + square = finat.TensorProductElement([finat.Lagrange(interval, 3)] * 2) + restricted = finat.RestrictedElement(square, restriction_domain="facet") + assert isinstance(restricted, finat.EnrichedElement) + + cube = finat.TensorProductElement([restricted, finat.Lagrange(interval, 3)]) + expanded = as_enriched(cube) + assert len(expanded.elements) == len(restricted.elements) > 1 + assert sum(element.space_dimension() for element in expanded.elements) \ + == cube.space_dimension() + check_nodal(cube) + @pytest.fixture(scope="module") def hexahedron(): @@ -113,3 +238,63 @@ def cubed(ps): values = nodal_values(element, evaluation) expected = numpy.einsum("...i,...i->...", values, values)[..., None] * values assert numpy.allclose(nodal_values(element, cubed), expected) + + +def test_direct_sum_must_be_the_first_factor(): + # A sum in a later factor interleaves with the factors before it, so its + # summands do not stack along the flat basis index. + interval = ufc_simplex(1) + line = finat.Lagrange(interval, 3) + restricted = finat.RestrictedElement( + finat.TensorProductElement([line] * 2), restriction_domain="facet") + with pytest.raises(NotImplementedError): + as_enriched(finat.TensorProductElement([line, restricted])) + + +def hdivcurl_cell(family): + if family.startswith("RTC"): + return ufl.quadrilateral + return ufl.TensorProductCell(ufl.quadrilateral, ufl.interval) + + +@pytest.mark.parametrize("family", ("RTCE", "RTCF", "NCE", "NCF")) +@pytest.mark.parametrize("degree", (1, 2, 3)) +def test_hdivcurl_dual_evaluation(family, degree): + # On a hexahedron one factor of a summand is itself a direct sum, which + # only stacks once the sum is brought out through the product and the + # pullback around it. + element = create_element(finat.ufl.FiniteElement(family, hdivcurl_cell(family), degree)) + check_nodal(element) + + +@pytest.mark.parametrize("family", ("RTCE", "RTCF", "NCE", "NCF")) +@pytest.mark.parametrize("domain", ("interior", "facet")) +def test_restricted_hdivcurl_dual_basis(family, domain): + # Restriction selects disjoint subsets of the DoFs, so a restricted + # H(div)/H(curl) element stays nodal even where the summands are not + # orthogonal to each other, as several of them map to the same component. + if family.startswith("RTC"): + cell = ufl.quadrilateral + else: + cell = ufl.TensorProductCell(ufl.quadrilateral, ufl.interval) + element = create_element(finat.ufl.FiniteElement(family, cell, 2)[domain]) + check_nodal(element) + + # Each summand has a dual basis on its own points, and together they + # account for every functional. Bringing the sum outermost rewrites one + # level at a time, so recurse to the summands that are not sums themselves. + def summands(e): + expanded = as_enriched(e) + if expanded is None: + return (e,) + return tuple(chain.from_iterable(map(summands, expanded.elements))) + + elements = summands(element) + assert len(elements) > 1 + assert sum(e.space_dimension() for e in elements) == element.space_dimension() + for e in elements: + Q, x = e.dual_basis + assert Q.shape == e.index_shape + e.value_shape + assert set(Q.free_indices) <= set(x.indices) + assert len(element.dual_basis[1].points) \ + == sum(len(e.dual_basis[1].points) for e in elements)