diff --git a/finat/finiteelementbase.py b/finat/finiteelementbase.py index 1a5ba6fb..310d25d3 100644 --- a/finat/finiteelementbase.py +++ b/finat/finiteelementbase.py @@ -4,7 +4,8 @@ import gem import numpy from gem.interpreter import evaluate -from gem.optimise import delta_elimination, sum_factorise, traverse_product +from gem.optimise import (delta_elimination, is_contraction, sum_factorise, + traverse_product) from gem.utils import cached_property from finat.quadrature import make_quadrature @@ -266,8 +267,10 @@ def dual_evaluation(self, fn, coordinate_mapping=None): expr = fn(x) # Apply targeted sum factorisation and delta elimination to - # the expression - sum_indices, factors = delta_elimination(*traverse_product(expr)) + # the expression, preserving contractions that fn already factorised + # FIXME: a single pass of gem.optimise.contraction will choke on too many indices + # This is a temporary workaround https://github.com/firedrakeproject/fiat/issues/283 + sum_indices, factors = delta_elimination(*traverse_product(expr, stop_at=is_contraction)) expr = sum_factorise(sum_indices, factors) # NOTE: any shape indices in the expression are because the # expression is tensor valued. @@ -277,11 +280,9 @@ def dual_evaluation(self, fn, coordinate_mapping=None): Qi = Q[basis_indices + shape_indices] expri = expr[shape_indices] evaluation = gem.IndexSum(Qi * expri, x.indices + shape_indices) - # Now we want to factorise over the new contraction with x, - # ignoring any shape indices to avoid hitting the sum- - # factorisation index limit (this is a bit of a hack). - # Really need to do a more targeted job here. - evaluation = gem.optimise.contraction(evaluation, shape_indices) + # Factorise over the new contraction with Qi, keeping whole the + # contractions that fn already factorised + evaluation = gem.optimise.contraction(evaluation, stop_at=is_contraction) return evaluation, basis_indices def dual_transformation(self, Q, coordinate_mapping=None): diff --git a/finat/tensorfiniteelement.py b/finat/tensorfiniteelement.py index 7a4fc2fb..bb690233 100644 --- a/finat/tensorfiniteelement.py +++ b/finat/tensorfiniteelement.py @@ -3,7 +3,8 @@ import numpy import gem -from gem.optimise import delta_elimination, sum_factorise, traverse_product +from gem.optimise import (delta_elimination, is_contraction, sum_factorise, + traverse_product) from gem.utils import cached_property from finat.finiteelementbase import FiniteElementBase @@ -177,8 +178,10 @@ def dual_evaluation(self, fn, coordinate_mapping=None): expr = fn(x) # Apply targeted sum factorisation and delta elimination to - # the expression - sum_indices, factors = delta_elimination(*traverse_product(expr)) + # the expression, preserving contractions that fn already factorised + # FIXME: a single pass of gem.optimise.contraction will choke on too many indices + # This is a temporary workaround https://github.com/firedrakeproject/fiat/issues/283 + sum_indices, factors = delta_elimination(*traverse_product(expr, stop_at=is_contraction)) expr = sum_factorise(sum_indices, factors) # NOTE: any shape indices in the expression are because the # expression is tensor valued. @@ -200,7 +203,7 @@ def dual_evaluation(self, fn, coordinate_mapping=None): # This doesn't work perfectly, the resulting code doesn't have # a minimal memory footprint, although the operation count # does appear to be minimal. - evaluation = gem.optimise.contraction(evaluation) + evaluation = gem.optimise.contraction(evaluation, stop_at=is_contraction) return evaluation, scalar_i + tensor_vi @property diff --git a/gem/optimise.py b/gem/optimise.py index 47782479..2e1af754 100644 --- a/gem/optimise.py +++ b/gem/optimise.py @@ -630,17 +630,41 @@ def traverse_sum(expression, stop_at=None): return result -def contraction(expression, ignore=None): +def is_contraction(expression: Node) -> bool: + """Test whether an expression is a tensor contraction. + + Parameters + ---------- + expression : + A GEM expression. + + Returns + ------- + bool + Whether the expression is a contraction. + + Notes + ----- + Pass this as ``stop_at`` to keep a contraction that is already sum + factorised out of a surrounding one. Flattening it discards its + factorisation, along with any subexpression it shares with another + factor, and inflates the number of indices to factorise over. + + """ + return isinstance(expression, IndexSum) + + +def contraction(expression, stop_at=None): """Optimise the contractions of the tensor product at the root of the expression, including: - IndexSum-Delta cancellation - Sum factorisation - :arg ignore: Optional set of indices to ignore when applying sum - factorisation (otherwise all summation indices will be - considered). Use this if your expression has many contraction - indices. + :arg stop_at: Optional predicate on GEM expressions that are not to + be broken into further factors, see :func:`traverse_product`. + The contraction at the root is always broken up, as that is the + one being optimised. This routine was designed with finite element coefficient evaluation in mind. @@ -654,18 +678,13 @@ def contraction(expression, ignore=None): # Flatten product tree, eliminate deltas, sum factorise def rebuild(expression): - sum_indices, factors = traverse_product(expression, index_replacer=index_replacer) + root = expression + sum_indices, factors = traverse_product( + expression, index_replacer=index_replacer, + stop_at=None if stop_at is None else lambda e: e is not root and stop_at(e)) sum_indices, factors = delta_elimination(sum_indices, factors, index_replacer=index_replacer) factors = [index_replacer(f, ()) for f in factors] - if ignore is not None: - # TODO: This is a really blunt instrument and one might - # plausibly want the ignored indices to be contracted on - # the inside rather than the outside. - extra = tuple(i for i in sum_indices if i in ignore) - to_factor = tuple(i for i in sum_indices if i not in ignore) - return IndexSum(sum_factorise(to_factor, factors), extra) - else: - return sum_factorise(sum_indices, factors) + return sum_factorise(sum_indices, factors) # Sometimes the value shape is composed as a ListTensor, which # could get in the way of decomposing factors. In particular, diff --git a/test/finat/test_dual_basis.py b/test/finat/test_dual_basis.py index 7cfe1ae6..98f6ea29 100644 --- a/test/finat/test_dual_basis.py +++ b/test/finat/test_dual_basis.py @@ -3,6 +3,7 @@ import finat import gem from FIAT import ufc_simplex +from gem.interpreter import evaluate @pytest.mark.parametrize("dim", (2, 3)) @@ -45,3 +46,69 @@ def test_enriched_element_dual_evaluation(): assert isinstance(expr.children[0], gem.Concatenate) assert len(indices) == 1 assert indices[0].extent == enriched.space_dimension() + + +@pytest.fixture(scope="module") +def hexahedron(): + line = finat.Lagrange(ufc_simplex(1), 1) + return finat.TensorProductElement([finat.TensorProductElement([line, line]), line]) + + +def coefficient_evaluation(element, ps, dofs): + """Evaluate a coefficient at a point set, sum factorised as TSFC does.""" + beta = element.get_indices() + zeta = element.get_value_indices() + dim = element.cell.get_spatial_dimension() + table = element.basis_evaluation(0, ps)[(0,) * dim] + dofs = gem.Literal(dofs.reshape([index.extent for index in beta])) + value = gem.Product(gem.Indexed(table, beta + zeta), gem.Indexed(dofs, beta)) + return gem.ComponentTensor(gem.optimise.contraction(gem.IndexSum(value, beta)), zeta) + + +def nodal_values(element, fn): + """Dual evaluate fn against a nodal element, giving its values at the nodes.""" + expression, indices = element.dual_evaluation(fn) + result, = evaluate([gem.ComponentTensor(expression, indices)]) + return result.arr + + +@pytest.mark.parametrize("power", (2, 3, 4)) +def test_dual_evaluation_of_powers(hexahedron, power): + # Each evaluation contracts over the three tensor-product directions, so + # a product of them carries more indices than one sum factorisation can + # search. The evaluations are already factorised, so keep them that way. + numpy.random.seed(0) + dofs = numpy.random.rand(hexahedron.space_dimension()) + + def evaluation(ps): + return coefficient_evaluation(hexahedron, ps, dofs) + + def monomial(ps): + expression = evaluation(ps) + for _ in range(power - 1): + expression = gem.Product(expression, evaluation(ps)) + return expression + + assert numpy.allclose(nodal_values(hexahedron, monomial), + nodal_values(hexahedron, evaluation) ** power) + + +def test_dual_evaluation_of_coupled_evaluations(hexahedron): + # Contracting the value indices of two evaluations couples them into a + # single contraction, which no ordering of the factors can break up. + element = finat.TensorFiniteElement(hexahedron, (3,)) + numpy.random.seed(0) + dofs = numpy.random.rand(element.space_dimension()) + + def evaluation(ps): + return coefficient_evaluation(element, ps, dofs) + + def cubed(ps): + u = evaluation(ps) + i, j = gem.Index(extent=3), gem.Index(extent=3) + square = gem.IndexSum(gem.Product(gem.Indexed(u, (i,)), gem.Indexed(u, (i,))), (i,)) + return gem.ComponentTensor(gem.Product(square, gem.Indexed(u, (j,))), (j,)) + + values = nodal_values(element, evaluation) + expected = numpy.einsum("...i,...i->...", values, values)[..., None] * values + assert numpy.allclose(nodal_values(element, cubed), expected) diff --git a/test/gem/test_sum_factorise.py b/test/gem/test_sum_factorise.py index fb81cd7d..b506d43c 100644 --- a/test/gem/test_sum_factorise.py +++ b/test/gem/test_sum_factorise.py @@ -3,6 +3,8 @@ import gem from gem.interpreter import evaluate +from gem import optimise +from gem.node import traversal from gem.optimise import sum_factorise @@ -50,3 +52,33 @@ def test_too_many_indices_in_one_contraction(): table = gem.Indexed(gem.Literal(numpy.ones((2,) * 7)), indices) with pytest.raises(NotImplementedError): sum_factorise(indices, [table]) + + +def test_contraction_preserves_factorised_contractions(): + # A dual evaluation contracts the weights with an expression whose + # coefficient evaluations are already sum factorised. Those carry the + # point index of the contraction, so flattening them yields a single + # connected contraction that is too large to factorise. + numpy.random.seed(0) + p, q = gem.Index(extent=4), gem.Index(extent=4) + ijk = tuple(gem.Index(extent=3) for _ in range(3)) + + table = gem.Indexed(gem.Literal(numpy.random.rand(3, 3, 3, 4)), ijk + (p,)) + dofs = numpy.random.rand(3, 3, 3) + evaluation = optimise.contraction(gem.IndexSum(gem.Product(table, gem.Indexed(gem.Literal(dofs), ijk)), ijk)) + assert optimise.is_contraction(evaluation) + + weights = numpy.random.rand(4, 4) + cubed = gem.Product(gem.Product(gem.Indexed(gem.Literal(weights), (q, p)), evaluation), + gem.Product(evaluation, evaluation)) + expression = gem.IndexSum(cubed, (p,)) + + with pytest.raises(NotImplementedError): + optimise.contraction(expression) + + optimised = optimise.contraction(expression, stop_at=optimise.is_contraction) + assert evaluation in set(traversal([optimised])) + + result, = evaluate([gem.ComponentTensor(optimised, (q,))]) + expected = weights.dot(numpy.einsum("ijkp,ijk->p", numpy.asarray(table.children[0].array), dofs) ** 3) + assert numpy.allclose(result.arr, expected)