From 808ca7ef38bd867959df233061b880ca775002c6 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 10 Sep 2026 15:59:55 +0100 Subject: [PATCH 1/3] Preserve direct-sum point indices in dual evaluation --- finat/enriched.py | 26 ++++++++++++++------------ test/finat/test_dual_basis.py | 26 ++++++++++++++++++++++---- 2 files changed, 36 insertions(+), 16 deletions(-) diff --git a/finat/enriched.py b/finat/enriched.py index f355b690..bb65d40c 100644 --- a/finat/enriched.py +++ b/finat/enriched.py @@ -245,29 +245,31 @@ def _dual_evaluation(self, fn, coordinate_mapping=None): 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. + returns. The summand point indices remain free, so the caller can + choose how to contract each direct-sum component. - 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. + The summands do not share their points, so their evaluations stack + along the basis index while retaining their own point indices. + Concatenating over a free basis index is what + :func:`~gem.unconcatenate.unconcatenate` splits downstream. """ if not self.is_nodal_enriched: raise NotImplementedError( f"Dual evaluation not defined for non-nodal {type(self).__name__}" ) - # Each summand contracts through its own dual_basis, so a non-nodal - # sum has to be refused here as well as in dual_basis: this path - # never asks self for one. + # Each summand uses its own dual_basis, so a non-nodal sum has to be + # refused here as well as in dual_basis: this path never asks self + # for one. evals = [] + point_indices = [] for element in self.elements: - expr, point_indices, indices = element.dual_evaluation( + expr, element_points, indices = element.dual_evaluation( fn, coordinate_mapping=coordinate_mapping) - evals.append(broadcast_tensor(gem.IndexSum(expr, point_indices), indices)) + evals.append(broadcast_tensor(expr, indices)) + point_indices.extend(element_points) beta = self.get_indices() - return gem.Indexed(gem.Concatenate(*evals), beta), (), beta + return gem.Indexed(gem.Concatenate(*evals), beta), tuple(dict.fromkeys(point_indices)), beta @singledispatch diff --git a/test/finat/test_dual_basis.py b/test/finat/test_dual_basis.py index 5fb58d98..bf1306f7 100644 --- a/test/finat/test_dual_basis.py +++ b/test/finat/test_dual_basis.py @@ -12,6 +12,7 @@ from finat.quadrature import QuadratureRule from finat.quadrature_element import QuadratureElement from gem.interpreter import evaluate +from gem.unconcatenate import unconcatenate from FIAT import ufc_simplex @@ -50,11 +51,28 @@ 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)]) + expr, _, indices = element.dual_evaluation(tabulate) n = element.space_dimension() + strides = tuple( + numpy.prod(tuple(index.extent for index in indices[offset + 1:]), dtype=int) + for offset in range(len(indices)) + ) + variable = gem.FlexiblyIndexed( + gem.Variable("A", (n,)), ((0, tuple(zip(indices, strides))),) + ) + blocks = [] + for variable, evaluation in unconcatenate([(variable, expr)]): + point_indices = tuple( + index for index in evaluation.free_indices + if index not in variable.free_indices and index not in j + ) + blocks.append(gem.ComponentTensor( + gem.IndexSum(evaluation, point_indices), + variable.index_ordering(), + )) + i = gem.Index(extent=n) + evaluation = gem.Indexed(gem.Concatenate(*blocks), (i,)) + result, = evaluate([gem.ComponentTensor(evaluation, (i,) + j)]) assert numpy.allclose(result.arr.reshape(n, n), numpy.eye(n)) From d5a558d4887870b1825d80b5f339755ebdc22f2f Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 11 Sep 2026 23:40:24 +0100 Subject: [PATCH 2/3] Fold tables of zeros in constant_fold_zero numpy.array_equal compares shapes before it compares values, so it returned False for every table that the pass gave it and folded nothing but a scalar zero. A table of zeros stayed a dense Literal. Dual evaluation between facet-restricted elements on a tensor product cell tabulates the interior basis functions of each direct-sum component at the boundary points of every other component. Those tables are exactly zero, so 30 of the 49 component pairs that a hexahedron produces contribute nothing. Folding the tables lets Indexed and Product drop those pairs, which returns the interpolation to the O(degree^(dim + 1)) cost that sum factorisation gives. The kernel of tests/tsfc/test_dual_evaluation.py::test_dual_argument_is_sum_factorised loses 72% of its flops at degree 16. Fold on the value of the table, and keep the dtype so that a table of integers does not become a floating point zero. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_013tWTW5ErfhCgV63nXai8HN --- gem/optimise.py | 6 +++--- test/gem/test_simplify.py | 27 +++++++++++++++++++++++++++ 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/gem/optimise.py b/gem/optimise.py index 494a55db..19dff5cb 100644 --- a/gem/optimise.py +++ b/gem/optimise.py @@ -194,9 +194,9 @@ def _constant_fold_zero(node, self): @_constant_fold_zero.register(Literal) def _constant_fold_zero_literal(node, self): - if numpy.array_equal(node.array, 0): - # All zeros, make symbolic zero - return Zero(node.shape) + if not node.array.any(): + # A table of any shape that holds only zeros is a symbolic zero. + return Zero(node.shape, dtype=node.dtype) else: return node diff --git a/test/gem/test_simplify.py b/test/gem/test_simplify.py index 94aeac3d..57b48a69 100644 --- a/test/gem/test_simplify.py +++ b/test/gem/test_simplify.py @@ -100,3 +100,30 @@ def test_flatten_indexsum(A): result = gem.IndexSum(gem.IndexSum(Aij, (i,)), (j,)) expected = gem.IndexSum(Aij, (i, j)) assert result == expected + + +@pytest.mark.parametrize("shape", [(), (3,), (7, 2), (2, 3, 4)]) +def test_constant_fold_zero_table(shape): + """A Literal of any shape that holds only zeros folds to a Zero.""" + zeros = gem.Literal(numpy.zeros(shape)) + folded, = gem.optimise.constant_fold_zero([zeros]) + assert isinstance(folded, gem.Zero) + assert folded.shape == shape + + +def test_constant_fold_zero_keeps_nonzero_table(): + """A Literal that holds a nonzero entry stays a Literal.""" + array = numpy.zeros((7, 2)) + array[3, 1] = 1.0 + literal = gem.Literal(array) + folded, = gem.optimise.constant_fold_zero([literal]) + assert folded == literal + + +def test_constant_fold_zero_removes_product(): + """Folding a zero table removes the expression that it multiplies.""" + i, j = gem.indices(2) + zeros = gem.Indexed(gem.Literal(numpy.zeros((7, 2))), (i, j)) + other = gem.Indexed(gem.Variable("v", (7, 2)), (i, j)) + folded, = gem.optimise.constant_fold_zero([gem.Product(zeros, other)]) + assert isinstance(folded, gem.Zero) From 8cd5f4e5a016462cc02a02f35629e9b38f4da05d Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sun, 13 Sep 2026 01:36:50 +0100 Subject: [PATCH 3/3] Block a direct sum's basis and its dual basis the same way EnrichedElement presented two decompositions of the same direct sum. basis_evaluation concatenated blocks of shape elem.index_shape over self.elements, giving (4,4,3), (24,4) for NCE3, while _dual_evaluation concatenated whatever the as_enriched rewriting returned, giving (4,4,3), (96,). A basis and its dual basis must be blocked alike, so nothing downstream could pair them up and contract them. Promote EnrichedElement's private _summands to a `summands` property on every element, and block basis_evaluation, point_evaluation, dual_basis and _dual_evaluation along it alike. as_enriched on a FlattenedDimensions dropped the wrapper and returned summands on the tensor product cell, which cannot tabulate against the entities of the quadrilateral or hexahedron they came from. Distribute the flattening over the sum instead, as the other wrappers already do. Add split_contraction, which carries the identity that a sum over a whole direct sum is the sum of the sums over its blocks. Unlike unconcatenate it needs no assignment variable to carry the concatenation index, because the sum itself is what the Concatenate splits against. split_group holds the part that the two now share. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01KWajdMc5VFPbPuB1cupu9F --- finat/enriched.py | 43 +++++++------ finat/finiteelementbase.py | 22 +++++++ gem/unconcatenate.py | 114 ++++++++++++++++++++++++++++------ test/finat/test_dual_basis.py | 6 +- 4 files changed, 144 insertions(+), 41 deletions(-) diff --git a/finat/enriched.py b/finat/enriched.py index bb65d40c..ed9b8c13 100644 --- a/finat/enriched.py +++ b/finat/enriched.py @@ -115,7 +115,7 @@ def merge(tables): tables = tuple(tables) zeta = self.get_value_indices() tensors = [] - for elem, table in zip(self.elements, tables): + for elem, table in zip(self.summands, tables): beta_i = elem.get_indices() tensors.append(gem.ComponentTensor( gem.Indexed(table, beta_i + zeta), @@ -138,7 +138,7 @@ def basis_evaluation(self, order, ps, entity=None, coordinate_mapping=None): :param entity: the cell entity on which to tabulate. ''' results = [element.basis_evaluation(order, ps, entity, coordinate_mapping=coordinate_mapping) - for element in self.elements] + for element in self.summands] return self._compose_evaluations(results) def point_evaluation(self, order, refcoords, entity=None, coordinate_mapping=None): @@ -153,7 +153,7 @@ def point_evaluation(self, order, refcoords, entity=None, coordinate_mapping=Non :param entity: the cell entity on which to tabulate. ''' results = [element.point_evaluation(order, refcoords, entity, coordinate_mapping) - for element in self.elements] + for element in self.summands] return self._compose_evaluations(results) @property @@ -166,20 +166,17 @@ def mapping(self): return result @cached_property - def _summands(self): + 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. + An element is brought out as a direct sum one level at a time. A summand + of :attr:`elements` may therefore be a direct sum in turn, and these are + the elements that remain once every level is brought out. They 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) + return tuple(chain.from_iterable(element.summands + for element in self.elements)) @property def dual_basis(self): @@ -205,12 +202,12 @@ def dual_basis(self): f"Dual basis not defined for non-nodal {type(self).__name__}" ) if any(type(e).dual_transformation is not FiniteElementBase.dual_transformation - for e in self._summands): + for e in self.summands): raise NotImplementedError( f"dual_basis not defined for {type(self).__name__} with a summand" " that has its own dual_transformation; use dual_evaluation instead" ) - duals = [element.dual_basis for element in self._summands] + duals = [element.dual_basis for element in self.summands] x = UnionPointSet([xk for _, xk in duals]) p, = x.indices zeta = self.get_value_indices() @@ -219,7 +216,7 @@ def dual_basis(self): 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)): + 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. @@ -262,7 +259,7 @@ def _dual_evaluation(self, fn, coordinate_mapping=None): # for one. evals = [] point_indices = [] - for element in self.elements: + for element in self.summands: expr, element_points, indices = element.dual_evaluation( fn, coordinate_mapping=coordinate_mapping) evals.append(broadcast_tensor(expr, indices)) @@ -292,7 +289,15 @@ def as_enriched_enriched(element): @as_enriched.register(FlattenedDimensions) def as_enriched_flattened(element): - return as_enriched(element.product) + """Distribute the flattening over the sum the product is. + + Each summand keeps the cell of the element that it came out of. The + summands therefore tabulate against the same entities as that element. + """ + summands = as_enriched(element.product) + if summands is None: + return None + return distribute_over_sum(FlattenedDimensions, summands) @as_enriched.register(DiscontinuousElement) diff --git a/finat/finiteelementbase.py b/finat/finiteelementbase.py index 6da8abd9..06574e3b 100644 --- a/finat/finiteelementbase.py +++ b/finat/finiteelementbase.py @@ -291,6 +291,28 @@ def dual_basis(self): f"Dual basis not defined for element {type(self).__name__}" ) + @cached_property + def summands(self): + """The direct summands whose bases stack into this element's basis. + + A direct sum blocks its tabulation and its dual basis along these + summands alike. A contraction of the one against the other therefore + splits into a sum over them. + + Returns + ------- + tuple + The elements, on this element's own cell and in basis order, whose + tabulations concatenate into this element's tabulation and whose + dual bases stack into its dual basis. An element that is not a + direct sum is its own only summand. + """ + from finat.enriched import as_enriched # Avoid circular import + summands = as_enriched(self) + if summands is None: + return (self,) + return summands.summands + def dual_evaluation(self, fn, coordinate_mapping=None): '''Get a GEM expression for performing the dual basis evaluation at the nodes of the reference element. Currently only works for flat diff --git a/gem/unconcatenate.py b/gem/unconcatenate.py index 3e20edf8..88399e8d 100644 --- a/gem/unconcatenate.py +++ b/gem/unconcatenate.py @@ -63,7 +63,7 @@ from gem.interpreter import evaluate -__all__ = ['flatten', 'unconcatenate'] +__all__ = ['flatten', 'split_contraction', 'unconcatenate'] def find_group(expressions, splittable_indices): @@ -175,16 +175,24 @@ def replace_node(expression, mapping, cut=None): return mapper(expression) -def _unconcatenate(cache, pairs): - # Tail-call recursive core of unconcatenate. - # Assumes that input has already been sanitised. - # 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 +def split_group(cache, concat_group): + """Splits a group of indexed Concatenate nodes into their blocks. + + Parameters + ---------- + cache + Index splitting cache :py:class:`dict`. + concat_group + A group of indexed :py:class:`Concatenate` nodes, as + :py:func:`find_group` returns. - # Get the index split + Returns + ------- + tuple + The index that the group shares, one multiindex for each block, and + one substitution for each block. A substitution replaces every node + of the group by that block of it. + """ concat_ref = next(iter(concat_group)) assert isinstance(concat_ref, Indexed) concat_expr, = concat_ref.children @@ -197,19 +205,31 @@ def _unconcatenate(cache, pairs): for child in concat_expr.children) cache[index] = multiindices - def cut(node): - """No need to rebuild expression of independent of the - relevant concatenation index.""" - return index not in node.free_indices - - # Build Concatenate node replacement mappings mappings = [{} for i in range(len(multiindices))] for concat_ref in concat_group: concat_expr, = concat_ref.children - for i in range(len(multiindices)): - sub_ref = Indexed(concat_expr.children[i], multiindices[i]) + for i, multiindex in enumerate(multiindices): + sub_ref = Indexed(concat_expr.children[i], multiindex) sub_ref, = remove_componenttensors((sub_ref,)) mappings[i][concat_ref] = sub_ref + return index, multiindices, mappings + + +def _unconcatenate(cache, pairs): + # Tail-call recursive core of unconcatenate. + # Assumes that input has already been sanitised. + # 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 + + index, multiindices, mappings = split_group(cache, concat_group) + + def cut(node): + """No need to rebuild expression of independent of the + relevant concatenation index.""" + return index not in node.free_indices # Finally, split assignment pairs split_pairs = [] @@ -224,6 +244,64 @@ def cut(node): return _unconcatenate(cache, split_pairs) +def _split_contraction(cache, expression, indices): + # Tail-call recursive core of split_contraction. + # Assumes that input has already been sanitised. + concat_group = find_group([expression], set(indices)) + if concat_group is None: + return [(expression, indices)] + + index, multiindices, mappings = split_group(cache, concat_group) + + def cut(node): + """No need to rebuild expression of independent of the + relevant concatenation index.""" + return index not in node.free_indices + + # Split the contraction, one block at a time + rest = tuple(i for i in indices if i != index) + terms = [] + for multiindex, mapping in zip(multiindices, mappings): + terms.extend(_split_contraction(cache, replace_node(expression, mapping, cut), + rest + multiindex)) + return terms + + +def split_contraction(expression, indices, cache=None): + """Splits a contraction along the :py:class:`Concatenate` nodes it sums over. + + No assignment variable need carry the concatenation index here. The sum + is what the Concatenate splits against. A sum over a whole direct sum is + the sum of the sums over its blocks: + + sum_j Indexed(Concatenate(A, B), (j,)) * Indexed(Concatenate(C, D), (j,)) + = sum_{ja} A_ja * C_ja + sum_{jb} B_jb * D_jb. + + Every Concatenate that one index indexes must concatenate the same blocks. + A FInAT element gives that guarantee: it blocks its tabulation and its dual + basis along the same summands. + + Parameters + ---------- + expression + A scalar GEM expression. + indices + The multiindex that ``expression`` is summed over. + cache + Index splitting cache :py:class:`dict` (optional). + + Returns + ------- + list + The (expression, multiindex) pairs whose index sums add up to the + index sum of ``expression`` over ``indices``. + """ + if cache is None: + cache = {} + expression, = remove_componenttensors([expression]) + return _split_contraction(cache, expression, tuple(indices)) + + def unconcatenate(pairs, cache=None): """Splits a list of (indexed variable, expression) pairs along :py:class:`Concatenate` nodes embedded in the expressions. diff --git a/test/finat/test_dual_basis.py b/test/finat/test_dual_basis.py index bf1306f7..cdefd63d 100644 --- a/test/finat/test_dual_basis.py +++ b/test/finat/test_dual_basis.py @@ -81,10 +81,8 @@ def check_dual_basis(element): 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) + assert len(x.points) == sum(len(e.dual_basis[1].points) + for e in element.summands) i = element.get_indices() j = element.get_indices()