Fix cell_avg + MixedFunctionSpace

tests/firedrake/regression/test_assemble.py

@@ -415,3 +415,41 @@ def test_assemble_tensor_empty_shape(mesh):
     v = Function(V).assign(1)
     expected = assemble(inner(v, v)*dx)
     assert np.allclose(result, expected)
+
+
+@pytest.mark.parametrize("coefficient", [False, True], ids=["Expr", "Function"])
+def test_cell_avg(coefficient):
+    mesh = UnitSquareMesh(3, 3)
+    x = SpatialCoordinate(mesh)
+    expr = dot(x, x) ** 2
+    if coefficient:
+        V = FunctionSpace(mesh, "CG", 4)
+        expr = Function(V).interpolate(expr)
+
+    result = assemble(inner(cell_avg(expr), expr) * dx)
+
+    Q = FunctionSpace(mesh, "DG", 0)
+    p = Function(Q)
+    p.project(expr)
+    expect = assemble(inner(p, expr) * dx)
+    assert np.isclose(result, expect)
+
+
+def test_cell_avg_mfs():
+    mesh = UnitSquareMesh(3, 3)
+    V = VectorFunctionSpace(mesh, "CG", 2)
+    Q = FunctionSpace(mesh, "DG", 0)
+    Z = MixedFunctionSpace([V, Q])
+    z = Function(Z)
+    usub, psub = z.subfunctions
+    usub.interpolate(SpatialCoordinate(mesh))
+    psub.interpolate(Constant(1))
+
+    expect = 6
+    result1 = assemble(inner(cell_avg(div(usub)), psub + div(usub))*dx)
+    assert np.isclose(result1, expect)
+
+    # This fails if do_replace_functions=True in entity_avg in tscf/ufl_utils.py
+    u, p = split(z)
+    result2 = assemble(inner(cell_avg(div(u)), p + div(u))*dx)
+    assert np.isclose(result2, expect)

tsfc/ufl_utils.py

@@ -43,6 +43,7 @@ def compute_form_data(form,
                       do_apply_default_restrictions=True,
                       do_apply_restrictions=True,
                       do_estimate_degrees=True,
+                      do_replace_functions=True,
                       coefficients_to_split=None,
                       complex_mode=False):
     """Preprocess UFL form in a format suitable for TSFC. Return
@@ -64,7 +65,7 @@ def compute_form_data(form,
         do_apply_default_restrictions=do_apply_default_restrictions,
         do_apply_restrictions=do_apply_restrictions,
         do_estimate_degrees=do_estimate_degrees,
-        do_replace_functions=True,
+        do_replace_functions=do_replace_functions,
         coefficients_to_split=coefficients_to_split,
         complex_mode=complex_mode,
         do_remove_component_tensors=do_remove_component_tensors,
@@ -112,7 +113,9 @@ def entity_avg(integrand, measure, argument_multiindices):
     degree = estimate_total_polynomial_degree(integrand)
     form = integrand * measure
     fd = compute_form_data(form, do_estimate_degrees=False,
-                           do_apply_function_pullbacks=False)
+                           do_apply_function_pullbacks=False,
+                           do_replace_functions=False,
+                           )
     itg_data, = fd.integral_data
     integral, = itg_data.integrals
     integrand = integral.integrand()