Test interpolation of products of coefficient evaluations

.github/actions/install/action.yml

@@ -163,6 +163,14 @@ runs:
         firedrake-clean
         pip list
 
+    - name: "DROP BEFORE MERGE: install FIAT from firedrakeproject/fiat#280"
+      shell: bash
+      run: |
+        . venv/bin/activate
+        pip install --verbose --no-build-isolation --no-deps --force-reinstall \
+          git+https://github.com/firedrakeproject/fiat.git@pbrubeck/atomic-contraction
+        pip list | grep -i fiat
+
     - name: Run firedrake-check
       shell: bash
       run: |

tests/firedrake/regression/test_interpolate.py

@@ -785,3 +785,46 @@ def test_interpolate_indexed():
     I1 = assemble(interpolate(u2, U), mat_type="nest")
     I1_block = assemble(interpolate(TrialFunction(U), U))
     assert np.allclose(I1.petscmat.getNestSubMatrix(0, 1)[:, :], I1_block.petscmat[:, :])
+
+
+@pytest.fixture
+def hexmesh():
+    return ExtrudedMesh(UnitSquareMesh(1, 1, quadrilateral=True), 1)
+
+
+@pytest.mark.parametrize("source,target,expr,expected", [
+    (FunctionSpace, FunctionSpace,
+     lambda f: f * f * f,
+     lambda f: f ** 3),
+    (FunctionSpace, FunctionSpace,
+     lambda f: f * f * f * f,
+     lambda f: f ** 4),
+    (VectorFunctionSpace, VectorFunctionSpace,
+     lambda f: dot(f, f) * f,
+     lambda f: np.einsum("...i,...i,...j->...j", f, f, f)),
+    (TensorFunctionSpace, TensorFunctionSpace,
+     lambda f: dot(f, f),
+     lambda f: np.einsum("...ij,...jk->...ik", f, f)),
+    (TensorFunctionSpace, FunctionSpace,
+     lambda f: inner(f, f),
+     lambda f: np.einsum("...ij,...ij->...", f, f)),
+], ids=["cube", "quartic", "vector", "tensor", "inner"])
+def test_interpolate_too_many_indices(hexmesh, source, target, expr, expected):
+    """Test that products of several coefficient evaluations
+    do not exceed the sum-factorisation index limit.
+
+    Internally, sum_factorise breaks the contraction into
+    independent subproblems to avoid a monolithic contraction
+    with too many indices.
+
+    Evaluations that a value index contracts together do not split
+    into independent subproblems, and are instead kept whole.
+    """
+    V = source(hexmesh, "CG", 1)
+    W = target(hexmesh, "CG", 1)
+
+    w = Function(V)
+    w.dat.data[...] = np.arange(1, w.dat.data.size + 1).reshape(w.dat.data.shape)
+    u = Function(W).interpolate(expr(w))
+
+    assert np.allclose(u.dat.data_ro, expected(w.dat.data_ro))