Keep factorised contractions atomic in dual evaluation

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):

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

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,

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)

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)