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,8 @@ 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
+        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 +278,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,8 @@ 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
+        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 +201,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

@@ -382,25 +382,60 @@ def count(pair):
     return result, flops
 
 
-def sum_factorise(sum_indices, factors):
-    """Optimise a tensor product through sum factorisation.
+def _independent_contractions(sum_indices, groups):
+    """Split a contraction into independent subproblems.
+
+    Two contraction indices only interact if some factor carries both of
+    them, so the factors and the indices form a graph whose connected
+    components can be contracted separately.
 
     :arg sum_indices: free indices for contractions
-    :arg factors: product factors
-    :returns: optimised GEM expression
+    :arg groups: product factors, grouped by free indices
+    :returns: a pair of the list of (indices, groups) subproblems and the
+              list of groups carrying no contraction index
     """
-    if len(factors) == 0 and len(sum_indices) == 0:
-        # Empty product
-        return one
+    # Union-find over the contraction indices
+    parent = {index: index for index in sum_indices}
+
+    def find(index):
+        if parent[index] == index:
+            return index
+
+        parent[index] = find(parent[index])
+        return parent[index]
+
+    index_set = set(sum_indices)
+    shared = [[i for i in group.free_indices if i in index_set] for group in groups]
+    for indices in shared:
+        for index in indices[1:]:
+            root, other = find(indices[0]), find(index)
+            if root != other:
+                parent[other] = root
+
+    subproblems = OrderedDict((find(index), ([], [])) for index in sum_indices)
+    for index in sum_indices:
+        subproblems[find(index)][0].append(index)
+
+    rest = []
+    for group, indices in zip(groups, shared):
+        if indices:
+            subproblems[find(indices[0])][1].append(group)
+        else:
+            rest.append(group)
+    return list(subproblems.values()), rest
+
+
+def _sum_factorise_connected(sum_indices, groups):
+    """Sum factorise a single connected contraction by exhaustive search.
 
+    :arg sum_indices: free indices for contractions, which must not split
+                      into independent subproblems
+    :arg groups: product factors, grouped by free indices
+    :returns: optimised GEM expression
+    """
     if len(sum_indices) > 6:
         raise NotImplementedError("Too many indices for sum factorisation!")
 
-    # Form groups by free indices
-    groups = groupby(factors, key=lambda f: f.free_indices)
-    groups = [Product(*terms) for _, terms in groups]
-
-    # Sum factorisation
     expression = None
     best_flops = numpy.inf
 
@@ -435,6 +470,33 @@ def sum_factorise(sum_indices, factors):
     return expression
 
 
+def sum_factorise(sum_indices, factors):
+    """Optimise a tensor product through sum factorisation.
+
+    :arg sum_indices: free indices for contractions
+    :arg factors: product factors
+    :returns: optimised GEM expression
+    """
+    if len(factors) == 0 and len(sum_indices) == 0:
+        # Empty product
+        return one
+
+    # Form groups by free indices
+    groups = groupby(factors, key=lambda f: f.free_indices)
+    groups = [Product(*terms) for _, terms in groups]
+
+    # Contractions that share no factor are independent of each other, so
+    # factorise them separately rather than searching the orderings that
+    # interleave them.
+    subproblems, terms = _independent_contractions(sum_indices, groups)
+    terms = terms + [_sum_factorise_connected(indices, subgroups)
+                     for indices, subgroups in subproblems]
+    if not terms:
+        return one
+    expression, _ = associate(Product, terms)
+    return expression
+
+
 def make_sum(summands):
     """Constructs an operation-minimal sum of GEM expressions."""
     groups = groupby(summands, key=lambda f: f.free_indices)
@@ -568,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.
@@ -592,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)