Allow Expressions with submeshes

ffcx/ir/representation.py

@@ -596,12 +596,10 @@ def _compute_expression_ir(
     points = expr[1]
     expr = expr[0]
 
-    try:
-        cell = ufl.domain.extract_unique_domain(expr).ufl_cell()
-    except AttributeError:
-        # This case corresponds to a spatially constant expression
-        # without any dependencies
-        cell = None
+    expr_domain = max(
+        ufl.domain.extract_domains(expr), default=None, key=lambda d: d.topological_dimension
+    )
+    cell = expr_domain.ufl_cell() if expr_domain else None
 
     # Prepare dimensions of all unique element in expression, including
     # elements for arguments, coefficients and coordinate mappings
@@ -684,8 +682,6 @@ def _compute_expression_ir(
         _offset += np.prod(constant.ufl_shape, dtype=int)
 
     base_ir["original_constant_offsets"] = original_constant_offsets
-
-    expr_domain = ufl.domain.extract_unique_domain(expr)
     base_ir["coordinate_element_hash"] = (
         expr_domain.ufl_coordinate_element().basix_hash() if expr_domain is not None else 0
     )

test/test_jit_expression.py

@@ -482,3 +482,70 @@ def test_coordinate_free_expression(compile_args):
     # Check that the expression evaluates to [1, 0] at all points
     expected = np.array([[1.0, 0.0], [1.0, 0.0], [1.0, 0.0]])
     assert np.allclose(A, expected)
+
+
+def test_mixed_mesh_expression(compile_args):
+    """Test facet expression containing quantities from parent and facet mesh."""
+    c_el = basix.ufl.element("Lagrange", "triangle", 1, shape=(2,))
+    mesh = ufl.Mesh(c_el)
+
+    f_c_el = basix.ufl.element("Lagrange", "interval", 1, shape=(2,))
+    facet_mesh = ufl.Mesh(f_c_el)
+    f_el = basix.ufl.element("P", "interval", 1)
+    V = ufl.FunctionSpace(facet_mesh, f_el)
+    c = ufl.Coefficient(V)
+
+    n = ufl.FacetNormal(mesh)
+    expr = c * n
+
+    dtype = np.float64
+    points = np.array([[0.3], [0.5], [0.8]], dtype=dtype)
+
+    obj, _, _ = ffcx.codegeneration.jit.compile_expressions(
+        [(expr, points)], cffi_extra_compile_args=compile_args
+    )
+
+    ffi = cffi.FFI()
+    expression = obj[0]
+
+    c_type = "double"
+    c_xtype = "double"
+
+    output = np.zeros(points.shape[0] * 2, dtype=dtype)
+
+    # Define mesh
+    coords = np.array([[0.1, 0.3, 0], [2, 0, 0.0], [0, 1, 0.0]], dtype=dtype)
+
+    # Define exact normals
+    face_normal = np.array([0.0, 0.0, 1.0])
+    e0 = coords[2] - coords[1]
+    e1 = coords[0] - coords[2]
+    e2 = coords[1] - coords[0]
+    edges = np.array([e0, e1, e2])
+    edge_normals = np.cross(edges, face_normal)[:, :2]
+    norms = np.linalg.norm(edge_normals, axis=1, keepdims=True)
+    edge_normals_normalized = edge_normals / norms
+
+    # Pack coefficients, constants, and other data for expression evaluation
+    u_coeffs = np.array([0.1, 0.5], dtype=dtype)
+    consts = np.array([], dtype=dtype)
+    entity_index = np.array([0], dtype=np.intc)
+    quad_perm = np.array([0], dtype=np.uint8)
+
+    ref_coeff = u_coeffs[0] + points[:, 0] * (u_coeffs[1] - u_coeffs[0])
+    for i, normal in enumerate(edge_normals_normalized):
+        # Tabulate facet normal
+        output[:] = 0
+        entity_index[0] = i
+        expression.tabulate_tensor_float64(
+            ffi.cast(f"{c_type} *", output.ctypes.data),
+            ffi.cast(f"{c_type} *", u_coeffs.ctypes.data),
+            ffi.cast(f"{c_type} *", consts.ctypes.data),
+            ffi.cast(f"{c_xtype} *", coords.ctypes.data),
+            ffi.cast("int *", entity_index.ctypes.data),
+            ffi.cast("uint8_t *", quad_perm.ctypes.data),
+            ffi.NULL,
+        )
+        ref_sol = ref_coeff.reshape(-1, 1) * normal
+
+        np.testing.assert_allclose(output, ref_sol.flatten())